hud-python 0.3.4__py3-none-any.whl → 0.4.0__py3-none-any.whl

hud/otel/collector.py

@@ -0,0 +1,142 @@
+"""Global span collector for building in-memory traces.
+
+This module provides a way to collect spans during execution
+and retrieve them as a Trace object, enabling replay functionality
+without modifying agent code.
+"""
+
+from __future__ import annotations
+
+import logging
+import threading
+from contextvars import ContextVar
+from typing import TYPE_CHECKING
+
+from opentelemetry import trace
+from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
+
+from hud.types import Trace
+
+if TYPE_CHECKING:
+    from opentelemetry.sdk.trace import ReadableSpan
+
+logger = logging.getLogger(__name__)
+
+# Global storage for collected spans by task_run_id
+_TRACE_STORAGE: dict[str, TraceCollector] = {}
+_LOCK = threading.Lock()
+
+# Context variable to track if collection is enabled
+_collecting_enabled: ContextVar[bool] = ContextVar("collecting_enabled", default=False)
+
+
+class TraceCollector:
+    """Collects spans for a single task run."""
+
+    def __init__(self, task_run_id: str) -> None:
+        self.task_run_id = task_run_id
+        self.spans: list[ReadableSpan] = []
+        self._lock = threading.Lock()
+
+    def add_span(self, span: ReadableSpan) -> None:
+        """Thread-safe span addition."""
+        with self._lock:
+            self.spans.append(span)
+
+    def to_trace(self) -> Trace:
+        """Convert collected spans to a Trace object."""
+        from .exporters import HudSpan, _span_to_dict
+
+        trace = Trace()
+
+        # Convert spans to TraceSteps
+        for span in self.spans:
+            try:
+                # Use the same conversion logic as the exporter
+                span_dict = _span_to_dict(span)
+                hud_span = HudSpan.model_validate(span_dict)
+
+                # The attributes field is already a TraceStep
+                step = hud_span.attributes
+                # Add timing from the span itself
+                step.start_timestamp = hud_span.start_time
+                step.end_timestamp = hud_span.end_time
+                trace.append(step)
+
+            except Exception as e:
+                # Log but don't fail the whole trace
+                logger.debug("Failed to convert span: %s", e)
+
+        return trace
+
+
+class CollectingSpanExporter(SpanExporter):
+    """A span exporter that collects spans in memory for replay."""
+
+    def export(self, spans: list[ReadableSpan]) -> SpanExportResult:
+        """Collect spans if collection is enabled."""
+        if not _collecting_enabled.get():
+            return SpanExportResult.SUCCESS
+
+        for span in spans:
+            # Extract task_run_id from span
+            task_run_id = span.attributes.get("hud.task_run_id") if span.attributes else None
+            if not task_run_id or not isinstance(task_run_id, str):
+                continue
+
+            # Get or create collector
+            with _LOCK:
+                if task_run_id not in _TRACE_STORAGE:
+                    _TRACE_STORAGE[task_run_id] = TraceCollector(task_run_id)
+                collector = _TRACE_STORAGE[task_run_id]
+
+            # Add span
+            collector.add_span(span)
+
+        return SpanExportResult.SUCCESS
+
+    def shutdown(self) -> None:
+        """Clean up resources."""
+        with _LOCK:
+            _TRACE_STORAGE.clear()
+
+
+def enable_trace_collection(enabled: bool = True) -> None:
+    """Enable or disable in-memory trace collection."""
+    _collecting_enabled.set(enabled)
+
+
+def get_trace(task_run_id: str) -> Trace | None:
+    """Retrieve collected trace for a task run ID.
+
+    Returns None if no trace was collected or collection was disabled.
+    """
+    with _LOCK:
+        collector = _TRACE_STORAGE.get(task_run_id)
+        if collector:
+            return collector.to_trace()
+    return None
+
+
+def clear_trace(task_run_id: str) -> None:
+    """Clear collected trace for a task run ID."""
+    with _LOCK:
+        _TRACE_STORAGE.pop(task_run_id, None)
+
+
+def install_collector() -> None:
+    """Install the collecting span exporter.
+
+    This should be called after configure_telemetry().
+    """
+    provider = trace.get_tracer_provider()
+    # Guard for SDK tracer providers only
+    if hasattr(provider, "add_span_processor"):
+        from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+
+        exporter = CollectingSpanExporter()
+        processor = SimpleSpanProcessor(exporter)
+        try:
+            provider.add_span_processor(processor)  # type: ignore[attr-defined]
+        except Exception:
+            logger.warning("Failed to add span processor")

hud/otel/config.py

