agenta 0.13.0a0__py3-none-any.whl → 0.13.0a1__py3-none-any.whl

agenta/client/backend/types/simple_evaluation_output.py

@@ -20,11 +20,19 @@ class SimpleEvaluationOutput(pydantic.BaseModel):
     evaluation_type: EvaluationType
 
     def json(self, **kwargs: typing.Any) -> str:
-        kwargs_with_defaults: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
+        kwargs_with_defaults: typing.Any = {
+            "by_alias": True,
+            "exclude_unset": True,
+            **kwargs,
+        }
         return super().json(**kwargs_with_defaults)
 
     def dict(self, **kwargs: typing.Any) -> typing.Dict[str, typing.Any]:
-        kwargs_with_defaults: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
+        kwargs_with_defaults: typing.Any = {
+            "by_alias": True,
+            "exclude_unset": True,
+            **kwargs,
+        }
         return super().dict(**kwargs_with_defaults)
 
     class Config:

agenta/client/backend/types/span.py

@@ -30,11 +30,19 @@ class Span(pydantic.BaseModel):
     children: typing.Optional[typing.List[Span]]
 
     def json(self, **kwargs: typing.Any) -> str:
-        kwargs_with_defaults: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
+        kwargs_with_defaults: typing.Any = {
+            "by_alias": True,
+            "exclude_unset": True,
+            **kwargs,
+        }
         return super().json(**kwargs_with_defaults)
 
     def dict(self, **kwargs: typing.Any) -> typing.Dict[str, typing.Any]:
-        kwargs_with_defaults: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
+        kwargs_with_defaults: typing.Any = {
+            "by_alias": True,
+            "exclude_unset": True,
+            **kwargs,
+        }
         return super().dict(**kwargs_with_defaults)
 
     class Config:

agenta/client/backend/types/span_detail.py

@@ -31,11 +31,19 @@ class SpanDetail(pydantic.BaseModel):
     config: typing.Optional[typing.Dict[str, typing.Any]]
 
     def json(self, **kwargs: typing.Any) -> str:
-        kwargs_with_defaults: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
+        kwargs_with_defaults: typing.Any = {
+            "by_alias": True,
+            "exclude_unset": True,
+            **kwargs,
+        }
         return super().json(**kwargs_with_defaults)
 
     def dict(self, **kwargs: typing.Any) -> typing.Dict[str, typing.Any]:
-        kwargs_with_defaults: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
+        kwargs_with_defaults: typing.Any = {
+            "by_alias": True,
+            "exclude_unset": True,
+            **kwargs,
+        }
         return super().dict(**kwargs_with_defaults)
 
     class Config:

agenta/client/backend/types/span_variant.py

@@ -17,11 +17,19 @@ class SpanVariant(pydantic.BaseModel):
     revision: int
 
     def json(self, **kwargs: typing.Any) -> str:
-        kwargs_with_defaults: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
+        kwargs_with_defaults: typing.Any = {
+            "by_alias": True,
+            "exclude_unset": True,
+            **kwargs,
+        }
         return super().json(**kwargs_with_defaults)
 
     def dict(self, **kwargs: typing.Any) -> typing.Dict[str, typing.Any]:
-        kwargs_with_defaults: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
+        kwargs_with_defaults: typing.Any = {
+            "by_alias": True,
+            "exclude_unset": True,
+            **kwargs,
+        }
         return super().dict(**kwargs_with_defaults)
 
     class Config:

agenta/client/backend/types/template.py

@@ -17,11 +17,19 @@ class Template(pydantic.BaseModel):
     image: TemplateImageInfo
 
     def json(self, **kwargs: typing.Any) -> str:
-        kwargs_with_defaults: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
+        kwargs_with_defaults: typing.Any = {
+            "by_alias": True,
+            "exclude_unset": True,
+            **kwargs,
+        }
         return super().json(**kwargs_with_defaults)
 
     def dict(self, **kwargs: typing.Any) -> typing.Dict[str, typing.Any]:
-        kwargs_with_defaults: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
+        kwargs_with_defaults: typing.Any = {
+            "by_alias": True,
+            "exclude_unset": True,
+            **kwargs,
+        }
         return super().dict(**kwargs_with_defaults)
 
     class Config:

agenta/client/backend/types/template_image_info.py

@@ -22,11 +22,19 @@ class TemplateImageInfo(pydantic.BaseModel):
     template_uri: typing.Optional[str]
 
     def json(self, **kwargs: typing.Any) -> str:
