agenta 0.12.2__py3-none-any.whl → 0.32.0a1__py3-none-any.whl

agenta/sdk/tracing/processors.py

@@ -0,0 +1,121 @@
+from typing import Optional, Dict, List
+
+from opentelemetry.baggage import get_all as get_baggage
+from opentelemetry.context import Context
+from opentelemetry.sdk.trace import Span
+from opentelemetry.sdk.trace.export import (
+    SpanExporter,
+    ReadableSpan,
+    BatchSpanProcessor,
+    _DEFAULT_MAX_QUEUE_SIZE,
+    _DEFAULT_MAX_EXPORT_BATCH_SIZE,
+)
+
+from agenta.sdk.utils.logging import log
+from agenta.sdk.tracing.conventions import Reference
+
+
+class TraceProcessor(BatchSpanProcessor):
+    def __init__(
+        self,
+        span_exporter: SpanExporter,
+        references: Dict[str, str] = None,
+        max_queue_size: int = None,
+        schedule_delay_millis: float = None,
+        max_export_batch_size: int = None,
+        export_timeout_millis: float = None,
+    ):
+        super().__init__(
+            span_exporter,
+            _DEFAULT_MAX_QUEUE_SIZE,
+            12 * 60 * 60 * 1000,  # 12 hours
+            _DEFAULT_MAX_EXPORT_BATCH_SIZE,
+            500,  # < 1 second (0.5 seconds)
+        )
+
+        self._registry = dict()
+        self._exporter = span_exporter
+        self.references = references or dict()
+        self.spans: Dict[int, List[ReadableSpan]] = dict()
+
+    def on_start(
+        self,
+        span: Span,
+        parent_context: Optional[Context] = None,
+    ) -> None:
+        baggage = get_baggage(parent_context)
+
+        for key in self.references.keys():
+            span.set_attribute(f"ag.refs.{key}", self.references[key])
+
+        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])
+
+        if span.context.trace_id not in self._registry:
+            self._registry[span.context.trace_id] = dict()
+
+        self._registry[span.context.trace_id][span.context.span_id] = True
+
+    def on_end(
+        self,
+        span: ReadableSpan,
+    ):
+        if self.done:
+            return
+
+        if span.context.trace_id not in self.spans:
+            self.spans[span.context.trace_id] = list()
+
+        self.spans[span.context.trace_id].append(span)
+
+        del self._registry[span.context.trace_id][span.context.span_id]
+
+        if len(self._registry[span.context.trace_id]) == 0:
+            self.export(span.context.trace_id)
+
+    def export(
+        self,
+        trace_id: int,
+    ):
+        spans = self.spans[trace_id]
+
+        for span in spans:
+            self.queue.appendleft(span)
+
+        with self.condition:
+            self.condition.notify()
+
+        del self.spans[trace_id]
+
+    def force_flush(
+        self,
+        timeout_millis: int = None,
+    ) -> bool:
+        ret = super().force_flush(timeout_millis)
+
+        if not ret:
+            log.warning("Agenta - Skipping export due to timeout.")
+
+    def is_ready(
+        self,
+        trace_id: Optional[int] = None,
+    ) -> bool:
+        is_ready = True
+
+        try:
+            is_ready = self._exporter.is_ready(trace_id)
+        except:  # pylint: disable=bare-except
+            pass
+
+        return is_ready
+
+    def fetch(
+        self,
+        trace_id: Optional[int] = None,
+    ) -> Dict[str, ReadableSpan]:
+        trace = self._exporter.fetch(trace_id)  # type: ignore
+
+        return trace

agenta/sdk/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.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,
+        )

agenta/sdk/tracing/tracing.py

