agenta 0.32.0a1__py3-none-any.whl → 0.33.0__py3-none-any.whl

agenta/sdk/decorators/routing.py

@@ -1,4 +1,3 @@
-from sre_parse import NOT_LITERAL_UNI_IGNORE
 from typing import Type, Any, Callable, Dict, Optional, Tuple, List
 from inspect import signature, iscoroutinefunction, Signature, Parameter
 from functools import wraps
@@ -6,6 +5,7 @@ from traceback import format_exception
 from asyncio import sleep
 from uuid import UUID
 from pydantic import BaseModel, HttpUrl, ValidationError
+from os import environ
 
 from fastapi import Body, FastAPI, HTTPException, Request
 
@@ -38,12 +38,18 @@ from agenta.sdk.types import (
 
 import agenta as ag
 
-
-app = FastAPI()
 log.setLevel("DEBUG")
 
+AGENTA_RUNTIME_PREFIX = environ.get("AGENTA_RUNTIME_PREFIX", "")
+
+app = FastAPI(
+    docs_url=f"{AGENTA_RUNTIME_PREFIX}/docs",  # Swagger UI
+    openapi_url=f"{AGENTA_RUNTIME_PREFIX}/openapi.json",  # OpenAPI schema
+)
 
-app.include_router(router, prefix="")
+app.include_router(router, prefix=AGENTA_RUNTIME_PREFIX)
+
+log.error("Agenta - Runtime Prefix:" + AGENTA_RUNTIME_PREFIX)
 
 
 class PathValidator(BaseModel):
@@ -114,13 +120,13 @@ class entrypoint:
     routes = list()
 
     _middleware = False
-    _run_path = "/run"
-    _test_path = "/test"
+    _run_path = f"{AGENTA_RUNTIME_PREFIX}/run"
+    _test_path = f"{AGENTA_RUNTIME_PREFIX}/test"
     _config_key = "ag_config"
     # LEGACY
-    _legacy_playground_run_path = "/playground/run"
-    _legacy_generate_path = "/generate"
-    _legacy_generate_deployed_path = "/generate_deployed"
+    _legacy_playground_run_path = f"{AGENTA_RUNTIME_PREFIX}/playground/run"
+    _legacy_generate_path = f"{AGENTA_RUNTIME_PREFIX}/generate"
+    _legacy_generate_deployed_path = f"{AGENTA_RUNTIME_PREFIX}/generate_deployed"
 
     def __init__(
         self,
@@ -142,8 +148,8 @@ class entrypoint:
             app.add_middleware(InlineMiddleware)
             app.add_middleware(VaultMiddleware)
             app.add_middleware(ConfigMiddleware)
-            app.add_middleware(OTelMiddleware)
             app.add_middleware(AuthMiddleware)
+            app.add_middleware(OTelMiddleware)
             app.add_middleware(CORSMiddleware)
         ### ------------------ #
 
@@ -175,14 +181,22 @@ class entrypoint:
         self.update_run_wrapper_signature(wrapper=run_wrapper)
 
         run_route = f"{entrypoint._run_path}{route_path}"
-        app.post(run_route, response_model=BaseResponse)(run_wrapper)
+        app.post(
+            run_route,
+            response_model=BaseResponse,
+            response_model_exclude_none=True,
+        )(run_wrapper)
 
         # LEGACY
         # TODO: Removing this implies breaking changes in :
         # - calls to /generate_deployed must be replaced with calls to /run
         if route_path == "":
             run_route = entrypoint._legacy_generate_deployed_path
-            app.post(run_route, response_model=BaseResponse)(run_wrapper)
+            app.post(
+                run_route,
+                response_model=BaseResponse,
+                response_model_exclude_none=True,
+            )(run_wrapper)
         # LEGACY
         ### ----------- #
 
@@ -203,21 +217,33 @@ class entrypoint:
         self.update_test_wrapper_signature(wrapper=test_wrapper, config_instance=config)
 
         test_route = f"{entrypoint._test_path}{route_path}"
-        app.post(test_route, response_model=BaseResponse)(test_wrapper)
+        app.post(
+            test_route,
+            response_model=BaseResponse,
+            response_model_exclude_none=True,
+        )(test_wrapper)
 
         # LEGACY
         # TODO: Removing this implies breaking changes in :
         # - calls to /generate must be replaced with calls to /test
         if route_path == "":
             test_route = entrypoint._legacy_generate_path
-            app.post(test_route, response_model=BaseResponse)(test_wrapper)
+            app.post(
+                test_route,
+                response_model=BaseResponse,
+                response_model_exclude_none=True,
+            )(test_wrapper)
         # LEGACY
 
         # LEGACY
         # TODO: Removing this implies no breaking changes
         if route_path == "":
             test_route = entrypoint._legacy_playground_run_path
-            app.post(test_route, response_model=BaseResponse)(test_wrapper)
+            app.post(
+                test_route,
+                response_model=BaseResponse,
+                response_model_exclude_none=True,
+            )(test_wrapper)
         # LEGACY
         ### ------------ #
 
@@ -249,17 +275,18 @@ class entrypoint:
             )
         # LEGACY
 
-        app.openapi_schema = None  # Forces FastAPI to re-generate the schema
-        openapi_schema = app.openapi()
-        openapi_schema["agenta_sdk"] = {"version": get_current_version()}
+        openapi_schema = self.openapi()
+
         for _route in entrypoint.routes:
             if _route["config"] is not None:
                 self.override_config_in_schema(
                     openapi_schema=openapi_schema,
                     func_name=_route["func"],
-                    endpoint=_route["endpoint"],
+                    endpoint=_route["endpoint"].replace(AGENTA_RUNTIME_PREFIX, ""),
                     config=_route["config"],
                 )
+
+        app.openapi_schema = openapi_schema
         ### --------------- #
 
     def parse_config(self) -> Tuple[Optional[Type[BaseModel]], Dict[str, Any]]:
@@ -369,7 +396,9 @@ class entrypoint:
     ):
         display_exception("Application Exception")
 
