agenta 0.44.3__tar.gz → 0.45.0__tar.gz

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

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

{agenta-0.44.3 → agenta-0.45.0}/agenta/sdk/decorators/routing.py

@@ -363,20 +363,44 @@ class entrypoint:
         tree = None
         content_type = "text/plain"
         tree_id = None
+        trace_id = None
+        span_id = None
 
         with suppress():
             if isinstance(result, (dict, list)):
                 content_type = "application/json"
             data = self.patch_result(result)
 
-            tree, tree_id = await self.fetch_inline_trace(inline)
+            (
+                tree,
+                tree_id,
+                trace_id,
+                span_id,
+            ) = await self.fetch_inline_trace(inline)
 
         try:
             return BaseResponse(
-                data=data, tree=tree, content_type=content_type, tree_id=tree_id
+                data=data,
+                tree=tree,
+                content_type=content_type,
+                tree_id=tree_id,
+                trace_id=trace_id,
+                span_id=span_id,
             )
         except:  # pylint: disable=bare-except
-            return BaseResponse(data=data, content_type=content_type)
+            try:
+                return BaseResponse(
+                    data=data,
+                    content_type=content_type,
+                    tree_id=tree_id,
+                    trace_id=trace_id,
+                    span_id=span_id,
+                )
+            except:  # pylint: disable=bare-except
+                return BaseResponse(
+                    data=data,
+                    content_type=content_type,
+                )
 
     def handle_failure(
         self,
@@ -444,8 +468,10 @@ class entrypoint:
 
         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
+        _trace_id = link.get("trace_id") if link else None  # in int format
+        tree_id = (
+            str(UUID(int=_trace_id)) if _trace_id else None
+        )  # in uuid_as_str format
 
         return tree_id
 
@@ -461,24 +487,28 @@ class entrypoint:
 
         link = context.link
 
+        _trace_id = link.get("trace_id") if link else None  # in int format
+        _span_id = link.get("span_id") if link else None  # in int format
+
         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 = str(UUID(int=_trace_id)) if _trace_id else None
+        trace_id = UUID(int=_trace_id).hex if _trace_id else None
+        span_id = UUID(int=_span_id).hex[16:] if _span_id else None
 
-        if _tree_id is not None:
+        if _trace_id is not None:
             if inline:
                 remaining_steps = NOFSTEPS
                 while (
-                    not ag.tracing.is_inline_trace_ready(_tree_id)
+                    not ag.tracing.is_inline_trace_ready(_trace_id)
                     and remaining_steps > 0
                 ):
                     await sleep(TIMESTEP)
 
                     remaining_steps -= 1
 
-                tree = ag.tracing.get_inline_trace(_tree_id)
+                tree = ag.tracing.get_inline_trace(_trace_id)
 
-        return tree, tree_id
+        return tree, tree_id, trace_id, span_id
 
     # --- OpenAPI --- #
 

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

@@ -131,8 +131,8 @@ class instrument:  # pylint: disable=invalid-name
 
         if not context.link:
             context.link = {
-                "tree_id": span.get_span_context().trace_id,
-                "node_id": span.get_span_context().span_id,
+                "trace_id": span.get_span_context().trace_id,
+                "span_id": span.get_span_context().span_id,
             }
 
             tracing_context.set(context)
@@ -297,9 +297,7 @@ 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.44.3 → agenta-0.45.0}/agenta/sdk/managers/apps.py

@@ -12,11 +12,6 @@ AppType = Literal["SERVICE:completion", "SERVICE:chat", "CUSTOM"]
 DEFAULT_APP_TYPE = "SERVICE:completion"
 
 
-AppType = Literal["SERVICE:completion", "SERVICE:chat", "CUSTOM"]
-
-DEFAULT_APP_TYPE = "SERVICE:completion"
-
-
 class AppManager:
     @classmethod
     @handle_exceptions()

{agenta-0.44.3 → agenta-0.45.0}/agenta/sdk/tracing/inline.py

@@ -160,6 +160,9 @@ class OTelExtraDTO(BaseModel):
 
 
 class SpanDTO(BaseModel):
+    trace_id: str
+    span_id: str
+
     scope: Optional[ProjectScopeDTO] = None
 
     lifecycle: Optional[LifecycleDTO] = None
