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

agenta/sdk/engines/tracing/processors.py

@@ -0,0 +1,190 @@
+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, SpanProcessor
+from opentelemetry.sdk.trace.export import (
+    SpanExporter,
+    ReadableSpan,
+    BatchSpanProcessor,
+)
+
+from agenta.sdk.utils.logging import get_module_logger
+from agenta.sdk.engines.tracing.conventions import Reference
+
+log = get_module_logger(__name__)
+
+
+class TraceProcessor(SpanProcessor):
+    def __init__(
+        self,
+        span_exporter: SpanExporter,
+        references: Dict[str, str] = None,
+        inline: bool = False,
+        max_queue_size: int = None,
+        schedule_delay_millis: float = None,
+        max_export_batch_size: int = None,
+        export_timeout_millis: float = None,
+    ):
+        self.references = references or dict()
+        self.inline = inline is True
+
+        self._registry = dict()
+        self._exporter = span_exporter
+        self._spans: Dict[int, List[ReadableSpan]] = dict()
+
+        # --- DISTRIBUTED
+        if not self.inline:
+            self._delegate = BatchSpanProcessor(
+                span_exporter,
+                max_queue_size,
+                schedule_delay_millis,
+                max_export_batch_size,
+                export_timeout_millis,
+            )
+        # --- DISTRIBUTED
+
+    def on_start(
+        self,
+        span: Span,
+        parent_context: Optional[Context] = None,
+    ) -> None:
+        for key in self.references.keys():
+            span.set_attribute(f"ag.refs.{key}", self.references[key])
+
+        baggage = get_baggage(parent_context)
+
+        for key in baggage.keys():
+            if key.startswith("ag.refs."):
+                _key = key.replace("ag.refs.", "")
+                if _key in [_.value for _ in Reference.__members__.values()]:
+                    span.set_attribute(key, baggage[key])
+
+        trace_id = span.context.trace_id
+        span_id = span.context.span_id
+
+        self._registry.setdefault(trace_id, {})
+        self._registry[trace_id][span_id] = True
+
+    def on_end(
+        self,
+        span: ReadableSpan,
+    ):
+        trace_id = span.context.trace_id
+        span_id = span.context.span_id
+
+        self._spans.setdefault(trace_id, []).append(span)
+        self._registry.setdefault(trace_id, {})
+        self._registry[trace_id].pop(span_id, None)
+
+        if not self._registry[trace_id]:
+            spans = self._spans.pop(trace_id, [])
+            self._registry.pop(trace_id, None)
+
+            # --- INLINE
+            if self.inline:
+                self._exporter.export(spans)
+            # --- INLINE
+
+            # --- DISTRIBUTED
+            else:
+                for span in spans:
+                    self._delegate.on_end(span)
+
+                self._delegate.force_flush()
+            # --- DISTRIBUTED
+
+    def force_flush(
+        self,
+        timeout_millis: int = None,
+    ) -> bool:
+        # --- INLINE
+        if self.inline:
+            try:
+                ret = self._exporter.force_flush(timeout_millis)
+            except:  # pylint: disable=bare-except
+                ret = True
+        # --- INLINE
+
+        # --- DISTRIBUTED
+        else:
+            ret = self._delegate.force_flush(timeout_millis)
+        # --- DISTRIBUTED
+
+        if not ret:
+            log.warning("Agenta - Skipping export due to timeout.")
+
+        return ret
+
+    def shutdown(self) -> None:
+        # --- INLINE
+        if self.inline:
+            self._exporter.shutdown()
+        # --- INLINE
+
+        # --- DISTRIBUTED
+        else:
+            self._delegate.shutdown()
+        # --- DISTRIBUTED
+
+    def is_ready(
+        self,
+        trace_id: Optional[int] = None,
+    ) -> bool:
+        is_ready = True
+
+        # --- INLINE
+        if self.inline:
+            try:
+                is_ready = self._exporter.is_ready(trace_id)
+            except:  # pylint: disable=bare-except
+                pass
+        # --- INLINE
+
+        return is_ready
+
+    def fetch(
+        self,
+        trace_id: Optional[int] = None,
+    ) -> Dict[str, ReadableSpan]:
+        trace = None
+
+        # --- INLINE
+        if self.inline:
+            try:
+                trace = self._exporter.fetch(trace_id)  # type: ignore
+            except:  # pylint: disable=bare-except
+                pass
+        # --- INLINE
+
+        return trace
+
+
+# Internal storage for the last ended span context
+_last_ended_span_context = None
+_lock = Lock()
+
+
+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
+
+
+def _get_last_ended():
+    """Get the last ended span context"""
+    with _lock:
+        return _last_ended_span_context
+
+
+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/sdk/engines/tracing/propagation.py