@@ -0,0 +1,237 @@
+from typing import Optional, Any, Dict, Callable
+from enum import Enum
+
+from httpx import get as check
+
+from opentelemetry.trace import (
+    get_current_span,
+    set_tracer_provider,
+    get_tracer_provider,
+    Status,
+    StatusCode,
+)
+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 log
+from agenta.sdk.tracing.processors import TraceProcessor
+from agenta.sdk.tracing.exporters import InlineExporter, OTLPExporter
+from agenta.sdk.tracing.spans import CustomSpan
+from agenta.sdk.tracing.inline import parse_inline_trace
+from agenta.sdk.tracing.conventions import Reference, is_valid_attribute_key
+
+
+class Tracing(metaclass=Singleton):
+    VERSION = "0.1.0"
+
+    Status = Status
+    StatusCode = StatusCode
+
+    def __init__(
+        self,
+        url: str,
+        redact: Optional[Callable[..., Any]] = None,
+        redact_on_error: Optional[bool] = True,
+    ) -> None:
+        # ENDPOINT (OTLP)
+        self.otlp_url = url
+        # HEADERS (OTLP)
+        self.headers: Dict[str, str] = dict()
+        # REFERENCES
+        self.references: Dict[str, str] = dict()
+        # CREDENTIALS
+        self.credentials: Dict[int, str] = dict()
+
+        # TRACER PROVIDER
+        self.tracer_provider: Optional[TracerProvider] = None
+        # TRACE PROCESSORS -- INLINE
+        self.inline: Optional[TraceProcessor] = None
+        # TRACER
+        self.tracer: Optional[Tracer] = None
+        # INLINE SPANS for INLINE TRACES (INLINE PROCESSOR)
+        self.inline_spans: Dict[str, Any] = dict()
+
+        # REDACT
+        self.redact = redact
+        self.redact_on_error = redact_on_error
+
+    # PUBLIC
+
+    def configure(
+        self,
+        api_key: Optional[str] = None,
+        service_id: Optional[str] = None,
+        # DEPRECATING
+        app_id: Optional[str] = None,
+    ):
+        # HEADERS (OTLP)
+        if api_key:
+            self.headers["Authorization"] = f"ApiKey {api_key}"
+        # REFERENCES
+        if service_id:
+            self.references["service.id"] = service_id
+        if app_id:
+            self.references["application.id"] = app_id
+
+        # TRACER PROVIDER
+        self.tracer_provider = TracerProvider(
+            resource=Resource(attributes={"service.name": "agenta-sdk"})
+        )
+        # TRACE PROCESSORS -- INLINE
+        self.inline = TraceProcessor(
+            InlineExporter(
+                registry=self.inline_spans,
+            ),
+            references=self.references,
+        )
+        self.tracer_provider.add_span_processor(self.inline)
+        # TRACE PROCESSORS -- OTLP
+        try:
+            log.info(
+                "Agenta - OLTP URL: %s",
+                self.otlp_url,
+            )
+            # check(
+            #     self.otlp_url,
+            #     headers=self.headers,
+            #     timeout=1,
+            # )
+
+            _otlp = TraceProcessor(
+                OTLPExporter(
+                    endpoint=self.otlp_url,
+                    headers=self.headers,
+                    credentials=self.credentials,
+                ),
+                references=self.references,
+            )
+
+            self.tracer_provider.add_span_processor(_otlp)
+        except:  # pylint: disable=bare-except
+            log.warning("Agenta - OLTP unreachable, skipping exports.")
+
+        # GLOBAL TRACER PROVIDER -- INSTRUMENTATION LIBRARIES
+        set_tracer_provider(self.tracer_provider)
+        # TRACER
+        self.tracer: Tracer = self.tracer_provider.get_tracer("agenta.tracer")
+
+    def get_current_span(self):
+        _span = None
+
+        with suppress():
+            _span = get_current_span()
+
+            if _span.is_recording():
+                return CustomSpan(_span)
+
+        return _span
+
+    def store_internals(
+        self,
+        attributes: Dict[str, Any],
+        span: Optional[Span] = None,
+    ):
+        with suppress():
+            if span is None:
+                span = self.get_current_span()
+
+            span.set_attributes(
+                attributes={"internals": attributes},
+                namespace="data",
+            )
+
+    def store_refs(
+        self,
+        refs: Dict[str, str],
+        span: Optional[Span] = None,
+    ):
+        with suppress():
+            if span is None:
+                span = self.get_current_span()
+
+            for key in refs.keys():
+                if key in [_.value for _ in Reference.__members__.values()]:
+                    # ADD REFERENCE TO THIS SPAN
+                    span.set_attribute(
+                        key.value if isinstance(key, Enum) else key,
+                        refs[key],
+                        namespace="refs",
+                    )
+
+                    # AND TO ALL SPANS CREATED AFTER THIS ONE
+                    self.references[key] = refs[key]
+                    # TODO: THIS SHOULD BE REPLACED BY A TRACE CONTEXT !!!
+
+    def store_meta(
+        self,
+        meta: Dict[str, Any],
+        span: Optional[Span] = None,
+    ):
+        with suppress():
+            if span is None:
+                span = self.get_current_span()
+
+            for key in meta.keys():
+                if is_valid_attribute_key(key):
+                    span.set_attribute(
+                        key,
+                        meta[key],
+                        namespace="meta",
+                    )
+
+    def store_metrics(
+        self,
+        metrics: Dict[str, Any],
+        span: Optional[Span] = None,
+    ):
+        with suppress():
+            if span is None:
+                span = self.get_current_span()
+
+            for key in metrics.keys():
+                if is_valid_attribute_key(key):
+                    span.set_attribute(
+                        key,
+                        metrics[key],
+                        namespace="metrics",
+                    )
+
+    def is_inline_trace_ready(
+        self,
+        trace_id: Optional[int] = None,
+    ) -> bool:
+        is_ready = True
+
+        with suppress():
+            if trace_id is not None:
+                is_ready = self.inline.is_ready(trace_id)
+
+        return is_ready
+
+    def get_inline_trace(
+        self,
+        trace_id: Optional[int] = None,
+    ) -> Dict[str, Any]:
+        _inline_trace = {}
+
+        with suppress():
+            is_ready = self.inline.is_ready(trace_id)
+
+            if is_ready is True:
+                otel_spans = self.inline.fetch(trace_id)
+
+                if otel_spans:
+                    _inline_trace = parse_inline_trace(otel_spans)
+
+        return _inline_trace
+
+
+def get_tracer(
+    tracing: Tracing,
+) -> Tracer:
+    if tracing is None or tracing.tracer is None or tracing.tracer_provider is None:
+        return get_tracer_provider().get_tracer("default.tracer")
+
+    return tracing.tracer