agenta 0.26.0a0__py3-none-any.whl → 0.27.0a0__py3-none-any.whl

agenta/sdk/tracing/attributes.py

@@ -0,0 +1,181 @@
+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 _encode_value(value: Any) -> Optional[Attribute]:
+    if value is None:
+        return None
+
+    if isinstance(value, (str, int, float, bool, bytes)):
+        return value
+
+    if isinstance(value, dict) or isinstance(value, list):
+        encoded = dumps(value)
+        value = "@ag.type=json:" + encoded
+        return value
+
+    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
+
+
+"""
+def _decode_key(namespace: Optional[str] = None, key: str = "") -> str:
+    if namespace is None:
+        return key
+
+    return key.replace(f"ag.{namespace}.", "")
+"""
+
+"""
+def _decode_value(value: Attribute) -> Any:
+    if isinstance(value, (int, float, bool, bytes)):
+        return value
+
+    if isinstance(value, str):
+        if value == "@ag.type=none:":
+            return None
+
+        if value.startswith("@ag.type=json:"):
+            encoded = value[len("@ag.type=json:") :]
+            value = loads(encoded)
+            return value
+
+        return value
+
+    return value
+"""
+
+"""
+def deserialize(
+    *,
+    namespace: str,
+    attributes: Dict[str, Any],
+    max_depth: Optional[int] = None,
+) -> Dict[str, Any]:
+    if not isinstance(attributes, dict):
+        return {}
+
+    return {
+        _decode_key(namespace, key): _decode_value(value)
+        for key, value in attributes
+        if key != _decode_key(namespace, key)
+    }
+"""

agenta/sdk/tracing/context.py

@@ -0,0 +1,21 @@
+from contextvars import ContextVar
+from contextlib import contextmanager
+from traceback import format_exc
+
+from agenta.sdk.utils.logging import log
+
+tracing_context = ContextVar("tracing_context", default={})
+
+
+@contextmanager
+def tracing_context_manager():
+    _tracing_context = {"health": {"status": "ok"}}
+
+    token = tracing_context.set(_tracing_context)
+    try:
+        yield
+    except Exception as e:
+        log.error(f"Error with tracing context: {_tracing_context}")
+        log.error(f"Exception: {format_exc()}")
+    finally:
+        tracing_context.reset(token)

agenta/sdk/tracing/conventions.py

@@ -0,0 +1,43 @@
+from typing import Literal
+
+from opentelemetry.trace import SpanKind
+
+Namespace = Literal[
+    "data.inputs",
+    "data.internals",
+    "data.outputs",
+    "metrics.scores",
+    "metrics.unit.costs",
+    "metrics.unit.tokens",
+    "meta.configuration",
+    "meta.version",
+    "tags",
+    "refs",
+]
+
+Code = Literal[
+    "OK",
+    "UNSET",
+    "ERROR",
+]
+
+
+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/tracing/exporters.py

@@ -0,0 +1,53 @@
+from typing import Sequence, Dict, List
+
+from opentelemetry.sdk.trace.export import (
+    ConsoleSpanExporter,
+    SpanExporter,
+    SpanExportResult,
+    ReadableSpan,
+)
+from opentelemetry.exporter.otlp.proto.http.trace_exporter import (
+    OTLPSpanExporter,
+)
+
+from agenta.sdk.utils.exceptions import suppress
+
+
+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 fetch(self, trace_id: int) -> List[ReadableSpan]:
+        trace = self._registry.get(trace_id, [])
+
+        del self._registry[trace_id]
+
+        return trace
+
+
+ConsoleExporter = ConsoleSpanExporter
+InlineExporter = InlineTraceExporter
+OTLPExporter = OTLPSpanExporter