aeri-python 4.0.0__py3-none-any.whl

aeri/_client/span_filter.py

@@ -0,0 +1,76 @@
+"""Span filter predicates for controlling OpenTelemetry span export.
+
+This module provides composable filter functions that determine which spans
+the AeriSpanProcessor forwards to the Aeri backend.
+"""
+
+from opentelemetry.sdk.trace import ReadableSpan
+
+from aeri._client.constants import AERI_TRACER_NAME
+
+KNOWN_LLM_INSTRUMENTATION_SCOPE_PREFIXES = frozenset(
+    {
+        AERI_TRACER_NAME,
+        "agent_framework",
+        "ai",
+        "haystack",
+        "langsmith",
+        "litellm",
+        "openinference",
+        "opentelemetry.instrumentation.anthropic",
+        "strands-agents",
+        "vllm",
+    }
+)
+"""Known instrumentation scope namespace prefixes.
+
+Prefix matching is boundary-aware:
+- exact match (``scope == prefix``)
+- direct descendant scopes (``scope.startswith(prefix + ".")``)
+
+Please create a Github issue in https://github.com/aeri/aeri if you'd like to expand this default allow list. 
+"""
+
+
+def is_aeri_span(span: ReadableSpan) -> bool:
+    """Return whether the span was created by the Aeri SDK tracer."""
+    return (
+        span.instrumentation_scope is not None
+        and span.instrumentation_scope.name == AERI_TRACER_NAME
+    )
+
+
+def is_genai_span(span: ReadableSpan) -> bool:
+    """Return whether the span has any ``gen_ai.*`` semantic convention attribute."""
+    if span.attributes is None:
+        return False
+
+    return any(
+        isinstance(key, str) and key.startswith("gen_ai")
+        for key in span.attributes.keys()
+    )
+
+
+def _matches_scope_prefix(scope_name: str, prefix: str) -> bool:
+    """Return whether a scope matches a prefix using namespace boundaries."""
+    return scope_name == prefix or scope_name.startswith(f"{prefix}.")
+
+
+def is_known_llm_instrumentor(span: ReadableSpan) -> bool:
+    """Return whether the span comes from a known LLM instrumentation scope."""
+    if span.instrumentation_scope is None:
+        return False
+
+    scope_name = span.instrumentation_scope.name
+
+    return any(
+        _matches_scope_prefix(scope_name, prefix)
+        for prefix in KNOWN_LLM_INSTRUMENTATION_SCOPE_PREFIXES
+    )
+
+
+def is_default_export_span(span: ReadableSpan) -> bool:
+    """Return whether a span should be exported by default."""
+    return (
+        is_aeri_span(span) or is_genai_span(span) or is_known_llm_instrumentor(span)
+    )

aeri/_client/span_processor.py