-        kwargs_with_defaults: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
+        kwargs_with_defaults: typing.Any = {
+            "by_alias": True,
+            "exclude_unset": True,
+            **kwargs,
+        }
         return super().json(**kwargs_with_defaults)
 
     def dict(self, **kwargs: typing.Any) -> typing.Dict[str, typing.Any]:
-        kwargs_with_defaults: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
+        kwargs_with_defaults: typing.Any = {
+            "by_alias": True,
+            "exclude_unset": True,
+            **kwargs,
+        }
         return super().dict(**kwargs_with_defaults)
 
     class Config:

agenta/client/backend/types/test_set_output_response.py

@@ -17,11 +17,19 @@ class TestSetOutputResponse(pydantic.BaseModel):
     created_at: dt.datetime
 
     def json(self, **kwargs: typing.Any) -> str:
-        kwargs_with_defaults: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
+        kwargs_with_defaults: typing.Any = {
+            "by_alias": True,
+            "exclude_unset": True,
+            **kwargs,
+        }
         return super().json(**kwargs_with_defaults)
 
     def dict(self, **kwargs: typing.Any) -> typing.Dict[str, typing.Any]:
-        kwargs_with_defaults: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
+        kwargs_with_defaults: typing.Any = {
+            "by_alias": True,
+            "exclude_unset": True,
+            **kwargs,
+        }
         return super().dict(**kwargs_with_defaults)
 
     class Config:

agenta/client/backend/types/test_set_simple_response.py

@@ -17,11 +17,19 @@ class TestSetSimpleResponse(pydantic.BaseModel):
     created_at: str
 
     def json(self, **kwargs: typing.Any) -> str:
-        kwargs_with_defaults: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
+        kwargs_with_defaults: typing.Any = {
+            "by_alias": True,
+            "exclude_unset": True,
+            **kwargs,
+        }
         return super().json(**kwargs_with_defaults)
 
     def dict(self, **kwargs: typing.Any) -> typing.Dict[str, typing.Any]:
-        kwargs_with_defaults: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
+        kwargs_with_defaults: typing.Any = {
+            "by_alias": True,
+            "exclude_unset": True,
+            **kwargs,
+        }
         return super().dict(**kwargs_with_defaults)
 
     class Config:

agenta/client/backend/types/trace_detail.py

@@ -31,11 +31,19 @@ class TraceDetail(pydantic.BaseModel):
     config: typing.Dict[str, typing.Any]
 
     def json(self, **kwargs: typing.Any) -> str:
-        kwargs_with_defaults: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
+        kwargs_with_defaults: typing.Any = {
+            "by_alias": True,
+            "exclude_unset": True,
+            **kwargs,
+        }
         return super().json(**kwargs_with_defaults)
 
     def dict(self, **kwargs: typing.Any) -> typing.Dict[str, typing.Any]:
-        kwargs_with_defaults: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
+        kwargs_with_defaults: typing.Any = {
+            "by_alias": True,
+            "exclude_unset": True,
+            **kwargs,
+        }
         return super().dict(**kwargs_with_defaults)
 
     class Config:

agenta/client/backend/types/uri.py

@@ -15,11 +15,19 @@ class Uri(pydantic.BaseModel):
     uri: str
 
     def json(self, **kwargs: typing.Any) -> str:
-        kwargs_with_defaults: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
+        kwargs_with_defaults: typing.Any = {
+            "by_alias": True,
+            "exclude_unset": True,
+            **kwargs,
+        }
         return super().json(**kwargs_with_defaults)
 
     def dict(self, **kwargs: typing.Any) -> typing.Dict[str, typing.Any]:
-        kwargs_with_defaults: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
+        kwargs_with_defaults: typing.Any = {
+            "by_alias": True,
+            "exclude_unset": True,
+            **kwargs,
+        }
         return super().dict(**kwargs_with_defaults)
 
     class Config:

agenta/client/backend/types/validation_error.py

@@ -18,11 +18,19 @@ class ValidationError(pydantic.BaseModel):
     type: str
 
     def json(self, **kwargs: typing.Any) -> str:
-        kwargs_with_defaults: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
+        kwargs_with_defaults: typing.Any = {
+            "by_alias": True,
+            "exclude_unset": True,
+            **kwargs,
+        }
         return super().json(**kwargs_with_defaults)
 
     def dict(self, **kwargs: typing.Any) -> typing.Dict[str, typing.Any]:
-        kwargs_with_defaults: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
+        kwargs_with_defaults: typing.Any = {
+            "by_alias": True,
+            "exclude_unset": True,
+            **kwargs,
+        }
         return super().dict(**kwargs_with_defaults)
 
     class Config:

agenta/client/backend/types/variant_action.py