-        status_code = 500
+        status_code = (
+            getattr(error, "status_code") if hasattr(error, "status_code") else 500
+        )
         stacktrace = format_exception(error, value=error, tb=error.__traceback__)  # type: ignore
 
         raise HTTPException(
@@ -430,8 +459,8 @@ class entrypoint:
         link = context.link
 
         tree = None
-        _tree_id = link.get("tree_id") if link else None # in int format
-        tree_id = str(UUID(int=_tree_id)) if _tree_id else None # in uuid_as_str format
+        _tree_id = link.get("tree_id") if link else None  # in int format
+        tree_id = str(UUID(int=_tree_id)) if _tree_id else None  # in uuid_as_str format
 
         if _tree_id is not None:
             if inline:
@@ -544,6 +573,82 @@ class entrypoint:
                 )
             )
 
+    def openapi(self):
+        app.openapi_schema = None  # Forces FastAPI to re-generate the schema
+
+        openapi_schema = app.openapi()
+
+        # ✅ Fix paths by removing the prefix
+        updated_paths = {}
+        for path, methods in openapi_schema["paths"].items():
+            new_path = (
+                path[len(AGENTA_RUNTIME_PREFIX) :]
+                if path.startswith(AGENTA_RUNTIME_PREFIX)
+                else path
+            )
+            updated_paths[new_path] = methods
+        openapi_schema["paths"] = updated_paths  # Replace paths
+
+        # ✅ Fix schema names and update `$ref` references
+        if "components" in openapi_schema and "schemas" in openapi_schema["components"]:
+            updated_schemas = {}
+            schema_name_map = {}  # Map old schema names to new schema names
+
+            for schema_name, schema_value in openapi_schema["components"][
+                "schemas"
+            ].items():
+                if AGENTA_RUNTIME_PREFIX and AGENTA_RUNTIME_PREFIX != "":
+                    new_schema_name = schema_name.replace(
+                        AGENTA_RUNTIME_PREFIX.lstrip("/").replace("/", "_") + "_", ""
+                    ).strip("_")
+                else:
+                    new_schema_name = schema_name
+                updated_schemas[new_schema_name] = schema_value
+                schema_name_map[schema_name] = new_schema_name  # Store mapping
+
+            # ✅ Fix `$ref` references
+            for path, methods in updated_paths.items():
+                for method in methods.values():
+                    if "requestBody" in method and "content" in method["requestBody"]:
+                        for content_type, content in method["requestBody"][
+                            "content"
+                        ].items():
+                            if "$ref" in content["schema"]:
+                                old_ref = content["schema"]["$ref"]
+                                schema_name = old_ref.split("/")[
+                                    -1
+                                ]  # Extract schema name
+                                if schema_name in schema_name_map:
+                                    content["schema"][
+                                        "$ref"
+                                    ] = f"#/components/schemas/{schema_name_map[schema_name]}"
+
+                    if "responses" in method:
+                        for status_code, response in method["responses"].items():
+                            if "content" in response:
+                                for content_type, content in response[
+                                    "content"
+                                ].items():
+                                    if "$ref" in content["schema"]:
+                                        old_ref = content["schema"]["$ref"]
+                                        schema_name = old_ref.split("/")[
+                                            -1
+                                        ]  # Extract schema name
+                                        if schema_name in schema_name_map:
+                                            content["schema"][
+                                                "$ref"
+                                            ] = f"#/components/schemas/{schema_name_map[schema_name]}"
+
+            # ✅ Update OpenAPI schema with fixed schemas
+            openapi_schema["components"]["schemas"] = updated_schemas
+
+            # ✅ Add Agenta SDK version info
+            openapi_schema["agenta_sdk"] = {"version": get_current_version()}
+
+            print(openapi_schema)
+
+        return openapi_schema
+
     def override_config_in_schema(
         self,
         openapi_schema: dict,
@@ -553,6 +658,7 @@ class entrypoint:
     ):
         """Override config in OpenAPI schema to add agenta-specific metadata."""
         endpoint = endpoint[1:].replace("/", "_")
