agenta 0.34.6__py3-none-any.whl → 0.35.0__py3-none-any.whl

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/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/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/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/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/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/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/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/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/sdk/middleware/otel.py

@@ -4,8 +4,10 @@ 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):
@@ -13,28 +15,11 @@ class OTelMiddleware(BaseHTTPMiddleware):
         super().__init__(app)
 
     async def dispatch(self, request: Request, call_next: Callable):
-        request.state.otel = {"baggage": {}}
+        request.state.otel = {"baggage": {}, "traceparent": None}
 
         with suppress():
-            baggage = await self._get_baggage(request)
+            _, traceparent, baggage = extract(request.headers)
 
-            request.state.otel = {"baggage": baggage}
+            request.state.otel = {"baggage": baggage, "traceparent": traceparent}
 
         return await call_next(request)
-
-    async def _get_baggage(
-        self,
-        request,
-    ):
-        _baggage = {"baggage": request.headers.get("Baggage", "")}
-
-        context = W3CBaggagePropagator().extract(_baggage)
-
-        baggage = {}
-
-        if context:
-            for partial in context.values():
-                for key, value in partial.items():
-                    baggage[key] = value
-
-        return baggage

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/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/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,

agenta/sdk/tracing/processors.py

@@ -8,7 +8,9 @@ from opentelemetry.sdk.trace.export import (
     ReadableSpan,
     BatchSpanProcessor,
     _DEFAULT_MAX_QUEUE_SIZE,
+    _DEFAULT_SCHEDULE_DELAY_MILLIS,
     _DEFAULT_MAX_EXPORT_BATCH_SIZE,
+    _DEFAULT_EXPORT_TIMEOUT_MILLIS,
 )
 
 from agenta.sdk.utils.logging import log
@@ -20,6 +22,7 @@ class TraceProcessor(BatchSpanProcessor):
         self,
         span_exporter: SpanExporter,
         references: Dict[str, str] = None,
+        inline: bool = False,
         max_queue_size: int = None,
         schedule_delay_millis: float = None,
         max_export_batch_size: int = None,
@@ -28,67 +31,86 @@ class TraceProcessor(BatchSpanProcessor):
         super().__init__(
             span_exporter,
             _DEFAULT_MAX_QUEUE_SIZE,
-            12 * 60 * 60 * 1000,  # 12 hours
+            60 * 60 * 1000 if inline else _DEFAULT_SCHEDULE_DELAY_MILLIS,
             _DEFAULT_MAX_EXPORT_BATCH_SIZE,
-            500,  # < 1 second (0.5 seconds)
+            500 if inline else _DEFAULT_EXPORT_TIMEOUT_MILLIS,  # < 1 second
         )
 
-        self._registry = dict()
-        self._exporter = span_exporter
         self.references = references or dict()
-        self.spans: Dict[int, List[ReadableSpan]] = dict()
+        self.inline = inline is True
+
+        # --- INLINE
+        if self.inline:
+            self._registry = dict()
+            self._exporter = span_exporter
+            self._spans: Dict[int, List[ReadableSpan]] = dict()
+        # --- INLINE
 
     def on_start(
         self,
         span: Span,
         parent_context: Optional[Context] = None,
     ) -> None:
-        baggage = get_baggage(parent_context)
-
         for key in self.references.keys():
             span.set_attribute(f"ag.refs.{key}", self.references[key])
 
+        baggage = get_baggage(parent_context)
+
         for key in baggage.keys():
             if key.startswith("ag.refs."):
                 _key = key.replace("ag.refs.", "")
                 if _key in [_.value for _ in Reference.__members__.values()]:
                     span.set_attribute(key, baggage[key])
 
-        if span.context.trace_id not in self._registry:
-            self._registry[span.context.trace_id] = dict()
+        # --- INLINE
+        if self.inline:
+            if span.context.trace_id not in self._registry:
+                self._registry[span.context.trace_id] = dict()
 
-        self._registry[span.context.trace_id][span.context.span_id] = True
+            self._registry[span.context.trace_id][span.context.span_id] = True
+        # --- INLINE
 
     def on_end(
         self,
         span: ReadableSpan,
     ):