@@ -792,6 +795,9 @@ def _parse_from_attributes(
 def parse_from_otel_span_dto(
     otel_span_dto: OTelSpanDTO,
 ) -> SpanDTO:
+    trace_id = str(otel_span_dto.context.trace_id[2:])
+    span_id = str(otel_span_dto.context.span_id[2:])
+
     lifecyle = LifecycleDTO(
         created_at=datetime.now(),
     )
@@ -800,7 +806,7 @@ def parse_from_otel_span_dto(
 
     types = _parse_from_types(otel_span_dto)
 
-    tree_id = UUID(otel_span_dto.context.trace_id[2:])
+    tree_id = UUID(trace_id)
 
     tree_type: str = types.get("tree")
 
@@ -809,7 +815,7 @@ def parse_from_otel_span_dto(
         type=tree_type.lower() if tree_type else None,
     )
 
-    node_id = UUID(tree_id.hex[16:] + otel_span_dto.context.span_id[2:])
+    node_id = UUID(trace_id[16:] + span_id)
 
     node_type = NodeType.TASK
     try:
@@ -871,6 +877,8 @@ def parse_from_otel_span_dto(
     )
 
     span_dto = SpanDTO(
+        trace_id=trace_id,
+        span_id=span_id,
         lifecycle=lifecyle,
         root=root,
         tree=tree,

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

@@ -1,8 +1,9 @@
 from typing import Optional, Dict, List
+from threading import Lock
 
 from opentelemetry.baggage import get_all as get_baggage
 from opentelemetry.context import Context
-from opentelemetry.sdk.trace import Span
+from opentelemetry.sdk.trace import Span, SpanProcessor
 from opentelemetry.sdk.trace.export import (
     SpanExporter,
     ReadableSpan,
@@ -154,3 +155,26 @@ class TraceProcessor(BatchSpanProcessor):
         # --- INLINE
 
         return trace
+
+
+invocation_link = {}
+invocation_link_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()
+
+            with invocation_link_lock:
+                invocation_link["trace_id"] = ctx.trace_id
+                invocation_link["span_id"] = ctx.span_id
+
+
+def use_invocation_link() -> Optional[Dict[str, str]]:
+    with invocation_link_lock:
+        if invocation_link:
+            return invocation_link.copy()
+        else:
+            return None

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

@@ -1,8 +1,11 @@
 from typing import Optional, Any, Dict, Callable
 from enum import Enum
+from uuid import UUID
 
+from pydantic import BaseModel
 from httpx import get as check
 
+
 from opentelemetry.trace import (
     get_current_span,
     set_tracer_provider,
@@ -10,13 +13,19 @@ from opentelemetry.trace import (
     Status,
     StatusCode,
 )
+from opentelemetry.sdk import trace
 from opentelemetry.sdk.trace import Span, Tracer, TracerProvider
 from opentelemetry.sdk.resources import Resource
 
+
 from agenta.sdk.utils.singleton import Singleton
 from agenta.sdk.utils.exceptions import suppress
 from agenta.sdk.utils.logging import get_module_logger
-from agenta.sdk.tracing.processors import TraceProcessor
+from agenta.sdk.tracing.processors import (
+    TraceProcessor,
+    InvocationLinkHook,
+    use_invocation_link,
+)
 from agenta.sdk.tracing.exporters import InlineExporter, OTLPExporter
 from agenta.sdk.tracing.spans import CustomSpan
 from agenta.sdk.tracing.inline import parse_inline_trace
@@ -24,9 +33,27 @@ 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
 
+from agenta.sdk.context.tracing import tracing_context
+
 log = get_module_logger(__name__)
 
 
+_original_init = trace.TracerProvider.__init__
+
+
+def patched_init(self, *args, **kwargs):
+    _original_init(self, *args, **kwargs)
+    self.add_span_processor(InvocationLinkHook())
+
+
+trace.TracerProvider.__init__ = patched_init
+
+
+class Link(BaseModel):
+    trace_id: str
+    span_id: str
+
+
 class Tracing(metaclass=Singleton):
     VERSION = "0.1.0"
 
@@ -242,6 +269,40 @@ 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,
+                    )
+
+        return None
+
 
 def get_tracer(
     tracing: Tracing,

{agenta-0.44.3 → agenta-0.45.0}/agenta/sdk/types.py

@@ -39,6 +39,8 @@ class BaseResponse(BaseModel):
     content_type: Optional[str] = "string"
     tree: Optional[AgentaNodesResponse] = None
     tree_id: Optional[str] = None
+    trace_id: Optional[str] = None
+    span_id: Optional[str] = None
 
     model_config = ConfigDict(use_enum_values=True, exclude_none=True)
 

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

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