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

agenta/sdk/engines/tracing/attributes.py

@@ -0,0 +1,185 @@
+from json import loads, dumps
+from typing import Optional, Union, Sequence, Any, Dict
+
+Primitive = Union[str, int, float, bool, bytes]
+PrimitivesSequence = Sequence[Primitive]
+Attribute = Union[Primitive, PrimitivesSequence]
+
+
+def _marshal(
+    unmarshalled: Dict[str, Any],
+    *,
+    parent_key: Optional[str] = "",
+    depth: Optional[int] = 0,
+    max_depth: Optional[int] = None,
+) -> Dict[str, Any]:
+    """
+    Marshals a dictionary of unmarshalled attributes into a flat dictionary
+
+    Example:
+    unmarshalled = {
+        "ag": {
+            "type": "tree",
+            "node": {
+                "name": "root",
+                "children": [
+                    {
+                        "name": "child1",
+                    },
+                    {
+                        "name": "child2",
+                    }
+                ]
+            }
+        }
+    }
+    marshalled = {
+        "ag.type": "tree",
+        "ag.node.name": "root",
+        "ag.node.children.0.name": "child1",
+        "ag.node.children.1.name": "child2"
+    }
+    """
+    marshalled = {}
+
+    # If max_depth is set and we've reached it,
+    # just return the unmarshalled attributes
+    if max_depth is not None and depth >= max_depth:
+        marshalled[parent_key] = unmarshalled
+        # MISSING ENCODING TO JSON IF NOT PRIMITIVE
+
+        return marshalled
+
+    # Otherwise,
+    # iterate over the unmarshalled attributes and marshall them
+    for key, value in unmarshalled.items():
+        child_key = f"{parent_key}.{key}" if parent_key else key
+
+        if isinstance(value, dict):
+            dict_key = child_key
+
+            marshalled.update(
+                _marshal(
+                    value,
+                    parent_key=dict_key,
+                    depth=depth + 1,
+                    max_depth=max_depth,
+                )
+            )
+        elif isinstance(value, list):
+            if max_depth is not None and depth + 1 >= max_depth:
+                marshalled[child_key] = value
+                # MISSING ENCODING TO JSON IF NOT PRIMITIVE
+            else:
+                for i, item in enumerate(value):
+                    list_key = f"{child_key}.{i}"
+
+                    if isinstance(item, (dict, list)):
+                        marshalled.update(
+                            _marshal(
+                                item,
+                                parent_key=list_key,
+                                depth=depth + 1,
+                                max_depth=max_depth,
+                            )
+                        )
+                    else:
+                        marshalled[list_key] = item
+                        # MISSING ENCODING TO JSON IF NOT PRIMITIVE
+        else:
+            marshalled[child_key] = value
+            # MISSING ENCODING TO JSON IF NOT PRIMITIVE
+
+    return marshalled
+
+
+def _encode_key(
+    namespace: Optional[str] = None,
+    key: str = "",
+) -> str:
+    if namespace is None:
+        return key
+
+    return f"ag.{namespace}.{key}"
+
+
+def _make_serializable(value: Any) -> Any:
+    """
+    Transform complex nested structures into JSON-serializable form.
+    Handles Pydantic models, nested dictionaries and lists recursively.
+    """
+    if value is None or isinstance(value, (str, int, float, bool, bytes)):
+        return value
+
+    # Handle Pydantic objects (prioritize v2 over v1 API)
+    if hasattr(value, "model_dump"):  # Pydantic v2
+        return value.model_dump()
+    elif hasattr(value, "dict"):  # Pydantic v1
+        return value.dict()
+
+    if isinstance(value, dict):
+        try:
+            # Test serialization without modifying - optimizes for already-serializable dicts
+            dumps(
+                value
+            )  # If serialization fails, we'll catch the exception and process deeply
+            return value  # Avoid unnecessary recursion for serializable dicts
+        except TypeError:
+            return {k: _make_serializable(v) for k, v in value.items()}
+    elif isinstance(value, list):
+        try:
+            # Test serialization without modifying - optimizes for already-serializable lists
+            dumps(
+                value
+            )  # If serialization fails, we'll catch the exception and process deeply
+            return value  # Avoid unnecessary recursion for serializable lists
+        except TypeError:
+            return [_make_serializable(item) for item in value]
+
+    return repr(value)
+
+
+def _encode_value(value: Any) -> Optional[Attribute]:
+    """
+    Encode values for tracing, ensuring proper JSON serialization.
+    Adds the @ag.type=json: prefix only to appropriate values.
+    """
+    if value is None:
+        return None
+
+    if isinstance(value, (str, int, float, bool, bytes)):
+        return value
+
+    try:
+        if (
+            isinstance(value, (dict, list))
+            or hasattr(value, "model_dump")
+            or hasattr(value, "dict")
+        ):
+            serializable_value = _make_serializable(value)
+            return "@ag.type=json:" + dumps(serializable_value)
+    except TypeError:
+        pass
+
+    return repr(value)
+
+
+def serialize(
+    *,
+    namespace: str,
+    attributes: Dict[str, Any],
+    max_depth: Optional[int] = None,
+) -> Dict[str, str]:
+    if not isinstance(attributes, dict):
+        return {}
+
+    _attributes = {
+        k: v
+        for k, v in {
+            _encode_key(namespace, key): _encode_value(value)
+            for key, value in _marshal(attributes, max_depth=max_depth).items()
+        }.items()
+        if v is not None
+    }
+
+    return _attributes

