agenta 0.32.0__py3-none-any.whl → 0.32.0a2__py3-none-any.whl

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

@@ -0,0 +1,27 @@
+from typing import Optional, Protocol, Any
+
+from agenta.sdk.litellm.mocks import MOCKS
+from agenta.sdk.context.routing import routing_context
+
+
+class LitellmProtocol(Protocol):
+    async def acompletion(self, *args: Any, **kwargs: Any) -> Any:
+        ...
+
+
+litellm: Optional[LitellmProtocol] = None  # pylint: disable=invalid-name
+
+
+async def acompletion(*args, **kwargs):
+    mock = routing_context.get().mock
+
+    if mock:
+        if mock not in MOCKS:
+            raise ValueError(f"Mock {mock} not found")
+
+        return MOCKS[mock](*args, **kwargs)
+
+    if not litellm:
+        raise ValueError("litellm not found")
+
+    return await litellm.acompletion(*args, **kwargs)

agenta/sdk/litellm/mocks/__init__.py

@@ -0,0 +1,32 @@
+from typing import Callable
+
+from pydantic import BaseModel
+
+
+class MockMessageModel(BaseModel):
+    content: str
+
+
+class MockChoiceModel(BaseModel):
+    message: MockMessageModel
+
+
+class MockResponseModel(BaseModel):
+    choices: list[MockChoiceModel]
+
+
+def hello_mock_response(*args, **kwargs) -> MockResponseModel:
+    return MockResponseModel(
+        choices=[
+            MockChoiceModel(
+                message=MockMessageModel(
+                    content="world",
+                )
+            )
+        ],
+    )
+
+
+MOCKS: dict[str, Callable[..., MockResponseModel]] = {
+    "hello": hello_mock_response,
+}

agenta/sdk/managers/vault.py

@@ -0,0 +1,16 @@
+from typing import Optional, Dict, Any
+
+from agenta.sdk.context.routing import routing_context
+
+
+class VaultManager:
+    @staticmethod
+    def get_from_route() -> Optional[Dict[str, Any]]:
+        context = routing_context.get()
+
+        secrets = context.secrets
+
+        if not secrets:
+            return None
+
+        return secrets

agenta/sdk/middleware/auth.py

@@ -98,7 +98,11 @@ class AuthMiddleware(BaseHTTPMiddleware):
 
             cookies = {"sAccessToken": access_token} if access_token else None
 
-            baggage = request.state.otel.get("baggage") if request.state.otel else {}
+            baggage = (
+                getattr(request.state.otel, "baggage")
+                if hasattr(request.state, "otel")
+                else {}
+            )
 
             project_id = (
                 # CLEANEST

agenta/sdk/middleware/config.py

@@ -40,7 +40,7 @@ class ConfigMiddleware(BaseHTTPMiddleware):
         request: Request,
         call_next: Callable,
     ):
-        request.state.config = {}
+        request.state.config = {"parameters": None, "references": None}
 
         with suppress():
             parameters, references = await self._get_config(request)
@@ -116,13 +116,14 @@ class ConfigMiddleware(BaseHTTPMiddleware):
 
         for ref_key in ["application_ref", "variant_ref", "environment_ref"]:
             refs = config.get(ref_key)
-            ref_prefix = ref_key.split("_", maxsplit=1)[0]
+            if refs:
+                ref_prefix = ref_key.split("_", maxsplit=1)[0]
 
-            for ref_part_key in ["id", "slug", "version"]:
-                ref_part = refs.get(ref_part_key)
+                for ref_part_key in ["id", "slug", "version"]:
+                    ref_part = refs.get(ref_part_key)
 
-                if ref_part:
-                    references[ref_prefix + "." + ref_part_key] = ref_part
+                    if ref_part:
+                        references[ref_prefix + "." + ref_part_key] = str(ref_part)
 
         _cache.put(_hash, {"parameters": parameters, "references": references})
 
@@ -159,12 +160,20 @@ class ConfigMiddleware(BaseHTTPMiddleware):
             or body.get("app")
         )
 
