agenta 0.45.0__tar.gz → 0.45.2__tar.gz

{agenta-0.45.0 → agenta-0.45.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: agenta
-Version: 0.45.0
+Version: 0.45.2
 Summary: The SDK for agenta is an open-source LLMOps platform.
 Keywords: LLMOps,LLM,evaluation,prompt engineering
 Author: Mahmoud Mabrouk

{agenta-0.45.0 → agenta-0.45.2}/agenta/sdk/decorators/tracing.py

@@ -297,7 +297,9 @@ class instrument:  # pylint: disable=invalid-name
             not in (
                 ignore
                 if isinstance(ignore, list)
-                else io.keys() if ignore is True else []
+                else io.keys()
+                if ignore is True
+                else []
             )
         }
 

{agenta-0.45.0 → agenta-0.45.2}/agenta/sdk/litellm/mockllm.py

@@ -1,23 +1,31 @@
 from typing import Optional, Protocol, Any
+from os import environ
+
+from agenta.sdk.utils.logging import get_module_logger
 
 from agenta.sdk.litellm.mocks import MOCKS
 from agenta.sdk.context.routing import routing_context
 
+AGENTA_LITELLM_MOCK = environ.get("AGENTA_LITELLM_MOCK") or None
+
+log = get_module_logger(__name__)
+
 
 class LitellmProtocol(Protocol):
-    async def acompletion(self, *args: Any, **kwargs: Any) -> Any:
-        ...
+    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
+    mock = AGENTA_LITELLM_MOCK or routing_context.get().mock
 
     if mock:
+        log.debug("Mocking litellm: %s.", mock)
+
         if mock not in MOCKS:
-            raise ValueError(f"Mock {mock} not found")
+            mock = "hello"
 
         return MOCKS[mock](*args, **kwargs)
 

{agenta-0.45.0 → agenta-0.45.2}/agenta/sdk/litellm/mocks/__init__.py

@@ -1,4 +1,5 @@
 from typing import Callable
+from asyncio import sleep
 
 from pydantic import BaseModel
 
@@ -30,7 +31,26 @@ def hello_mock_response(*args, **kwargs) -> MockResponseModel:
 def chat_mock_response(*args, **kwargs) -> MockResponseModel:
     return MockResponseModel(
         choices=[
-            MockChoiceModel(message=MockMessageModel(content="world", role="assistant"))
+            MockChoiceModel(
+                message=MockMessageModel(
+                    content="world",
+                    role="assistant",
+                )
+            )
+        ],
+    )
+
+
+def delay_mock_response(*args, **kwargs) -> MockResponseModel:
+    sleep(2)
+
+    return MockResponseModel(
+        choices=[
+            MockChoiceModel(
+                message=MockMessageModel(
+                    content="delay",
+                )
+            )
         ],
     )
 
@@ -38,4 +58,5 @@ def chat_mock_response(*args, **kwargs) -> MockResponseModel:
 MOCKS: dict[str, Callable[..., MockResponseModel]] = {
     "hello": hello_mock_response,
     "chat": chat_mock_response,
+    "delay": delay_mock_response,
 }

{agenta-0.45.0 → agenta-0.45.2}/agenta/sdk/middleware/config.py

@@ -92,15 +92,21 @@ class ConfigMiddleware(BaseHTTPMiddleware):
                 return parameters, references
 
         config = {}
