agenta 0.52.6__py3-none-any.whl → 0.63.2__py3-none-any.whl

agenta/sdk/tracing/attributes.py

@@ -6,7 +6,7 @@ PrimitivesSequence = Sequence[Primitive]
 Attribute = Union[Primitive, PrimitivesSequence]
 
 
-def _marshal(
+def _marshall(
     unmarshalled: Dict[str, Any],
     *,
     parent_key: Optional[str] = "",
@@ -59,7 +59,7 @@ def _marshal(
             dict_key = child_key
 
             marshalled.update(
-                _marshal(
+                _marshall(
                     value,
                     parent_key=dict_key,
                     depth=depth + 1,
@@ -76,7 +76,7 @@ def _marshal(
 
                     if isinstance(item, (dict, list)):
                         marshalled.update(
-                            _marshal(
+                            _marshall(
                                 item,
                                 parent_key=list_key,
                                 depth=depth + 1,
@@ -177,7 +177,7 @@ def serialize(
         k: v
         for k, v in {
             _encode_key(namespace, key): _encode_value(value)
-            for key, value in _marshal(attributes, max_depth=max_depth).items()
+            for key, value in _marshall(attributes, max_depth=max_depth).items()
         }.items()
         if v is not None
     }

agenta/sdk/tracing/exporters.py

@@ -1,4 +1,7 @@
-from typing import Sequence, Dict, List, Optional
+from typing import Sequence, Dict, List, Optional, Any
+from threading import Thread
+from os import environ
+from uuid import UUID
 
 from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
 from opentelemetry.sdk.trace.export import (
@@ -8,18 +11,21 @@ from opentelemetry.sdk.trace.export import (
     ReadableSpan,
 )
 
+from agenta.sdk.utils.constants import TRUTHY
 from agenta.sdk.utils.logging import get_module_logger
 from agenta.sdk.utils.exceptions import suppress
 from agenta.sdk.utils.cache import TTLLRUCache
-from agenta.sdk.context.tracing import (
-    tracing_exporter_context_manager,
-    tracing_exporter_context,
-    TracingExporterContext,
+from agenta.sdk.contexts.tracing import (
+    otlp_context_manager,
+    otlp_context,
+    OTLPContext,
 )
 
 
 log = get_module_logger(__name__)
 
+_ASYNC_EXPORT = environ.get("AGENTA_OTLP_ASYNC_EXPORT", "true").lower() in TRUTHY
+
 
 class InlineTraceExporter(SpanExporter):
     def __init__(
@@ -45,6 +51,8 @@ class InlineTraceExporter(SpanExporter):
 
                 self._registry[trace_id].append(span)
 
+            return
+
     def shutdown(self) -> None:
         self._shutdown = True
 
@@ -84,7 +92,7 @@ class OTLPExporter(OTLPSpanExporter):
         self.credentials = credentials
 
     def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
-        grouped_spans: Dict[str, List[str]] = {}
+        grouped_spans: Dict[Optional[str], List[ReadableSpan]] = dict()
 
         for span in spans:
             trace_id = span.get_span_context().trace_id
@@ -94,18 +102,28 @@ class OTLPExporter(OTLPSpanExporter):
                 credentials = self.credentials.get(trace_id)
 
             if credentials not in grouped_spans:
-                grouped_spans[credentials] = []
+                grouped_spans[credentials] = list()
 
             grouped_spans[credentials].append(span)
 
         serialized_spans = []
 
         for credentials, _spans in grouped_spans.items():
-            with tracing_exporter_context_manager(
-                context=TracingExporterContext(
+            with otlp_context_manager(
+                context=OTLPContext(
                     credentials=credentials,
                 )
             ):
+                for _span in _spans:
+                    trace_id = _span.get_span_context().trace_id
+                    span_id = _span.get_span_context().span_id
+
+                    # log.debug(
+                    #     "[SPAN]  [EXPORT]",
+                    #     trace_id=UUID(int=trace_id).hex,
+                    #     span_id=UUID(int=span_id).hex[-16:],
+                    # )
+
                 serialized_spans.append(super().export(_spans))
 
         if all(serialized_spans):
@@ -114,16 +132,48 @@ class OTLPExporter(OTLPSpanExporter):
             return SpanExportResult.FAILURE
 
     def _export(self, serialized_data: bytes, timeout_sec: Optional[float] = None):
-        credentials = tracing_exporter_context.get().credentials
+        try:
+            credentials = otlp_context.get().credentials
+
+            if credentials:
+                self._session.headers.update({"Authorization": credentials})
+
+            def __export():
+                with suppress():
+                    resp = None
+                    if timeout_sec is not None:
+                        resp = super(OTLPExporter, self)._export(
+                            serialized_data,
+                            timeout_sec,
+                        )
+                    else:
+                        resp = super(OTLPExporter, self)._export(
+                            serialized_data,
+                        )
+
+                    # log.debug(
+                    #     "[SPAN] [_EXPORT]",
+                    #     data=serialized_data,
+                    #     resp=resp,
+                    # )
+
+            if _ASYNC_EXPORT is True:
+                # log.debug("[SPAN] [ASYNC.X]", credentials=(credentials is not None))
+                thread = Thread(target=__export, daemon=False)
+                thread.start()
+            else:
+                # log.debug("[SPAN] [ SYNC.X]", credentials=(credentials is not None))
+                return __export()
 
-        if credentials:
-            self._session.headers.update({"Authorization": credentials})
+        except Exception as e:
+            log.error(f"Export failed with error: {e}", exc_info=True)
 
-        with suppress():
-            if timeout_sec is not None:
-                return super()._export(serialized_data, timeout_sec)
-            else:
-                return super()._export(serialized_data)
+        finally:
+
+            class Response:
+                ok = True
+
+            return Response()
 
 
 ConsoleExporter = ConsoleSpanExporter

agenta/sdk/tracing/inline.py

@@ -532,11 +532,11 @@ from json import loads, JSONDecodeError, dumps
 from copy import copy
 
 
-def _unmarshal_attributes(
+def _unmarshall_attributes(
     marshalled: Dict[str, Any],
 ) -> Dict[str, Any]:
     """
-    Unmarshals a dictionary of marshalled attributes into a nested dictionary
+    Unmarshalls a dictionary of marshalled attributes into a nested dictionary
 
     Example:
     marshalled = {
@@ -566,42 +566,34 @@ def _unmarshal_attributes(
 
     for key, value in marshalled.items():
         keys = key.split(".")
-
-        level = unmarshalled
-
-        for i, part in enumerate(keys[:-1]):
-            if part.isdigit():
-                part = int(part)
-
-                if not isinstance(level, list):
-                    level = []
-
-                while len(level) <= part:
-                    level.append({})
-
-                level = level[part]
-
+        current = unmarshalled
+
+        for i, key in enumerate(keys):
+            is_last = i == len(keys) - 1
+            next_key = keys[i + 1] if not is_last else None
+            is_index = key.isdigit()
+            key = int(key) if is_index else key
+
+            if is_last:
+                if isinstance(current, list) and isinstance(key, int):
+                    while len(current) <= key:
+                        current.append(None)
+                    current[key] = value
+                elif isinstance(current, dict):
+                    current[key] = value
             else:
-                if part not in level:
-                    level[part] = {} if not keys[i + 1].isdigit() else []
-
-                level = level[part]
-
-        last_key = keys[-1]
-
-        if last_key.isdigit():
-            last_key = int(last_key)
-
-            if not isinstance(level, list):
-                level = []
-
-            while len(level) <= last_key:
-                level.append(None)
-
-            level[last_key] = value
-
-        else:
-            level[last_key] = value
+                next_is_index = next_key.isdigit() if next_key else False
+
+                if isinstance(current, list) and isinstance(key, int):
+                    while len(current) <= key:
+                        current.append([] if next_is_index else {})
+                    if current[key] is None:
+                        current[key] = [] if next_is_index else {}
+                    current = current[key]
+                elif isinstance(current, dict):
+                    if key not in current:
+                        current[key] = [] if next_is_index else {}
+                    current = current[key]
 
     return unmarshalled
 
@@ -750,7 +742,7 @@ def _parse_from_attributes(
     for key in _data.keys():
         del otel_span_dto.attributes[_encode_key("data", key)]
 
-    # _data = _unmarshal_attributes(_data)
+    # _data = _unmarshall_attributes(_data)
     _data = _data if _data else None
 
     # METRICS
@@ -759,7 +751,7 @@ def _parse_from_attributes(
     for key in _metrics.keys():
         del otel_span_dto.attributes[_encode_key("metrics", key)]
 
-    # _metrics = _unmarshal_attributes(_metrics)
+    # _metrics = _unmarshall_attributes(_metrics)
     _metrics = _metrics if _metrics else None
 
     # META
@@ -768,7 +760,7 @@ def _parse_from_attributes(
     for key in _meta.keys():
         del otel_span_dto.attributes[_encode_key("meta", key)]
 
-    # _meta = _unmarshal_attributes(_meta)
+    # _meta = _unmarshall_attributes(_meta)
     _meta = _meta if _meta else None
 
     # TAGS
@@ -904,7 +896,7 @@ def parse_to_agenta_span_dto(
 ) -> SpanDTO:
     # DATA
     if span_dto.data:
-        span_dto.data = _unmarshal_attributes(span_dto.data)
+        span_dto.data = _unmarshall_attributes(span_dto.data)
 
         if "outputs" in span_dto.data:
             if "__default__" in span_dto.data["outputs"]:
@@ -912,19 +904,19 @@ def parse_to_agenta_span_dto(
 
     # METRICS
     if span_dto.metrics:
-        span_dto.metrics = _unmarshal_attributes(span_dto.metrics)
+        span_dto.metrics = _unmarshall_attributes(span_dto.metrics)
 
     # META
     if span_dto.meta:
-        span_dto.meta = _unmarshal_attributes(span_dto.meta)
+        span_dto.meta = _unmarshall_attributes(span_dto.meta)
 
     # TAGS
     if span_dto.tags:
-        span_dto.tags = _unmarshal_attributes(span_dto.tags)
+        span_dto.tags = _unmarshall_attributes(span_dto.tags)
 
     # REFS
     if span_dto.refs:
-        span_dto.refs = _unmarshal_attributes(span_dto.refs)
+        span_dto.refs = _unmarshall_attributes(span_dto.refs)
 
     if isinstance(span_dto.links, list):
         for link in span_dto.links:

agenta/sdk/tracing/processors.py

@@ -1,5 +1,7 @@
 from typing import Optional, Dict, List
 from threading import Lock
+from json import dumps
+from uuid import UUID
 
 from opentelemetry.baggage import get_all as get_baggage
 from opentelemetry.context import Context
@@ -9,10 +11,13 @@ from opentelemetry.sdk.trace.export import (
     ReadableSpan,
     BatchSpanProcessor,
 )
+from opentelemetry.trace import SpanContext
 
 from agenta.sdk.utils.logging import get_module_logger
 from agenta.sdk.tracing.conventions import Reference
 
+from agenta.sdk.contexts.tracing import TracingContext
+
 log = get_module_logger(__name__)
 
 
@@ -50,6 +55,15 @@ class TraceProcessor(SpanProcessor):
         span: Span,
         parent_context: Optional[Context] = None,
     ) -> None:
+        trace_id = span.context.trace_id
+        span_id = span.context.span_id
+
+        # log.debug(
+        #     "[SPAN] [START] ",
+        #     trace_id=UUID(int=trace_id).hex,
+        #     span_id=UUID(int=span_id).hex[-16:],
+        # )
+
         for key in self.references.keys():
             span.set_attribute(f"ag.refs.{key}", self.references[key])
 
@@ -61,6 +75,83 @@ class TraceProcessor(SpanProcessor):
                 if _key in [_.value for _ in Reference.__members__.values()]:
                     span.set_attribute(key, baggage[key])
 
+        context = TracingContext.get()
+
+        trace_type = span.attributes.get("trace_type") if span.attributes else None
+
+        context.annotate = (
+            context.annotate
+            or (context.type == "annotation")
+            or (trace_type == "annotation")
+        )
+        context.type = (
+            (str(trace_type) if trace_type else None)
+            or context.type
+            or ("annotation" if context.annotate else "invocation")
+        )
+
+        span.set_attribute("ag.type.tree", context.type)
+
+        if context.flags:
+            for key in context.flags.keys():
+                span.set_attribute(f"ag.flags.{key}", context.flags[key])
+        # if context.tags:
+        #     for key in context.tags.keys():
+        #         span.set_attribute(f"ag.tags.{key}", context.tags[key])
+        # if context.meta:
+        #     span.set_attribute(f"ag.meta.", dumps(context.meta))
+
+        # --- DISTRIBUTED
+        if not self.inline:
+            if context.links:
+                for key, link in context.links.items():
+                    try:
+                        link = link.model_dump(mode="json", exclude_none=True)
+                    except:  # pylint: disable=bare-except
+                        pass
+                    if not isinstance(link, dict):
+                        continue
+                    if not link.get("trace_id") or not link.get("span_id"):
+                        continue
+
+                    span.add_link(
+                        context=SpanContext(
+                            trace_id=int(str(link.get("trace_id")), 16),
+                            span_id=int(str(link.get("span_id")), 16),
+                            is_remote=True,
+                        ),
+                        attributes=dict(
+                            key=str(key),
+                        ),
+                    )
+
+        if context.references:
+            for key, ref in context.references.items():
+                try:
+                    ref = ref.model_dump(mode="json", exclude_none=True)
+                except:  # pylint: disable=bare-except
+                    pass
+                if not isinstance(ref, dict):
+                    continue
+                if not ref.get("id") and not ref.get("slug") and not ref.get("version"):
+                    continue
+
+                if ref.get("id"):
+                    span.set_attribute(
+                        f"ag.refs.{key}.id",
+                        str(ref.get("id")),
+                    )
+                if ref.get("slug"):
+                    span.set_attribute(
+                        f"ag.refs.{key}.slug",
+                        str(ref.get("slug")),
+                    )
+                if ref.get("version"):
+                    span.set_attribute(
+                        f"ag.refs.{key}.version",
+                        str(ref.get("version")),
+                    )
+
         trace_id = span.context.trace_id
         span_id = span.context.span_id
 
@@ -74,6 +165,12 @@ class TraceProcessor(SpanProcessor):
         trace_id = span.context.trace_id
         span_id = span.context.span_id
 
+        # log.debug(
+        #     "[SPAN] [END]   ",
+        #     trace_id=UUID(int=trace_id).hex,
+        #     span_id=UUID(int=span_id).hex[-16:],
+        # )
+
         self._spans.setdefault(trace_id, []).append(span)
         self._registry.setdefault(trace_id, {})
         self._registry[trace_id].pop(span_id, None)

agenta/sdk/tracing/propagation.py

@@ -6,6 +6,8 @@ from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapProp
 from opentelemetry.baggage import set_baggage
 from opentelemetry.context import get_current
 
+from agenta.sdk.contexts.tracing import TracingContext
+
 import agenta as ag
 
 
@@ -72,7 +74,7 @@ def inject(
 
     _context = get_current()
 
-    ctx = ag.sdk.context.tracing.tracing_context.get()
+    ctx = TracingContext.get()
 
     # --- Inject traceparent --- #
     try:

agenta/sdk/tracing/spans.py

@@ -35,6 +35,10 @@ class CustomSpan(Span):  # INHERITANCE FOR TYPING ONLY
 
     ## --- PROXY METHODS --- ##
 
+    @property
+    def name(self) -> str:
+        return self._span.name
+
     def get_span_context(self):
         return self._span.get_span_context()
 

agenta/sdk/tracing/tracing.py

@@ -52,8 +52,6 @@ class Link(BaseModel):
 
 
 class Tracing(metaclass=Singleton):
-    VERSION = "0.1.0"
-
     Status = Status
     StatusCode = StatusCode
 
@@ -101,19 +99,6 @@ class Tracing(metaclass=Singleton):
             resource=Resource(attributes={"service.name": "agenta-sdk"})
         )
 
-        # --- 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)
@@ -131,6 +116,19 @@ class Tracing(metaclass=Singleton):
         except:  # pylint: disable=bare-except
             log.warning("Agenta - OLTP unreachable, skipping exports.")
 
+        # --- 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
+
         # GLOBAL TRACER PROVIDER -- INSTRUMENTATION LIBRARIES
         set_tracer_provider(self.tracer_provider)
         # TRACER