agenta/sdk/engines/tracing/conventions.py

@@ -0,0 +1,49 @@
+from enum import Enum
+from re import fullmatch
+
+from opentelemetry.trace import SpanKind
+
+
+class Reference(str, Enum):
+    #
+    VARIANT_ID = "variant.id"
+    VARIANT_SLUG = "variant.slug"
+    VARIANT_VERSION = "variant.version"
+    #
+    ENVIRONMENT_ID = "environment.id"
+    ENVIRONMENT_SLUG = "environment.slug"
+    ENVIRONMENT_VERSION = "environment.version"
+    #
+    APPLICATION_ID = "application.id"
+    APPLICATION_SLUG = "application.slug"
+    #
+
+
+_PATTERN = r"[A-Za-z0-9._-]+"
+
+
+def is_valid_attribute_key(
+    string: str,
+):
+    return bool(fullmatch(_PATTERN, string))
+
+
+def parse_span_kind(type: str) -> SpanKind:
+    kind = SpanKind.INTERNAL
+    if type in [
+        "agent",
+        "chain",
+        "workflow",
+    ]:
+        kind = SpanKind.SERVER
+    elif type in [
+        "tool",
+        "embedding",
+        "query",
+        "completion",
+        "chat",
+        "rerank",
+    ]:
+        kind = SpanKind.CLIENT
+
+    return kind

agenta/sdk/engines/tracing/exporters.py

@@ -0,0 +1,130 @@
+from typing import Sequence, Dict, List, Optional
+
+from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+from opentelemetry.sdk.trace.export import (
+    ConsoleSpanExporter,
+    SpanExporter,
+    SpanExportResult,
+    ReadableSpan,
+)
+
+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.contexts.tracing import (
+    otlp_context_manager,
+    otlp_context,
+    OTLPContext,
+)
+
+
+log = get_module_logger(__name__)
+
+
+class InlineTraceExporter(SpanExporter):
+    def __init__(
+        self,
+        registry: Dict[str, List[ReadableSpan]],
+    ):
+        self._shutdown = False
+        self._registry = registry
+
+    def export(
+        self,
+        spans: Sequence[ReadableSpan],
+    ) -> SpanExportResult:
+        if self._shutdown:
+            return
+
+        with suppress():
+            for span in spans:
+                trace_id = span.get_span_context().trace_id
+
+                if trace_id not in self._registry:
+                    self._registry[trace_id] = []
+
+                self._registry[trace_id].append(span)
+
+    def shutdown(self) -> None:
+        self._shutdown = True
+
+    def force_flush(self, timeout_millis: int = 30000) -> bool:
+        return True
+
+    def is_ready(
+        self,
+        trace_id: int,
+    ) -> bool:
+        is_ready = trace_id in self._registry
+        return is_ready
+
+    def fetch(
+        self,
+        trace_id: int,
+    ) -> List[ReadableSpan]:
+        trace = self._registry.get(trace_id, [])
+
+        if trace_id in self._registry:
+            del self._registry[trace_id]
+
+        return trace
+
+
+class OTLPExporter(OTLPSpanExporter):
+    _MAX_RETRY_TIMEOUT = 2
+
+    def __init__(
+        self,
+        *args,
+        credentials: Optional[TTLLRUCache] = None,
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+
+        self.credentials = credentials
+
+    def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
+        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] = []
+
+            grouped_spans[credentials].append(span)
+
+        serialized_spans = []
+
+        for credentials, _spans in grouped_spans.items():
+            with otlp_context_manager(
+                context=OTLPContext(
+                    credentials=credentials,
+                )
+            ):
+                serialized_spans.append(super().export(_spans))
+
+        if all(serialized_spans):
+            return SpanExportResult.SUCCESS
+        else:
+            return SpanExportResult.FAILURE
+
+    def _export(self, serialized_data: bytes, timeout_sec: Optional[float] = None):
+        credentials = otlp_context.get().credentials
+
+        if credentials:
+            self._session.headers.update({"Authorization": credentials})
+
+        with suppress():
+            if timeout_sec is not None:
+                return super()._export(serialized_data, timeout_sec)
+            else:
+                return super()._export(serialized_data)
+
+
+ConsoleExporter = ConsoleSpanExporter
+InlineExporter = InlineTraceExporter