-        if not any([application_id, application_slug]):
+        application_version = (
+            # CLEANEST
+            baggage.get("application_version")
+            # ALTERNATIVE
+            or request.query_params.get("application_version")
+        )
+
+        if not any([application_id, application_slug, application_version]):
             return None
 
         return Reference(
             id=application_id,
             slug=application_slug,
+            version=application_version,
         )
 
     async def _parse_variant_ref(

agenta/sdk/middleware/inline.py

@@ -0,0 +1,38 @@
+from typing import Callable
+
+from starlette.middleware.base import BaseHTTPMiddleware
+from fastapi import Request, FastAPI
+
+
+from agenta.sdk.utils.exceptions import suppress
+
+from agenta.sdk.utils.constants import TRUTHY
+
+
+class InlineMiddleware(BaseHTTPMiddleware):
+    def __init__(self, app: FastAPI):
+        super().__init__(app)
+
+    async def dispatch(
+        self,
+        request: Request,
+        call_next: Callable,
+    ):
+        request.state.inline = False
+
+        with suppress():
+            baggage = request.state.otel.get("baggage") if request.state.otel else {}
+
+            inline = (
+                str(
+                    # CLEANEST
+                    baggage.get("inline")
+                    # ALTERNATIVE
+                    or request.query_params.get("inline")
+                )
+                in TRUTHY
+            )
+
+            request.state.inline = inline
+
+        return await call_next(request)

agenta/sdk/middleware/mock.py

@@ -0,0 +1,33 @@
+from typing import Callable
+
+from starlette.middleware.base import BaseHTTPMiddleware
+from fastapi import Request, FastAPI
+
+
+from agenta.sdk.utils.exceptions import suppress
+
+
+class MockMiddleware(BaseHTTPMiddleware):
+    def __init__(self, app: FastAPI):
+        super().__init__(app)
+
+    async def dispatch(
+        self,
+        request: Request,
+        call_next: Callable,
+    ):
+        request.state.mock = None
+
+        with suppress():
+            baggage = request.state.otel.get("baggage") if request.state.otel else {}
+
+            mock = (
+                # CLEANEST
+                baggage.get("mock")
+                # ALTERNATIVE
+                or request.query_params.get("mock")
+            )
+
+            request.state.mock = mock
+
+        return await call_next(request)

agenta/sdk/middleware/vault.py

@@ -1,6 +1,6 @@
 from os import getenv
 from json import dumps
-from typing import Callable, Dict, Optional, List, Any
+from typing import Callable, Dict, Optional, List, Any, get_args
 
 import httpx
 from fastapi import FastAPI, Request
@@ -18,25 +18,11 @@ from agenta.sdk.middleware.cache import TTLLRUCache, CACHE_CAPACITY, CACHE_TTL
 import agenta as ag
 
 
-# ProviderKind (agenta.client.backend.types.provider_kind import ProviderKind) defines a type hint that allows \
-# for a fixed set of string literals representing various provider names, alongside `typing.Any`.
-PROVIDER_KINDS = []
+_PROVIDER_KINDS = []
 
-# Rationale behind the following:
-# -------------------------------
-# You cannot loop directly over the values in `typing.Literal` because:
-# - `Literal` is not iterable.
-# - `ProviderKind.__args__` includes `Literal` and `Any`, but the actual string values
-#   are nested within the `Literal`'s own `__args__` attribute.
-
-# To solve this, we programmatically extract the values from `Literal` while retaining
-# the structure of ProviderKind. This ensures:
-# 1. We don't modify the original `ProviderKind` type definition.
-# 2. We dynamically access the literal values for use at runtime when necessary.
 for arg in ProviderKind.__args__:  # type: ignore
     if hasattr(arg, "__args__"):
-        PROVIDER_KINDS.extend(arg.__args__)
-
+        _PROVIDER_KINDS.extend(arg.__args__)
 
 _CACHE_ENABLED = getenv("AGENTA_MIDDLEWARE_CACHE_ENABLED", "true").lower() in TRUTHY
 
@@ -100,7 +86,7 @@ class VaultMiddleware(BaseHTTPMiddleware):
         local_secrets: List[SecretDTO] = []
 
         try:
-            for provider_kind in PROVIDER_KINDS:
+            for provider_kind in _PROVIDER_KINDS:
                 provider = provider_kind
                 key_name = f"{provider.upper()}_API_KEY"
                 key = getenv(key_name)
@@ -108,7 +94,8 @@ class VaultMiddleware(BaseHTTPMiddleware):
                 if not key:
                     continue
 
-                secret = SecretDTO(  # 'kind' attribute in SecretDTO defaults to 'provider_kind'
+                secret = SecretDTO(
+                    # kind=...  # defaults to 'provider_kind'
                     data=ProviderKeyDTO(
                         provider=provider,
                         key=key,

agenta/sdk/tracing/exporters.py

@@ -48,7 +48,6 @@ class InlineTraceExporter(SpanExporter):
         trace_id: int,
     ) -> bool:
         is_ready = trace_id in self._registry
-
         return is_ready
 
     def fetch(

agenta/sdk/tracing/inline.py

@@ -701,15 +701,6 @@ def _parse_from_semconv(
 def _parse_from_links(
     otel_span_dto: OTelSpanDTO,
 ) -> dict:
-    # TESTING
-    otel_span_dto.links = [
-        OTelLinkDTO(
-            context=otel_span_dto.context,
-            attributes={"ag.type.link": "testcase"},
-        )
-    ]
-    # -------
-
     # LINKS
     links = None
     otel_links = None
@@ -926,8 +917,9 @@ def parse_to_agenta_span_dto(
     if span_dto.refs:
         span_dto.refs = _unmarshal_attributes(span_dto.refs)
 
-    for link in span_dto.links:
-        link.tree_id = None
+    if isinstance(span_dto.links, list):
+        for link in span_dto.links:
+            link.tree_id = None
 
     if span_dto.nodes:
         for v in span_dto.nodes.values():
@@ -1030,6 +1022,24 @@ def _parse_readable_spans(
     otel_span_dtos = list()
 
     for span in spans:
+        otel_events = [
+            OTelEventDTO(
+                name=event.name,
+                timestamp=_timestamp_ns_to_datetime(event.timestamp),
+                attributes=event.attributes,
+            )
+            for event in span.events
+        ]
+        otel_links = [
+            OTelLinkDTO(
+                context=OTelContextDTO(
+                    trace_id=_int_to_hex(link.context.trace_id, 128),
+                    span_id=_int_to_hex(link.context.span_id, 64),
+                ),
+                attributes=link.attributes,
+            )
+            for link in span.links
+        ]
         otel_span_dto = OTelSpanDTO(
             context=OTelContextDTO(
                 trace_id=_int_to_hex(span.get_span_context().trace_id, 128),
@@ -1045,14 +1055,7 @@ def _parse_readable_spans(
             status_code=OTelStatusCode("STATUS_CODE_" + span.status.status_code.name),
             status_message=span.status.description,
             attributes=span.attributes,
-            events=[
-                OTelEventDTO(
-                    name=event.name,
-                    timestamp=_timestamp_ns_to_datetime(event.timestamp),
-                    attributes=event.attributes,
-                )
-                for event in span.events
-            ],
+            events=otel_events if len(otel_events) > 0 else None,
             parent=(
                 OTelContextDTO(
                     trace_id=_int_to_hex(span.parent.trace_id, 128),
@@ -1061,16 +1064,7 @@ def _parse_readable_spans(
                 if span.parent
                 else None
             ),
-            links=[
-                OTelLinkDTO(
-                    context=OTelContextDTO(
-                        trace_id=_int_to_hex(link.context.trace_id, 128),
-                        span_id=_int_to_hex(link.context.span_id, 64),
-                    ),
-                    attributes=link.attributes,
-                )
-                for link in span.links
-            ],
+            links=otel_links if len(otel_links) > 0 else None,
         )
 
         otel_span_dtos.append(otel_span_dto)