-        async with httpx.AsyncClient() as client:
-            response = await client.post(
-                f"{self.host}/api/variants/configs/fetch",
-                headers=headers,
-                json=refs,
-            )
-
-            if response.status_code == 200:
-                config = response.json()
+
+        is_test_path = request.url.path.endswith("/test")
+        are_refs_missing = not variant_ref and not environment_ref
+        should_fetch = not is_test_path or not are_refs_missing
+
+        if should_fetch:
+            async with httpx.AsyncClient() as client:
+                response = await client.post(
+                    f"{self.host}/api/variants/configs/fetch",
+                    headers=headers,
+                    json=refs,
+                )
+
+                if response.status_code == 200:
+                    config = response.json()
 
         if not config:
             config["application_ref"] = refs[

{agenta-0.45.0 → agenta-0.45.2}/agenta/sdk/tracing/processors.py

@@ -157,24 +157,30 @@ class TraceProcessor(BatchSpanProcessor):
         return trace
 
 
-invocation_link = {}
-invocation_link_lock = Lock()
+# Internal storage for the last ended span context
+_last_ended_span_context = None
+_lock = Lock()
 
 
-class InvocationLinkHook(SpanProcessor):
-    def on_end(self, span):
-        # Only capture root spans if you want
-        if span.parent is None:
-            ctx = span.get_span_context()
+def _set_last_ended(span_ctx) -> None:
+    """Set the last ended span context"""
+    with _lock:
+        global _last_ended_span_context
+        _last_ended_span_context = span_ctx
 
-            with invocation_link_lock:
-                invocation_link["trace_id"] = ctx.trace_id
-                invocation_link["span_id"] = ctx.span_id
 
+def _get_last_ended():
+    """Get the last ended span context"""
+    with _lock:
+        return _last_ended_span_context
 
-def use_invocation_link() -> Optional[Dict[str, str]]:
-    with invocation_link_lock:
-        if invocation_link:
-            return invocation_link.copy()
-        else:
-            return None
+
+class EndedSpanRecorder(SpanProcessor):
+    """Records the last ended span context for later reference.
+
+    This allows accessing span information even after the span has been ended,
+    which is useful for linking annotations to auto-instrumented spans.
+    """
+
+    def on_end(self, span):
+        _set_last_ended(span.get_span_context())

{agenta-0.45.0 → agenta-0.45.2}/agenta/sdk/tracing/tracing.py

@@ -23,8 +23,8 @@ from agenta.sdk.utils.exceptions import suppress
 from agenta.sdk.utils.logging import get_module_logger
 from agenta.sdk.tracing.processors import (
     TraceProcessor,
-    InvocationLinkHook,
-    use_invocation_link,
+    EndedSpanRecorder,
+    _get_last_ended,
 )
 from agenta.sdk.tracing.exporters import InlineExporter, OTLPExporter
 from agenta.sdk.tracing.spans import CustomSpan
@@ -43,7 +43,7 @@ _original_init = trace.TracerProvider.__init__
 
 def patched_init(self, *args, **kwargs):
     _original_init(self, *args, **kwargs)
-    self.add_span_processor(InvocationLinkHook())
+    self.add_span_processor(EndedSpanRecorder())
 
 
 trace.TracerProvider.__init__ = patched_init
@@ -269,37 +269,58 @@ class Tracing(metaclass=Singleton):
     ):
         return inject(*args, **kwargs)
 
-    def get_invocation_link(self) -> Optional[Link]:
-        with suppress():
-            link = tracing_context.get().link
-
-            trace_id = None
-            span_id = None
-
-            if link:
-                trace_id = UUID(int=link.get("trace_id")).hex if link else None
-                span_id = UUID(int=link.get("span_id")).hex[16:] if link else None
-
-            if trace_id and span_id:
-                return Link(
-                    trace_id=trace_id,
-                    span_id=span_id,
-                )
-                
-            link = use_invocation_link()
-
-            trace_id = None
-            span_id = None
-
-            if link:
-                trace_id = UUID(int=link.get("trace_id")).hex if link else None
-                span_id = UUID(int=link.get("span_id")).hex[16:] if link else None
-
-                if trace_id and span_id:
-                    return Link(
-                        trace_id=trace_id,
-                        span_id=span_id,
-                    )
+    def get_current_span_context(self):
+        """Get the current active span context if available.
+
+        Returns:
+            SpanContext or None if no active span
+        """
+        span = get_current_span()
+        ctx = span.get_span_context()
+        return ctx if ctx and ctx.is_valid else None
+
+    def get_last_span_context(self):
+        """Get the last closed span context if available.
+
+        This is useful for accessing span information after a span has closed,
+        particularly with auto-instrumentation libraries.
+
+        Returns:
+            SpanContext or None if no spans have been closed
+        """
+        return _get_last_ended()
+
+    def get_span_context(self):
+        """Get the most relevant span context.
+
+        First tries to get the current active span context.
+        If no active span exists, falls back to the last closed span.
+
+        Returns:
+            SpanContext or None if no relevant span context is available
+        """
+        return self.get_current_span_context() or self.get_last_span_context()
+
+    def build_invocation_link(self, span_ctx=None) -> Optional[Link]:
+        """
+        Builds a Link object containing the hex-formatted trace_id and span_id
+        from the current (or fallback last ended) span context.
+        Useful to link annotations to spans.
+
+        Args:
+            span_ctx: Optional SpanContext to convert to a Link
+
+        Returns:
+            Link object with trace_id and span_id or None if no valid context
+        """
+        if span_ctx is None:
+            span_ctx = self.get_span_context()
+
+        if span_ctx and span_ctx.is_valid:
+            return Link(
+                trace_id=f"{span_ctx.trace_id:032x}",
+                span_id=f"{span_ctx.span_id:016x}",
+            )
 
         return None
 

{agenta-0.45.0 → agenta-0.45.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "agenta"
-version = "0.45.0"
+version = "0.45.2"
 description = "The SDK for agenta is an open-source LLMOps platform."
 readme = "README.md"
 authors = ["Mahmoud Mabrouk <mahmoud@agenta.ai>"]