akadmin/headroom
1
0
Fork 0
mirror of https://github.com/headroomlabs-ai/headroom.git synced 2026-08-27 14:17:10 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

feat: add OTEL observability core

Browse source 
Add a shared observability layer for OTEL metrics and Langfuse tracing.
Instrument the shared compression pipeline once so SDK and proxy paths can
reuse the same metrics and spans without scattering telemetry concerns.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This commit is contained in:
JerrettDavis 2026-04-09 21:20:34 -05:00
parent e4f72569c7
commit ebd482c0a0
9 changed files with 1291 additions and 121 deletions
Download patch file Download diff file Expand all files Collapse all files

27
headroom/__init__.py
View file

@ -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",

15
headroom/compress.py
View file

@ -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

41
headroom/observability/__init__.py Normal file
View file

@ -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",
]

511
headroom/observability/metrics.py Normal file
View file

@ -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()

228
headroom/observability/tracing.py Normal file
View file

@ -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()

309
headroom/transforms/pipeline.py
View file

@ -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(

10
pyproject.toml
View file

@ -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]

200
tests/test_observability_metrics.py Normal file
View file

@ -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()

71
tests/test_observability_tracing.py Normal file
View file

@ -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