diff --git a/headroom/__init__.py b/headroom/__init__.py index 026602d16..e52c4eba0 100644 --- a/headroom/__init__.py +++ b/headroom/__init__.py @@ -131,6 +131,20 @@ except ImportError: ScopeLevel = None # type: ignore[assignment,misc] with_memory = None # type: ignore[assignment] +from .observability import ( + HeadroomOtelMetrics, + HeadroomTracer, + LangfuseTracingConfig, + OTelMetricsConfig, + configure_langfuse_tracing, + configure_otel_metrics, + get_headroom_tracer, + get_langfuse_tracing_status, + get_otel_metrics, + get_otel_metrics_status, + reset_headroom_tracing, + reset_otel_metrics, +) from .providers import AnthropicProvider, OpenAIProvider, Provider, TokenCounter # Relevance scoring - BM25 always available, embedding requires sentence-transformers @@ -223,6 +237,19 @@ __all__ = [ "count_tokens_text", "count_tokens_messages", "generate_report", + # Observability + "HeadroomOtelMetrics", + "HeadroomTracer", + "LangfuseTracingConfig", + "OTelMetricsConfig", + "configure_otel_metrics", + "configure_langfuse_tracing", + "get_headroom_tracer", + "get_langfuse_tracing_status", + "get_otel_metrics", + "get_otel_metrics_status", + "reset_headroom_tracing", + "reset_otel_metrics", # Memory - hierarchical memory system "with_memory", # Main user-facing API "Memory", diff --git a/headroom/compress.py b/headroom/compress.py index 256a1237e..71808a9c9 100644 --- a/headroom/compress.py +++ b/headroom/compress.py @@ -57,9 +57,11 @@ Examples: from __future__ import annotations import logging +import threading from dataclasses import dataclass, field from typing import Any +from .observability import get_otel_metrics from .utils import extract_user_query as _extract_user_query logger = logging.getLogger(__name__) @@ -67,7 +69,7 @@ logger = logging.getLogger(__name__) # Lazy-initialized singleton pipeline _pipeline = None -_pipeline_lock = None +_pipeline_lock = threading.Lock() @dataclass @@ -171,6 +173,11 @@ def compress( ) except Exception as e: + get_otel_metrics().record_compression_failure( + model=model, + operation="compress", + error_type=type(e).__name__, + ) logger.warning("Compression failed, returning original messages: %s", e) return CompressResult( messages=messages, @@ -188,12 +195,6 @@ def _get_pipeline() -> Any: if _pipeline is not None: return _pipeline - import threading - - global _pipeline_lock - if _pipeline_lock is None: - _pipeline_lock = threading.Lock() - with _pipeline_lock: if _pipeline is not None: return _pipeline diff --git a/headroom/observability/__init__.py b/headroom/observability/__init__.py new file mode 100644 index 000000000..44eea15a5 --- /dev/null +++ b/headroom/observability/__init__.py @@ -0,0 +1,41 @@ +"""Operational observability helpers for Headroom.""" + +from .metrics import ( + HeadroomOtelMetrics, + OTelMetricsConfig, + configure_otel_metrics, + get_otel_metrics, + get_otel_metrics_status, + reset_otel_metrics, + set_otel_metrics, + shutdown_otel_metrics, +) +from .tracing import ( + HeadroomTracer, + LangfuseTracingConfig, + configure_langfuse_tracing, + get_headroom_tracer, + get_langfuse_tracing_status, + reset_headroom_tracing, + set_headroom_tracer, + shutdown_headroom_tracing, +) + +__all__ = [ + "HeadroomOtelMetrics", + "OTelMetricsConfig", + "configure_otel_metrics", + "get_otel_metrics", + "get_otel_metrics_status", + "HeadroomTracer", + "LangfuseTracingConfig", + "configure_langfuse_tracing", + "get_headroom_tracer", + "get_langfuse_tracing_status", + "reset_otel_metrics", + "reset_headroom_tracing", + "set_otel_metrics", + "set_headroom_tracer", + "shutdown_headroom_tracing", + "shutdown_otel_metrics", +] diff --git a/headroom/observability/metrics.py b/headroom/observability/metrics.py new file mode 100644 index 000000000..0e7135aad --- /dev/null +++ b/headroom/observability/metrics.py @@ -0,0 +1,511 @@ +"""OpenTelemetry-backed operational metrics for Headroom.""" + +from __future__ import annotations + +import logging +import os +from dataclasses import dataclass, field +from importlib.metadata import PackageNotFoundError +from importlib.metadata import version as package_version +from threading import Lock +from typing import Any, Literal + +from opentelemetry import metrics + +logger = logging.getLogger(__name__) + +MetricExporter = Literal["console", "otlp_http"] + +_SCOPE_NAME = "headroom" +_DEFAULT_EXPORT_INTERVAL_MS = 10000 +_MILLISECONDS_TO_SECONDS = 1000.0 + +_metrics_lock = Lock() +_global_metrics: HeadroomOtelMetrics | None = None +_owned_meter_provider: Any | None = None +_owned_metrics_config: OTelMetricsConfig | None = None + + +def _headroom_version() -> str: + try: + return package_version("headroom-ai") + except PackageNotFoundError: + return "unknown" + + +def _parse_bool(raw: str | None, default: bool = False) -> bool: + if raw is None: + return default + value = raw.strip().lower() + if value in {"1", "true", "yes", "on"}: + return True + if value in {"0", "false", "no", "off"}: + return False + return default + + +def _parse_int(raw: str | None, default: int) -> int: + if raw is None: + return default + try: + value = int(raw.strip()) + except ValueError: + return default + return value if value > 0 else default + + +def _parse_key_value_pairs(raw: str | None) -> dict[str, str]: + if raw is None: + return {} + + pairs: dict[str, str] = {} + for item in raw.split(","): + part = item.strip() + if not part or "=" not in part: + continue + key, value = part.split("=", 1) + key = key.strip() + value = value.strip() + if key and value: + pairs[key] = value + return pairs + + +@dataclass(slots=True) +class OTelMetricsConfig: + """Configuration for Headroom-managed OTEL metric export.""" + + enabled: bool = False + service_name: str = "headroom" + exporter: MetricExporter = "otlp_http" + endpoint: str | None = None + headers: dict[str, str] = field(default_factory=dict) + export_interval_millis: int = _DEFAULT_EXPORT_INTERVAL_MS + resource_attributes: dict[str, str] = field(default_factory=dict) + + @classmethod + def from_env(cls, *, default_service_name: str = "headroom") -> OTelMetricsConfig: + exporter_raw = ( + os.environ.get("HEADROOM_OTEL_METRICS_EXPORTER", "otlp_http") + .strip() + .lower() + .replace("-", "_") + ) + if exporter_raw not in {"console", "otlp_http"}: + logger.warning( + "Unknown HEADROOM_OTEL_METRICS_EXPORTER=%s; falling back to otlp_http", + exporter_raw, + ) + exporter_raw = "otlp_http" + + return cls( + enabled=_parse_bool(os.environ.get("HEADROOM_OTEL_METRICS_ENABLED"), default=False), + service_name=os.environ.get("HEADROOM_OTEL_SERVICE_NAME", default_service_name).strip() + or default_service_name, + exporter=exporter_raw, # type: ignore[arg-type] + endpoint=os.environ.get("HEADROOM_OTEL_METRICS_ENDPOINT") or None, + headers=_parse_key_value_pairs(os.environ.get("HEADROOM_OTEL_METRICS_HEADERS")), + export_interval_millis=_parse_int( + os.environ.get("HEADROOM_OTEL_METRICS_EXPORT_INTERVAL_MS"), + _DEFAULT_EXPORT_INTERVAL_MS, + ), + resource_attributes=_parse_key_value_pairs( + os.environ.get("HEADROOM_OTEL_RESOURCE_ATTRIBUTES") + ), + ) + + def status(self) -> dict[str, Any]: + return { + "configured": True, + "enabled": self.enabled, + "service_name": self.service_name, + "exporter": self.exporter, + "endpoint": self.endpoint, + "resource_attributes": dict(self.resource_attributes), + } + + +class HeadroomOtelMetrics: + """Shared OTEL metrics facade for Headroom operations.""" + + def __init__(self, meter_provider: Any | None = None): + if meter_provider is None: + self._meter = metrics.get_meter(_SCOPE_NAME, _headroom_version()) + else: + self._meter = meter_provider.get_meter(_SCOPE_NAME, _headroom_version()) + + self._proxy_requests = self._meter.create_counter( + "headroom.proxy.requests", + description="Proxy requests handled by Headroom.", + unit="1", + ) + self._proxy_cached_requests = self._meter.create_counter( + "headroom.proxy.requests.cached", + description="Proxy requests served with provider cache participation.", + unit="1", + ) + self._proxy_failed_requests = self._meter.create_counter( + "headroom.proxy.requests.failed", + description="Proxy requests that failed.", + unit="1", + ) + self._proxy_rate_limited_requests = self._meter.create_counter( + "headroom.proxy.requests.rate_limited", + description="Proxy requests rejected by rate limiting.", + unit="1", + ) + self._proxy_input_tokens = self._meter.create_counter( + "headroom.proxy.tokens.input", + description="Input tokens received by the proxy.", + unit="1", + ) + self._proxy_output_tokens = self._meter.create_counter( + "headroom.proxy.tokens.output", + description="Output tokens returned by upstream providers.", + unit="1", + ) + self._proxy_saved_tokens = self._meter.create_counter( + "headroom.proxy.tokens.saved", + description="Input tokens saved by Headroom compression.", + unit="1", + ) + self._proxy_cache_read_tokens = self._meter.create_counter( + "headroom.proxy.cache.read_tokens", + description="Provider cache read tokens observed by the proxy.", + unit="1", + ) + self._proxy_cache_write_tokens = self._meter.create_counter( + "headroom.proxy.cache.write_tokens", + description="Provider cache write tokens observed by the proxy.", + unit="1", + ) + self._proxy_cache_write_ttl_tokens = self._meter.create_counter( + "headroom.proxy.cache.write_ttl_tokens", + description="Provider cache write tokens by observed TTL bucket.", + unit="1", + ) + self._proxy_uncached_input_tokens = self._meter.create_counter( + "headroom.proxy.cache.uncached_input_tokens", + description="Proxy input tokens not served from provider cache.", + unit="1", + ) + self._proxy_cache_busts = self._meter.create_counter( + "headroom.proxy.cache.busts", + description="Requests that lost provider cache efficiency.", + unit="1", + ) + self._proxy_cache_bust_tokens_lost = self._meter.create_counter( + "headroom.proxy.cache.bust_tokens_lost", + description="Tokens that lost provider cache discount because of compression.", + unit="1", + ) + self._proxy_latency = self._meter.create_histogram( + "headroom.proxy.request.duration", + description="End-to-end proxy request duration.", + unit="s", + ) + self._proxy_overhead = self._meter.create_histogram( + "headroom.proxy.overhead.duration", + description="Time spent inside Headroom optimization logic.", + unit="s", + ) + self._proxy_ttfb = self._meter.create_histogram( + "headroom.proxy.ttfb.duration", + description="Upstream time to first byte observed by Headroom.", + unit="s", + ) + self._compression_runs = self._meter.create_counter( + "headroom.compression.runs", + description="Compression pipeline runs executed by Headroom.", + unit="1", + ) + self._compression_failures = self._meter.create_counter( + "headroom.compression.failures", + description="Compression operations that failed before producing a result.", + unit="1", + ) + self._compression_input_tokens = self._meter.create_counter( + "headroom.compression.tokens.input", + description="Input tokens analyzed by Headroom compression.", + unit="1", + ) + self._compression_output_tokens = self._meter.create_counter( + "headroom.compression.tokens.output", + description="Output tokens produced by Headroom compression.", + unit="1", + ) + self._compression_saved_tokens = self._meter.create_counter( + "headroom.compression.tokens.saved", + description="Tokens removed by Headroom compression.", + unit="1", + ) + self._compression_duration = self._meter.create_histogram( + "headroom.compression.pipeline.duration", + description="Compression pipeline execution duration.", + unit="s", + ) + self._compression_stage_duration = self._meter.create_histogram( + "headroom.compression.stage.duration", + description="Per-stage compression timing emitted by the pipeline.", + unit="s", + ) + self._compression_transforms = self._meter.create_counter( + "headroom.compression.transforms", + description="Transforms applied during compression.", + unit="1", + ) + self._waste_signal_tokens = self._meter.create_counter( + "headroom.compression.waste.tokens", + description="Waste tokens detected in compressed inputs.", + unit="1", + ) + + @staticmethod + def _attrs(**attrs: Any) -> dict[str, Any]: + filtered: dict[str, Any] = {} + for key, value in attrs.items(): + if value is None or value == "": + continue + filtered[key] = value + return filtered + + def record_proxy_request( + self, + *, + provider: str, + model: str, + input_tokens: int, + output_tokens: int, + tokens_saved: int, + latency_ms: float, + cached: bool = False, + overhead_ms: float = 0.0, + ttfb_ms: float = 0.0, + cache_read_tokens: int = 0, + cache_write_tokens: int = 0, + cache_write_5m_tokens: int = 0, + cache_write_1h_tokens: int = 0, + uncached_input_tokens: int = 0, + ) -> None: + attrs = self._attrs(provider=provider, model=model, cached=cached) + + self._proxy_requests.add(1, attrs) + if cached: + self._proxy_cached_requests.add(1, attrs) + + self._proxy_input_tokens.add(max(input_tokens, 0), attrs) + self._proxy_output_tokens.add(max(output_tokens, 0), attrs) + self._proxy_saved_tokens.add(max(tokens_saved, 0), attrs) + self._proxy_latency.record(max(latency_ms, 0.0) / _MILLISECONDS_TO_SECONDS, attrs) + + if overhead_ms > 0: + self._proxy_overhead.record(overhead_ms / _MILLISECONDS_TO_SECONDS, attrs) + if ttfb_ms > 0: + self._proxy_ttfb.record(ttfb_ms / _MILLISECONDS_TO_SECONDS, attrs) + if cache_read_tokens > 0: + self._proxy_cache_read_tokens.add(cache_read_tokens, attrs) + if cache_write_tokens > 0: + self._proxy_cache_write_tokens.add(cache_write_tokens, attrs) + if uncached_input_tokens > 0: + self._proxy_uncached_input_tokens.add(uncached_input_tokens, attrs) + if cache_write_5m_tokens > 0: + self._proxy_cache_write_ttl_tokens.add( + cache_write_5m_tokens, + self._attrs(provider=provider, model=model, ttl="5m"), + ) + if cache_write_1h_tokens > 0: + self._proxy_cache_write_ttl_tokens.add( + cache_write_1h_tokens, + self._attrs(provider=provider, model=model, ttl="1h"), + ) + + def record_proxy_failed(self, *, provider: str | None = None, model: str | None = None) -> None: + self._proxy_failed_requests.add(1, self._attrs(provider=provider, model=model)) + + def record_proxy_rate_limited( + self, + *, + provider: str | None = None, + model: str | None = None, + ) -> None: + self._proxy_rate_limited_requests.add(1, self._attrs(provider=provider, model=model)) + + def record_proxy_cache_bust(self, *, tokens_lost: int) -> None: + self._proxy_cache_busts.add(1) + self._proxy_cache_bust_tokens_lost.add(max(tokens_lost, 0)) + + def record_pipeline_run( + self, + *, + model: str, + provider: str | None, + tokens_before: int, + tokens_after: int, + duration_ms: float, + timing: dict[str, float] | None = None, + transforms_applied: list[str] | None = None, + waste_signals: dict[str, int] | None = None, + ) -> None: + attrs = self._attrs(model=model, provider=provider) + tokens_saved = max(tokens_before - tokens_after, 0) + + self._compression_runs.add(1, attrs) + self._compression_input_tokens.add(max(tokens_before, 0), attrs) + self._compression_output_tokens.add(max(tokens_after, 0), attrs) + self._compression_saved_tokens.add(tokens_saved, attrs) + self._compression_duration.record(max(duration_ms, 0.0) / _MILLISECONDS_TO_SECONDS, attrs) + + if transforms_applied: + for transform in transforms_applied: + self._compression_transforms.add( + 1, self._attrs(model=model, provider=provider, transform=transform) + ) + + if timing: + for stage, stage_ms in timing.items(): + if stage == "pipeline_total" or stage.startswith("_"): + continue + self._compression_stage_duration.record( + max(stage_ms, 0.0) / _MILLISECONDS_TO_SECONDS, + self._attrs(model=model, provider=provider, stage=stage), + ) + + if waste_signals: + for signal_name, token_count in waste_signals.items(): + if token_count > 0: + self._waste_signal_tokens.add( + token_count, + self._attrs(model=model, provider=provider, signal=signal_name), + ) + + def record_compression_failure( + self, + *, + model: str, + operation: str, + error_type: str, + ) -> None: + self._compression_failures.add( + 1, + self._attrs(model=model, operation=operation, error_type=error_type), + ) + + +def get_otel_metrics() -> HeadroomOtelMetrics: + global _global_metrics + + if _global_metrics is None: + with _metrics_lock: + if _global_metrics is None: + _global_metrics = HeadroomOtelMetrics() + + return _global_metrics + + +def set_otel_metrics(otel_metrics: HeadroomOtelMetrics) -> HeadroomOtelMetrics: + global _global_metrics + with _metrics_lock: + _global_metrics = otel_metrics + return otel_metrics + + +def configure_otel_metrics(config: OTelMetricsConfig | None = None) -> HeadroomOtelMetrics: + global _global_metrics + global _owned_meter_provider + global _owned_metrics_config + + resolved = config or OTelMetricsConfig() + if not resolved.enabled: + return get_otel_metrics() + + try: + from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter + from opentelemetry.sdk.metrics import MeterProvider + from opentelemetry.sdk.metrics.export import ( + ConsoleMetricExporter, + PeriodicExportingMetricReader, + ) + from opentelemetry.sdk.resources import SERVICE_NAME, SERVICE_VERSION, Resource + except ImportError: + logger.warning( + "OpenTelemetry SDK/exporter packages are not installed. " + "Install headroom-ai[otel] to enable managed OTEL metric export." + ) + return get_otel_metrics() + + exporter: Any + if resolved.exporter == "console": + exporter = ConsoleMetricExporter() + else: + exporter_kwargs: dict[str, Any] = {} + if resolved.endpoint is not None: + exporter_kwargs["endpoint"] = resolved.endpoint + if resolved.headers: + exporter_kwargs["headers"] = resolved.headers + exporter = OTLPMetricExporter(**exporter_kwargs) + + reader = PeriodicExportingMetricReader( + exporter, + export_interval_millis=resolved.export_interval_millis, + ) + resource = Resource.create( + { + SERVICE_NAME: resolved.service_name, + SERVICE_VERSION: _headroom_version(), + **resolved.resource_attributes, + } + ) + meter_provider = MeterProvider(resource=resource, metric_readers=[reader]) + otel_metrics = HeadroomOtelMetrics(meter_provider=meter_provider) + + previous_provider = None + with _metrics_lock: + previous_provider = _owned_meter_provider + _owned_meter_provider = meter_provider + _owned_metrics_config = resolved + _global_metrics = otel_metrics + + if previous_provider is not None: + try: + previous_provider.shutdown() + except Exception: + logger.debug("Failed to shut down previous OTEL metrics provider", exc_info=True) + + return otel_metrics + + +def get_otel_metrics_status() -> dict[str, Any]: + with _metrics_lock: + if _owned_metrics_config is None: + return { + "configured": False, + "enabled": False, + "service_name": None, + "exporter": None, + "endpoint": None, + "resource_attributes": {}, + } + return _owned_metrics_config.status() + + +def shutdown_otel_metrics() -> None: + global _global_metrics + global _owned_meter_provider + global _owned_metrics_config + + provider = None + with _metrics_lock: + provider = _owned_meter_provider + _owned_meter_provider = None + _owned_metrics_config = None + _global_metrics = None + + if provider is not None: + try: + provider.shutdown() + except Exception: + logger.debug("Failed to shut down OTEL metrics provider", exc_info=True) + + +def reset_otel_metrics() -> None: + shutdown_otel_metrics() diff --git a/headroom/observability/tracing.py b/headroom/observability/tracing.py new file mode 100644 index 000000000..76eb9caf6 --- /dev/null +++ b/headroom/observability/tracing.py @@ -0,0 +1,228 @@ +"""OTEL tracing helpers for Headroom and Langfuse.""" + +from __future__ import annotations + +import base64 +import logging +import os +from dataclasses import dataclass, field +from threading import Lock +from typing import Any + +from opentelemetry import trace + +from .metrics import _headroom_version, _parse_bool, _parse_key_value_pairs + +logger = logging.getLogger(__name__) + +_SCOPE_NAME = "headroom" + +_tracing_lock = Lock() +_global_tracer: HeadroomTracer | None = None +_owned_tracer_provider: Any | None = None +_owned_langfuse_config: LangfuseTracingConfig | None = None + + +@dataclass(slots=True) +class LangfuseTracingConfig: + """Configuration for Headroom-managed Langfuse OTLP trace export.""" + + enabled: bool = False + public_key: str = field(default="", repr=False) + secret_key: str = field(default="", repr=False) + base_url: str = "https://cloud.langfuse.com" + service_name: str = "headroom" + resource_attributes: dict[str, str] = field(default_factory=dict) + + @property + def endpoint(self) -> str: + return f"{self.base_url.rstrip('/')}/api/public/otel/v1/traces" + + @property + def auth_header(self) -> str: + encoded = base64.b64encode(f"{self.public_key}:{self.secret_key}".encode()).decode() + return f"Basic {encoded}" + + @property + def headers(self) -> dict[str, str]: + return { + "Authorization": self.auth_header, + "x-langfuse-ingestion-version": "4", + } + + @classmethod + def from_env(cls, *, default_service_name: str = "headroom") -> LangfuseTracingConfig: + public_key = os.environ.get("LANGFUSE_PUBLIC_KEY", "").strip() + secret_key = os.environ.get("LANGFUSE_SECRET_KEY", "").strip() + + return cls( + enabled=_parse_bool( + os.environ.get("HEADROOM_LANGFUSE_ENABLED"), + default=False, + ), + public_key=public_key, + secret_key=secret_key, + base_url=( + os.environ.get("LANGFUSE_BASE_URL") + or os.environ.get("LANGFUSE_OTEL_HOST") + or "https://cloud.langfuse.com" + ).strip(), + service_name=os.environ.get( + "HEADROOM_LANGFUSE_SERVICE_NAME", default_service_name + ).strip() + or default_service_name, + resource_attributes=_parse_key_value_pairs( + os.environ.get("HEADROOM_LANGFUSE_RESOURCE_ATTRIBUTES") + ), + ) + + def is_complete(self) -> bool: + return bool(self.public_key and self.secret_key) + + def status(self) -> dict[str, Any]: + return { + "configured": True, + "enabled": self.enabled, + "service_name": self.service_name, + "base_url": self.base_url, + "endpoint": self.endpoint, + } + + +class HeadroomTracer: + """Tracer facade used by shared Headroom compression paths.""" + + def __init__(self, tracer_provider: Any | None = None): + if tracer_provider is None: + self._tracer = trace.get_tracer(_SCOPE_NAME, _headroom_version()) + else: + self._tracer = tracer_provider.get_tracer(_SCOPE_NAME, _headroom_version()) + + def start_as_current_span( + self, + name: str, + *, + attributes: dict[str, Any] | None = None, + ) -> Any: + return self._tracer.start_as_current_span( + name, + attributes=attributes, + record_exception=True, + set_status_on_exception=True, + ) + + +def get_headroom_tracer() -> HeadroomTracer: + global _global_tracer + + if _global_tracer is None: + with _tracing_lock: + if _global_tracer is None: + _global_tracer = HeadroomTracer() + + return _global_tracer + + +def set_headroom_tracer(headroom_tracer: HeadroomTracer) -> HeadroomTracer: + global _global_tracer + with _tracing_lock: + _global_tracer = headroom_tracer + return headroom_tracer + + +def configure_langfuse_tracing( + config: LangfuseTracingConfig | None = None, +) -> HeadroomTracer: + global _global_tracer + global _owned_tracer_provider + global _owned_langfuse_config + + resolved = config or LangfuseTracingConfig() + if not resolved.enabled: + return get_headroom_tracer() + if not resolved.is_complete(): + logger.warning( + "Langfuse tracing is enabled but LANGFUSE_PUBLIC_KEY/LANGFUSE_SECRET_KEY are missing." + ) + return get_headroom_tracer() + + try: + from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter + from opentelemetry.sdk.resources import SERVICE_NAME, SERVICE_VERSION, Resource + from opentelemetry.sdk.trace import TracerProvider + from opentelemetry.sdk.trace.export import BatchSpanProcessor + except ImportError: + logger.warning( + "OpenTelemetry SDK/exporter packages are not installed. " + "Install headroom-ai[otel] to enable Langfuse OTLP tracing." + ) + return get_headroom_tracer() + + resource = Resource.create( + { + SERVICE_NAME: resolved.service_name, + SERVICE_VERSION: _headroom_version(), + **resolved.resource_attributes, + } + ) + tracer_provider = TracerProvider(resource=resource) + tracer_provider.add_span_processor( + BatchSpanProcessor( + OTLPSpanExporter( + endpoint=resolved.endpoint, + headers=resolved.headers, + ) + ) + ) + headroom_tracer = HeadroomTracer(tracer_provider=tracer_provider) + + previous_provider = None + with _tracing_lock: + previous_provider = _owned_tracer_provider + _owned_tracer_provider = tracer_provider + _owned_langfuse_config = resolved + _global_tracer = headroom_tracer + + if previous_provider is not None: + try: + previous_provider.shutdown() + except Exception: + logger.debug("Failed to shut down previous Langfuse tracer provider", exc_info=True) + + return headroom_tracer + + +def get_langfuse_tracing_status() -> dict[str, Any]: + with _tracing_lock: + if _owned_langfuse_config is None: + return { + "configured": False, + "enabled": False, + "service_name": None, + "base_url": None, + "endpoint": None, + } + return _owned_langfuse_config.status() + + +def shutdown_headroom_tracing() -> None: + global _global_tracer + global _owned_tracer_provider + global _owned_langfuse_config + + provider = None + with _tracing_lock: + provider = _owned_tracer_provider + _owned_tracer_provider = None + _owned_langfuse_config = None + _global_tracer = None + + if provider is not None: + try: + provider.shutdown() + except Exception: + logger.debug("Failed to shut down Langfuse tracer provider", exc_info=True) + + +def reset_headroom_tracing() -> None: + shutdown_headroom_tracing() diff --git a/headroom/transforms/pipeline.py b/headroom/transforms/pipeline.py index 5c826429a..06cbda913 100644 --- a/headroom/transforms/pipeline.py +++ b/headroom/transforms/pipeline.py @@ -4,6 +4,7 @@ from __future__ import annotations import logging import time +from contextlib import nullcontext from typing import TYPE_CHECKING, Any from ..config import ( @@ -17,6 +18,7 @@ from ..config import ( TransformResult, WasteSignals, ) +from ..observability import get_headroom_tracer, get_otel_metrics from ..tokenizer import Tokenizer from ..utils import deep_copy_messages from .base import Transform @@ -145,6 +147,16 @@ class TransformPipeline: return Tokenizer(get_tokenizer(model), model) # type: ignore[arg-type] + def _provider_name(self) -> str | None: + if self._provider is None: + return None + + name = getattr(self._provider, "provider_name", None) + if isinstance(name, str) and name: + return name + + return self._provider.__class__.__name__.removesuffix("Provider").lower() + def apply( self, messages: list[dict[str, Any]], @@ -166,7 +178,9 @@ class TransformPipeline: Returns: Combined TransformResult. """ + record_metrics = kwargs.pop("record_metrics", True) tokenizer = self._get_tokenizer(model) + provider_name = self._provider_name() # Get model limit from kwargs (should be set by client) model_limit = kwargs.get("model_limit") @@ -188,135 +202,204 @@ class TransformPipeline: model, ) - # Track all transforms applied - all_transforms: list[str] = [] - all_markers: list[str] = [] - all_warnings: list[str] = [] - all_timing: dict[str, float] = {} # transform_name → ms - - # Track transform diffs if enabled - transform_diffs: list[TransformDiff] = [] - generate_diff = self.config.generate_diff_artifact - - t_copy = time.perf_counter() - current_messages = deep_copy_messages(messages) - copy_ms = (time.perf_counter() - t_copy) * 1000 - - all_timing["_deep_copy"] = copy_ms - all_timing["_initial_token_count"] = count_ms - - pipeline_start = time.perf_counter() - - frozen_count = kwargs.get("frozen_message_count", 0) - if frozen_count > 0: - logger.info( - "Pipeline: freezing first %d/%d messages (prefix cached by provider)", - frozen_count, - len(messages), + tracer = get_headroom_tracer() + span_attributes = { + "headroom.model": model, + "headroom.provider": provider_name or "unknown", + "headroom.message_count": len(messages), + "headroom.tokens.before": tokens_before, + } + pipeline_span_context = ( + tracer.start_as_current_span( + "headroom.compression.pipeline", + attributes=span_attributes, ) + if record_metrics + else nullcontext() + ) - for transform in self.transforms: - # Check if transform should run - if not transform.should_apply(current_messages, tokenizer, **kwargs): - continue + with pipeline_span_context as pipeline_span: + # Track all transforms applied + all_transforms: list[str] = [] + all_markers: list[str] = [] + all_warnings: list[str] = [] + all_timing: dict[str, float] = {} # transform_name → ms - # Time the transform - t0 = time.perf_counter() - result = transform.apply(current_messages, tokenizer, **kwargs) - duration_ms = (time.perf_counter() - t0) * 1000 + # Track transform diffs if enabled + transform_diffs: list[TransformDiff] = [] + generate_diff = self.config.generate_diff_artifact - # Update messages for next transform - current_messages = result.messages + t_copy = time.perf_counter() + current_messages = deep_copy_messages(messages) + copy_ms = (time.perf_counter() - t_copy) * 1000 - # Use token counts reported by the transform itself — avoids - # redundant O(N) recount of the full message list after each step. - tokens_before_transform = result.tokens_before - tokens_after_transform = result.tokens_after + all_timing["_deep_copy"] = copy_ms + all_timing["_initial_token_count"] = count_ms - # Accumulate results - all_transforms.extend(result.transforms_applied) - all_markers.extend(result.markers_inserted) - all_warnings.extend(result.warnings) - all_timing[transform.name] = duration_ms + pipeline_start = time.perf_counter() - # Merge sub-transform timing (e.g. ContentRouter's per-compressor breakdown) - if result.timing: - all_timing.update(result.timing) - - # Log transform results - if result.transforms_applied: + frozen_count = kwargs.get("frozen_message_count", 0) + if frozen_count > 0: logger.info( - "Transform %s: %d -> %d tokens (saved %d) [%.1fms]", - transform.name, - tokens_before_transform, - tokens_after_transform, - tokens_before_transform - tokens_after_transform, - duration_ms, + "Pipeline: freezing first %d/%d messages (prefix cached by provider)", + frozen_count, + len(messages), + ) + + for transform in self.transforms: + # Check if transform should run + if not transform.should_apply(current_messages, tokenizer, **kwargs): + continue + + transform_span_context = ( + tracer.start_as_current_span( + "headroom.compression.transform", + attributes={ + "headroom.model": model, + "headroom.provider": provider_name or "unknown", + "headroom.transform": transform.name, + }, + ) + if record_metrics + else nullcontext() + ) + + with transform_span_context as transform_span: + # Time the transform + t0 = time.perf_counter() + result = transform.apply(current_messages, tokenizer, **kwargs) + duration_ms = (time.perf_counter() - t0) * 1000 + + # Update messages for next transform + current_messages = result.messages + + # Use token counts reported by the transform itself — avoids + # redundant O(N) recount of the full message list after each step. + tokens_before_transform = result.tokens_before + tokens_after_transform = result.tokens_after + + if transform_span is not None and transform_span.is_recording(): + transform_span.set_attribute( + "headroom.tokens.before", tokens_before_transform + ) + transform_span.set_attribute( + "headroom.tokens.after", tokens_after_transform + ) + transform_span.set_attribute( + "headroom.tokens.saved", + tokens_before_transform - tokens_after_transform, + ) + transform_span.set_attribute("headroom.duration_ms", duration_ms) + transform_span.set_attribute( + "headroom.transforms_applied", + len(result.transforms_applied), + ) + + # Accumulate results + all_transforms.extend(result.transforms_applied) + all_markers.extend(result.markers_inserted) + all_warnings.extend(result.warnings) + all_timing[transform.name] = duration_ms + + # Merge sub-transform timing (e.g. ContentRouter's per-compressor breakdown) + if result.timing: + all_timing.update(result.timing) + + # Log transform results + if result.transforms_applied: + logger.info( + "Transform %s: %d -> %d tokens (saved %d) [%.1fms]", + transform.name, + tokens_before_transform, + tokens_after_transform, + tokens_before_transform - tokens_after_transform, + duration_ms, + ) + else: + logger.debug( + "Transform %s: no changes [%.1fms]", transform.name, duration_ms + ) + + # Record diff if enabled + if generate_diff: + transform_diffs.append( + TransformDiff( + transform_name=transform.name, + tokens_before=tokens_before_transform, + tokens_after=tokens_after_transform, + tokens_saved=tokens_before_transform - tokens_after_transform, + details=", ".join(result.transforms_applied) + if result.transforms_applied + else "", + duration_ms=duration_ms, + ) + ) + + # Single final token count — the only full recount in the pipeline. + # Earlier per-transform counts come from each transform's own result. + t_final_count = time.perf_counter() + tokens_after = tokenizer.count_messages(current_messages) + all_timing["_final_token_count"] = (time.perf_counter() - t_final_count) * 1000 + + pipeline_ms = (time.perf_counter() - pipeline_start) * 1000 + all_timing["pipeline_total"] = pipeline_ms + + # Log pipeline summary + total_saved = tokens_before - tokens_after + timing_parts = " ".join(f"{k}={v:.0f}ms" for k, v in all_timing.items()) + if total_saved > 0: + logger.info( + "Pipeline complete: %d -> %d tokens (saved %d, %.1f%% reduction) [%s]", + tokens_before, + tokens_after, + total_saved, + (total_saved / tokens_before * 100) if tokens_before > 0 else 0, + timing_parts, ) else: - logger.debug("Transform %s: no changes [%.1fms]", transform.name, duration_ms) + logger.debug("Pipeline complete: no token savings [%s]", timing_parts) - # Record diff if enabled + # Build diff artifact if enabled + diff_artifact = None if generate_diff: - transform_diffs.append( - TransformDiff( - transform_name=transform.name, - tokens_before=tokens_before_transform, - tokens_after=tokens_after_transform, - tokens_saved=tokens_before_transform - tokens_after_transform, - details=", ".join(result.transforms_applied) - if result.transforms_applied - else "", - duration_ms=duration_ms, - ) + diff_artifact = DiffArtifact( + request_id=kwargs.get("request_id", ""), + original_tokens=tokens_before, + optimized_tokens=tokens_after, + total_tokens_saved=tokens_before - tokens_after, + transforms=transform_diffs, ) - # Single final token count — the only full recount in the pipeline. - # Earlier per-transform counts come from each transform's own result. - t_final_count = time.perf_counter() - tokens_after = tokenizer.count_messages(current_messages) - all_timing["_final_token_count"] = (time.perf_counter() - t_final_count) * 1000 + # Detect waste signals in original messages (only when significant compression) + waste_signals: WasteSignals | None = None + if tokens_before > tokens_after and (tokens_before - tokens_after) > 100: + try: + from ..parser import parse_messages - pipeline_ms = (time.perf_counter() - pipeline_start) * 1000 - all_timing["pipeline_total"] = pipeline_ms + _, _, waste_signals = parse_messages(messages, tokenizer) + if waste_signals.total() == 0: + waste_signals = None + except Exception: + pass - # Log pipeline summary - total_saved = tokens_before - tokens_after - timing_parts = " ".join(f"{k}={v:.0f}ms" for k, v in all_timing.items()) - if total_saved > 0: - logger.info( - "Pipeline complete: %d -> %d tokens (saved %d, %.1f%% reduction) [%s]", - tokens_before, - tokens_after, - total_saved, - (total_saved / tokens_before * 100) if tokens_before > 0 else 0, - timing_parts, - ) - else: - logger.debug("Pipeline complete: no token savings [%s]", timing_parts) + if pipeline_span is not None and pipeline_span.is_recording(): + pipeline_span.set_attribute("headroom.tokens.after", tokens_after) + pipeline_span.set_attribute("headroom.tokens.saved", total_saved) + pipeline_span.set_attribute("headroom.duration_ms", pipeline_ms) + pipeline_span.set_attribute("headroom.transforms_applied", len(all_transforms)) + pipeline_span.set_attribute("headroom.warnings", len(all_warnings)) - # Build diff artifact if enabled - diff_artifact = None - if generate_diff: - diff_artifact = DiffArtifact( - request_id=kwargs.get("request_id", ""), - original_tokens=tokens_before, - optimized_tokens=tokens_after, - total_tokens_saved=tokens_before - tokens_after, - transforms=transform_diffs, - ) - - # Detect waste signals in original messages (only when significant compression) - waste_signals: WasteSignals | None = None - if tokens_before > tokens_after and (tokens_before - tokens_after) > 100: - try: - from ..parser import parse_messages - - _, _, waste_signals = parse_messages(messages, tokenizer) - if waste_signals.total() == 0: - waste_signals = None - except Exception: - pass + if record_metrics: + get_otel_metrics().record_pipeline_run( + model=model, + provider=provider_name, + tokens_before=tokens_before, + tokens_after=tokens_after, + duration_ms=pipeline_ms, + timing=all_timing, + transforms_applied=all_transforms, + waste_signals=waste_signals.to_dict() if waste_signals is not None else None, + ) return TransformResult( messages=current_messages, @@ -350,7 +433,7 @@ class TransformPipeline: TransformResult with simulated changes. """ # apply() already works on a copy, so this is safe - return self.apply(messages, model, **kwargs) + return self.apply(messages, model, record_metrics=False, **kwargs) def create_pipeline( diff --git a/pyproject.toml b/pyproject.toml index d8b0e037e..84f2fc792 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -50,6 +50,7 @@ dependencies = [ "litellm==1.82.3", # Model registry, pricing, and provider support "click>=8.1.0", # CLI framework "rich>=13.0.0", # Rich terminal output + "opentelemetry-api>=1.24.0", # Safe no-op OTEL API for instrumentation ] [project.optional-dependencies] @@ -102,6 +103,11 @@ image = [ reports = [ "jinja2>=3.0.0", ] +# OpenTelemetry metrics export +otel = [ + "opentelemetry-sdk>=1.24.0", + "opentelemetry-exporter-otlp-proto-http>=1.24.0", +] # any-llm multi-provider backend (requires Python 3.11+) anyllm = [ "any-llm-sdk>=1.0.0; python_version >= '3.11'", @@ -173,6 +179,8 @@ dev = [ "fastapi>=0.100.0", "uvicorn>=0.23.0", "httpx[http2]>=0.24.0", + "opentelemetry-sdk>=1.24.0", + "opentelemetry-exporter-otlp-proto-http>=1.24.0", "ollama>=0.4.0", "langchain-ollama>=0.2.0", "hnswlib>=0.8.0", @@ -182,7 +190,7 @@ dev = [ ] # All optional dependencies (everything you need) all = [ - "headroom-ai[proxy,code,ml,memory,relevance,image,reports,evals,voice,html,benchmark,mcp]", + "headroom-ai[proxy,code,ml,memory,relevance,image,reports,otel,evals,voice,html,benchmark,mcp]", ] [project.scripts] diff --git a/tests/test_observability_metrics.py b/tests/test_observability_metrics.py new file mode 100644 index 000000000..f49c3b71c --- /dev/null +++ b/tests/test_observability_metrics.py @@ -0,0 +1,200 @@ +"""Tests for OTEL-backed operational observability.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +import pytest +from opentelemetry.sdk.metrics import MeterProvider +from opentelemetry.sdk.metrics.export import InMemoryMetricReader + +from headroom.observability import HeadroomOtelMetrics, reset_otel_metrics, set_otel_metrics +from headroom.proxy.prometheus_metrics import PrometheusMetrics +from headroom.transforms.pipeline import TransformPipeline + + +def _collect_metrics(reader: InMemoryMetricReader) -> dict[str, Any]: + data = reader.get_metrics_data() + collected: dict[str, Any] = {} + + for resource_metric in data.resource_metrics: + for scope_metric in resource_metric.scope_metrics: + for metric in scope_metric.metrics: + collected[metric.name] = metric + + return collected + + +def _find_point(metric: Any, **expected_attributes: Any) -> Any: + for point in metric.data.data_points: + if all(point.attributes.get(key) == value for key, value in expected_attributes.items()): + return point + raise AssertionError(f"No datapoint matched attributes: {expected_attributes}") + + +def test_headroom_otel_metrics_records_proxy_and_pipeline_metrics() -> None: + reader = InMemoryMetricReader() + provider = MeterProvider(metric_readers=[reader]) + otel_metrics = HeadroomOtelMetrics(meter_provider=provider) + + otel_metrics.record_proxy_request( + provider="anthropic", + model="claude-opus-4-6", + input_tokens=120, + output_tokens=30, + tokens_saved=45, + latency_ms=18.5, + cached=True, + overhead_ms=4.0, + ttfb_ms=12.0, + cache_read_tokens=25, + cache_write_tokens=35, + cache_write_5m_tokens=10, + cache_write_1h_tokens=25, + uncached_input_tokens=60, + ) + otel_metrics.record_proxy_cache_bust(tokens_lost=7) + otel_metrics.record_pipeline_run( + model="claude-opus-4-6", + provider="anthropic", + tokens_before=120, + tokens_after=75, + duration_ms=6.5, + timing={"_deep_copy": 0.2, "router": 3.5, "pipeline_total": 6.5}, + transforms_applied=["router:smart_crusher:0.35"], + waste_signals={"json_bloat": 12}, + ) + + metrics = _collect_metrics(reader) + + requests = metrics["headroom.proxy.requests"] + request_point = _find_point( + requests, + provider="anthropic", + model="claude-opus-4-6", + cached=True, + ) + assert request_point.value == 1 + + latency = metrics["headroom.proxy.request.duration"] + latency_point = _find_point( + latency, + provider="anthropic", + model="claude-opus-4-6", + cached=True, + ) + assert latency_point.count == 1 + assert latency_point.sum == pytest.approx(0.0185) + + ttl_tokens = metrics["headroom.proxy.cache.write_ttl_tokens"] + five_minute_ttl = _find_point( + ttl_tokens, + provider="anthropic", + model="claude-opus-4-6", + ttl="5m", + ) + assert five_minute_ttl.value == 10 + + compression_runs = metrics["headroom.compression.runs"] + compression_point = _find_point( + compression_runs, + provider="anthropic", + model="claude-opus-4-6", + ) + assert compression_point.value == 1 + + stage_duration = metrics["headroom.compression.stage.duration"] + router_stage = _find_point( + stage_duration, + provider="anthropic", + model="claude-opus-4-6", + stage="router", + ) + assert router_stage.count == 1 + assert router_stage.sum == pytest.approx(0.0035) + + assert len(stage_duration.data.data_points) == 1 + + waste_tokens = metrics["headroom.compression.waste.tokens"] + waste_point = _find_point( + waste_tokens, + provider="anthropic", + model="claude-opus-4-6", + signal="json_bloat", + ) + assert waste_point.value == 12 + + +@dataclass +class _SpyMetrics: + pipeline_calls: list[dict[str, Any]] = field(default_factory=list) + + def record_pipeline_run(self, **kwargs: Any) -> None: + self.pipeline_calls.append(kwargs) + + +@dataclass +class _SpyProxyMetrics: + failed_calls: list[dict[str, Any]] = field(default_factory=list) + rate_limited_calls: list[dict[str, Any]] = field(default_factory=list) + + def record_proxy_failed(self, **kwargs: Any) -> None: + self.failed_calls.append(kwargs) + + def record_proxy_rate_limited(self, **kwargs: Any) -> None: + self.rate_limited_calls.append(kwargs) + + +def test_transform_pipeline_simulate_skips_metric_recording() -> None: + spy = _SpyMetrics() + set_otel_metrics(spy) # type: ignore[arg-type] + + try: + pipeline = TransformPipeline(transforms=[]) + messages = [{"role": "user", "content": "hello world"}] + + pipeline.apply(messages, model="gpt-4o", model_limit=1024) + assert len(spy.pipeline_calls) == 1 + + pipeline.simulate(messages, model="gpt-4o", model_limit=1024) + assert len(spy.pipeline_calls) == 1 + finally: + reset_otel_metrics() + + +def test_proxy_failure_and_rate_limit_metrics_include_provider_labels() -> None: + reader = InMemoryMetricReader() + provider = MeterProvider(metric_readers=[reader]) + otel_metrics = HeadroomOtelMetrics(meter_provider=provider) + + otel_metrics.record_proxy_failed(provider="openai") + otel_metrics.record_proxy_rate_limited(provider="anthropic", model="claude-sonnet") + + metrics = _collect_metrics(reader) + + failed_point = _find_point(metrics["headroom.proxy.requests.failed"], provider="openai") + assert failed_point.value == 1 + + rate_limited_point = _find_point( + metrics["headroom.proxy.requests.rate_limited"], + provider="anthropic", + model="claude-sonnet", + ) + assert rate_limited_point.value == 1 + + +@pytest.mark.asyncio +async def test_prometheus_metrics_reads_late_configured_otel_metrics() -> None: + spy = _SpyProxyMetrics() + metrics = PrometheusMetrics() + set_otel_metrics(spy) # type: ignore[arg-type] + + try: + await metrics.record_failed(provider="openai") + await metrics.record_rate_limited(provider="anthropic", model="claude-sonnet") + + assert spy.failed_calls == [{"provider": "openai", "model": None}] + assert spy.rate_limited_calls == [{"provider": "anthropic", "model": "claude-sonnet"}] + finally: + reset_otel_metrics() diff --git a/tests/test_observability_tracing.py b/tests/test_observability_tracing.py new file mode 100644 index 000000000..d10054631 --- /dev/null +++ b/tests/test_observability_tracing.py @@ -0,0 +1,71 @@ +"""Tests for Langfuse/OTEL tracing helpers.""" + +from __future__ import annotations + +import pytest +from opentelemetry.sdk.resources import Resource +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + +from headroom.observability import ( + HeadroomTracer, + LangfuseTracingConfig, + get_langfuse_tracing_status, + reset_headroom_tracing, + set_headroom_tracer, +) +from headroom.transforms.pipeline import TransformPipeline + + +def test_langfuse_tracing_config_builds_trace_endpoint() -> None: + config = LangfuseTracingConfig( + enabled=True, + public_key="pk-lf-test", + secret_key="sk-lf-test", + base_url="https://cloud.langfuse.com", + service_name="headroom-proxy", + ) + + assert config.endpoint == "https://cloud.langfuse.com/api/public/otel/v1/traces" + assert config.headers["x-langfuse-ingestion-version"] == "4" + assert config.headers["Authorization"].startswith("Basic ") + assert "sk-lf-test" not in repr(config) + + +def test_transform_pipeline_emits_trace_spans() -> None: + exporter = InMemorySpanExporter() + provider = TracerProvider(resource=Resource.create({"service.name": "headroom-test"})) + provider.add_span_processor(SimpleSpanProcessor(exporter)) + set_headroom_tracer(HeadroomTracer(tracer_provider=provider)) + + try: + pipeline = TransformPipeline(transforms=[]) + messages = [{"role": "user", "content": "hello world"}] + pipeline.apply(messages, model="gpt-4o", model_limit=1024) + + spans = exporter.get_finished_spans() + assert len(spans) == 1 + span = spans[0] + assert span.name == "headroom.compression.pipeline" + assert span.attributes["headroom.model"] == "gpt-4o" + assert span.attributes["headroom.tokens.before"] >= 1 + assert span.attributes["headroom.tokens.after"] >= 1 + finally: + reset_headroom_tracing() + + +def test_langfuse_tracing_status_defaults_to_unconfigured() -> None: + reset_headroom_tracing() + status = get_langfuse_tracing_status() + assert status["configured"] is False + assert status["enabled"] is False + + +def test_langfuse_tracing_requires_explicit_enable(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LANGFUSE_PUBLIC_KEY", "pk-lf-test") + monkeypatch.setenv("LANGFUSE_SECRET_KEY", "sk-lf-test") + + config = LangfuseTracingConfig.from_env(default_service_name="headroom-proxy") + + assert config.enabled is False