-        if self.done:
-            return
+        # --- INLINE
+        if self.inline:
+            if self.done:
+                return
+
+            if span.context.trace_id not in self._spans:
+                self._spans[span.context.trace_id] = list()
 
-        if span.context.trace_id not in self.spans:
-            self.spans[span.context.trace_id] = list()
+            self._spans[span.context.trace_id].append(span)
 
-        self.spans[span.context.trace_id].append(span)
+            del self._registry[span.context.trace_id][span.context.span_id]
 
-        del self._registry[span.context.trace_id][span.context.span_id]
+            if len(self._registry[span.context.trace_id]) == 0:
+                self.export(span.context.trace_id)
+        # --- INLINE
 
-        if len(self._registry[span.context.trace_id]) == 0:
-            self.export(span.context.trace_id)
+        # --- DISTRIBUTED
+        else:
+            super().on_end(span)
+        # --- DISTRIBUTED
 
     def export(
         self,
         trace_id: int,
     ):
-        spans = self.spans[trace_id]
+        # --- INLINE
+        if self.inline:
+            spans = self._spans[trace_id]
 
-        for span in spans:
-            self.queue.appendleft(span)
+            for span in spans:
+                self.queue.appendleft(span)
 
-        with self.condition:
-            self.condition.notify()
+            with self.condition:
+                self.condition.notify()
 
-        del self.spans[trace_id]
+            del self._spans[trace_id]
+        # --- INLINE
 
     def force_flush(
         self,
@@ -105,10 +127,13 @@ class TraceProcessor(BatchSpanProcessor):
     ) -> bool:
         is_ready = True
 
-        try:
-            is_ready = self._exporter.is_ready(trace_id)
-        except:  # pylint: disable=bare-except
-            pass
+        # --- INLINE
+        if self.inline:
+            try:
+                is_ready = self._exporter.is_ready(trace_id)
+            except:  # pylint: disable=bare-except
+                pass
+        # --- INLINE
 
         return is_ready
 
@@ -116,6 +141,14 @@ class TraceProcessor(BatchSpanProcessor):
         self,
         trace_id: Optional[int] = None,
     ) -> Dict[str, ReadableSpan]:
-        trace = self._exporter.fetch(trace_id)  # type: ignore
+        trace = None
+
+        # --- INLINE
+        if self.inline:
+            try:
+                trace = self._exporter.fetch(trace_id)  # type: ignore
+            except:  # pylint: disable=bare-except
+                pass
+        # --- INLINE
 
         return trace

agenta/sdk/tracing/propagation.py

@@ -0,0 +1,100 @@
+from typing import Tuple, Optional, Dict, Any
+
+from opentelemetry.trace import Span, set_span_in_context, get_current_span
+from opentelemetry.baggage.propagation import W3CBaggagePropagator
+from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
+from opentelemetry.baggage import set_baggage
+from opentelemetry.context import get_current
+
+import agenta as ag
+
+
+def extract(
+    headers: Dict[str, str],
+) -> Tuple[Optional[str], Optional[Any], Dict[str, str]]:
+    # --- Extract credentials --- #
+    credentials = None
+
+    try:
+        credentials = (
+            headers.get("Authorization")  # Uppercase
+            or headers.get("authorization")  # Lowercase
+            or None
+        )
+
+    except:  # pylint: disable=bare-except
+        pass
+
+    # --- Extract traceparent --- #
+    traceparent = None
+
+    try:
+        _carrier = {
+            "traceparent": headers.get("Traceparent")  # Uppercase
+            or headers.get("traceparent")  # Lowercase
+            or "",
+        }
+
+        _context = TraceContextTextMapPropagator().extract(_carrier)
+
+        traceparent = _context
+    except:  # pylint: disable=bare-except
+        pass
+
+    # --- Extract baggage --- #
+    baggage = {}
+
+    try:
+        _carrier = {
+            "baggage": headers.get("Baggage")  # Uppercase
+            or headers.get("baggage")  # Lowercase
+            or "",
+        }
+
+        _context = W3CBaggagePropagator().extract(_carrier)
+
+        if _context:
+            for partial in _context.values():
+                for key, value in partial.items():
+                    baggage[key] = value
+
+    except:  # pylint: disable=bare-except
+        pass
+
+    # --- #
+    return credentials, traceparent, baggage
+
+
+def inject(
+    headers: Optional[Dict[str, str]] = None,
+) -> Dict[str, str]:
+    headers = headers or {}
+
+    _context = get_current()
+
+    ctx = ag.sdk.context.tracing.tracing_context.get()
+
+    # --- Inject traceparent --- #
+    try:
+        TraceContextTextMapPropagator().inject(headers, context=_context)
+
+    except:  # pylint: disable=bare-except
+        pass
+
+    # --- Inject baggage --- #
+    try:
+        if ctx.baggage:
+            for key, value in ctx.baggage.items():
+                _context = set_baggage(key, value, context=_context)
+
+        W3CBaggagePropagator().inject(headers, context=_context)
+
+    except:  # pylint: disable=bare-except
+        pass
+
+    # --- Inject credentials --- #
+    if ctx.credentials:
+        headers["Authorization"] = ctx.credentials
+
+    # --- #
+    return headers

