agenta 0.34.7__tar.gz → 0.35.0__tar.gz

{agenta-0.34.7 → agenta-0.35.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: agenta
-Version: 0.34.7
+Version: 0.35.0
 Summary: The SDK for agenta is an open-source LLMOps platform.
 Keywords: LLMOps,LLM,evaluation,prompt engineering
 Author: Mahmoud Mabrouk
@@ -22,9 +22,9 @@ Requires-Dist: fastapi (>=0.100.0)
 Requires-Dist: httpx (>=0.24,<0.28)
 Requires-Dist: importlib-metadata (>=8.0.0,<9.0)
 Requires-Dist: litellm (>=1.48.0,<2.0.0)
-Requires-Dist: opentelemetry-api (>=1.27.0,<2.0.0)
-Requires-Dist: opentelemetry-exporter-otlp (>=1.27.0,<2.0.0)
-Requires-Dist: opentelemetry-sdk (>=1.27.0,<2.0.0)
+Requires-Dist: opentelemetry-api (==1.27.0)
+Requires-Dist: opentelemetry-exporter-otlp (==1.27.0)
+Requires-Dist: opentelemetry-sdk (==1.27.0)
 Requires-Dist: posthog (>=3.1.0,<4.0.0)
 Requires-Dist: pydantic (>=2)
 Requires-Dist: python-dotenv (>=1.0.0,<2.0.0)

{agenta-0.34.7 → agenta-0.35.0}/agenta/sdk/agenta_init.py

@@ -89,11 +89,6 @@ class AgentaSingleton:
 
         self.base_id = getenv("AGENTA_BASE_ID")
 
-        self.service_id = getenv("AGENTA_SERVICE_ID") or self.base_id or None
-
-        log.info("Agenta - Service ID: %s", self.service_id)
-        log.info("Agenta - Application ID: %s", self.app_id)
-
         self.tracing = Tracing(
             url=f"{self.host}/api/observability/v1/otlp/traces",  # type: ignore
             redact=redact,
@@ -102,9 +97,6 @@ class AgentaSingleton:
 
         self.tracing.configure(
             api_key=self.api_key,
-            service_id=self.service_id,
-            # DEPRECATING
-            app_id=self.app_id,
         )
 
         self.api = AgentaApi(

{agenta-0.34.7 → agenta-0.35.0}/agenta/sdk/context/tracing.py

@@ -7,6 +7,8 @@ from pydantic import BaseModel
 
 
 class TracingContext(BaseModel):
+    traceparent: Optional[Dict[str, Any]] = None
+    baggage: Optional[Dict[str, Any]] = None
     credentials: Optional[str] = None
     parameters: Optional[Dict[str, Any]] = None
     references: Optional[Dict[str, Any]] = None

{agenta-0.34.7 → agenta-0.35.0}/agenta/sdk/decorators/routing.py

@@ -330,6 +330,8 @@ class entrypoint:
             raise HTTPException(status_code=500, detail="Missing 'request'.")
 
         state = request.state
+        traceparent = state.otel.get("traceparent")
+        baggage = state.otel.get("baggage")
         credentials = state.auth.get("credentials")
         parameters = state.config.get("parameters")
         references = state.config.get("references")
@@ -346,6 +348,8 @@ class entrypoint:
         ):
             with tracing_context_manager(
                 context=TracingContext(
+                    traceparent=traceparent,
+                    baggage=baggage,
                     credentials=credentials,
                     parameters=parameters,
                     references=references,
@@ -378,14 +382,13 @@ class entrypoint:
                 content_type = "application/json"
             data = self.patch_result(result)
 
-            if inline:
-                tree, tree_id = await self.fetch_inline_trace(inline)
+            tree, tree_id = await self.fetch_inline_trace(inline)
 
         try:
             return BaseResponse(
                 data=data, tree=tree, content_type=content_type, tree_id=tree_id
             )
-        except:
+        except:  # pylint: disable=bare-except
             return BaseResponse(data=data, content_type=content_type)
 
     def handle_failure(
@@ -397,6 +400,9 @@ class entrypoint:
         status_code = (
             getattr(error, "status_code") if hasattr(error, "status_code") else 500
         )
+        if status_code in [401, 403]:  # Reserved HTTP codes for auth middleware
+            status_code = 424  # Proxy Authentication Required
+
         stacktrace = format_exception(error, value=error, tb=error.__traceback__)  # type: ignore
 
         raise HTTPException(
@@ -444,6 +450,18 @@ class entrypoint:
 
         return data
 
+    async def fetch_inline_trace_id(
+        self,
+    ):
+        context = tracing_context.get()
+
+        link = context.link
+
+        _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
+
+        return tree_id
+
     async def fetch_inline_trace(
         self,
         inline: bool,
@@ -472,6 +490,7 @@ class entrypoint:
                     remaining_steps -= 1
 
                 tree = ag.tracing.get_inline_trace(_tree_id)
+
         return tree, tree_id
 
     # --- OpenAPI --- #

{agenta-0.34.7 → agenta-0.35.0}/agenta/sdk/decorators/tracing.py

@@ -4,8 +4,9 @@ from functools import wraps
 from itertools import chain
 from inspect import iscoroutinefunction, getfullargspec
 
-from opentelemetry import baggage as baggage
-from opentelemetry.context import attach, detach
+from opentelemetry import baggage
+from opentelemetry.context import attach, detach, get_current
+from opentelemetry.baggage import set_baggage, get_all
 
 from agenta.sdk.utils.exceptions import suppress
 from agenta.sdk.context.tracing import tracing_context
@@ -49,7 +50,15 @@ class instrument:  # pylint: disable=invalid-name
 
                 token = self._attach_baggage()
 
-                with ag.tracer.start_as_current_span(func.__name__, kind=self.kind):
+                ctx = self._get_traceparent()
+
+                with ag.tracer.start_as_current_span(
+                    name=func.__name__,
+                    kind=self.kind,
+                    context=ctx,
+                ):
+                    self._set_link()
+
                     self._pre_instrument(func, *args, **kwargs)
 
                     result = await func(*args, **kwargs)
@@ -69,7 +78,15 @@ class instrument:  # pylint: disable=invalid-name
 
                 token = self._attach_baggage()
 
-                with ag.tracer.start_as_current_span(func.__name__, kind=self.kind):
+                ctx = self._get_traceparent()
+
+                with ag.tracer.start_as_current_span(
+                    name=func.__name__,
+                    kind=self.kind,
+                    context=ctx,
+                ):
+                    self._set_link()
+
                     self._pre_instrument(func, *args, **kwargs)
 
                     result = func(*args, **kwargs)
@@ -90,6 +107,30 @@ class instrument:  # pylint: disable=invalid-name
 
         self.kind = parse_span_kind(self.type)
 
+    def _get_traceparent(self):
+        context = tracing_context.get()
+
+        traceparent = context.traceparent
+
+        if not context.link:
+            for key, value in get_all(get_current()).items():
+                traceparent = set_baggage(name=key, value=value, context=traceparent)
+
+            return traceparent
+
+    def _set_link(self):
+        span = ag.tracing.get_current_span()
+
+        context = tracing_context.get()
+
+        if not context.link:
+            context.link = {
+                "tree_id": span.get_span_context().trace_id,
+                "node_id": span.get_span_context().span_id,
+            }
+
+            tracing_context.set(context)
+
     def _attach_baggage(self):
         context = tracing_context.get()
 
@@ -122,7 +163,7 @@ class instrument:  # pylint: disable=invalid-name
         with suppress():
             trace_id = span.context.trace_id
 
-            ag.tracing.credentials[trace_id] = context.credentials
+            ag.tracing.credentials.put(trace_id, context.credentials)
 
             span.set_attributes(
                 attributes={"node": self.type},
@@ -153,6 +194,7 @@ class instrument:  # pylint: disable=invalid-name
         result,
     ):
         span = ag.tracing.get_current_span()
+
         with suppress():
             cost = None
             usage = {}
@@ -206,15 +248,6 @@ class instrument:  # pylint: disable=invalid-name
 
             span.set_status("OK")
 
-        with suppress():
-            if hasattr(span, "parent") and span.parent is None:
-                context = tracing_context.get()
-                context.link = {
-                    "tree_id": span.get_span_context().trace_id,
-                    "node_id": span.get_span_context().span_id,
-                }
-                tracing_context.set(context)
-
     def _parse(
         self,
         func,

{agenta-0.34.7 → agenta-0.35.0}/agenta/sdk/litellm/mocks/__init__.py

@@ -26,7 +26,20 @@ def hello_mock_response(*args, **kwargs) -> MockResponseModel:
         ],
     )
 
+def chat_mock_response(*args, **kwargs) -> MockResponseModel:
+    return MockResponseModel(
+        choices=[
+            MockChoiceModel(
+                message=MockMessageModel(
+                    content="world",
+                    role="assistant"
+                )
+            )
+        ],
+    )
+
 
 MOCKS: dict[str, Callable[..., MockResponseModel]] = {
     "hello": hello_mock_response,
+    "chat": chat_mock_response
 }

{agenta-0.34.7 → agenta-0.35.0}/agenta/sdk/middleware/auth.py

@@ -8,7 +8,7 @@ from starlette.middleware.base import BaseHTTPMiddleware
 from fastapi import FastAPI, Request
 from fastapi.responses import JSONResponse
 
-from agenta.sdk.middleware.cache import TTLLRUCache, CACHE_CAPACITY, CACHE_TTL
+from agenta.sdk.utils.cache import TTLLRUCache
 from agenta.sdk.utils.constants import TRUTHY
 from agenta.sdk.utils.exceptions import display_exception
 from agenta.sdk.utils.logging import log
@@ -22,7 +22,7 @@ _CACHE_ENABLED = getenv("AGENTA_MIDDLEWARE_CACHE_ENABLED", "false").lower() in T
 
 _ALWAYS_ALLOW_LIST = [f"{AGENTA_RUNTIME_PREFIX}/health"]
 
-_cache = TTLLRUCache(capacity=CACHE_CAPACITY, ttl=CACHE_TTL)
+_cache = TTLLRUCache()
 
 
 class DenyResponse(JSONResponse):
@@ -54,7 +54,6 @@ class AuthMiddleware(BaseHTTPMiddleware):
         super().__init__(app)
 
         self.host = ag.DEFAULT_AGENTA_SINGLETON_INSTANCE.host
-        self.resource_id = ag.DEFAULT_AGENTA_SINGLETON_INSTANCE.service_id
 
     async def dispatch(self, request: Request, call_next: Callable):
         try:
@@ -111,9 +110,6 @@ class AuthMiddleware(BaseHTTPMiddleware):
 
             params = {"action": "run_service", "resource_type": "service"}
 
-            if self.resource_id:
-                params["resource_id"] = self.resource_id
-
             if project_id:
                 params["project_id"] = project_id
 

{agenta-0.34.7 → agenta-0.35.0}/agenta/sdk/middleware/config.py

@@ -10,7 +10,7 @@ from fastapi import Request, FastAPI
 
 import httpx
 
-from agenta.sdk.middleware.cache import TTLLRUCache, CACHE_CAPACITY, CACHE_TTL
+from agenta.sdk.utils.cache import TTLLRUCache
 from agenta.sdk.utils.constants import TRUTHY
 from agenta.sdk.utils.exceptions import suppress
 
@@ -19,7 +19,7 @@ import agenta as ag
 
 _CACHE_ENABLED = getenv("AGENTA_MIDDLEWARE_CACHE_ENABLED", "false").lower() in TRUTHY
 
-_cache = TTLLRUCache(capacity=CACHE_CAPACITY, ttl=CACHE_TTL)
+_cache = TTLLRUCache()
 
 
 class Reference(BaseModel):
@@ -147,37 +147,32 @@ class ConfigMiddleware(BaseHTTPMiddleware):
 
         application_id = (
             # CLEANEST
-            baggage.get("application_id")
+            baggage.get("ag.refs.application.id")
             # ALTERNATIVE
             or request.query_params.get("application_id")
             # LEGACY
+            or baggage.get("application_id")
             or request.query_params.get("app_id")
             or self.application_id
         )
         application_slug = (
             # CLEANEST
-            baggage.get("application_slug")
+            baggage.get("ag.refs.application.slug")
             # ALTERNATIVE
             or request.query_params.get("application_slug")
             # LEGACY
+            or baggage.get("application_slug")
             or request.query_params.get("app_slug")
             or body.get("app")
         )
 
-        application_version = (
-            # CLEANEST
-            baggage.get("application_version")
-            # ALTERNATIVE
-            or request.query_params.get("application_version")
-        )
-
-        if not any([application_id, application_slug, application_version]):
+        if not any([application_id, application_slug, None]):
             return None
 
         return Reference(
             id=application_id,
             slug=application_slug,
-            version=application_version,
+            version=None,
         )
 
     async def _parse_variant_ref(
@@ -194,24 +189,29 @@ class ConfigMiddleware(BaseHTTPMiddleware):
 
         variant_id = (
             # CLEANEST
-            baggage.get("variant_id")
+            baggage.get("ag.refs.variant.id")
             # ALTERNATIVE
             or request.query_params.get("variant_id")
+            # LEGACY
+            or baggage.get("variant_id")
         )
         variant_slug = (
             # CLEANEST
-            baggage.get("variant_slug")
+            baggage.get("ag.refs.variant.slug")
             # ALTERNATIVE
             or request.query_params.get("variant_slug")
             # LEGACY
+            or baggage.get("variant_slug")
             or request.query_params.get("config")
             or body.get("config")
         )
         variant_version = (
             # CLEANEST
-            baggage.get("variant_version")
+            baggage.get("ag.refs.variant.version")
             # ALTERNATIVE
             or request.query_params.get("variant_version")
+            # LEGACY
+            or baggage.get("variant_version")
         )
 
         if not any([variant_id, variant_slug, variant_version]):
@@ -237,24 +237,29 @@ class ConfigMiddleware(BaseHTTPMiddleware):
 
         environment_id = (
             # CLEANEST
-            baggage.get("environment_id")
+            baggage.get("ag.refs.environment.id")
             # ALTERNATIVE
             or request.query_params.get("environment_id")
+            # LEGACY
+            or baggage.get("environment_id")
         )
         environment_slug = (
             # CLEANEST
-            baggage.get("environment_slug")
+            baggage.get("ag.refs.environment.slug")
             # ALTERNATIVE
             or request.query_params.get("environment_slug")
             # LEGACY
+            or baggage.get("environment_slug")
             or request.query_params.get("environment")
             or body.get("environment")
         )
         environment_version = (
             # CLEANEST
-            baggage.get("environment_version")
+            baggage.get("ag.refs.environment.version")
             # ALTERNATIVE
             or request.query_params.get("environment_version")
+            # LEGACY
+            or baggage.get("environment_version")
         )
 
         if not any([environment_id, environment_slug, environment_version]):

{agenta-0.34.7 → agenta-0.35.0}/agenta/sdk/middleware/inline.py

@@ -3,9 +3,7 @@ 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
 
 

{agenta-0.34.7 → agenta-0.35.0}/agenta/sdk/middleware/mock.py

@@ -3,7 +3,6 @@ from typing import Callable
 from starlette.middleware.base import BaseHTTPMiddleware
 from fastapi import Request, FastAPI
 
-
 from agenta.sdk.utils.exceptions import suppress
 
 

agenta-0.35.0/agenta/sdk/middleware/otel.py

@@ -0,0 +1,25 @@
+from typing import Callable
+
+from starlette.middleware.base import BaseHTTPMiddleware
+from fastapi import Request, FastAPI
+
+from opentelemetry.baggage.propagation import W3CBaggagePropagator
+from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
+
+from agenta.sdk.utils.exceptions import suppress
+from agenta.sdk.tracing.propagation import extract
+
+
+class OTelMiddleware(BaseHTTPMiddleware):
+    def __init__(self, app: FastAPI):
+        super().__init__(app)
+
+    async def dispatch(self, request: Request, call_next: Callable):
+        request.state.otel = {"baggage": {}, "traceparent": None}
+
+        with suppress():
+            _, traceparent, baggage = extract(request.headers)
+
+            request.state.otel = {"baggage": baggage, "traceparent": traceparent}
+
+        return await call_next(request)

{agenta-0.34.7 → agenta-0.35.0}/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, get_args
+from typing import Callable, Dict, Optional, List, Any
 
 import httpx
 from fastapi import FastAPI, Request
@@ -13,7 +13,7 @@ from agenta.client.backend.types.secret_dto import SecretDto as SecretDTO
 from agenta.client.backend.types.provider_key_dto import (
     ProviderKeyDto as ProviderKeyDTO,
 )
-from agenta.sdk.middleware.cache import TTLLRUCache, CACHE_CAPACITY, CACHE_TTL
+from agenta.sdk.utils.cache import TTLLRUCache
 
 import agenta as ag
 
@@ -26,7 +26,7 @@ for arg in ProviderKind.__args__:  # type: ignore
 
 _CACHE_ENABLED = getenv("AGENTA_MIDDLEWARE_CACHE_ENABLED", "false").lower() in TRUTHY
 
-_cache = TTLLRUCache(capacity=CACHE_CAPACITY, ttl=CACHE_TTL)
+_cache = TTLLRUCache()
 
 
 class VaultMiddleware(BaseHTTPMiddleware):

{agenta-0.34.7 → agenta-0.35.0}/agenta/sdk/tracing/exporters.py

@@ -1,4 +1,4 @@
-from typing import Sequence, Dict, List
+from typing import Sequence, Dict, List, Optional
 
 from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
 from opentelemetry.sdk.trace.export import (
@@ -14,10 +14,14 @@ from agenta.sdk.context.exporting import (
     exporting_context,
     ExportingContext,
 )
+from agenta.sdk.utils.cache import TTLLRUCache
 
 
 class InlineTraceExporter(SpanExporter):
-    def __init__(self, registry: Dict[str, List[ReadableSpan]]):
+    def __init__(
+        self,
+        registry: Dict[str, List[ReadableSpan]],
+    ):
         self._shutdown = False
         self._registry = registry
 
@@ -65,29 +69,45 @@ class InlineTraceExporter(SpanExporter):
 class OTLPExporter(OTLPSpanExporter):
     _MAX_RETRY_TIMEOUT = 2
 
-    def __init__(self, *args, credentials: Dict[int, str] = None, **kwargs):
+    def __init__(
+        self,
+        *args,
+        credentials: Optional[TTLLRUCache] = None,
+        **kwargs,
+    ):
         super().__init__(*args, **kwargs)
 
         self.credentials = credentials
 
     def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
-        credentials = None
+        grouped_spans: Dict[str, List[str]] = {}
+
+        for span in spans:
+            trace_id = span.get_span_context().trace_id
+
+            credentials = None
+            if self.credentials:
+                credentials = self.credentials.get(trace_id)
+
+            if credentials not in grouped_spans:
+                grouped_spans[credentials] = []
 
-        if self.credentials:
-            trace_ids = set(span.get_span_context().trace_id for span in spans)
+            grouped_spans[credentials].append(span)
 
-            if len(trace_ids) == 1:
-                trace_id = trace_ids.pop()
+        serialized_spans = []
 
-                if trace_id in self.credentials:
-                    credentials = self.credentials.pop(trace_id)
+        for credentials, _spans in grouped_spans.items():
+            with exporting_context_manager(
+                context=ExportingContext(
+                    credentials=credentials,
+                )
+            ):
+                serialized_spans.append(super().export(_spans))
 
-        with exporting_context_manager(
-            context=ExportingContext(
-                credentials=credentials,
-            )
-        ):
-            return super().export(spans)
+        if all(serialized_spans):
+            return SpanExportResult.SUCCESS
+        else:
+            return SpanExportResult.FAILURE
 
     def _export(self, serialized_data: bytes):
         credentials = exporting_context.get().credentials

{agenta-0.34.7 → agenta-0.35.0}/agenta/sdk/tracing/inline.py

@@ -1061,7 +1061,7 @@ def _parse_readable_spans(
                     trace_id=_int_to_hex(span.parent.trace_id, 128),
                     span_id=_int_to_hex(span.parent.span_id, 64),
                 )
-                if span.parent
+                if span.parent and not span.parent.is_remote
                 else None
             ),
             links=otel_links if len(otel_links) > 0 else None,