@@ -0,0 +1,206 @@
+"""Span processor for Aeri OpenTelemetry integration.
+
+This module defines the AeriSpanProcessor class, which extends OpenTelemetry's
+BatchSpanProcessor with Aeri-specific functionality. It handles exporting
+spans to the Aeri API with proper authentication and filtering.
+
+Key features:
+- HTTP-based span export to Aeri API
+- Basic authentication with Aeri API keys
+- Configurable batch processing behavior
+- Project-scoped span filtering to prevent cross-project data leakage
+"""
+
+import base64
+import os
+from typing import Callable, Dict, List, Optional
+
+from opentelemetry import context as context_api
+from opentelemetry.context import Context
+from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+from opentelemetry.sdk.trace import ReadableSpan, Span
+from opentelemetry.sdk.trace.export import BatchSpanProcessor
+from opentelemetry.trace import format_span_id
+
+from aeri._client.environment_variables import (
+    AERI_FLUSH_AT,
+    AERI_FLUSH_INTERVAL,
+    AERI_OTEL_TRACES_EXPORT_PATH,
+)
+from aeri._client.propagation import _get_propagated_attributes_from_context
+from aeri._client.span_filter import is_default_export_span, is_aeri_span
+from aeri._client.utils import span_formatter
+from aeri.logger import aeri_logger
+from aeri.version import __version__ as aeri_version
+
+
+class AeriSpanProcessor(BatchSpanProcessor):
+    """OpenTelemetry span processor that exports spans to the Aeri API.
+
+    This processor extends OpenTelemetry's BatchSpanProcessor with Aeri-specific functionality:
+    1. Project-scoped span filtering to prevent cross-project data leakage
+    2. Instrumentation scope filtering to block spans from specific libraries/frameworks
+    3. Configurable batch processing parameters for optimal performance
+    4. HTTP-based span export to the Aeri OTLP endpoint
+    5. Debug logging for span processing operations
+    6. Authentication with Aeri API using Basic Auth
+
+    The processor is designed to efficiently handle large volumes of spans with
+    minimal overhead, while ensuring spans are only sent to the correct project.
+    It integrates with OpenTelemetry's standard span lifecycle, adding Aeri-specific
+    filtering and export capabilities.
+    """
+
+    def __init__(
+        self,
+        *,
+        public_key: str,
+        secret_key: str,
+        base_url: str,
+        timeout: Optional[int] = None,
+        flush_at: Optional[int] = None,
+        flush_interval: Optional[float] = None,
+        blocked_instrumentation_scopes: Optional[List[str]] = None,
+        should_export_span: Optional[Callable[[ReadableSpan], bool]] = None,
+        additional_headers: Optional[Dict[str, str]] = None,
+    ):
+        self.public_key = public_key
+        self.blocked_instrumentation_scopes = (
+            blocked_instrumentation_scopes
+            if blocked_instrumentation_scopes is not None
+            else []
+        )
+        self._should_export_span = should_export_span or is_default_export_span
+
+        env_flush_at = os.environ.get(AERI_FLUSH_AT, None)
+        flush_at = flush_at or int(env_flush_at) if env_flush_at is not None else None
+
+        env_flush_interval = os.environ.get(AERI_FLUSH_INTERVAL, None)
+        flush_interval = (
+            flush_interval or float(env_flush_interval)
+            if env_flush_interval is not None
+            else None
+        )
+
+        basic_auth_header = "Basic " + base64.b64encode(
+            f"{public_key}:{secret_key}".encode("utf-8")
+        ).decode("ascii")
+
+        # Prepare default headers
+        default_headers = {
+            "Authorization": basic_auth_header,
+            "x-aeri-sdk-name": "python",
+            "x-aeri-sdk-version": aeri_version,
+            "x-aeri-public-key": public_key,
+        }
+
+        # Merge additional headers if provided
+        headers = {**default_headers, **(additional_headers or {})}
+
+        traces_export_path = os.environ.get(AERI_OTEL_TRACES_EXPORT_PATH, None)
+
+        endpoint = (
+            f"{base_url}/{traces_export_path}"
+            if traces_export_path
+            else f"{base_url}/api/public/otel/v1/traces"
+        )
+
+        aeri_span_exporter = OTLPSpanExporter(
+            endpoint=endpoint,
+            headers=headers,
+            timeout=timeout,
+        )
+
+        super().__init__(
+            span_exporter=aeri_span_exporter,
+            export_timeout_millis=timeout * 1_000 if timeout else None,
+            max_export_batch_size=flush_at,
+            schedule_delay_millis=flush_interval * 1_000
+            if flush_interval is not None
+            else None,
+        )
+
+    def on_start(self, span: Span, parent_context: Optional[Context] = None) -> None:
+        context = parent_context or context_api.get_current()
+        propagated_attributes = _get_propagated_attributes_from_context(context)
+
+        if propagated_attributes:
+            span.set_attributes(propagated_attributes)
+
+            aeri_logger.debug(
+                f"Propagated {len(propagated_attributes)} attributes to span '{format_span_id(span.context.span_id)}': {propagated_attributes}"
+            )
+
+        return super().on_start(span, parent_context)
+
+    def on_end(self, span: ReadableSpan) -> None:
+        # Only export spans that belong to the scoped project
+        # This is important to not send spans to wrong project in multi-project setups
+        if is_aeri_span(span) and not self._is_aeri_project_span(span):
+            aeri_logger.debug(
+                f"Security: Span rejected - belongs to project '{span.instrumentation_scope.attributes.get('public_key') if span.instrumentation_scope and span.instrumentation_scope.attributes else None}' but processor is for '{self.public_key}'. "
+                f"This prevents cross-project data leakage in multi-project environments."
+            )
+            return
+
+        # Do not export spans from blocked instrumentation scopes
+        if self._is_blocked_instrumentation_scope(span):
+            aeri_logger.debug(
+                "Trace: Dropping span due to blocked instrumentation scope | "
+                f"span_name='{span.name}' | "
+                f"instrumentation_scope='{self._get_scope_name(span)}'"
+            )
+            return
+
+        # Apply custom or default span filter
+        try:
+            should_export = self._should_export_span(span)
+        except Exception as error:
+            aeri_logger.error(
+                "Trace: should_export_span callback raised an error. "
+                f"Dropping span name='{span.name}' scope='{self._get_scope_name(span)}'. "
+                f"Error: {error}"
+            )
+            return
+
+        if not should_export:
+            aeri_logger.debug(
+                "Trace: Dropping span due to should_export_span filter | "
+                f"span_name='{span.name}' | "
+                f"instrumentation_scope='{self._get_scope_name(span)}'"
+            )
+            return
+
+        aeri_logger.debug(
+            f"Trace: Processing span name='{span._name}' | Full details:\n{span_formatter(span)}"
+        )
+
+        super().on_end(span)
+
+    def _is_blocked_instrumentation_scope(self, span: ReadableSpan) -> bool:
+        return (
+            span.instrumentation_scope is not None
+            and span.instrumentation_scope.name in self.blocked_instrumentation_scopes
+        )
+
+    def _is_aeri_project_span(self, span: ReadableSpan) -> bool:
+        if not is_aeri_span(span):
+            return False
+
+        if span.instrumentation_scope is not None:
+            public_key_on_span = (
+                span.instrumentation_scope.attributes.get("public_key", None)
+                if span.instrumentation_scope.attributes
+                else None
+            )
+
+            return public_key_on_span == self.public_key
+
+        return False
+
+    @staticmethod
+    def _get_scope_name(span: ReadableSpan) -> Optional[str]:
+        if span.instrumentation_scope is None:
+            return None
+
+        return span.instrumentation_scope.name