+
         schema_key = f"Body_{func_name}_{endpoint}_post"
         schema_to_override = openapi_schema["components"]["schemas"][schema_key]
 
@@ -560,7 +666,7 @@ class entrypoint:
         config_class_name = type(config).__name__
         config_schema = openapi_schema["components"]["schemas"][config_class_name]
         # Process each field in the config class
-        for field_name, field in config.__class__.__fields__.items():
+        for field_name, field in config.__class__.model_fields.items():
             # Check if field has Annotated metadata for MultipleChoice
             if hasattr(field, "metadata") and field.metadata:
                 for meta in field.metadata:

agenta/sdk/decorators/tracing.py

@@ -165,15 +165,27 @@ class instrument:  # pylint: disable=invalid-name
                 usage = {"total_tokens": usage}
 
             span.set_attributes(
-                attributes={"total": cost},
+                attributes={"total": float(cost) if cost else None},
                 namespace="metrics.unit.costs",
             )
             span.set_attributes(
                 attributes=(
                     {
-                        "prompt": usage.get("prompt_tokens", None),
-                        "completion": usage.get("completion_tokens", None),
-                        "total": usage.get("total_tokens", None),
+                        "prompt": (
+                            float(usage.get("prompt_tokens"))
+                            if usage.get("prompt_tokens", None)
+                            else None
+                        ),
+                        "completion": (
+                            float(usage.get("completion_tokens"))
+                            if usage.get("completion_tokens", None)
+                            else None
+                        ),
+                        "total": (
+                            float(usage.get("total_tokens", None))
+                            if usage.get("total_tokens", None)
+                            else None
+                        ),
                     }
                 ),
                 namespace="metrics.unit.tokens",

agenta/sdk/litellm/litellm.py

@@ -154,16 +154,34 @@ def litellm_handler():
                 pass
 
             span.set_attributes(
-                attributes={"total": kwargs.get("response_cost")},
+                attributes={
+                    "total": (
+                        float(kwargs.get("response_cost"))
+                        if kwargs.get("response_cost")
+                        else None
+                    )
+                },
                 namespace="metrics.unit.costs",
             )
 
             span.set_attributes(
                 attributes=(
                     {
-                        "prompt": response_obj.usage.prompt_tokens,
-                        "completion": response_obj.usage.completion_tokens,
-                        "total": response_obj.usage.total_tokens,
+                        "prompt": (
+                            float(response_obj.usage.prompt_tokens)
+                            if response_obj.usage.prompt_tokens
+                            else None
+                        ),
+                        "completion": (
+                            float(response_obj.usage.completion_tokens)
+                            if response_obj.usage.completion_tokens
+                            else None
+                        ),
+                        "total": (
+                            float(response_obj.usage.total_tokens)
+                            if response_obj.usage.total_tokens
+                            else None
+                        ),
                     }
                 ),
                 namespace="metrics.unit.tokens",
@@ -264,16 +282,34 @@ def litellm_handler():
                 pass
 
             span.set_attributes(
-                attributes={"total": kwargs.get("response_cost")},
+                attributes={
+                    "total": (
+                        float(kwargs.get("response_cost"))
+                        if kwargs.get("response_cost")
+                        else None
+                    )
+                },
                 namespace="metrics.unit.costs",
             )
 
             span.set_attributes(
                 attributes=(
                     {
-                        "prompt": response_obj.usage.prompt_tokens,
-                        "completion": response_obj.usage.completion_tokens,
-                        "total": response_obj.usage.total_tokens,
+                        "prompt": (
+                            float(response_obj.usage.prompt_tokens)
+                            if response_obj.usage.prompt_tokens
+                            else None
+                        ),
+                        "completion": (
+                            float(response_obj.usage.completion_tokens)
+                            if response_obj.usage.completion_tokens
+                            else None
+                        ),
+                        "total": (
+                            float(response_obj.usage.total_tokens)
+                            if response_obj.usage.total_tokens
+                            else None
+                        ),
                     }
                 ),
                 namespace="metrics.unit.tokens",

agenta/sdk/litellm/mockllm.py

@@ -1,7 +1,7 @@
 from typing import Optional, Protocol, Any
 
-from agenta.sdk.context.routing import routing_context
 from agenta.sdk.litellm.mocks import MOCKS
+from agenta.sdk.context.routing import routing_context
 
 
 class LitellmProtocol(Protocol):
@@ -19,7 +19,7 @@ async def acompletion(*args, **kwargs):
         if mock not in MOCKS:
             raise ValueError(f"Mock {mock} not found")
 
-        return MOCKS[mock]
+        return MOCKS[mock](*args, **kwargs)
 
     if not litellm:
         raise ValueError("litellm not found")

agenta/sdk/litellm/mocks/__init__.py

@@ -1,3 +1,5 @@
+from typing import Callable
+
 from pydantic import BaseModel
 
 
@@ -13,8 +15,8 @@ class MockResponseModel(BaseModel):
     choices: list[MockChoiceModel]
 
 
-MOCKS = {
-    "hello": MockResponseModel(
+def hello_mock_response(*args, **kwargs) -> MockResponseModel:
+    return MockResponseModel(
         choices=[
             MockChoiceModel(
                 message=MockMessageModel(
@@ -22,5 +24,9 @@ MOCKS = {
                 )
             )
         ],
-    ),
+    )
+
+
+MOCKS: dict[str, Callable[..., MockResponseModel]] = {
+    "hello": hello_mock_response,
 }

agenta/sdk/managers/apps.py

@@ -0,0 +1,64 @@
+import logging
+
+from agenta.sdk.utils.exceptions import handle_exceptions
+
+import agenta as ag
+
+
+logger = logging.getLogger(__name__)
+
+
+class AppManager:
+    @classmethod
+    @handle_exceptions()
+    def create(cls, *, app_slug: str, template_key: str):
+        app_response = ag.api.apps.create_app(
+            app_name=app_slug, template_key=template_key
+        )
+        return app_response
+
+    @classmethod
+    @handle_exceptions()
+    async def acreate(cls, *, app_slug: str, template_key: str):
+        app_response = await ag.async_api.apps.create_app(
+            app_name=app_slug, template_key=template_key
+        )
+        return app_response
+
+    @classmethod
+    @handle_exceptions()
+    def list(cls):
+        apps_response = ag.api.apps.list_apps()
+        return apps_response
+
+    @classmethod
+    @handle_exceptions()
+    async def alist(cls):
+        apps_response = await ag.async_api.apps.list_apps()
+        return apps_response
+
+    @classmethod
+    @handle_exceptions()
+    def update(cls, *, app_id: str, app_slug: str):
+        app_response = ag.api.apps.update_app(app_id=app_id, app_name=app_slug)
+        return app_response
+
+    @classmethod
+    @handle_exceptions()
+    async def aupdate(cls, *, app_id: str, app_slug: str):
+        app_response = await ag.async_api.apps.update_app(
+            app_id=app_id, app_name=app_slug
+        )
+        return app_response
+
+    @classmethod
+    @handle_exceptions()
+    def delete(cls, *, app_id: str):
+        ag.api.apps.remove_app(app_id=app_id)
+        return None
+
+    @classmethod
+    @handle_exceptions()
+    async def adelete(cls, *, app_id: str):
+        await ag.async_api.apps.remove_app(app_id=app_id)
+        return None

agenta/sdk/managers/shared.py

@@ -505,12 +505,12 @@ class SharedManager:
                     slug=variant_slug,
                     version=None,
                     id=None,
-                ),
+                ).model_dump(),
                 application_ref=SharedManager._ref_or_none(  # type: ignore  # type: ignore
                     slug=app_slug,
                     version=None,
                     id=app_id,
-                ),
+                ).model_dump(),
             )
         )