245 lines
8.1 KiB
Python
245 lines
8.1 KiB
Python
"""Request correlation, OpenTelemetry spans and stage timing.
|
|
|
|
This module deliberately has a no-op default. The RAG domain remains runnable
|
|
without the OpenTelemetry packages or collector, while a configured runtime
|
|
gets one trace context shared by FastAPI's async middleware and its sync
|
|
thread-pool endpoint.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import base64
|
|
import re
|
|
import uuid
|
|
from contextlib import contextmanager
|
|
from contextvars import ContextVar
|
|
from functools import wraps
|
|
from time import perf_counter
|
|
from typing import Any, Iterator, Mapping
|
|
|
|
from .metrics import Metrics, NullMetrics, STAGE_DURATION
|
|
|
|
_CORRELATION_ID: ContextVar[str | None] = ContextVar(
|
|
"duocthu_correlation_id", default=None
|
|
)
|
|
_STAGE: ContextVar[str | None] = ContextVar("duocthu_stage", default=None)
|
|
_SAFE_CORRELATION = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._:-]{0,127}$")
|
|
_metrics: Metrics = NullMetrics()
|
|
|
|
try:
|
|
from opentelemetry import trace
|
|
|
|
_tracer = trace.get_tracer("duocthu.ai-service")
|
|
except ImportError: # pragma: no cover - exercised only in minimal installs
|
|
trace = None
|
|
_tracer = None
|
|
|
|
|
|
def configure_telemetry(settings: Any, metrics: Metrics | None = None) -> bool:
|
|
"""Configure one OTLP tracer provider; return whether tracing is active.
|
|
|
|
Metrics are configured independently and continue to work when tracing is
|
|
disabled or optional OpenTelemetry packages are absent.
|
|
"""
|
|
global _metrics, _tracer
|
|
_metrics = metrics or NullMetrics()
|
|
if not getattr(settings, "otel_enabled", False):
|
|
return False
|
|
|
|
try:
|
|
from opentelemetry import trace as otel_trace
|
|
from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
|
|
from opentelemetry.sdk.resources import Resource
|
|
from opentelemetry.sdk.trace import TracerProvider
|
|
from opentelemetry.sdk.trace.export import BatchSpanProcessor
|
|
from opentelemetry.sdk.trace.sampling import ParentBased, TraceIdRatioBased
|
|
except ImportError:
|
|
return False
|
|
|
|
current = otel_trace.get_tracer_provider()
|
|
# OpenTelemetry only permits setting the global provider once. Reuse a
|
|
# provider installed by a host/test instead of replacing it and warning.
|
|
if isinstance(current, TracerProvider):
|
|
_tracer = current.get_tracer("duocthu.ai-service")
|
|
return True
|
|
|
|
provider = TracerProvider(
|
|
resource=Resource.create(
|
|
{
|
|
"service.name": getattr(settings, "otel_service_name", "ai-service"),
|
|
"service.version": "0.1.0",
|
|
"deployment.environment": getattr(settings, "environment", "local"),
|
|
}
|
|
),
|
|
sampler=ParentBased(
|
|
TraceIdRatioBased(float(getattr(settings, "otel_sample_ratio", 1.0)))
|
|
),
|
|
)
|
|
exporter = OTLPSpanExporter(endpoint=settings.otel_exporter_otlp_endpoint)
|
|
provider.add_span_processor(BatchSpanProcessor(exporter))
|
|
for extra in _langfuse_exporters(settings, OTLPSpanExporter):
|
|
provider.add_span_processor(BatchSpanProcessor(extra))
|
|
otel_trace.set_tracer_provider(provider)
|
|
_tracer = provider.get_tracer("duocthu.ai-service")
|
|
return True
|
|
|
|
|
|
def _langfuse_exporters(settings: Any, exporter_cls: Any) -> list[Any]:
|
|
"""Zero or one extra OTLP exporter pointed at Langfuse.
|
|
|
|
Langfuse accepts standard OTLP/HTTP on `/api/public/otel/v1/traces`,
|
|
authenticated with HTTP Basic using the project's public key as the
|
|
username and the secret key as the password. Returning a list keeps the
|
|
caller unchanged when Langfuse is not configured -- which is the default,
|
|
and the case for every local run without a `.env`.
|
|
"""
|
|
base_url = str(getattr(settings, "langfuse_base_url", "") or "").rstrip("/")
|
|
public_key = str(getattr(settings, "langfuse_public_key", "") or "")
|
|
secret_key = str(getattr(settings, "langfuse_secret_key", "") or "")
|
|
if not (base_url and public_key and secret_key):
|
|
return []
|
|
|
|
token = base64.b64encode(f"{public_key}:{secret_key}".encode()).decode()
|
|
return [
|
|
exporter_cls(
|
|
endpoint=f"{base_url}/api/public/otel/v1/traces",
|
|
headers={"Authorization": f"Basic {token}"},
|
|
)
|
|
]
|
|
|
|
|
|
def normalize_correlation_id(candidate: str | None) -> str:
|
|
value = (candidate or "").strip()
|
|
return value if _SAFE_CORRELATION.fullmatch(value) else str(uuid.uuid4())
|
|
|
|
|
|
@contextmanager
|
|
def correlation_context(candidate: str | None) -> Iterator[str]:
|
|
correlation_id = normalize_correlation_id(candidate)
|
|
token = _CORRELATION_ID.set(correlation_id)
|
|
try:
|
|
yield correlation_id
|
|
finally:
|
|
_CORRELATION_ID.reset(token)
|
|
|
|
|
|
def current_correlation_id() -> str:
|
|
value = _CORRELATION_ID.get()
|
|
return value or str(uuid.uuid4())
|
|
|
|
|
|
def current_trace_id() -> str | None:
|
|
if trace is None:
|
|
return None
|
|
context = trace.get_current_span().get_span_context()
|
|
if not context.is_valid:
|
|
return None
|
|
return f"{context.trace_id:032x}"
|
|
|
|
|
|
def current_stage() -> str | None:
|
|
return _STAGE.get()
|
|
|
|
|
|
@contextmanager
|
|
def request_span(
|
|
method: str, route: str, headers: Mapping[str, str]
|
|
) -> Iterator[Any]:
|
|
"""Start a server span and extract an incoming W3C ``traceparent``."""
|
|
if _tracer is None:
|
|
yield None
|
|
return
|
|
try:
|
|
from opentelemetry import propagate
|
|
from opentelemetry.trace import SpanKind
|
|
except ImportError: # pragma: no cover - minimal install fallback
|
|
yield None
|
|
return
|
|
|
|
parent = propagate.extract(headers)
|
|
with _tracer.start_as_current_span(
|
|
f"{method} {route}",
|
|
context=parent,
|
|
kind=SpanKind.SERVER,
|
|
attributes={
|
|
"http.request.method": method,
|
|
"http.route": route,
|
|
"duocthu.correlation_id": current_correlation_id(),
|
|
},
|
|
) as span:
|
|
yield span
|
|
|
|
|
|
@contextmanager
|
|
def dependency_span(provider: str, operation: str) -> Iterator[Any]:
|
|
"""Trace one external dependency call without creating another histogram."""
|
|
if _tracer is None:
|
|
yield None
|
|
return
|
|
with _tracer.start_as_current_span(
|
|
f"provider.{provider}.{operation}",
|
|
attributes={
|
|
"duocthu.provider": provider,
|
|
"duocthu.operation": operation,
|
|
},
|
|
) as span:
|
|
yield span
|
|
|
|
|
|
@contextmanager
|
|
def stage(name: str, **attributes: Any) -> Iterator[Any]:
|
|
"""Create a child span and observe a low-cardinality stage histogram."""
|
|
started = perf_counter()
|
|
outcome = "ok"
|
|
span = None
|
|
token = _STAGE.set(name)
|
|
try:
|
|
if _tracer is None:
|
|
yield None
|
|
else:
|
|
safe_attributes = {
|
|
key: value
|
|
for key, value in attributes.items()
|
|
if value is not None and isinstance(value, (str, bool, int, float))
|
|
}
|
|
safe_attributes["duocthu.stage"] = name
|
|
with _tracer.start_as_current_span(
|
|
f"rag.stage.{name}", attributes=safe_attributes
|
|
) as span:
|
|
yield span
|
|
except BaseException as exc:
|
|
outcome = "cancelled" if exc.__class__.__name__ == "CancelledError" else "error"
|
|
if span is not None:
|
|
span.set_attribute("duocthu.outcome", outcome)
|
|
raise
|
|
finally:
|
|
elapsed = perf_counter() - started
|
|
if span is not None:
|
|
span.set_attribute("duocthu.outcome", outcome)
|
|
span.set_attribute("duocthu.duration_ms", elapsed * 1000)
|
|
_metrics.observe(STAGE_DURATION, elapsed, stage=name, outcome=outcome)
|
|
_STAGE.reset(token)
|
|
|
|
|
|
def annotate_current_span(**attributes: Any) -> None:
|
|
if trace is None:
|
|
return
|
|
span = trace.get_current_span()
|
|
if not span.is_recording():
|
|
return
|
|
for key, value in attributes.items():
|
|
if value is not None and isinstance(value, (str, bool, int, float)):
|
|
span.set_attribute(key, value)
|
|
|
|
|
|
def traced_stage(name: str):
|
|
"""Decorator form for pure functions that should remain otherwise untouched."""
|
|
def decorate(function):
|
|
@wraps(function)
|
|
def wrapped(*args, **kwargs):
|
|
with stage(name):
|
|
return function(*args, **kwargs)
|
|
|
|
return wrapped
|
|
|
|
return decorate
|