agenta/sdk/tracing/tracing.py

@@ -21,6 +21,8 @@ from agenta.sdk.tracing.exporters import InlineExporter, OTLPExporter
 from agenta.sdk.tracing.spans import CustomSpan
 from agenta.sdk.tracing.inline import parse_inline_trace
 from agenta.sdk.tracing.conventions import Reference, is_valid_attribute_key
+from agenta.sdk.tracing.propagation import extract, inject
+from agenta.sdk.utils.cache import TTLLRUCache
 
 
 class Tracing(metaclass=Singleton):
@@ -42,7 +44,7 @@ class Tracing(metaclass=Singleton):
         # REFERENCES
         self.references: Dict[str, str] = dict()
         # CREDENTIALS
-        self.credentials: Dict[int, str] = dict()
+        self.credentials: TTLLRUCache = TTLLRUCache(ttl=(60 * 60))  # 1 hour x 512 keys
 
         # TRACER PROVIDER
         self.tracer_provider: Optional[TracerProvider] = None
@@ -62,37 +64,33 @@ class Tracing(metaclass=Singleton):
     def configure(
         self,
         api_key: Optional[str] = None,
-        service_id: Optional[str] = None,
-        # DEPRECATING
-        app_id: Optional[str] = None,
+        inline: Optional[bool] = True,
     ):
         # HEADERS (OTLP)
         if api_key:
             self.headers["Authorization"] = f"ApiKey {api_key}"
-        # REFERENCES
-        if service_id:
-            self.references["service.id"] = service_id
-        if app_id:
-            self.references["application.id"] = app_id
 
         # TRACER PROVIDER
         self.tracer_provider = TracerProvider(
             resource=Resource(attributes={"service.name": "agenta-sdk"})
         )
-        # TRACE PROCESSORS -- INLINE
-        self.inline = TraceProcessor(
-            InlineExporter(
-                registry=self.inline_spans,
-            ),
-            references=self.references,
-        )
-        self.tracer_provider.add_span_processor(self.inline)
+
+        # --- INLINE
+        if inline:
+            # TRACE PROCESSORS -- INLINE
+            self.inline = TraceProcessor(
+                InlineExporter(
+                    registry=self.inline_spans,
+                ),
+                references=self.references,
+                inline=inline,
+            )
+            self.tracer_provider.add_span_processor(self.inline)
+        # --- INLINE
+
         # TRACE PROCESSORS -- OTLP
         try:
-            log.info(
-                "Agenta - OLTP URL: %s",
-                self.otlp_url,
-            )
+            log.info("Agenta - OLTP URL: %s", self.otlp_url)
             # check(
             #     self.otlp_url,
             #     headers=self.headers,
@@ -205,7 +203,7 @@ class Tracing(metaclass=Singleton):
         is_ready = True
 
         with suppress():