@@ -16,11 +16,19 @@ class VariantAction(pydantic.BaseModel):
     action: VariantActionEnum
 
     def json(self, **kwargs: typing.Any) -> str:
-        kwargs_with_defaults: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
+        kwargs_with_defaults: typing.Any = {
+            "by_alias": True,
+            "exclude_unset": True,
+            **kwargs,
+        }
         return super().json(**kwargs_with_defaults)
 
     def dict(self, **kwargs: typing.Any) -> typing.Dict[str, typing.Any]:
-        kwargs_with_defaults: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
+        kwargs_with_defaults: typing.Any = {
+            "by_alias": True,
+            "exclude_unset": True,
+            **kwargs,
+        }
         return super().dict(**kwargs_with_defaults)
 
     class Config:

agenta/client/backend/types/variant_action_enum.py

@@ -14,7 +14,9 @@ class VariantActionEnum(str, enum.Enum):
     START = "START"
     STOP = "STOP"
 
-    def visit(self, start: typing.Callable[[], T_Result], stop: typing.Callable[[], T_Result]) -> T_Result:
+    def visit(
+        self, start: typing.Callable[[], T_Result], stop: typing.Callable[[], T_Result]
+    ) -> T_Result:
         if self is VariantActionEnum.START:
             return start()
         if self is VariantActionEnum.STOP:

agenta/client/backend/types/with_pagination.py

@@ -18,11 +18,19 @@ class WithPagination(pydantic.BaseModel):
     page_size: int = pydantic.Field(alias="pageSize")
 
     def json(self, **kwargs: typing.Any) -> str:
-        kwargs_with_defaults: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
+        kwargs_with_defaults: typing.Any = {
+            "by_alias": True,
+            "exclude_unset": True,
+            **kwargs,
+        }
         return super().json(**kwargs_with_defaults)
 
     def dict(self, **kwargs: typing.Any) -> typing.Dict[str, typing.Any]:
-        kwargs_with_defaults: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
+        kwargs_with_defaults: typing.Any = {
+            "by_alias": True,
+            "exclude_unset": True,
+            **kwargs,
+        }
         return super().dict(**kwargs_with_defaults)
 
     class Config:

agenta/client/backend/types/workspace_member_response.py

@@ -17,11 +17,19 @@ class WorkspaceMemberResponse(pydantic.BaseModel):
     roles: typing.List[WorkspacePermission]
 
     def json(self, **kwargs: typing.Any) -> str:
-        kwargs_with_defaults: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
+        kwargs_with_defaults: typing.Any = {
+            "by_alias": True,
+            "exclude_unset": True,
+            **kwargs,
+        }
         return super().json(**kwargs_with_defaults)
 
     def dict(self, **kwargs: typing.Any) -> typing.Dict[str, typing.Any]:
-        kwargs_with_defaults: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
+        kwargs_with_defaults: typing.Any = {
+            "by_alias": True,
+            "exclude_unset": True,
+            **kwargs,
+        }
         return super().dict(**kwargs_with_defaults)
 
     class Config:

agenta/client/backend/types/workspace_permission.py

@@ -19,11 +19,19 @@ class WorkspacePermission(pydantic.BaseModel):
     permissions: typing.Optional[typing.List[Permission]]
 
     def json(self, **kwargs: typing.Any) -> str:
-        kwargs_with_defaults: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
+        kwargs_with_defaults: typing.Any = {
+            "by_alias": True,
+            "exclude_unset": True,
+            **kwargs,
+        }
         return super().json(**kwargs_with_defaults)
 
     def dict(self, **kwargs: typing.Any) -> typing.Dict[str, typing.Any]:
-        kwargs_with_defaults: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
+        kwargs_with_defaults: typing.Any = {
+            "by_alias": True,
+            "exclude_unset": True,
+            **kwargs,
+        }
         return super().dict(**kwargs_with_defaults)
 
     class Config:

agenta/client/backend/types/workspace_response.py

@@ -23,11 +23,19 @@ class WorkspaceResponse(pydantic.BaseModel):
     members: typing.Optional[typing.List[WorkspaceMemberResponse]]
 
     def json(self, **kwargs: typing.Any) -> str:
-        kwargs_with_defaults: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
+        kwargs_with_defaults: typing.Any = {
+            "by_alias": True,
+            "exclude_unset": True,
+            **kwargs,
+        }
         return super().json(**kwargs_with_defaults)
 
     def dict(self, **kwargs: typing.Any) -> typing.Dict[str, typing.Any]:
