mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-13 12:14:00 +00:00
Export traces as OpenTelemetry GenAI spans and metrics
Finished traces are replayed with their recorded timestamps and explicit parents under the request's server span, using gen_ai.* attribute names. Content is exported only when OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT opts in. gen_ai.client token-usage and operation-duration histograms are recorded per model call.
This commit is contained in:
1 parent
3f45d019b0
commit
b5e9257659
3 files changed
+609
No files matched your search
@@ -0,0 +1,274 @@
|
||||
"""OpenTelemetry export for execution traces (GenAI semantic conventions).
|
||||
|
||||
A finished :class:`~docsgpt.tracing.core.Trace` is *replayed*: every recorded
|
||||
span is re-created with its original start and end timestamps and explicit
|
||||
parents, under the OTel context captured when the trace started (normally the
|
||||
HTTP server span). Replaying after the fact means no OTel context is ever
|
||||
attached across the agent loop's generator yields. The cost is that GenAI
|
||||
spans reach the backend only when the request ends, and auto-instrumented
|
||||
HTTP/DB spans made during a step do not nest under it.
|
||||
|
||||
Attribute names follow the (still incubating) ``gen_ai.*`` conventions and are
|
||||
written as literals rather than imported from the semconv package. Prompt and
|
||||
tool content is exported only when
|
||||
``OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT`` opts in.
|
||||
|
||||
Nothing here runs unless an OTel SDK is installed as the global provider (for
|
||||
example by launching under ``opentelemetry-instrument``).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import threading
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
from docsgpt.core.settings import settings
|
||||
from docsgpt.tracing.core import (
|
||||
KIND_EMBEDDING,
|
||||
KIND_LLM,
|
||||
KIND_TOOL,
|
||||
STATUS_ERROR,
|
||||
STATUS_OK,
|
||||
Trace,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_INSTRUMENTATION_NAME = "docsgpt.tracing"
|
||||
|
||||
#: DocsGPT provider names that differ from the ``gen_ai.provider.name`` values.
|
||||
_PROVIDER_NAMES = {
|
||||
"google": "gcp.gen_ai",
|
||||
"azure_openai": "azure.ai.openai",
|
||||
"aws_bedrock": "aws.bedrock",
|
||||
"mistral": "mistral_ai",
|
||||
}
|
||||
|
||||
#: Preview keys with a dedicated GenAI attribute; others export as ``docsgpt.preview.<key>``.
|
||||
_TOOL_PREVIEW_ATTRIBUTES = {
|
||||
"arguments": "gen_ai.tool.call.arguments",
|
||||
"result": "gen_ai.tool.call.result",
|
||||
}
|
||||
|
||||
_CAPTURE_VALUES = {"span_only", "span_and_event", "true"}
|
||||
|
||||
# Advisory buckets from the GenAI metrics conventions.
|
||||
_DURATION_BUCKETS = [0.01, 0.02, 0.04, 0.08, 0.16, 0.32, 0.64, 1.28, 2.56, 5.12, 10.24, 20.48, 40.96, 81.92]
|
||||
_TOKEN_BUCKETS = [1, 4, 16, 64, 256, 1024, 4096, 16384, 65536, 262144, 1048576, 4194304, 16777216, 67108864]
|
||||
|
||||
|
||||
def provider_name(internal: Optional[str]) -> str:
|
||||
"""Map a DocsGPT ``provider_name`` to the ``gen_ai.provider.name`` value."""
|
||||
if not internal:
|
||||
return "unknown"
|
||||
return _PROVIDER_NAMES.get(internal, internal)
|
||||
|
||||
|
||||
def content_capture_enabled() -> bool:
|
||||
"""Whether the standard GenAI content-capture variable opts into span content."""
|
||||
value = os.getenv("OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT", "")
|
||||
return value.strip().lower() in _CAPTURE_VALUES
|
||||
|
||||
|
||||
def _is_noop(provider: Any) -> bool:
|
||||
from opentelemetry import trace as ot_trace
|
||||
|
||||
return isinstance(provider, (ot_trace.ProxyTracerProvider, ot_trace.NoOpTracerProvider))
|
||||
|
||||
|
||||
def _attr_value(value: Any) -> Any:
|
||||
"""Coerce a recorded attribute into an OTel-legal value (or ``None`` to skip)."""
|
||||
if value is None:
|
||||
return None
|
||||
if isinstance(value, (str, bool, int, float)):
|
||||
return value
|
||||
if isinstance(value, (list, tuple)):
|
||||
items = [v for v in value if v is not None]
|
||||
if items and all(isinstance(v, str) for v in items):
|
||||
return list(items)
|
||||
if items and all(isinstance(v, bool) for v in items):
|
||||
return list(items)
|
||||
if items and all(isinstance(v, (int, float)) and not isinstance(v, bool) for v in items):
|
||||
return list(items)
|
||||
if not items:
|
||||
return None
|
||||
try:
|
||||
return json.dumps(value, ensure_ascii=False, default=str)
|
||||
except (TypeError, ValueError):
|
||||
return str(value)
|
||||
|
||||
|
||||
def _attributes(raw: Dict[str, Any]) -> Dict[str, Any]:
|
||||
out = {}
|
||||
for key, value in raw.items():
|
||||
coerced = _attr_value(value)
|
||||
if coerced is not None:
|
||||
out[key] = coerced
|
||||
return out
|
||||
|
||||
|
||||
def _preview_attributes(kind: str, previews: Dict[str, Any]) -> Dict[str, Any]:
|
||||
out = {}
|
||||
for key, value in previews.items():
|
||||
name = _TOOL_PREVIEW_ATTRIBUTES.get(key) if kind == KIND_TOOL else None
|
||||
out[name or f"docsgpt.preview.{key}"] = value if isinstance(value, str) else _attr_value(value)
|
||||
return out
|
||||
|
||||
|
||||
def _root_attributes(trace: Trace) -> Dict[str, Any]:
|
||||
return _attributes(
|
||||
{
|
||||
"docsgpt.trace.id": trace.id,
|
||||
"docsgpt.source": trace.source,
|
||||
"docsgpt.request_id": trace.request_id,
|
||||
"docsgpt.message_id": trace.message_id,
|
||||
"docsgpt.activity_id": trace.activity_id,
|
||||
"docsgpt.workflow_run_id": trace.workflow_run_id,
|
||||
"gen_ai.conversation.id": trace.conversation_id,
|
||||
"gen_ai.agent.id": trace.agent_id,
|
||||
"docsgpt.status": trace.status,
|
||||
"docsgpt.dropped_spans": trace.dropped_spans or None,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def export_trace(trace: Trace, tracer_provider: Any = None) -> Optional[str]:
|
||||
"""Replay ``trace`` as OTel spans; returns the OTel trace id (hex) or ``None``.
|
||||
|
||||
Args:
|
||||
trace: A finished trace.
|
||||
tracer_provider: Override the global provider (tests).
|
||||
|
||||
Returns:
|
||||
The 32-char hex OTel trace id, or ``None`` when export is disabled,
|
||||
no SDK is configured, or the trace is empty.
|
||||
"""
|
||||
if not settings.TRACES_OTEL_EXPORT or not trace.spans:
|
||||
return None
|
||||
from opentelemetry import trace as ot_trace
|
||||
from opentelemetry.trace import SpanKind, Status, StatusCode
|
||||
|
||||
provider = tracer_provider or ot_trace.get_tracer_provider()
|
||||
if _is_noop(provider):
|
||||
return None
|
||||
tracer = provider.get_tracer(_INSTRUMENTATION_NAME)
|
||||
capture = content_capture_enabled() and not trace.content_blocked
|
||||
|
||||
root = tracer.start_span(
|
||||
f"docsgpt {trace.source}",
|
||||
context=trace.otel_context,
|
||||
kind=SpanKind.INTERNAL,
|
||||
attributes=_root_attributes(trace),
|
||||
start_time=trace.start_ns,
|
||||
)
|
||||
if trace.status == STATUS_ERROR:
|
||||
root.set_status(Status(StatusCode.ERROR))
|
||||
exported: Dict[str, Any] = {}
|
||||
for span in trace.spans:
|
||||
parent = exported.get(span.parent_id, root)
|
||||
attributes = dict(span.attributes)
|
||||
if trace.conversation_id and span.kind in (KIND_LLM, KIND_TOOL):
|
||||
attributes.setdefault("gen_ai.conversation.id", trace.conversation_id)
|
||||
if span.status and span.status not in (STATUS_OK, STATUS_ERROR):
|
||||
attributes["docsgpt.status"] = span.status
|
||||
attributes = _attributes(attributes)
|
||||
if capture and span.previews:
|
||||
attributes.update(_preview_attributes(span.kind, span.previews))
|
||||
otel_span = tracer.start_span(
|
||||
span.name,
|
||||
context=ot_trace.set_span_in_context(parent),
|
||||
kind=SpanKind.CLIENT if span.kind in (KIND_LLM, KIND_EMBEDDING) else SpanKind.INTERNAL,
|
||||
attributes=attributes,
|
||||
start_time=trace.span_start_ns(span),
|
||||
)
|
||||
if span.status == STATUS_ERROR:
|
||||
otel_span.set_status(Status(StatusCode.ERROR, span.error))
|
||||
exported[span.id] = otel_span
|
||||
for span in trace.spans:
|
||||
exported[span.id].end(end_time=trace.span_end_ns(span))
|
||||
root.end(end_time=trace.start_ns + int((trace.duration_ms or 0) * 1e6))
|
||||
context = root.get_span_context()
|
||||
return format(context.trace_id, "032x") if context.trace_id else None
|
||||
|
||||
|
||||
# -- metrics --------------------------------------------------------------------
|
||||
|
||||
_instruments_lock = threading.Lock()
|
||||
_global_instruments: Optional[tuple] = None
|
||||
|
||||
|
||||
def _make_instruments(meter_provider: Any = None) -> tuple:
|
||||
from opentelemetry import metrics
|
||||
|
||||
meter = (
|
||||
meter_provider.get_meter(_INSTRUMENTATION_NAME)
|
||||
if meter_provider is not None
|
||||
else metrics.get_meter(_INSTRUMENTATION_NAME)
|
||||
)
|
||||
|
||||
def _histogram(name: str, unit: str, description: str, buckets: list):
|
||||
try:
|
||||
return meter.create_histogram(
|
||||
name, unit=unit, description=description,
|
||||
explicit_bucket_boundaries_advisory=buckets,
|
||||
)
|
||||
except TypeError: # API older than bucket advisories
|
||||
return meter.create_histogram(name, unit=unit, description=description)
|
||||
|
||||
duration = _histogram(
|
||||
"gen_ai.client.operation.duration", "s", "GenAI operation duration.", _DURATION_BUCKETS
|
||||
)
|
||||
tokens = _histogram(
|
||||
"gen_ai.client.token.usage", "{token}", "Number of input and output tokens used.", _TOKEN_BUCKETS
|
||||
)
|
||||
return duration, tokens
|
||||
|
||||
|
||||
def _instruments(meter_provider: Any = None) -> tuple:
|
||||
global _global_instruments
|
||||
if meter_provider is not None:
|
||||
return _make_instruments(meter_provider)
|
||||
with _instruments_lock:
|
||||
if _global_instruments is None:
|
||||
_global_instruments = _make_instruments()
|
||||
return _global_instruments
|
||||
|
||||
|
||||
def record_llm_metrics(
|
||||
*,
|
||||
provider: Optional[str],
|
||||
model: Optional[str],
|
||||
input_tokens: int,
|
||||
output_tokens: int,
|
||||
duration_s: float,
|
||||
error_type: Optional[str],
|
||||
operation: str = "chat",
|
||||
meter_provider: Any = None,
|
||||
) -> None:
|
||||
"""Record the two GenAI client metrics for one model call.
|
||||
|
||||
Token usage is recorded only for successful calls, as the conventions
|
||||
require; duration is recorded for every call, with ``error.type`` on
|
||||
failures. Without a configured SDK the meter is a no-op.
|
||||
"""
|
||||
if not settings.TRACES_OTEL_EXPORT:
|
||||
return
|
||||
try:
|
||||
duration, tokens = _instruments(meter_provider)
|
||||
attributes = {
|
||||
"gen_ai.operation.name": operation,
|
||||
"gen_ai.provider.name": provider_name(provider),
|
||||
}
|
||||
if model:
|
||||
attributes["gen_ai.request.model"] = str(model)
|
||||
if error_type:
|
||||
duration.record(duration_s, {**attributes, "error.type": error_type})
|
||||
return
|
||||
duration.record(duration_s, attributes)
|
||||
tokens.record(int(input_tokens or 0), {**attributes, "gen_ai.token.type": "input"})
|
||||
tokens.record(int(output_tokens or 0), {**attributes, "gen_ai.token.type": "output"})
|
||||
except Exception: # noqa: BLE001 - metrics must never break a call
|
||||
logger.debug("Failed to record GenAI metrics", exc_info=True)
|
||||
@@ -0,0 +1,235 @@
|
||||
"""Tests for replaying a finished trace as OpenTelemetry GenAI spans."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
from opentelemetry import trace as ot_trace
|
||||
from opentelemetry.sdk.metrics import MeterProvider
|
||||
from opentelemetry.sdk.metrics.export import InMemoryMetricReader
|
||||
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 opentelemetry.trace import SpanKind, StatusCode
|
||||
|
||||
from docsgpt import tracing
|
||||
from docsgpt.core.settings import settings
|
||||
from docsgpt.tracing import otel as trace_otel
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _settings(monkeypatch):
|
||||
monkeypatch.setattr(settings, "TRACES_ENABLED", True)
|
||||
monkeypatch.setattr(settings, "TRACES_OTEL_EXPORT", True)
|
||||
monkeypatch.setattr(settings, "TRACES_CAPTURE_CONTENT", True)
|
||||
monkeypatch.delenv("OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT", raising=False)
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def provider():
|
||||
exporter = InMemorySpanExporter()
|
||||
tp = TracerProvider()
|
||||
tp.add_span_processor(SimpleSpanProcessor(exporter))
|
||||
tp.exporter = exporter
|
||||
return tp
|
||||
|
||||
|
||||
def _sample_trace(otel_context=None):
|
||||
trace = tracing.start_trace(
|
||||
source="stream",
|
||||
request_id="req-1",
|
||||
message_id="msg-1",
|
||||
conversation_id="conv-1",
|
||||
agent_id="agent-1",
|
||||
capture_otel_context=False,
|
||||
)
|
||||
trace.otel_context = otel_context
|
||||
with tracing.activate(trace):
|
||||
with tracing.span(
|
||||
tracing.KIND_AGENT,
|
||||
"invoke_agent Support",
|
||||
attributes={"gen_ai.operation.name": "invoke_agent", "gen_ai.agent.name": "Support"},
|
||||
):
|
||||
with tracing.span(
|
||||
tracing.KIND_LLM,
|
||||
"chat gpt-4o",
|
||||
attributes={
|
||||
"gen_ai.operation.name": "chat",
|
||||
"gen_ai.provider.name": "openai",
|
||||
"gen_ai.request.model": "gpt-4o",
|
||||
"gen_ai.usage.input_tokens": 12,
|
||||
"gen_ai.usage.output_tokens": 3,
|
||||
"docsgpt.sources": ["a", "b"],
|
||||
"docsgpt.meta": {"nested": True},
|
||||
},
|
||||
) as llm:
|
||||
llm.preview("output", "hello")
|
||||
with tracing.span(
|
||||
tracing.KIND_TOOL,
|
||||
"execute_tool search",
|
||||
attributes={"gen_ai.tool.name": "search"},
|
||||
) as tool:
|
||||
tool.preview("arguments", {"q": "x"})
|
||||
tool.preview("result", "found")
|
||||
tool.end(error=RuntimeError("tool broke"))
|
||||
trace.finish()
|
||||
return trace
|
||||
|
||||
|
||||
class TestReplay:
|
||||
def test_spans_are_parented_and_timed(self, provider):
|
||||
trace = _sample_trace()
|
||||
otel_trace_id = trace_otel.export_trace(trace, tracer_provider=provider)
|
||||
spans = {s.name: s for s in provider.exporter.get_finished_spans()}
|
||||
root = spans["docsgpt stream"]
|
||||
agent = spans["invoke_agent Support"]
|
||||
llm = spans["chat gpt-4o"]
|
||||
tool = spans["execute_tool search"]
|
||||
assert agent.parent.span_id == root.context.span_id
|
||||
assert llm.parent.span_id == agent.context.span_id
|
||||
assert tool.parent.span_id == agent.context.span_id
|
||||
assert otel_trace_id == format(root.context.trace_id, "032x")
|
||||
recorded = {s.name: s for s in trace.spans}
|
||||
assert llm.start_time == trace.span_start_ns(recorded["chat gpt-4o"])
|
||||
assert llm.end_time == trace.span_end_ns(recorded["chat gpt-4o"])
|
||||
assert root.start_time == trace.start_ns
|
||||
assert root.end_time >= agent.end_time
|
||||
|
||||
def test_semconv_attributes_and_kinds(self, provider):
|
||||
trace_otel.export_trace(_sample_trace(), tracer_provider=provider)
|
||||
spans = {s.name: s for s in provider.exporter.get_finished_spans()}
|
||||
llm = spans["chat gpt-4o"]
|
||||
assert llm.kind == SpanKind.CLIENT
|
||||
assert llm.attributes["gen_ai.usage.input_tokens"] == 12
|
||||
assert llm.attributes["gen_ai.conversation.id"] == "conv-1"
|
||||
assert tuple(llm.attributes["docsgpt.sources"]) == ("a", "b")
|
||||
assert llm.attributes["docsgpt.meta"] == '{"nested": true}'
|
||||
root = spans["docsgpt stream"]
|
||||
assert root.attributes["docsgpt.request_id"] == "req-1"
|
||||
assert root.attributes["docsgpt.message_id"] == "msg-1"
|
||||
|
||||
def test_error_status(self, provider):
|
||||
trace_otel.export_trace(_sample_trace(), tracer_provider=provider)
|
||||
tool = {s.name: s for s in provider.exporter.get_finished_spans()}["execute_tool search"]
|
||||
assert tool.status.status_code == StatusCode.ERROR
|
||||
assert tool.attributes["error.type"] == "RuntimeError"
|
||||
|
||||
def test_no_content_by_default(self, provider):
|
||||
trace_otel.export_trace(_sample_trace(), tracer_provider=provider)
|
||||
for span in provider.exporter.get_finished_spans():
|
||||
assert not any("preview" in k or "call.arguments" in k for k in span.attributes)
|
||||
|
||||
def test_content_when_opted_in(self, provider, monkeypatch):
|
||||
monkeypatch.setenv("OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT", "span_only")
|
||||
trace_otel.export_trace(_sample_trace(), tracer_provider=provider)
|
||||
spans = {s.name: s for s in provider.exporter.get_finished_spans()}
|
||||
tool = spans["execute_tool search"]
|
||||
assert tool.attributes["gen_ai.tool.call.arguments"] == '{"q": "x"}'
|
||||
assert tool.attributes["gen_ai.tool.call.result"] == "found"
|
||||
assert spans["chat gpt-4o"].attributes["docsgpt.preview.output"] == "hello"
|
||||
|
||||
def test_content_blocked_trace_never_exports_content(self, provider, monkeypatch):
|
||||
monkeypatch.setenv("OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT", "SPAN_ONLY")
|
||||
trace = _sample_trace()
|
||||
trace.content_blocked = True
|
||||
trace_otel.export_trace(trace, tracer_provider=provider)
|
||||
for span in provider.exporter.get_finished_spans():
|
||||
assert not any("preview" in k or "call.result" in k for k in span.attributes)
|
||||
|
||||
def test_root_is_parented_to_captured_context(self, provider):
|
||||
server = provider.get_tracer("test").start_span("GET /stream")
|
||||
context = ot_trace.set_span_in_context(server)
|
||||
server.end()
|
||||
trace_otel.export_trace(_sample_trace(otel_context=context), tracer_provider=provider)
|
||||
root = {s.name: s for s in provider.exporter.get_finished_spans()}["docsgpt stream"]
|
||||
assert root.parent.span_id == server.get_span_context().span_id
|
||||
assert root.context.trace_id == server.get_span_context().trace_id
|
||||
|
||||
def test_skipped_without_sdk_provider(self):
|
||||
assert trace_otel.export_trace(_sample_trace(), tracer_provider=ot_trace.NoOpTracerProvider()) is None
|
||||
|
||||
def test_skipped_when_disabled(self, provider, monkeypatch):
|
||||
monkeypatch.setattr(settings, "TRACES_OTEL_EXPORT", False)
|
||||
assert trace_otel.export_trace(_sample_trace(), tracer_provider=provider) is None
|
||||
assert provider.exporter.get_finished_spans() == ()
|
||||
|
||||
|
||||
class TestProviderName:
|
||||
@pytest.mark.parametrize(
|
||||
"internal, expected",
|
||||
[
|
||||
("openai", "openai"),
|
||||
("google", "gcp.gen_ai"),
|
||||
("anthropic", "anthropic"),
|
||||
("azure_openai", "azure.ai.openai"),
|
||||
("groq", "groq"),
|
||||
(None, "unknown"),
|
||||
],
|
||||
)
|
||||
def test_mapping(self, internal, expected):
|
||||
assert trace_otel.provider_name(internal) == expected
|
||||
|
||||
|
||||
class TestMetrics:
|
||||
def test_llm_metrics_recorded(self):
|
||||
reader = InMemoryMetricReader()
|
||||
mp = MeterProvider(metric_readers=[reader])
|
||||
trace_otel.record_llm_metrics(
|
||||
provider="openai",
|
||||
model="gpt-4o",
|
||||
input_tokens=10,
|
||||
output_tokens=5,
|
||||
duration_s=0.25,
|
||||
error_type=None,
|
||||
meter_provider=mp,
|
||||
)
|
||||
data = reader.get_metrics_data()
|
||||
metrics = {
|
||||
m.name: m
|
||||
for rm in data.resource_metrics
|
||||
for sm in rm.scope_metrics
|
||||
for m in sm.metrics
|
||||
}
|
||||
usage = metrics["gen_ai.client.token.usage"]
|
||||
by_type = {
|
||||
p.attributes["gen_ai.token.type"]: p.sum for p in usage.data.data_points
|
||||
}
|
||||
assert by_type == {"input": 10, "output": 5}
|
||||
point = usage.data.data_points[0]
|
||||
assert point.attributes["gen_ai.provider.name"] == "openai"
|
||||
assert point.attributes["gen_ai.request.model"] == "gpt-4o"
|
||||
assert point.attributes["gen_ai.operation.name"] == "chat"
|
||||
duration = metrics["gen_ai.client.operation.duration"].data.data_points[0]
|
||||
assert duration.sum == pytest.approx(0.25)
|
||||
|
||||
def test_error_type_on_duration_only(self):
|
||||
reader = InMemoryMetricReader()
|
||||
mp = MeterProvider(metric_readers=[reader])
|
||||
trace_otel.record_llm_metrics(
|
||||
provider="openai",
|
||||
model="m",
|
||||
input_tokens=0,
|
||||
output_tokens=0,
|
||||
duration_s=0.1,
|
||||
error_type="Timeout",
|
||||
meter_provider=mp,
|
||||
)
|
||||
metrics = {
|
||||
m.name: m
|
||||
for rm in reader.get_metrics_data().resource_metrics
|
||||
for sm in rm.scope_metrics
|
||||
for m in sm.metrics
|
||||
}
|
||||
assert "gen_ai.client.token.usage" not in metrics
|
||||
point = metrics["gen_ai.client.operation.duration"].data.data_points[0]
|
||||
assert point.attributes["error.type"] == "Timeout"
|
||||
|
||||
def test_metrics_disabled(self, monkeypatch):
|
||||
monkeypatch.setattr(settings, "TRACES_OTEL_EXPORT", False)
|
||||
reader = InMemoryMetricReader()
|
||||
mp = MeterProvider(metric_readers=[reader])
|
||||
trace_otel.record_llm_metrics(
|
||||
provider="openai", model="m", input_tokens=1, output_tokens=1,
|
||||
duration_s=0.1, error_type=None, meter_provider=mp,
|
||||
)
|
||||
data = reader.get_metrics_data()
|
||||
assert data is None or not data.resource_metrics
|
||||
@@ -0,0 +1,100 @@
|
||||
"""Tests for writing a finished trace to Postgres and OTel."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from contextlib import contextmanager
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from docsgpt import tracing
|
||||
from docsgpt.core.settings import settings
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _settings(monkeypatch):
|
||||
monkeypatch.setattr(settings, "TRACES_ENABLED", True)
|
||||
|
||||
|
||||
def _trace_with_span():
|
||||
trace = tracing.start_trace(source="stream", request_id="r1", user_id="u1")
|
||||
with tracing.activate(trace):
|
||||
with tracing.span(tracing.KIND_LLM, "chat m"):
|
||||
pass
|
||||
return trace
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _patched_store():
|
||||
repo = MagicMock()
|
||||
repo.insert.return_value = True
|
||||
|
||||
@contextmanager
|
||||
def _session():
|
||||
yield MagicMock()
|
||||
|
||||
with patch("docsgpt.storage.db.session.db_session", _session), patch(
|
||||
"docsgpt.storage.db.repositories.request_traces.RequestTracesRepository",
|
||||
return_value=repo,
|
||||
):
|
||||
yield repo
|
||||
|
||||
|
||||
class TestFlush:
|
||||
def test_flush_writes_once(self):
|
||||
trace = _trace_with_span()
|
||||
with _patched_store() as repo:
|
||||
tracing.flush(trace)
|
||||
tracing.flush(trace)
|
||||
assert repo.insert.call_count == 1
|
||||
record = repo.insert.call_args[0][0]
|
||||
assert record["request_id"] == "r1"
|
||||
assert record["status"] == "ok"
|
||||
assert trace.finished and trace.flushed
|
||||
|
||||
def test_flush_status_override(self):
|
||||
trace = _trace_with_span()
|
||||
with _patched_store() as repo:
|
||||
tracing.flush(trace, status="paused")
|
||||
assert repo.insert.call_args[0][0]["status"] == "paused"
|
||||
|
||||
def test_empty_trace_is_not_stored(self):
|
||||
trace = tracing.start_trace(source="stream")
|
||||
with _patched_store() as repo:
|
||||
tracing.flush(trace)
|
||||
repo.insert.assert_not_called()
|
||||
|
||||
def test_store_failure_is_swallowed(self):
|
||||
trace = _trace_with_span()
|
||||
with _patched_store() as repo:
|
||||
repo.insert.side_effect = RuntimeError("db down")
|
||||
tracing.flush(trace) # must not raise
|
||||
|
||||
def test_otel_failure_does_not_block_store(self):
|
||||
trace = _trace_with_span()
|
||||
with _patched_store() as repo, patch(
|
||||
"docsgpt.tracing.otel.export_trace", side_effect=RuntimeError("otel")
|
||||
):
|
||||
tracing.flush(trace)
|
||||
repo.insert.assert_called_once()
|
||||
|
||||
def test_otel_trace_id_is_stored(self):
|
||||
trace = _trace_with_span()
|
||||
with _patched_store() as repo, patch(
|
||||
"docsgpt.tracing.otel.export_trace", return_value="ab" * 16
|
||||
):
|
||||
tracing.flush(trace)
|
||||
assert repo.insert.call_args[0][0]["otel_trace_id"] == "ab" * 16
|
||||
|
||||
def test_none_is_ignored(self):
|
||||
tracing.flush(None)
|
||||
tracing.discard(None)
|
||||
|
||||
|
||||
class TestDiscard:
|
||||
def test_discard_never_writes(self):
|
||||
trace = _trace_with_span()
|
||||
with _patched_store() as repo:
|
||||
tracing.discard(trace)
|
||||
tracing.flush(trace)
|
||||
repo.insert.assert_not_called()
|
||||
Reference in new issue
Block a user