@@ -0,0 +1,102 @@
+from typing import Tuple, Optional, Dict, Any
+
+from opentelemetry.trace import Span, set_span_in_context, get_current_span
+from opentelemetry.baggage.propagation import W3CBaggagePropagator
+from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
+from opentelemetry.baggage import set_baggage
+from opentelemetry.context import get_current
+
+from agenta.sdk.contexts.tracing import TracingContext
+
+import agenta as ag
+
+
+def extract(
+    headers: Dict[str, str],
+) -> Tuple[Optional[str], Optional[Any], Dict[str, str]]:
+    # --- Extract credentials --- #
+    credentials = None
+
+    try:
+        credentials = (
+            headers.get("Authorization")  # Uppercase
+            or headers.get("authorization")  # Lowercase
+            or None
+        )
+
+    except:  # pylint: disable=bare-except
+        pass
+
+    # --- Extract traceparent --- #
+    traceparent = None
+
+    try:
+        _carrier = {
+            "traceparent": headers.get("Traceparent")  # Uppercase
+            or headers.get("traceparent")  # Lowercase
+            or "",
+        }
+
+        _context = TraceContextTextMapPropagator().extract(_carrier)
+
+        traceparent = _context
+    except:  # pylint: disable=bare-except
+        pass
+
+    # --- Extract baggage --- #
+    baggage = {}
+
+    try:
+        _carrier = {
+            "baggage": headers.get("Baggage")  # Uppercase
+            or headers.get("baggage")  # Lowercase
+            or "",
+        }
+
+        _context = W3CBaggagePropagator().extract(_carrier)
+
+        if _context:
+            for partial in _context.values():
+                for key, value in partial.items():
+                    baggage[key] = value
+
+    except:  # pylint: disable=bare-except
+        pass
+
+    # --- #
+    return credentials, traceparent, baggage
+
+
+def inject(
+    headers: Optional[Dict[str, str]] = None,
+) -> Dict[str, str]:
+    headers = headers or {}
+
+    _context = get_current()
+
+    ctx = TracingContext.get()
+
+    # --- Inject traceparent --- #
+    try:
+        TraceContextTextMapPropagator().inject(headers, context=_context)
+
+    except:  # pylint: disable=bare-except
+        pass
+
+    # --- Inject baggage --- #
+    try:
+        if ctx.baggage:
+            for key, value in ctx.baggage.items():
+                _context = set_baggage(key, value, context=_context)
+
+        W3CBaggagePropagator().inject(headers, context=_context)
+
+    except:  # pylint: disable=bare-except
+        pass
+
+    # --- Inject credentials --- #
+    if ctx.credentials:
+        headers["Authorization"] = ctx.credentials
+
+    # --- #
+    return headers

agenta/sdk/engines/tracing/spans.py

@@ -0,0 +1,136 @@
+from typing import Optional, Union, Any, Dict
+
+from opentelemetry.trace import SpanContext
+from opentelemetry.trace.status import Status, StatusCode
+from opentelemetry.sdk.trace import Span
+
+from agenta.sdk.engines.tracing.attributes import serialize
+
+
+class CustomSpan(Span):  # INHERITANCE FOR TYPING ONLY
+    def __init__(
+        self,
+        span: Span,
+    ) -> None:
+        super().__init__(  # INHERITANCE FOR TYPING ONLY
+            name=span.name,
+            context=span.context,
+            parent=span.parent,
+            sampler=span._sampler,
+            trace_config=span._trace_config,
+            resource=span.resource,
+            attributes=span.attributes,
+            events=span.events,
+            links=span.links,
+            kind=span.kind,
+            span_processor=span._span_processor,
+            instrumentation_info=span.instrumentation_info,
+            record_exception=span._record_exception,
+            set_status_on_exception=span._set_status_on_exception,
+            limits=span._limits,
+            instrumentation_scope=span.instrumentation_scope,
+        )
+
+        self._span = span
+
+    ## --- PROXY METHODS --- ##
+
+    def get_span_context(self):
+        return self._span.get_span_context()
+
+    def is_recording(self) -> bool:
+        return self._span.is_recording()
+
+    def update_name(
+        self,
+        name: str,
+    ) -> None:
+        self._span.update_name(name)
+
+    def set_status(
+        self,
+        status: Union[Status, StatusCode],
+        description: Optional[str] = None,
+    ) -> None:
+        self._span.set_status(
+            status=status,
+            description=description,
+        )
+
+    def end(self) -> None:
+        self._span.end()
+
+    ## --- CUSTOM METHODS W/ ATTRIBUTES SERALIZATION --- ##
+
+    def set_attributes(
+        self,
+        attributes: Dict[str, Any],
+        namespace: Optional[str] = None,
+        max_depth: Optional[int] = None,
+    ) -> None:
+        self._span.set_attributes(
+            attributes=serialize(
+                namespace=namespace,
+                attributes=attributes,
+                max_depth=max_depth,
+            )
+        )
+
+    def set_attribute(
+        self,
+        key: str,
+        value: Any,
+        namespace: Optional[str] = None,
+    ) -> None:
+        self.set_attributes(
+            attributes={key: value},
+            namespace=namespace,
+        )
+
+    def add_event(
+        self,
+        name: str,
+        attributes: Optional[Dict[str, Any]] = None,
+        timestamp: Optional[int] = None,
+        namespace: Optional[str] = None,
+    ) -> None:
+        self._span.add_event(
+            name=name,
+            attributes=serialize(
+                namespace=namespace,
+                attributes=attributes,
+            ),
+            timestamp=timestamp,
+        )
+
+    def add_link(
+        self,
+        context: SpanContext,
+        attributes: Optional[Dict[str, Any]] = None,
+        namespace: Optional[str] = None,
+    ) -> None:
+        self._span.add_link(
+            context=context,
+            attributes=serialize(
+                namespace=namespace,
+                attributes=attributes,
+            ),
+        )
+
+    def record_exception(
+        self,
+        exception: BaseException,
+        attributes: Optional[Dict[str, Any]] = None,
+        timestamp: Optional[int] = None,
+        escaped: bool = False,
+        namespace: Optional[str] = None,
+    ) -> None:
+        self._span.record_exception(
+            exception=exception,
+            attributes=serialize(
+                namespace=namespace,
+                attributes=attributes,
+            ),
+            timestamp=timestamp,
+            escaped=escaped,
+        )