-        kwargs_with_defaults: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
+        kwargs_with_defaults: typing.Any = {
+            "by_alias": True,
+            "exclude_unset": True,
+            **kwargs,
+        }
         return super().dict(**kwargs_with_defaults)
 
     class Config:

agenta/client/backend/types/workspace_role_response.py

@@ -17,11 +17,19 @@ class WorkspaceRoleResponse(pydantic.BaseModel):
     role_description: typing.Optional[str]
 
     def json(self, **kwargs: typing.Any) -> str:
-        kwargs_with_defaults: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
+        kwargs_with_defaults: typing.Any = {
+            "by_alias": True,
+            "exclude_unset": True,
+            **kwargs,
+        }
         return super().json(**kwargs_with_defaults)
 
     def dict(self, **kwargs: typing.Any) -> typing.Dict[str, typing.Any]:
-        kwargs_with_defaults: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
+        kwargs_with_defaults: typing.Any = {
+            "by_alias": True,
+            "exclude_unset": True,
+            **kwargs,
+        }
         return super().dict(**kwargs_with_defaults)
 
     class Config:

agenta/sdk/__init__.py

@@ -9,6 +9,7 @@ from .types import (
     InFile,
     IntParam,
     MultipleChoiceParam,
+    GroupedMultipleChoiceParam,
     TextParam,
     MessagesInput,
     FileInputURL,

agenta/sdk/agenta_decorator.py

@@ -25,6 +25,7 @@ from .types import (
     InFile,
     IntParam,
     MultipleChoiceParam,
+    GroupedMultipleChoiceParam,
     TextParam,
     MessagesInput,
     FileInputURL,
@@ -83,7 +84,8 @@ def entrypoint(func: Callable[..., Any]) -> Callable[..., Any]:
         tracing.start_parent_span(
             name=func.__name__,
             inputs=func_params,
-            config=api_config_params,
+            config=config_params,
+            environment="playground",  # type: ignore #NOTE: wrapper is only called in playground
         )
 
         # Ingest files, prepare configurations and run llm app
@@ -97,7 +99,6 @@ def entrypoint(func: Callable[..., Any]) -> Callable[..., Any]:
         tracing.end_recording(
             outputs=llm_result.dict(),
             span=tracing.active_trace,
-            environment="playground",  # type: ignore #NOTE: wrapper is only called in playground
         )
         return llm_result
 
@@ -120,6 +121,7 @@ def entrypoint(func: Callable[..., Any]) -> Callable[..., Any]:
             name=func.__name__,
             inputs=func_params,
             config=config,
+            environment="playground",  # type: ignore #NOTE: wrapper is only called in playground
         )
 
         llm_result = await execute_function(
@@ -130,7 +132,6 @@ def entrypoint(func: Callable[..., Any]) -> Callable[..., Any]:
         tracing.end_recording(
             outputs=llm_result.dict(),
             span=tracing.active_trace,
-            environment="playground",  # type: ignore #NOTE: wrapper is only called in playground
         )
         return llm_result
 
@@ -446,6 +447,13 @@ def override_schema(openapi_schema: dict, func_name: str, endpoint: str, params:
         f"Body_{func_name}_{endpoint}_post"
     ]["properties"]
     for param_name, param_val in params.items():
+        if isinstance(param_val, GroupedMultipleChoiceParam):
+            subschema = find_in_schema(schema_to_override, param_name, "grouped_choice")
+            assert (
+                subschema
+            ), f"GroupedMultipleChoiceParam '{param_name}' is in the parameters but could not be found in the openapi.json"
+            subschema["choices"] = param_val.choices
+            subschema["default"] = param_val.default
         if isinstance(param_val, MultipleChoiceParam):
             subschema = find_in_schema(schema_to_override, param_name, "choice")
             default = str(param_val)

agenta/sdk/tracing/llm_tracing.py

@@ -92,7 +92,7 @@ class Tracing(object):
         self.tags.extend(tags)
 
     def start_parent_span(
-        self, name: str, inputs: Dict[str, Any], config: Dict[str, Any]
+        self, name: str, inputs: Dict[str, Any], config: Dict[str, Any], **kwargs
     ):
         trace_id = self._create_trace_id()
         span_id = self._create_span_id()
@@ -104,6 +104,7 @@ class Tracing(object):
             inputs=inputs,
             name=name,
             config=config,
+            environment=kwargs.get("environment"),
             spankind=SpanKind.WORKFLOW.value,
             status=SpanStatusCode.UNSET.value,
             start_time=datetime.now(timezone.utc),
@@ -111,7 +112,9 @@ class Tracing(object):
         self.active_trace = span
         self.recording_trace_id = trace_id
         self.parent_span_id = span.id
-        self.llm_logger.info(f"Recorded active_trace and parent_span_id: {span.id}")
+        self.llm_logger.info(
+            f"Recorded active_trace and setting parent_span_id: {span.id}"
+        )
 
     def start_span(
         self,
@@ -126,7 +129,10 @@ class Tracing(object):
             id=span_id,
             inputs=input,
             name=name,
+            app_id=self.app_id,
+            variant_id=self.variant_id,
             config=config,
+            environment=self.active_trace.environment,
             parent_span_id=self.parent_span_id,
             spankind=spankind.upper(),
             attributes={},
@@ -137,7 +143,9 @@ class Tracing(object):
         self.active_span = span
         self.span_dict[span.id] = span
         self.parent_span_id = span.id
-        self.llm_logger.info(f"Recorded active_span and parent_span_id: {span.id}")
+        self.llm_logger.info(
+            f"Recorded active_span and setting parent_span_id: {span.id}"
+        )
         return span
 
     def update_span_status(self, span: CreateSpan, value: str):
@@ -150,39 +158,34 @@ class Tracing(object):
             end_time=datetime.now(timezone.utc),
             outputs=[outputs["message"]],
             cost=outputs.get("cost", None),
-            environment=kwargs.get("environment"),
             tokens=outputs.get("usage"),
         )
 
         # Push span to list of recorded spans
         self.recorded_spans.append(updated_span)
-        self.llm_logger.info(f"Pushed {updated_span.spankind} span {updated_span.id} to recorded spans.")
-
-    def end_recording(self, outputs: Dict[str, Any], span: CreateSpan, **kwargs):
-        updated_span = CreateSpan(
-            **span.dict(),
-            end_time=datetime.now(timezone.utc),
-            outputs=[outputs["message"]],
-            cost=outputs.get("cost", None),
-            environment=kwargs.get("environment"),
-            tokens=outputs.get("usage"),
-        )
-        self.recorded_spans.append(updated_span)
         self.llm_logger.info(
-            f"Pushed workflow span {updated_span.id} to recorded spans."
+            f"Pushed {updated_span.spankind} span {updated_span.id} to recorded spans."
         )
+
+    def end_recording(self, outputs: Dict[str, Any], span: CreateSpan, **kwargs):
+        self.end_span(outputs=outputs, span=span, **kwargs)
+        if self.api_key == "":
+            return
+
         self.llm_logger.info(f"Preparing to send recorded spans for processing.")
+        self.llm_logger.info(f"Recorded spans => {len(self.recorded_spans)}")
         self.tasks_manager.add_task(
             self.active_trace.id,
             "trace",
             self.client.create_traces(
-                trace=self.recording_trace_id, spans=self.recorded_spans
+                trace=self.recording_trace_id, spans=self.recorded_spans  # type: ignore
             ),
             self.client,
         )
         self.llm_logger.info(
             f"Tracing for {span.id} recorded successfully and sent for processing."
         )
+        self._clear_recorded_spans()
 
     def _create_trace_id(self) -> str:
         """Creates a unique mongo id for the trace object.
@@ -201,3 +204,13 @@ class Tracing(object):
         """
 
         return str(ObjectId())
+
+    def _clear_recorded_spans(self) -> None:
+        """
+        Clear the list of recorded spans to prepare for next batch processing.
+        """
+
+        self.recorded_spans = []
+        self.llm_logger.info(
+            f"Cleared all recorded spans from batch: {self.recorded_spans}"
+        )

agenta/sdk/types.py

@@ -127,6 +127,44 @@ class MultipleChoiceParam(str):
         )
 
 
+class GroupedMultipleChoiceParam(str):
+    def __new__(cls, default: str = None, choices: Dict[str, List[str]] = None):
+        if choices is None:
+            choices = {}
+
+        if default and not any(default in choices for choices in choices.values()):
+            if not choices:
+                print(
+                    f"Warning: Default value {default} provided but choices are empty."
+                )
+            else:
+                raise ValueError(
+                    f"Default value {default} is not in the provided choices"
+                )
+
+        if not default:
+            for choices in choices.values():
+                if choices:
+                    default = choices[0]
+                    break
+
+        instance = super().__new__(cls, default)
+        instance.choices = choices
+        instance.default = default
+        return instance
+
+    @classmethod
+    def __modify_schema__(cls, field_schema: dict[str, Any], **kwargs):
+        choices = kwargs.get("choices", {})
+        field_schema.update(
+            {
+                "x-parameter": "grouped_choice",
+                "type": "string",
+                "choices": choices,
+            }
+        )
+
+
 class Message(BaseModel):
     role: str
     content: str