aeri/_client/utils.py

@@ -0,0 +1,132 @@
+"""Utility functions for Aeri OpenTelemetry integration.
+
+This module provides utility functions for working with OpenTelemetry spans,
+including formatting and serialization of span data, and async execution helpers.
+"""
+
+import asyncio
+import json
+import threading
+from hashlib import sha256
+from typing import Any, Coroutine
+
+from opentelemetry import trace as otel_trace_api
+from opentelemetry.sdk import util
+from opentelemetry.sdk.trace import ReadableSpan
+
+
+def span_formatter(span: ReadableSpan) -> str:
+    parent_id = (
+        otel_trace_api.format_span_id(span.parent.span_id) if span.parent else None
+    )
+    start_time = util.ns_to_iso_str(span._start_time) if span._start_time else None
+    end_time = util.ns_to_iso_str(span._end_time) if span._end_time else None
+    status = {
+        "status_code": str(span._status.status_code.name),
+    }
+
+    if span._status.description:
+        status["description"] = span._status.description
+
+    context = (
+        {
+            "trace_id": otel_trace_api.format_trace_id(span._context.trace_id),
+            "span_id": otel_trace_api.format_span_id(span._context.span_id),
+            "trace_state": repr(span._context.trace_state),
+        }
+        if span._context
+        else None
+    )
+
+    instrumentationScope = json.loads(
+        span._instrumentation_scope.to_json() if span._instrumentation_scope else "{}"
+    )
+
+    return (
+        json.dumps(
+            {
+                "name": span._name,
+                "context": context,
+                "kind": str(span.kind),
+                "parent_id": parent_id,
+                "start_time": start_time,
+                "end_time": end_time,
+                "status": status,
+                "attributes": span._format_attributes(span._attributes),
+                "events": span._format_events(span._events),
+                "links": span._format_links(span._links),
+                "resource": json.loads(span.resource.to_json()),
+                "instrumentationScope": instrumentationScope,
+            },
+            indent=2,
+        )
+        + "\n"
+    )
+
+
+class _RunAsyncThread(threading.Thread):
+    """Helper thread class for running async coroutines in a separate thread."""
+
+    def __init__(self, coro: Coroutine[Any, Any, Any]) -> None:
+        self.coro = coro
+        self.result: Any = None
+        self.exception: Exception | None = None
+        super().__init__()
+
+    def run(self) -> None:
+        try:
+            self.result = asyncio.run(self.coro)
+        except Exception as e:
+            self.exception = e
+
+
+def run_async_safely(coro: Coroutine[Any, Any, Any]) -> Any:
+    """Safely run an async coroutine, handling existing event loops.
+
+    This function detects if there's already a running event loop and uses
+    a separate thread if needed to avoid the "asyncio.run() cannot be called
+    from a running event loop" error. This is particularly useful in environments
+    like Jupyter notebooks, FastAPI applications, or other async frameworks.
+
+    Args:
+        coro: The coroutine to run
+
+    Returns:
+        The result of the coroutine
+
+    Raises:
+        Any exception raised by the coroutine
+
+    Example:
+        ```python
+        # Works in both sync and async contexts
+        async def my_async_function():
+            await asyncio.sleep(1)
+            return "done"
+
+        result = run_async_safely(my_async_function())
+        ```
+    """
+    try:
+        # Check if there's already a running event loop
+        loop = asyncio.get_running_loop()
+    except RuntimeError:
+        # No running loop, safe to use asyncio.run()
+        return asyncio.run(coro)
+
+    if loop and loop.is_running():
+        # There's a running loop, use a separate thread
+        thread = _RunAsyncThread(coro)
+        thread.start()
+        thread.join()
+
+        if thread.exception:
+            raise thread.exception
+        return thread.result
+    else:
+        # Loop exists but not running, safe to use asyncio.run()
+        return asyncio.run(coro)
+
+
+def get_sha256_hash_hex(value: Any) -> str:
+    return sha256(value.encode("utf-8")).digest().hex()