hud-python 0.4.1__py3-none-any.whl → 0.4.3__py3-none-any.whl

hud/otel/collector.py

@@ -1,142 +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")
+"""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

@@ -1,164 +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")
+"""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")