@@ -0,0 +1,164 @@
+"""Central configuration for OpenTelemetry inside HUD SDK.
+
+This file is responsible for
+1. creating the global ``TracerProvider``
+2. attaching span processors (HUD enrichment, batch + exporter)
+3. activating the community MCP instrumentation so that *every* MCP
+   request/response/notification is traced automatically.
+
+It is *idempotent*: calling :func:`configure_telemetry` more than once
+returns the same provider and does nothing.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+from opentelemetry import trace
+from opentelemetry.sdk.resources import Resource
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import BatchSpanProcessor
+
+from hud.settings import settings
+
+from .collector import enable_trace_collection, install_collector
+from .exporters import HudSpanExporter
+from .instrumentation import install_mcp_instrumentation
+from .processors import HudEnrichmentProcessor
+
+logger = logging.getLogger(__name__)
+
+# Global singleton provider so multiple calls do not create duplicates
+_TRACER_PROVIDER: TracerProvider | None = None
+
+
+def is_telemetry_configured() -> bool:
+    """Check if telemetry has been configured."""
+    return _TRACER_PROVIDER is not None
+
+
+# ---------------------------------------------------------------------------
+# Public API
+# ---------------------------------------------------------------------------
+
+
+def configure_telemetry(
+    *,
+    service_name: str = "hud-sdk",
+    service_version: str | None = None,
+    environment: str | None = None,
+    extra_resource_attributes: dict[str, Any] | None = None,
+    enable_otlp: bool = False,
+    otlp_endpoint: str | None = None,
+    otlp_headers: dict[str, str] | None = None,
+    enable_collection: bool = True,
+) -> TracerProvider:
+    """Initialise OpenTelemetry for the current Python process.
+
+    It is safe to call this in every entry-point; the provider will only
+    be created once.
+    """
+    global _TRACER_PROVIDER
+
+    if _TRACER_PROVIDER is not None:
+        return _TRACER_PROVIDER
+
+    # ------------------------------------------------------------------
+    # 1. Resource (identity of this service)
+    # ------------------------------------------------------------------
+    res_attrs: dict[str, Any] = {
+        "service.name": service_name,
+        "telemetry.sdk.name": "hud-otel",
+        "telemetry.sdk.language": "python",
+    }
+    if service_version:
+        res_attrs["service.version"] = service_version
+    if environment:
+        res_attrs["deployment.environment"] = environment
+    if extra_resource_attributes:
+        res_attrs.update(extra_resource_attributes)
+
+    resource = Resource.create(res_attrs)
+
+    # ------------------------------------------------------------------
+    # 2. Provider
+    # ------------------------------------------------------------------
+    provider = TracerProvider(resource=resource)
+    _TRACER_PROVIDER = provider
+
+    # ------------------------------------------------------------------
+    # 3. Processors / exporters
+    # ------------------------------------------------------------------
+    provider.add_span_processor(HudEnrichmentProcessor())
+
+    # HUD exporter (only if enabled and API key is available)
+    if settings.telemetry_enabled and settings.api_key:
+        exporter = HudSpanExporter(
+            telemetry_url=settings.hud_telemetry_url, api_key=settings.api_key
+        )
+        provider.add_span_processor(BatchSpanProcessor(exporter))
+    elif settings.telemetry_enabled and not settings.api_key and not enable_otlp:
+        # Error if no exporters are configured
+        raise ValueError(
+            "No telemetry backend configured. Either:\n"
+            "1. Set HUD_API_KEY environment variable for HUD telemetry\n"
+            "2. Use enable_otlp=True with configure_telemetry() for alternative backends (e.g., Jaeger)\n"  # noqa: E501
+        )
+    elif not settings.telemetry_enabled:
+        logger.info("HUD telemetry disabled via HUD_TELEMETRY_ENABLED=false")
+
+    # OTLP exporter (optional - for standard OTel viewers)
+    if enable_otlp:
+        try:
+            from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+
+            otlp_config = {}
+            if otlp_endpoint:
+                otlp_config["endpoint"] = otlp_endpoint
+                # Default to HTTP endpoint if not specified
+                if not otlp_endpoint.startswith(("http://", "https://")):
+                    otlp_config["endpoint"] = f"http://{otlp_endpoint}/v1/traces"
+            else:
+                # Default HTTP endpoint
+                otlp_config["endpoint"] = "http://localhost:4318/v1/traces"
+
+            if otlp_headers:
+                otlp_config["headers"] = otlp_headers
+
+            otlp_exporter = OTLPSpanExporter(**otlp_config)
+            provider.add_span_processor(BatchSpanProcessor(otlp_exporter))
+            logger.info("OTLP HTTP exporter enabled - endpoint: %s", otlp_config["endpoint"])
+        except ImportError:
+            logger.warning(
+                "OTLP export requested but opentelemetry-exporter-otlp-proto-http not installed. "
+                "Install with: pip install 'hud-python[agent]'"
+            )
+
+    # ------------------------------------------------------------------
+    # 4. Activate provider and instrumentation
+    # ------------------------------------------------------------------
+    trace.set_tracer_provider(provider)
+    install_mcp_instrumentation(provider)
+
+    # Install in-memory collector if requested
+    if enable_collection:
+        install_collector()
+        enable_trace_collection(True)
+        logger.debug("In-memory trace collection enabled")
+
+    # Agent instrumentation now handled by @hud.instrument decorators
+    logger.debug("OpenTelemetry configuration completed")
+
+    logger.debug("OpenTelemetry configured (provider id=%s)", id(provider))
+    return provider
+
+
+def shutdown_telemetry() -> None:
+    """Flush and shutdown the global provider (if configured)."""
+    global _TRACER_PROVIDER
+    if _TRACER_PROVIDER is None:
+        return
+    _TRACER_PROVIDER.shutdown()  # type: ignore[arg-type]
+    _TRACER_PROVIDER = None
+    logger.debug("OpenTelemetry shutdown complete")