-            if trace_id is not None:
+            if self.inline and trace_id:
                 is_ready = self.inline.is_ready(trace_id)
 
         return is_ready
@@ -217,16 +215,31 @@ class Tracing(metaclass=Singleton):
         _inline_trace = {}
 
         with suppress():
-            is_ready = self.inline.is_ready(trace_id)
+            if self.inline and trace_id:
+                is_ready = self.inline.is_ready(trace_id)
 
-            if is_ready is True:
-                otel_spans = self.inline.fetch(trace_id)
+                if is_ready is True:
+                    otel_spans = self.inline.fetch(trace_id)
 
-                if otel_spans:
-                    _inline_trace = parse_inline_trace(otel_spans)
+                    if otel_spans:
+                        _inline_trace = parse_inline_trace(otel_spans)
 
         return _inline_trace
 
+    def extract(
+        self,
+        *args,
+        **kwargs,
+    ):
+        return extract(*args, **kwargs)
+
+    def inject(
+        self,
+        *args,
+        **kwargs,
+    ):
+        return inject(*args, **kwargs)
+
 
 def get_tracer(
     tracing: Tracing,

agenta/sdk/utils/cache.py

@@ -0,0 +1,53 @@
+from typing import Optional
+from os import getenv
+from time import time
+from collections import OrderedDict
+from threading import Lock
+
+CACHE_CAPACITY = int(getenv("AGENTA_MIDDLEWARE_CACHE_CAPACITY", "512"))
+CACHE_TTL = int(getenv("AGENTA_MIDDLEWARE_CACHE_TTL", str(5 * 60)))  # 5 minutes
+
+
+class TTLLRUCache:
+    def __init__(
+        self,
+        capacity: Optional[int] = CACHE_CAPACITY,
+        ttl: Optional[int] = CACHE_TTL,
+    ):
+        self.cache = OrderedDict()
+        self.capacity = capacity
+        self.ttl = ttl
+        self.lock = Lock()
+
+    def get(self, key):
+        with self.lock:
+            # Get
+            value, expiry = self.cache.get(key, (None, None))
+
+            # Null check
+            if value is None:
+                return None
+
+            # TTL check
+            if time() > expiry:
+                del self.cache[key]
+                return None
+
+            # LRU update
+            self.cache.move_to_end(key)
+
+            return value
+
+    def put(self, key, value):
+        with self.lock:
+            try:
+                # LRU update
+                self.cache.move_to_end(key)
+
+            except KeyError:
+                # Capacity check
+                if len(self.cache) >= self.capacity:
+                    self.cache.popitem(last=False)
+
+            # Put
+            self.cache[key] = (value, time() + self.ttl)

{agenta-0.34.6.dist-info → agenta-0.35.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: agenta
-Version: 0.34.6
+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.6.dist-info → agenta-0.35.0.dist-info}/RECORD

@@ -197,20 +197,20 @@ agenta/docker/docker-assets/lambda_function.py,sha256=h4UZSSfqwpfsCgERv6frqwm_4J
 agenta/docker/docker-assets/main.py,sha256=7MI-21n81U7N7A0GxebNi0cmGWtJKcR2sPB6FcH2QfA,251
 agenta/docker/docker_utils.py,sha256=rKCSb3fDPizX8zwRBaLR0oQ0TnAmyJlzBhflICcwynE,3577
 agenta/sdk/__init__.py,sha256=WFYRfWh6IoYPSzMG2WF2Xz5amLQtzfWVscT0Q9OLpFY,2109
-agenta/sdk/agenta_init.py,sha256=4BkwWvLOGd82BElGfv60iDqr32OKvJqivRWSpAXjRGs,8019
+agenta/sdk/agenta_init.py,sha256=VmjOfzS5ioaUjsFxmSB1k4M4lDgFq73HbvqcId0IbuI,7719
 agenta/sdk/assets.py,sha256=4FN1IfZfBx9g6jW8KzT9ei4BQei9-7MasxHHD9UdeMg,7752
 agenta/sdk/client.py,sha256=trKyBOYFZRk0v5Eptxvh87yPf50Y9CqY6Qgv4Fy-VH4,2142
 agenta/sdk/context/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 agenta/sdk/context/exporting.py,sha256=16X8fgMhl58gehSlqANX97FiKxx4TkGiG4d2B0-7ZX0,516
 agenta/sdk/context/routing.py,sha256=FEsjw8EttI1SMyUo96ptcUsvHJnhoKwdr1szlkxxJNU,598
-agenta/sdk/context/tracing.py,sha256=zp7T_wLVkR-V1c0k7UAN69rwH9VV7MhoZD_IdNu-_RE,649
+agenta/sdk/context/tracing.py,sha256=xjErrXP1Nq1AfL-Cif1l-lNEfs12eQ3v_VCRgoKe7nY,743
 agenta/sdk/decorators/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-agenta/sdk/decorators/routing.py,sha256=FMRc6dSCc1lMblr8Eykj3A2tVI_N-Jjm6IJ6HxTY4CY,24730
-agenta/sdk/decorators/tracing.py,sha256=Wf7KWR1NknXlQEdB2sjMFrcSI-tSSvO9XubN-Ro5tkU,9216
+agenta/sdk/decorators/routing.py,sha256=a6llLTQTs-eGZ1023hyboNVJBeYcqI4zKEShRLvhTXk,25374
+agenta/sdk/decorators/tracing.py,sha256=tIi5gpwzvXd26pFW-7kYQtTFO8w1zzwFeAPY2Ldq4P4,9952
 agenta/sdk/litellm/__init__.py,sha256=Bpz1gfHQc0MN1yolWcjifLWznv6GjHggvRGQSpxpihM,37
 agenta/sdk/litellm/litellm.py,sha256=Xx_YJu05BYkmIe6uObjS6DwzjgwfNytGWf807Zh0vcU,10153
 agenta/sdk/litellm/mockllm.py,sha256=8V6dqdv8eA4P-VoXIwHNYlIjHG189P14POSfSfluVw0,678
-agenta/sdk/litellm/mocks/__init__.py,sha256=c-afSm0YkiHyyaLegvtFs6WuWZTl2Q7Fq_iUoxqBbQc,616
+agenta/sdk/litellm/mocks/__init__.py,sha256=UvunEgoDcVobBZXrqTnHt9ZEjJEFtqFHGiSCKfNiuHU,952
 agenta/sdk/managers/__init__.py,sha256=SN-LRwG0pRRDV3u2Q4JiiSTigN3-mYpzGNM35RzT4mc,238
 agenta/sdk/managers/apps.py,sha256=RSE5ER-a0IHo7iZrcQLshQZMemPTfAVsBfR8H1J-ryQ,1713
 agenta/sdk/managers/config.py,sha256=2LFl7aBF7FKd6UKhqcbnDXzlP8Rsj7hnCu5wBllVQzc,7753
@@ -220,25 +220,26 @@ agenta/sdk/managers/shared.py,sha256=vUhO2thObHxCIowyiM_KFsxPY2q6Rv_lHr2Up8WvFjg
 agenta/sdk/managers/variant.py,sha256=A5ga3mq3b0weUTXa9HO72MGaspthGcu1uK9K5OnP738,4172
 agenta/sdk/managers/vault.py,sha256=054ce9X_xKa2M4NtQWz-GugO6q_pYVWCP3IxbAJJcRw,337
 agenta/sdk/middleware/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-agenta/sdk/middleware/auth.py,sha256=XC0EKRu0H0a0SQTrUsDmkhHGfNFcKMNjL9Btna6ON_U,9835
-agenta/sdk/middleware/cache.py,sha256=-_e3_6f4cPfpuUpGRK_ZQvUx_54ez68db4SnBwfuESk,1078
-agenta/sdk/middleware/config.py,sha256=Mckj-xwIWj06fpMLrz43Ookn8Q7MuvqkF1mDmG17l20,7614
+agenta/sdk/middleware/auth.py,sha256=jdv9GIsfaWb_MsrzwR2iLhHY0YwKgIqqGQbczTWpC2k,9599
+agenta/sdk/middleware/config.py,sha256=-THnZp-4svzduei-ioRdgyk0ZBRWK10_6pLF-4QKXlE,7824
 agenta/sdk/middleware/cors.py,sha256=q3r7lGkrIdMcT_vuhsburMcjG7pyl7w0ycxrIrGJ2e8,921
-agenta/sdk/middleware/inline.py,sha256=utPuSn34WI_yfLK46xRYlJhRdvKvNA1Z1lWhhOLXoa8,903
-agenta/sdk/middleware/mock.py,sha256=kMG-yefb4N39Ds7VuniwjW5_IPWbnauhiLFyC1Bkyak,760
-agenta/sdk/middleware/otel.py,sha256=fXAYxKOdrCuWdDAYyWSCX4VLh5aNCWicAFi2C-p7wZY,1043
-agenta/sdk/middleware/vault.py,sha256=AsRa1mo9cNAh7p162U3NUhMlJpRzxOnojacqXL8nYnA,4313
+agenta/sdk/middleware/inline.py,sha256=ee8E4XBGcRSrHTvblqX1yRXuTN_sxLm7lY1jnywrBG8,901
+agenta/sdk/middleware/mock.py,sha256=bCUN9iJBxePyN9MBwBpJs-_iCNkUQeUjIIu3WElS1oQ,759
+agenta/sdk/middleware/otel.py,sha256=ShaB7iVaUjdq6lOXf0vmZthQsDUjAgIrseyPeJzVkEU,840
+agenta/sdk/middleware/vault.py,sha256=R4uLcGWuS0oigAjF-80jvqCScrF1_fpGpdYmA2D5A5A,4233
 agenta/sdk/router.py,sha256=mOguvtOwl2wmyAgOuWTsf98pQwpNiUILKIo67W_hR3A,119
 agenta/sdk/tracing/__init__.py,sha256=rQNe5-zT5Kt7_CDhq-lnUIi1EYTBVzVf_MbfcIxVD98,41
 agenta/sdk/tracing/attributes.py,sha256=zh8JQZSeYCLBeIRSopKJx6QQ-WEgw08Cr64DS_WOcT8,3833
 agenta/sdk/tracing/conventions.py,sha256=JBtznBXZ3aRkGKkLl7cPwdMNh3w1G-H2Ta2YrAxbr38,950
-agenta/sdk/tracing/exporters.py,sha256=PuSbamuLa4e9eGBM6tp7Smh5R41Jz07VBoBtMnkbtz8,2687
-agenta/sdk/tracing/inline.py,sha256=BQecJkZDlR85aodfvye-LuDhouee3RjLuUl7GMpfsSw,31282
-agenta/sdk/tracing/processors.py,sha256=d7MvJ_DRAP0RAlp8V8XrrY-m8pJ03nLxKoq15RnjedA,3334
+agenta/sdk/tracing/exporters.py,sha256=mnqZXbien7uzEdwfm06BFOZnC8JVJvNWI4ODSyMUeSc,3160
+agenta/sdk/tracing/inline.py,sha256=tY9mH1zMlJMqf2WWKJdakLtMHJKNeYRaynqvsPeqpTY,31312
+agenta/sdk/tracing/processors.py,sha256=HR7pSzHtV8WWr2qEeAki0yznWWQCKCW3HGm1nbVgmEM,4252
+agenta/sdk/tracing/propagation.py,sha256=EeOqDMqnh_MoEhGd1do_vy_tQBYUcoC8kpLqVoZeqg0,2561
 agenta/sdk/tracing/spans.py,sha256=nqUOjjirBxB8Eacv8Qj4Ra_6rknGi3lbJdNyKmk5ODQ,3707
-agenta/sdk/tracing/tracing.py,sha256=iZCIggUkR5f2zVcq0o7T8-to7pcFO349gUEBYXnBMzg,6888
+agenta/sdk/tracing/tracing.py,sha256=kWEt9AdUnTetgq7mR-fnap2isgL3NBg44cG7WFhZ_GU,7162
 agenta/sdk/types.py,sha256=HvFok4lEFxWow6WT71OrJ23WhTsj_FqqMyzlKByAqYk,19120
 agenta/sdk/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+agenta/sdk/utils/cache.py,sha256=69pQCtmh10DiYoOdKv4t2fCYDQCuWZtH4uy36kXewpw,1372
 agenta/sdk/utils/constants.py,sha256=zW3R4rjXOo2L5lz6q84l_zYuOM9u4mpPRHw_B1Dr_hI,67
 agenta/sdk/utils/costs.py,sha256=i8C7ud__pThLS55XkN4YW8czXtGeXr2mx7jjcOFeiXg,5955
 agenta/sdk/utils/exceptions.py,sha256=8FQP05ecHLpq6zSj3MbfdNp9SqP0a7iC_dA5uAe0vkE,1455
@@ -248,7 +249,7 @@ agenta/sdk/utils/logging.py,sha256=eFzEFuYpggfIhEKv09JZRqcDzkmZ482a_E2G-X0FK7Y,4
 agenta/sdk/utils/preinit.py,sha256=YlJL7RLfel0R7DFp-jK7OV-z4ZIQJM0oupYlk7g8b5o,1278
 agenta/sdk/utils/singleton.py,sha256=17Ph7LGnnV8HkPjImruKita2ni03Ari5jr0jqm__4sc,312
 agenta/sdk/utils/timing.py,sha256=rmBPSBuUnIu-epocUCVk0KcM2r36HuDoxkFqOZgfPhc,1507
-agenta-0.34.6.dist-info/METADATA,sha256=htEwbEb29csfhOVG0QBzRQLfuKkfw_WmIgv7XuI_ppY,29848
-agenta-0.34.6.dist-info/WHEEL,sha256=XbeZDeTWKc1w7CSIyre5aMDU_-PohRwTQceYnisIYYY,88
-agenta-0.34.6.dist-info/entry_points.txt,sha256=PDiu8_8AsL7ibU9v4iNoOKR1S7F2rdxjlEprjM9QOgo,46
-agenta-0.34.6.dist-info/RECORD,,
+agenta-0.35.0.dist-info/METADATA,sha256=xmUP8deXxTXgcdI1_2n4ROYik563Smi2D99mAmBwgd4,29827
+agenta-0.35.0.dist-info/WHEEL,sha256=XbeZDeTWKc1w7CSIyre5aMDU_-PohRwTQceYnisIYYY,88
+agenta-0.35.0.dist-info/entry_points.txt,sha256=PDiu8_8AsL7ibU9v4iNoOKR1S7F2rdxjlEprjM9QOgo,46
+agenta-0.35.0.dist-info/RECORD,,

agenta/sdk/middleware/cache.py

@@ -1,47 +0,0 @@
-from os import getenv
-from time import time
-from collections import OrderedDict
-
-CACHE_CAPACITY = int(getenv("AGENTA_MIDDLEWARE_CACHE_CAPACITY", "512"))
-CACHE_TTL = int(getenv("AGENTA_MIDDLEWARE_CACHE_TTL", str(5 * 60)))  # 5 minutes
-
-
-class TTLLRUCache:
-    def __init__(self, capacity: int, ttl: int):
-        self.cache = OrderedDict()
-        self.capacity = capacity
-        self.ttl = ttl
-
-    def get(self, key):
-        # CACHE
-        if key not in self.cache:
-            return None
-
-        value, expiry = self.cache[key]
-        # -----
-
-        # TTL
-        if time() > expiry:
-            del self.cache[key]
-
-            return None
-        # ---
-
-        # LRU
-        self.cache.move_to_end(key)
-        # ---
-
-        return value
-
-    def put(self, key, value):
-        # CACHE
-        if key in self.cache:
-            del self.cache[key]
-        # CACHE & LRU
-        elif len(self.cache) >= self.capacity:
-            self.cache.popitem(last=False)
-        # -----------
-
-        # TTL
-        self.cache[key] = (value, time() + self.ttl)
-        # ---