From 8db8185a0a1131131fff65cc5dff200df2735ebf Mon Sep 17 00:00:00 2001 From: kamilbenkirane Date: Tue, 5 May 2026 12:26:39 +0200 Subject: [PATCH 1/5] feat(telemetry): widen usage attrs, add metrics histograms, opt-in content events MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit V2 telemetry expansion bundling three improvements on top of #270/#272: 1. Span attributes: output_attributes() now emits gen_ai.usage.total_tokens, reasoning_tokens, and cached_input_tokens when the typed Usage carries them. Off-spec modality fields fall through to celeste.usage.. Adds the missing cached_tokens field to TextUsage so the data already produced by anthropic / openai / cohere / deepseek / openresponses providers actually reaches the span (previously dropped silently). Extends Google Gemini and chatcompletions provider mixins to surface cached-prompt tokens too. 2. Metrics: registers two GenAI semconv histograms — gen_ai.client.token.usage (one record per token category via gen_ai.token.type dimension) and gen_ai.client.operation.duration (with error.type on failures). Sampling- resilient token-rate and latency dashboards no longer require span scanning. Wired into _predict and _TracedStream._finalize. 3. Content events: opt-in via OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT env flag, default off for PII safety. When enabled, spans carry gen_ai.input.messages and gen_ai.output.messages events with the semconv {role, parts: [...]} shape. Text + reasoning + tool calls inline; multimodal artifacts (image/audio/video/document) by URL reference, never inline bytes. 19 new tests across unit_tests/test_telemetry_metrics.py and test_telemetry_content_events.py, plus extended attribute coverage in test_telemetry_streaming.py. 630 unit tests pass. Closes #271. --- src/celeste/client.py | 110 +++++--- src/celeste/modalities/text/io.py | 1 + .../protocols/chatcompletions/client.py | 2 + .../google/generate_content/client.py | 1 + src/celeste/telemetry.py | 248 +++++++++++++++- .../test_telemetry_content_events.py | 265 ++++++++++++++++++ tests/unit_tests/test_telemetry_metrics.py | 192 +++++++++++++ tests/unit_tests/test_telemetry_streaming.py | 59 +++- 8 files changed, 821 insertions(+), 57 deletions(-) create mode 100644 tests/unit_tests/test_telemetry_content_events.py create mode 100644 tests/unit_tests/test_telemetry_metrics.py diff --git a/src/celeste/client.py b/src/celeste/client.py index fbc6621..27e5536 100644 --- a/src/celeste/client.py +++ b/src/celeste/client.py @@ -1,5 +1,6 @@ """Base client for modality-specific AI operations.""" +import time import warnings from abc import ABC, abstractmethod from collections.abc import AsyncIterator @@ -213,44 +214,62 @@ async def _predict( Returns: Output of the parameterized type. """ + request_attrs = telemetry.request_attributes( + model=self.model, + provider=self.provider, + protocol=self.protocol, + modality=self.modality, + ) + started = time.monotonic() with telemetry.tracer.start_as_current_span( telemetry.span_name(self.modality, self.model), - attributes=telemetry.request_attributes( - model=self.model, - provider=self.provider, - protocol=self.protocol, - modality=self.modality, - ), + attributes=request_attrs, ) as span: - inputs, parameters = self._validate_artifacts(inputs, **parameters) - request_body = self._build_request( - inputs, extra_body=extra_body, **parameters - ) - response_data = await self._make_request( - request_body, - endpoint=endpoint, - extra_headers=extra_headers, - **parameters, - ) - content = self._parse_content(response_data) - content = self._transform_output(content, **parameters) - tool_calls = self._parse_tool_calls(response_data) - reasoning, signature = self._parse_reasoning(response_data) - kwargs: dict[str, Any] = {} - if reasoning is not None: - kwargs["reasoning"] = reasoning - if signature: - kwargs["signature"] = signature - output = self._output_class()( - content=content, - usage=self._get_usage(response_data), - finish_reason=self._get_finish_reason(response_data), - metadata=self._build_metadata(response_data), - tool_calls=tool_calls, - **kwargs, - ) - span.set_attributes(telemetry.output_attributes(output)) - return output + try: + inputs, parameters = self._validate_artifacts(inputs, **parameters) + input_event = telemetry._input_messages_event(inputs) + if input_event is not None: + span.add_event("gen_ai.input.messages", attributes=input_event) + request_body = self._build_request( + inputs, extra_body=extra_body, **parameters + ) + response_data = await self._make_request( + request_body, + endpoint=endpoint, + extra_headers=extra_headers, + **parameters, + ) + content = self._parse_content(response_data) + content = self._transform_output(content, **parameters) + tool_calls = self._parse_tool_calls(response_data) + reasoning, signature = self._parse_reasoning(response_data) + kwargs: dict[str, Any] = {} + if reasoning is not None: + kwargs["reasoning"] = reasoning + if signature: + kwargs["signature"] = signature + output = self._output_class()( + content=content, + usage=self._get_usage(response_data), + finish_reason=self._get_finish_reason(response_data), + metadata=self._build_metadata(response_data), + tool_calls=tool_calls, + **kwargs, + ) + span.set_attributes(telemetry.output_attributes(output)) + output_event = telemetry._output_messages_event(output) + if output_event is not None: + span.add_event("gen_ai.output.messages", attributes=output_event) + telemetry.record_token_usage(output.usage, request_attrs) + telemetry.record_operation_duration( + time.monotonic() - started, request_attrs + ) + return output + except BaseException as exc: + telemetry.record_operation_duration( + time.monotonic() - started, request_attrs, error=exc + ) + raise def _parse_tool_calls(self, response_data: dict[str, Any]) -> list[ToolCall]: """Parse tool calls from response. Override in providers that support tools.""" @@ -297,18 +316,19 @@ def _stream( request_body = self._build_request( inputs, extra_body=extra_body, streaming=True, **parameters ) + request_attrs = telemetry.request_attributes( + model=self.model, + provider=self.provider, + protocol=self.protocol, + modality=self.modality, + ) span = telemetry.tracer.start_span( telemetry.span_name(self.modality, self.model), - attributes={ - **telemetry.request_attributes( - model=self.model, - provider=self.provider, - protocol=self.protocol, - modality=self.modality, - ), - "gen_ai.request.stream": True, - }, + attributes={**request_attrs, "gen_ai.request.stream": True}, ) + input_event = telemetry._input_messages_event(inputs) + if input_event is not None: + span.add_event("gen_ai.input.messages", attributes=input_event) sse_iterator = self._make_stream_request( request_body, endpoint=endpoint, @@ -326,7 +346,7 @@ def _stream( }, **parameters, ) - return telemetry.trace_stream(stream, span) # type: ignore[return-value] + return telemetry.trace_stream(stream, span, metric_attributes=request_attrs) # type: ignore[return-value] @classmethod @abstractmethod diff --git a/src/celeste/modalities/text/io.py b/src/celeste/modalities/text/io.py index 95e7feb..e2951d4 100644 --- a/src/celeste/modalities/text/io.py +++ b/src/celeste/modalities/text/io.py @@ -50,6 +50,7 @@ class TextUsage(Usage): input_tokens: int | None = None output_tokens: int | None = None reasoning_tokens: int | None = None + cached_tokens: int | None = None class TextOutput(Output[TextContent]): diff --git a/src/celeste/protocols/chatcompletions/client.py b/src/celeste/protocols/chatcompletions/client.py index ea6353e..0c5dbde 100644 --- a/src/celeste/protocols/chatcompletions/client.py +++ b/src/celeste/protocols/chatcompletions/client.py @@ -117,10 +117,12 @@ def map_usage_fields(usage_data: dict[str, Any]) -> dict[str, int | float | None Shared by client and streaming across all capabilities. """ + prompt_tokens_details = usage_data.get("prompt_tokens_details") or {} return { UsageField.INPUT_TOKENS: usage_data.get("prompt_tokens"), UsageField.OUTPUT_TOKENS: usage_data.get("completion_tokens"), UsageField.TOTAL_TOKENS: usage_data.get("total_tokens"), + UsageField.CACHED_TOKENS: prompt_tokens_details.get("cached_tokens"), } def _parse_usage( diff --git a/src/celeste/providers/google/generate_content/client.py b/src/celeste/providers/google/generate_content/client.py index 54ba62f..02d20f0 100644 --- a/src/celeste/providers/google/generate_content/client.py +++ b/src/celeste/providers/google/generate_content/client.py @@ -110,6 +110,7 @@ def map_usage_fields(usage_data: dict[str, Any]) -> dict[str, int | float | None UsageField.OUTPUT_TOKENS: usage_data.get("candidatesTokenCount"), UsageField.TOTAL_TOKENS: usage_data.get("totalTokenCount"), UsageField.REASONING_TOKENS: usage_data.get("thoughtsTokenCount"), + UsageField.CACHED_TOKENS: usage_data.get("cachedContentTokenCount"), } def _parse_usage( diff --git a/src/celeste/telemetry.py b/src/celeste/telemetry.py index 7bfd363..49d8965 100644 --- a/src/celeste/telemetry.py +++ b/src/celeste/telemetry.py @@ -1,17 +1,22 @@ """OpenTelemetry GenAI telemetry for Celeste.""" import asyncio +import json +import os import time from collections.abc import Iterator from contextlib import contextmanager from types import TracebackType from typing import Any +from celeste.artifacts import Artifact from celeste.core import Modality, Protocol, Provider from celeste.exceptions import StreamNotExhaustedError -from celeste.io import Output +from celeste.io import Input, Output, Usage from celeste.models import Model from celeste.streaming import Stream +from celeste.tools import ToolCall +from celeste.types import Message _PROVIDER_NAME_MAP: dict[Provider, str] = { Provider.OPENAI: "openai", @@ -46,6 +51,9 @@ def set_status(self, *args: Any, **kwargs: Any) -> None: def record_exception(self, *args: Any, **kwargs: Any) -> None: """Discard the exception.""" + def add_event(self, *args: Any, **kwargs: Any) -> None: + """Discard the span event.""" + def end(self) -> None: """Discard the end signal.""" @@ -82,7 +90,25 @@ class _NoOpStatusCode: ERROR = "ERROR" +class _NoOpHistogram: + """Histogram stand-in used when ``opentelemetry-api`` is not installed.""" + + def record(self, value: float, attributes: dict[str, Any] | None = None) -> None: + """Discard the metric record.""" + + +class _NoOpMeter: + """Meter stand-in used when ``opentelemetry-api`` is not installed.""" + + def create_histogram( + self, name: str, unit: str = "", description: str = "" + ) -> _NoOpHistogram: + """Return a no-op histogram.""" + return _NoOpHistogram() + + try: + from opentelemetry import metrics as _otel_metrics from opentelemetry import trace as _otel_trace from opentelemetry.trace import Status as _OtelStatus from opentelemetry.trace import StatusCode as _OtelStatusCode @@ -91,11 +117,25 @@ class _NoOpStatusCode: use_span: Any = _otel_trace.use_span Status: Any = _OtelStatus StatusCode: Any = _OtelStatusCode + meter: Any = _otel_metrics.get_meter("celeste") except ImportError: tracer = _NoOpTracer() use_span = _noop_use_span Status = _NoOpStatus StatusCode = _NoOpStatusCode + meter = _NoOpMeter() + + +_token_usage_histogram: Any = meter.create_histogram( + name="gen_ai.client.token.usage", + unit="{token}", + description="Tokens used per GenAI call, sliced by gen_ai.token.type.", +) +_operation_duration_histogram: Any = meter.create_histogram( + name="gen_ai.client.operation.duration", + unit="s", + description="Wall-clock duration of GenAI calls.", +) def request_attributes( @@ -145,30 +185,202 @@ def span_name(modality: Modality, model: Model) -> str: return f"celeste.{modality.value} {model.id}" +# Maps `Usage` field names to GenAI semconv span attribute keys. +# Fields not in this map fall through to `celeste.usage.` (off-spec). +_GEN_AI_USAGE_FIELDS: dict[str, str] = { + "input_tokens": "gen_ai.usage.input_tokens", + "output_tokens": "gen_ai.usage.output_tokens", + "total_tokens": "gen_ai.usage.total_tokens", + "reasoning_tokens": "gen_ai.usage.reasoning_tokens", + "cached_tokens": "gen_ai.usage.cached_input_tokens", +} + +# Maps `Usage` field names to `gen_ai.token.type` dimension values for metrics. +# Only fields representing token *categories* belong here (input/output/cached/reasoning). +# `total_tokens` is omitted to avoid double-counting in histogram aggregations. +_GEN_AI_TOKEN_TYPES: dict[str, str] = { + "input_tokens": "input", + "output_tokens": "output", + "reasoning_tokens": "reasoning", + "cached_tokens": "cached_input", +} + + def output_attributes(output: Output[Any]) -> dict[str, Any]: """Extract GenAI response attributes from a typed Output.""" attrs: dict[str, Any] = {} - input_tokens = getattr(output.usage, "input_tokens", None) - output_tokens = getattr(output.usage, "output_tokens", None) - if isinstance(input_tokens, int): - attrs["gen_ai.usage.input_tokens"] = input_tokens - if isinstance(output_tokens, int): - attrs["gen_ai.usage.output_tokens"] = output_tokens + usage = output.usage + usage_fields = ( + type(usage).model_fields if hasattr(type(usage), "model_fields") else {} + ) + for field_name in usage_fields: + value = getattr(usage, field_name, None) + if not isinstance(value, int | float): + continue + semconv_key = _GEN_AI_USAGE_FIELDS.get(field_name) + attrs[semconv_key or f"celeste.usage.{field_name}"] = value if output.finish_reason is not None and output.finish_reason.reason is not None: attrs["gen_ai.response.finish_reasons"] = (output.finish_reason.reason,) return attrs +def record_token_usage(usage: Usage, attributes: dict[str, Any]) -> None: + """Record one ``gen_ai.client.token.usage`` observation per token category.""" + usage_fields = ( + type(usage).model_fields if hasattr(type(usage), "model_fields") else {} + ) + for field_name in usage_fields: + token_type = _GEN_AI_TOKEN_TYPES.get(field_name) + if token_type is None: + continue + value = getattr(usage, field_name, None) + if not isinstance(value, int | float): + continue + _token_usage_histogram.record( + value, + attributes={**attributes, "gen_ai.token.type": token_type}, + ) + + +def record_operation_duration( + duration_seconds: float, + attributes: dict[str, Any], + error: BaseException | None = None, +) -> None: + """Record one ``gen_ai.client.operation.duration`` observation.""" + attrs = {**attributes} + if error is not None: + attrs["error.type"] = type(error).__name__ + _operation_duration_histogram.record(duration_seconds, attributes=attrs) + + +# Read the semconv-standard env flag once at import. Per the GenAI spec, +# content capture is opt-in because prompts/responses often contain user data. +_CAPTURE_CONTENT: bool = ( + os.environ.get("OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT", "") + .strip() + .lower() + == "true" +) + +# Maps celeste artifact classes to semconv part type names. +_ARTIFACT_PART_TYPES: dict[str, str] = { + "ImageArtifact": "image", + "AudioArtifact": "audio", + "VideoArtifact": "video", + "DocumentArtifact": "document", +} + + +def _artifact_part(artifact: Artifact) -> dict[str, Any]: + """Convert a celeste artifact to a semconv part by URL reference.""" + part: dict[str, Any] = { + "type": _ARTIFACT_PART_TYPES.get(type(artifact).__name__, "media"), + } + if artifact.url is not None: + part["uri"] = artifact.url + if artifact.mime_type is not None: + part["mime_type"] = str(artifact.mime_type) + return part + + +def _content_to_parts(content: Any) -> list[dict[str, Any]]: + """Convert a celeste content value to a list of semconv parts.""" + if content is None: + return [] + if isinstance(content, list): + parts: list[dict[str, Any]] = [] + for item in content: + parts.extend(_content_to_parts(item)) + return parts + if isinstance(content, Artifact): + return [_artifact_part(content)] + if isinstance(content, str): + return [{"type": "text", "content": content}] + return [{"type": "text", "content": json.dumps(content, default=str)}] + + +def _tool_call_part(tool_call: ToolCall) -> dict[str, Any]: + """Convert a ToolCall into a semconv tool_call part.""" + return { + "type": "tool_call", + "id": tool_call.id, + "name": tool_call.name, + "arguments": json.dumps(tool_call.arguments, default=str), + } + + +def _message_to_dict(message: Message) -> dict[str, Any]: + """Convert a celeste Message into a semconv ``{role, parts}`` dict.""" + parts = _content_to_parts(message.content) + if message.reasoning is not None: + parts.append({"type": "reasoning", "content": message.reasoning}) + if message.tool_calls: + parts.extend(_tool_call_part(call) for call in message.tool_calls) + return {"role": message.role.value, "parts": parts} + + +def _input_messages_event(inputs: Input) -> dict[str, Any] | None: + """Build the ``gen_ai.input.messages`` event attributes from inputs. + + Returns ``None`` when content capture is disabled. + """ + if not _CAPTURE_CONTENT: + return None + messages: list[dict[str, Any]] = [] + for message in getattr(inputs, "messages", None) or []: + if isinstance(message, Message): + messages.append(_message_to_dict(message)) + prompt = getattr(inputs, "prompt", None) + if prompt is not None: + parts: list[dict[str, Any]] = [{"type": "text", "content": str(prompt)}] + for media_field in ("image", "video", "audio", "document"): + media = getattr(inputs, media_field, None) + if media is not None: + parts.extend(_content_to_parts(media)) + messages.append({"role": "user", "parts": parts}) + if not messages: + return None + return {"messages": json.dumps(messages, default=str)} + + +def _output_messages_event(output: Output[Any]) -> dict[str, Any] | None: + """Build the ``gen_ai.output.messages`` event attributes from an Output. + + Returns ``None`` when content capture is disabled. + """ + if not _CAPTURE_CONTENT: + return None + parts = _content_to_parts(output.content) + reasoning = getattr(output, "reasoning", None) + if reasoning is not None: + parts.append({"type": "reasoning", "content": reasoning}) + if output.tool_calls: + parts.extend(_tool_call_part(call) for call in output.tool_calls) + if not parts: + return None + return { + "messages": json.dumps([{"role": "assistant", "parts": parts}], default=str) + } + + class _TracedStream: """Stream-shaped wrapper that emits GenAI telemetry against an OTel span.""" - def __init__(self, inner: Stream[Any, Any, Any], span: Any) -> None: + def __init__( + self, + inner: Stream[Any, Any, Any], + span: Any, + metric_attributes: dict[str, Any] | None = None, + ) -> None: """Initialize wrapper around a constructed Stream and a detached span.""" self._inner = inner self._span = span + self._metric_attributes = metric_attributes or {} self._started = time.monotonic() self._seen_first = False self._ended = False + self._error: BaseException | None = None def __aiter__(self) -> "_TracedStream": """Return self as async iterator.""" @@ -183,6 +395,7 @@ async def __anext__(self) -> Any: self._finalize() raise except Exception as exc: + self._error = exc self._span.record_exception(exc) self._span.set_status(Status(StatusCode.ERROR, str(exc))) self._finalize() @@ -242,28 +455,41 @@ async def aclose(self) -> None: self._finalize() def _finalize(self) -> None: - """End the span exactly once; emit usage attrs if Output is populated.""" + """End the span exactly once; emit usage, metrics, and output event.""" if self._ended: return self._ended = True + duration = time.monotonic() - self._started try: output = self._inner.output except StreamNotExhaustedError: output = None if output is not None: self._span.set_attributes(output_attributes(output)) + output_event = _output_messages_event(output) + if output_event is not None: + self._span.add_event("gen_ai.output.messages", attributes=output_event) + record_token_usage(output.usage, self._metric_attributes) + record_operation_duration(duration, self._metric_attributes, error=self._error) self._span.end() -def trace_stream(stream: Stream[Any, Any, Any], span: Any) -> _TracedStream: +def trace_stream( + stream: Stream[Any, Any, Any], + span: Any, + metric_attributes: dict[str, Any] | None = None, +) -> _TracedStream: """Wrap a Stream to emit GenAI telemetry against ``span``.""" - return _TracedStream(stream, span) + return _TracedStream(stream, span, metric_attributes) __all__ = [ "Status", "StatusCode", + "meter", "output_attributes", + "record_operation_duration", + "record_token_usage", "request_attributes", "span_name", "trace_stream", diff --git a/tests/unit_tests/test_telemetry_content_events.py b/tests/unit_tests/test_telemetry_content_events.py new file mode 100644 index 0000000..f198274 --- /dev/null +++ b/tests/unit_tests/test_telemetry_content_events.py @@ -0,0 +1,265 @@ +"""Tests for opt-in GenAI content events (`gen_ai.input.messages` / `output.messages`).""" + +import json +from collections.abc import AsyncIterator +from typing import Any, ClassVar + +import pytest +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 Span + +from celeste import telemetry +from celeste.artifacts import ImageArtifact +from celeste.io import Chunk, Output, Usage +from celeste.mime_types import ImageMimeType +from celeste.modalities.text.io import TextInput, TextOutput, TextUsage +from celeste.parameters import Parameters +from celeste.streaming import Stream +from celeste.tools import ToolCall +from celeste.types import Message, Role + + +@pytest.fixture +def capture_enabled(monkeypatch: pytest.MonkeyPatch) -> None: + """Enable content capture by patching the module-level flag.""" + monkeypatch.setattr(telemetry, "_CAPTURE_CONTENT", True) + + +@pytest.fixture +def capture_disabled(monkeypatch: pytest.MonkeyPatch) -> None: + """Force-disable content capture (default behavior).""" + monkeypatch.setattr(telemetry, "_CAPTURE_CONTENT", False) + + +@pytest.fixture +def exporter() -> tuple[InMemorySpanExporter, TracerProvider]: + """Wire an in-memory span exporter.""" + span_exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(span_exporter)) + return span_exporter, provider + + +def _start_span(provider: TracerProvider, name: str = "test") -> Span: + """Start a detached span on the given provider's tracer.""" + return provider.get_tracer("celeste-test").start_span(name) + + +class _TestStreamUsage(Usage): + """Usage with token fields for streaming tests.""" + + input_tokens: int | None = None + output_tokens: int | None = None + + +class _TestStreamOutput(Output[str]): + """Output for streaming tests.""" + + pass + + +class _TestStream(Stream[_TestStreamOutput, Parameters, Chunk]): + """Concrete Stream that aggregates content + usage.""" + + _usage_class: ClassVar[type[Usage]] = _TestStreamUsage + _chunk_class: ClassVar[type[Chunk]] = Chunk + _output_class: ClassVar[type[Output]] = _TestStreamOutput + _empty_content: ClassVar[str] = "" + + def _aggregate_content(self, chunks: list[Chunk]) -> str: + return "".join(chunk.content for chunk in chunks) + + def _parse_chunk(self, event: dict[str, Any]) -> Chunk | None: + content = event.get("delta") + if not content and "usage" not in event: + return None + usage = _TestStreamUsage(**event["usage"]) if "usage" in event else None + return Chunk(content=content or "", finish_reason=None, usage=usage) + + +async def _async_iter(events: list[dict[str, Any]]) -> AsyncIterator[dict[str, Any]]: + for event in events: + yield event + + +class TestInputMessagesEvent: + """`_input_messages_event` builds the input messages payload from typed inputs.""" + + def test_capture_disabled_returns_none(self, capture_disabled: None) -> None: + """With the env flag off, the helper returns None — no event emitted.""" + inputs = TextInput(prompt="hello") + assert telemetry._input_messages_event(inputs) is None + + def test_text_prompt_only(self, capture_enabled: None) -> None: + """Plain text prompt produces a single user message with one text part.""" + inputs = TextInput(prompt="hello") + + result = telemetry._input_messages_event(inputs) + + assert result is not None + messages = json.loads(result["messages"]) + assert messages == [ + {"role": "user", "parts": [{"type": "text", "content": "hello"}]} + ] + + def test_messages_field_emitted(self, capture_enabled: None) -> None: + """When `inputs.messages` is set, those messages are serialized verbatim.""" + inputs = TextInput( + messages=[ + Message(role=Role.SYSTEM, content="be helpful"), + Message(role=Role.USER, content="hi"), + ] + ) + + result = telemetry._input_messages_event(inputs) + + assert result is not None + messages = json.loads(result["messages"]) + assert messages[0]["role"] == "system" + assert messages[0]["parts"] == [{"type": "text", "content": "be helpful"}] + assert messages[1]["role"] == "user" + + def test_image_input_emitted_as_reference(self, capture_enabled: None) -> None: + """Images attached to a prompt are emitted as URL references, not bytes.""" + inputs = TextInput( + prompt="describe this", + image=ImageArtifact( + url="https://example.com/img.png", + mime_type=ImageMimeType.PNG, + ), + ) + + result = telemetry._input_messages_event(inputs) + + assert result is not None + messages = json.loads(result["messages"]) + parts = messages[0]["parts"] + assert parts[0]["type"] == "text" + image_parts = [p for p in parts if p["type"] == "image"] + assert len(image_parts) == 1 + assert image_parts[0]["uri"] == "https://example.com/img.png" + assert image_parts[0]["mime_type"] == "image/png" + + +class TestOutputMessagesEvent: + """`_output_messages_event` builds the assistant payload from a typed Output.""" + + def test_capture_disabled_returns_none(self, capture_disabled: None) -> None: + """With the env flag off, the helper returns None.""" + output = TextOutput(content="reply", usage=TextUsage()) + assert telemetry._output_messages_event(output) is None + + def test_text_only(self, capture_enabled: None) -> None: + """Plain text output produces a single assistant message with one text part.""" + output = TextOutput(content="reply", usage=TextUsage()) + + result = telemetry._output_messages_event(output) + + assert result is not None + messages = json.loads(result["messages"]) + assert messages == [ + {"role": "assistant", "parts": [{"type": "text", "content": "reply"}]} + ] + + def test_reasoning_appended(self, capture_enabled: None) -> None: + """Reasoning content is appended as a `reasoning` part.""" + output = TextOutput( + content="answer", usage=TextUsage(), reasoning="thinking..." + ) + + result = telemetry._output_messages_event(output) + + assert result is not None + parts = json.loads(result["messages"])[0]["parts"] + types = [p["type"] for p in parts] + assert "text" in types and "reasoning" in types + reasoning_part = next(p for p in parts if p["type"] == "reasoning") + assert reasoning_part["content"] == "thinking..." + + def test_tool_calls_emitted_as_tool_call_parts(self, capture_enabled: None) -> None: + """Tool calls become semconv `tool_call` parts with id/name/arguments.""" + output = TextOutput( + content="", + usage=TextUsage(), + tool_calls=[ + ToolCall(id="call_1", name="get_weather", arguments={"city": "Paris"}) + ], + ) + + result = telemetry._output_messages_event(output) + + assert result is not None + parts = json.loads(result["messages"])[0]["parts"] + tool_parts = [p for p in parts if p["type"] == "tool_call"] + assert len(tool_parts) == 1 + assert tool_parts[0]["id"] == "call_1" + assert tool_parts[0]["name"] == "get_weather" + assert json.loads(tool_parts[0]["arguments"]) == {"city": "Paris"} + + def test_image_artifact_output_by_reference(self, capture_enabled: None) -> None: + """Image generation outputs emit URL references, never inline bytes.""" + output = TextOutput( + content=ImageArtifact( + url="https://example.com/out.png", mime_type=ImageMimeType.PNG + ), + usage=TextUsage(), + ) + + result = telemetry._output_messages_event(output) + + assert result is not None + parts = json.loads(result["messages"])[0]["parts"] + assert parts[0]["type"] == "image" + assert parts[0]["uri"] == "https://example.com/out.png" + assert parts[0]["mime_type"] == "image/png" + + +class TestStreamingContentEvents: + """Streaming flows emit content events at finalize when output is built.""" + + async def test_natural_exhaustion_emits_output_event( + self, + capture_enabled: None, + exporter: tuple[InMemorySpanExporter, TracerProvider], + ) -> None: + """After exhaustion, the span carries the `gen_ai.output.messages` event.""" + in_memory, provider = exporter + events = [ + {"delta": "Hello"}, + {"delta": " world", "usage": {"input_tokens": 1, "output_tokens": 2}}, + ] + stream = _TestStream(_async_iter(events)) + span = _start_span(provider) + wrapped = telemetry.trace_stream(stream, span) + + async for _ in wrapped: + pass + + finished = in_memory.get_finished_spans() + assert len(finished) == 1 + event_names = [e.name for e in finished[0].events] + assert "gen_ai.output.messages" in event_names + + async def test_aclose_before_output_skips_output_event( + self, + capture_enabled: None, + exporter: tuple[InMemorySpanExporter, TracerProvider], + ) -> None: + """When _output is None at finalize, the output event is skipped.""" + in_memory, provider = exporter + events = [{"delta": "Partial"}, {"delta": "..."}] + stream = _TestStream(_async_iter(events)) + span = _start_span(provider) + wrapped = telemetry.trace_stream(stream, span) + + await wrapped.__anext__() + await wrapped.aclose() + + finished = in_memory.get_finished_spans() + assert len(finished) == 1 + event_names = [e.name for e in finished[0].events] + assert "gen_ai.output.messages" not in event_names diff --git a/tests/unit_tests/test_telemetry_metrics.py b/tests/unit_tests/test_telemetry_metrics.py new file mode 100644 index 0000000..e2d7f22 --- /dev/null +++ b/tests/unit_tests/test_telemetry_metrics.py @@ -0,0 +1,192 @@ +"""Tests for GenAI metric histogram emission.""" + +from collections.abc import AsyncIterator +from typing import Any, ClassVar + +import pytest +from opentelemetry.sdk.metrics import MeterProvider +from opentelemetry.sdk.metrics.export import InMemoryMetricReader + +from celeste import telemetry +from celeste.io import Chunk, Output, Usage +from celeste.parameters import Parameters +from celeste.streaming import Stream + + +class _TestUsage(Usage): + """Usage with input/output/cached/reasoning fields for metric tests.""" + + input_tokens: int | None = None + output_tokens: int | None = None + cached_tokens: int | None = None + reasoning_tokens: int | None = None + + +class _TestOutput(Output[str]): + """Output for tests.""" + + pass + + +class _TestStream(Stream[_TestOutput, Parameters, Chunk]): + """Concrete Stream that aggregates usage from per-chunk metadata.""" + + _usage_class: ClassVar[type[Usage]] = _TestUsage + _chunk_class: ClassVar[type[Chunk]] = Chunk + _output_class: ClassVar[type[Output]] = _TestOutput + _empty_content: ClassVar[str] = "" + + def _aggregate_content(self, chunks: list[Chunk]) -> str: + return "".join(chunk.content for chunk in chunks) + + def _parse_chunk(self, event: dict[str, Any]) -> Chunk | None: + content = event.get("delta") + if not content and "usage" not in event: + return None + usage = _TestUsage(**event["usage"]) if "usage" in event else None + return Chunk(content=content or "", finish_reason=None, usage=usage) + + +async def _async_iter(events: list[dict[str, Any]]) -> AsyncIterator[dict[str, Any]]: + for event in events: + yield event + + +@pytest.fixture +def meter_setup(monkeypatch: pytest.MonkeyPatch) -> InMemoryMetricReader: + """Wire fresh histograms backed by InMemoryMetricReader for assertions.""" + reader = InMemoryMetricReader() + provider = MeterProvider(metric_readers=[reader]) + test_meter = provider.get_meter("celeste-test") + monkeypatch.setattr( + telemetry, + "_token_usage_histogram", + test_meter.create_histogram( + name="gen_ai.client.token.usage", + unit="{token}", + ), + ) + monkeypatch.setattr( + telemetry, + "_operation_duration_histogram", + test_meter.create_histogram( + name="gen_ai.client.operation.duration", + unit="s", + ), + ) + return reader + + +def _collect_metric(reader: InMemoryMetricReader, name: str) -> list[Any]: + """Pull metric data points by name from the reader.""" + metrics_data = reader.get_metrics_data() + if metrics_data is None: + return [] + points: list[Any] = [] + for resource_metric in metrics_data.resource_metrics: + for scope_metric in resource_metric.scope_metrics: + for metric in scope_metric.metrics: + if metric.name == name: + points.extend(metric.data.data_points) + return points + + +class TestRecordTokenUsage: + """`record_token_usage` emits one record per token category.""" + + def test_input_and_output_recorded_separately( + self, meter_setup: InMemoryMetricReader + ) -> None: + """input + output usage produce two records with different `gen_ai.token.type`.""" + usage = _TestUsage(input_tokens=100, output_tokens=50) + attrs = {"gen_ai.request.model": "test-model"} + + telemetry.record_token_usage(usage, attrs) + + points = _collect_metric(meter_setup, "gen_ai.client.token.usage") + token_types = sorted(p.attributes["gen_ai.token.type"] for p in points) + assert token_types == ["input", "output"] + by_type = {p.attributes["gen_ai.token.type"]: p.sum for p in points} + assert by_type["input"] == 100 + assert by_type["output"] == 50 + + def test_cached_and_reasoning_recorded_with_dedicated_token_type( + self, meter_setup: InMemoryMetricReader + ) -> None: + """cached_tokens and reasoning_tokens get distinct token.type values.""" + usage = _TestUsage( + input_tokens=10, output_tokens=5, cached_tokens=8, reasoning_tokens=3 + ) + telemetry.record_token_usage(usage, {}) + + points = _collect_metric(meter_setup, "gen_ai.client.token.usage") + token_types = sorted(p.attributes["gen_ai.token.type"] for p in points) + assert token_types == ["cached_input", "input", "output", "reasoning"] + + def test_none_fields_are_skipped(self, meter_setup: InMemoryMetricReader) -> None: + """Fields set to None do not produce records.""" + usage = _TestUsage(input_tokens=10) + telemetry.record_token_usage(usage, {}) + + points = _collect_metric(meter_setup, "gen_ai.client.token.usage") + assert len(points) == 1 + assert points[0].attributes["gen_ai.token.type"] == "input" + + +class TestRecordOperationDuration: + """`record_operation_duration` records latency with optional error.type.""" + + def test_success_records_duration_without_error_type( + self, meter_setup: InMemoryMetricReader + ) -> None: + """Successful call: duration recorded, no error.type attribute.""" + telemetry.record_operation_duration(1.5, {"gen_ai.request.model": "test-model"}) + + points = _collect_metric(meter_setup, "gen_ai.client.operation.duration") + assert len(points) == 1 + assert "error.type" not in points[0].attributes + assert points[0].sum == pytest.approx(1.5) + + def test_failure_records_duration_with_error_type( + self, meter_setup: InMemoryMetricReader + ) -> None: + """Failed call: duration recorded with error.type set to exception class name.""" + telemetry.record_operation_duration(0.2, {}, error=ValueError("bad input")) + + points = _collect_metric(meter_setup, "gen_ai.client.operation.duration") + assert len(points) == 1 + assert points[0].attributes["error.type"] == "ValueError" + + +class TestStreamingMetrics: + """Streaming flows record duration covering full consumption.""" + + async def test_streaming_records_duration_and_token_usage( + self, meter_setup: InMemoryMetricReader + ) -> None: + """Natural exhaustion records both metrics with correct values.""" + events = [ + {"delta": "Hello"}, + {"delta": " world", "usage": {"input_tokens": 12, "output_tokens": 34}}, + ] + stream = _TestStream(_async_iter(events)) + provider = ( + MeterProvider() + ) # noop meter for span; the spans go to default tracer + del provider + span = telemetry.tracer.start_span("test") + wrapped = telemetry.trace_stream( + stream, span, metric_attributes={"gen_ai.request.model": "test-model"} + ) + + async for _ in wrapped: + pass + + token_points = _collect_metric(meter_setup, "gen_ai.client.token.usage") + duration_points = _collect_metric( + meter_setup, "gen_ai.client.operation.duration" + ) + assert len(token_points) == 2 + assert len(duration_points) == 1 + assert duration_points[0].sum > 0 + assert "error.type" not in duration_points[0].attributes diff --git a/tests/unit_tests/test_telemetry_streaming.py b/tests/unit_tests/test_telemetry_streaming.py index 2026563..d67b29f 100644 --- a/tests/unit_tests/test_telemetry_streaming.py +++ b/tests/unit_tests/test_telemetry_streaming.py @@ -19,10 +19,20 @@ class _TestUsage(Usage): - """Usage with input/output token fields for telemetry assertions.""" + """Usage with the full GenAI semconv field set for telemetry assertions.""" input_tokens: int | None = None output_tokens: int | None = None + total_tokens: int | None = None + reasoning_tokens: int | None = None + cached_tokens: int | None = None + + +class _OffSpecUsage(Usage): + """Modality-specific usage with fields outside the GenAI semconv set.""" + + images_generated: int | None = None + audio_seconds: float | None = None class _TestOutput(Output[str]): @@ -234,3 +244,50 @@ async def test_natural_exhaustion_then_aclose_ends_span_once( attrs = finished[0].attributes or {} assert attrs["gen_ai.usage.input_tokens"] == 2 assert attrs["gen_ai.usage.output_tokens"] == 3 + + +class TestExtendedUsageAttributes: + """V2 widens `output_attributes` to emit total/reasoning/cached_input tokens.""" + + async def test_total_reasoning_cached_tokens_emitted( + self, exporter: tuple[InMemorySpanExporter, TracerProvider] + ) -> None: + """When the typed Usage carries the full set, all attributes appear on the span.""" + in_memory, provider = exporter + events = [ + { + "delta": "ok", + "usage": { + "input_tokens": 10, + "output_tokens": 5, + "total_tokens": 15, + "reasoning_tokens": 3, + "cached_tokens": 8, + }, + } + ] + stream = _TestStream(_async_iter(events)) + span = _start_span(provider) + wrapped = telemetry.trace_stream(stream, span) + + async for _ in wrapped: + pass + + finished = in_memory.get_finished_spans() + attrs = finished[0].attributes or {} + assert attrs["gen_ai.usage.input_tokens"] == 10 + assert attrs["gen_ai.usage.output_tokens"] == 5 + assert attrs["gen_ai.usage.total_tokens"] == 15 + assert attrs["gen_ai.usage.reasoning_tokens"] == 3 + assert attrs["gen_ai.usage.cached_input_tokens"] == 8 + + def test_off_spec_usage_emitted_under_celeste_namespace(self) -> None: + """Modality-specific Usage fields fall through to `celeste.usage.`.""" + usage = _OffSpecUsage(images_generated=4, audio_seconds=1.5) + output = _TestOutput(content="", usage=usage) + + attrs = telemetry.output_attributes(output) + + assert attrs["celeste.usage.images_generated"] == 4 + assert attrs["celeste.usage.audio_seconds"] == 1.5 + assert "gen_ai.usage.images_generated" not in attrs From b13aa8c48d72ef6f2f195a5a91ebe89933bb126b Mon Sep 17 00:00:00 2001 From: kamilbenkirane Date: Tue, 5 May 2026 12:34:01 +0200 Subject: [PATCH 2/5] refactor(telemetry): UsageField enum keys, public event helpers, shared finalize MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three small cleanups on top of the V2 telemetry expansion, no behavior change: 1. `_GEN_AI_USAGE_FIELDS` and `_GEN_AI_TOKEN_TYPES` now key on `UsageField` enum members instead of bare strings. Catches typos at import; aligns with the vocabulary providers already populate. 2. New public `add_input_event(span, inputs)` and `add_output_event(span, output)` helpers replace the cross-module calls to private `_input_messages_event` / `_output_messages_event`. Underscore-prefixed names stay as the dict-builder primitives; the public helpers handle `if event is not None: span.add_event(...)`. 3. New `record_output(span, output, attributes, duration, error=None)` extracts the four-step finalize sequence (set output_attributes + emit content event + record token usage + record duration) shared by `_predict` and `_TracedStream._finalize`. Also: drop dead `hasattr(type(usage), "model_fields")` guards (Usage is always Pydantic BaseModel by io.py contract); skip `bool` in the numeric usage iter (bool is a subclass of int — paranoid guard against ever populating a Usage boolean by accident). 630 tests still pass. --- src/celeste/client.py | 17 ++----- src/celeste/telemetry.py | 96 ++++++++++++++++++++++++++-------------- 2 files changed, 67 insertions(+), 46 deletions(-) diff --git a/src/celeste/client.py b/src/celeste/client.py index 27e5536..c3c061c 100644 --- a/src/celeste/client.py +++ b/src/celeste/client.py @@ -227,9 +227,7 @@ async def _predict( ) as span: try: inputs, parameters = self._validate_artifacts(inputs, **parameters) - input_event = telemetry._input_messages_event(inputs) - if input_event is not None: - span.add_event("gen_ai.input.messages", attributes=input_event) + telemetry.add_input_event(span, inputs) request_body = self._build_request( inputs, extra_body=extra_body, **parameters ) @@ -256,13 +254,8 @@ async def _predict( tool_calls=tool_calls, **kwargs, ) - span.set_attributes(telemetry.output_attributes(output)) - output_event = telemetry._output_messages_event(output) - if output_event is not None: - span.add_event("gen_ai.output.messages", attributes=output_event) - telemetry.record_token_usage(output.usage, request_attrs) - telemetry.record_operation_duration( - time.monotonic() - started, request_attrs + telemetry.record_output( + span, output, request_attrs, time.monotonic() - started ) return output except BaseException as exc: @@ -326,9 +319,7 @@ def _stream( telemetry.span_name(self.modality, self.model), attributes={**request_attrs, "gen_ai.request.stream": True}, ) - input_event = telemetry._input_messages_event(inputs) - if input_event is not None: - span.add_event("gen_ai.input.messages", attributes=input_event) + telemetry.add_input_event(span, inputs) sse_iterator = self._make_stream_request( request_body, endpoint=endpoint, diff --git a/src/celeste/telemetry.py b/src/celeste/telemetry.py index 49d8965..2dd9302 100644 --- a/src/celeste/telemetry.py +++ b/src/celeste/telemetry.py @@ -10,7 +10,7 @@ from typing import Any from celeste.artifacts import Artifact -from celeste.core import Modality, Protocol, Provider +from celeste.core import Modality, Protocol, Provider, UsageField from celeste.exceptions import StreamNotExhaustedError from celeste.io import Input, Output, Usage from celeste.models import Model @@ -185,38 +185,42 @@ def span_name(modality: Modality, model: Model) -> str: return f"celeste.{modality.value} {model.id}" -# Maps `Usage` field names to GenAI semconv span attribute keys. +# Maps `UsageField` enum members to GenAI semconv span attribute keys. # Fields not in this map fall through to `celeste.usage.` (off-spec). _GEN_AI_USAGE_FIELDS: dict[str, str] = { - "input_tokens": "gen_ai.usage.input_tokens", - "output_tokens": "gen_ai.usage.output_tokens", - "total_tokens": "gen_ai.usage.total_tokens", - "reasoning_tokens": "gen_ai.usage.reasoning_tokens", - "cached_tokens": "gen_ai.usage.cached_input_tokens", + UsageField.INPUT_TOKENS: "gen_ai.usage.input_tokens", + UsageField.OUTPUT_TOKENS: "gen_ai.usage.output_tokens", + UsageField.TOTAL_TOKENS: "gen_ai.usage.total_tokens", + UsageField.REASONING_TOKENS: "gen_ai.usage.reasoning_tokens", + UsageField.CACHED_TOKENS: "gen_ai.usage.cached_input_tokens", } -# Maps `Usage` field names to `gen_ai.token.type` dimension values for metrics. +# Maps `UsageField` enum members to `gen_ai.token.type` dimension values for metrics. # Only fields representing token *categories* belong here (input/output/cached/reasoning). # `total_tokens` is omitted to avoid double-counting in histogram aggregations. _GEN_AI_TOKEN_TYPES: dict[str, str] = { - "input_tokens": "input", - "output_tokens": "output", - "reasoning_tokens": "reasoning", - "cached_tokens": "cached_input", + UsageField.INPUT_TOKENS: "input", + UsageField.OUTPUT_TOKENS: "output", + UsageField.REASONING_TOKENS: "reasoning", + UsageField.CACHED_TOKENS: "cached_input", } -def output_attributes(output: Output[Any]) -> dict[str, Any]: - """Extract GenAI response attributes from a typed Output.""" - attrs: dict[str, Any] = {} - usage = output.usage - usage_fields = ( - type(usage).model_fields if hasattr(type(usage), "model_fields") else {} - ) - for field_name in usage_fields: +def _iter_usage_numeric(usage: Usage) -> Iterator[tuple[str, int | float]]: + """Yield (field_name, value) for each numeric `Usage` field that is set.""" + for field_name in type(usage).model_fields: value = getattr(usage, field_name, None) + if isinstance(value, bool): + continue if not isinstance(value, int | float): continue + yield field_name, value + + +def output_attributes(output: Output[Any]) -> dict[str, Any]: + """Extract GenAI response attributes from a typed Output.""" + attrs: dict[str, Any] = {} + for field_name, value in _iter_usage_numeric(output.usage): semconv_key = _GEN_AI_USAGE_FIELDS.get(field_name) attrs[semconv_key or f"celeste.usage.{field_name}"] = value if output.finish_reason is not None and output.finish_reason.reason is not None: @@ -226,16 +230,10 @@ def output_attributes(output: Output[Any]) -> dict[str, Any]: def record_token_usage(usage: Usage, attributes: dict[str, Any]) -> None: """Record one ``gen_ai.client.token.usage`` observation per token category.""" - usage_fields = ( - type(usage).model_fields if hasattr(type(usage), "model_fields") else {} - ) - for field_name in usage_fields: + for field_name, value in _iter_usage_numeric(usage): token_type = _GEN_AI_TOKEN_TYPES.get(field_name) if token_type is None: continue - value = getattr(usage, field_name, None) - if not isinstance(value, int | float): - continue _token_usage_histogram.record( value, attributes={**attributes, "gen_ai.token.type": token_type}, @@ -364,6 +362,34 @@ def _output_messages_event(output: Output[Any]) -> dict[str, Any] | None: } +def add_input_event(span: Any, inputs: Input) -> None: + """Add the ``gen_ai.input.messages`` event to ``span`` when capture is enabled.""" + event = _input_messages_event(inputs) + if event is not None: + span.add_event("gen_ai.input.messages", attributes=event) + + +def add_output_event(span: Any, output: Output[Any]) -> None: + """Add the ``gen_ai.output.messages`` event to ``span`` when capture is enabled.""" + event = _output_messages_event(output) + if event is not None: + span.add_event("gen_ai.output.messages", attributes=event) + + +def record_output( + span: Any, + output: Output[Any], + metric_attributes: dict[str, Any], + duration_seconds: float, + error: BaseException | None = None, +) -> None: + """Emit span attrs, content event, and metrics for a successful Output.""" + span.set_attributes(output_attributes(output)) + add_output_event(span, output) + record_token_usage(output.usage, metric_attributes) + record_operation_duration(duration_seconds, metric_attributes, error=error) + + class _TracedStream: """Stream-shaped wrapper that emits GenAI telemetry against an OTel span.""" @@ -465,12 +491,13 @@ def _finalize(self) -> None: except StreamNotExhaustedError: output = None if output is not None: - self._span.set_attributes(output_attributes(output)) - output_event = _output_messages_event(output) - if output_event is not None: - self._span.add_event("gen_ai.output.messages", attributes=output_event) - record_token_usage(output.usage, self._metric_attributes) - record_operation_duration(duration, self._metric_attributes, error=self._error) + record_output( + self._span, output, self._metric_attributes, duration, error=self._error + ) + else: + record_operation_duration( + duration, self._metric_attributes, error=self._error + ) self._span.end() @@ -486,9 +513,12 @@ def trace_stream( __all__ = [ "Status", "StatusCode", + "add_input_event", + "add_output_event", "meter", "output_attributes", "record_operation_duration", + "record_output", "record_token_usage", "request_attributes", "span_name", From 2f7f4b67f5c1c01ae25462171a0e30b7e3c578e1 Mon Sep 17 00:00:00 2001 From: kamilbenkirane Date: Tue, 5 May 2026 14:28:14 +0200 Subject: [PATCH 3/5] fix(telemetry): sync iteration drives the wrapper's __anext__, not the inner Stream's MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The previous `_TracedStream.__iter__` delegated to `iter(self._inner)`, which makes Python iterate the inner Stream directly via its blocking-portal `__iter__`. That bypasses the wrapper's `__anext__` — so for sync streaming consumers, the span never finalized, TTFC was never recorded, and metrics were never emitted. Fix: `_TracedStream.__iter__` now spins its own portal and drives `self.__anext__` (the wrapper's), mirroring `Stream.__iter__`. On exhaustion, exception, or generator-close (consumer break / GC), the finally block calls `self.aclose()` which finalizes the span via the existing async path. Caught by smoke-testing `celeste.text.sync.stream.generate(...)` against a real Groq llama-3.1-8b-instant call: previously the span was missing from finished spans entirely; now it appears alongside the async-streaming span with the same attribute set. --- src/celeste/telemetry.py | 20 +++++++++++++++++--- 1 file changed, 17 insertions(+), 3 deletions(-) diff --git a/src/celeste/telemetry.py b/src/celeste/telemetry.py index 2dd9302..49a8465 100644 --- a/src/celeste/telemetry.py +++ b/src/celeste/telemetry.py @@ -5,10 +5,12 @@ import os import time from collections.abc import Iterator -from contextlib import contextmanager +from contextlib import contextmanager, suppress from types import TracebackType from typing import Any +from anyio.from_thread import start_blocking_portal + from celeste.artifacts import Artifact from celeste.core import Modality, Protocol, Provider, UsageField from celeste.exceptions import StreamNotExhaustedError @@ -449,8 +451,20 @@ async def __aexit__( return False def __iter__(self) -> Iterator[Any]: - """Delegate sync iteration to the inner Stream.""" - return iter(self._inner) + """Sync iterator using a blocking portal — drives this wrapper's __anext__.""" + portal_cm = start_blocking_portal() + portal = portal_cm.__enter__() + try: + while True: + try: + yield portal.call(self.__anext__) + except StopAsyncIteration: + return + finally: + if not self._ended: + with suppress(RuntimeError): + portal.call(self.aclose) + portal_cm.__exit__(None, None, None) def __enter__(self) -> "_TracedStream": """Enter sync context — delegate to inner Stream.""" From fcc48af92480071b41d241e375751027676dadeb Mon Sep 17 00:00:00 2001 From: kamilbenkirane Date: Tue, 5 May 2026 14:58:46 +0200 Subject: [PATCH 4/5] refactor: dedupe telemetry test fixtures, fix Windows time.monotonic flake MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three review agents flagged the +800 LOC PR for: triplicated test stream classes, redundant fixtures, parametrizable repetition, and one Windows-fragile timing assertion. Applied: - New `tests/unit_tests/_telemetry_helpers.py` holds the canonical `TelemetryUsage` / `TelemetryOutput` / `TelemetryStream` / `async_iter` used by all three telemetry test files. - New `tests/unit_tests/conftest.py` exposes the shared `exporter` fixture + `start_test_span` helper. - `test_telemetry_metrics.py` parametrizes `record_token_usage` (3 cases) and `record_operation_duration` (success / failure). Drops the redundant `test_input_and_output_recorded_separately` (subsumed by the parametrized `all_token_types` case). - `test_telemetry_metrics.py:191` now asserts `>= 0` on duration sum — Windows `time.monotonic()` resolution can return 0 for sub-millisecond in-memory streams; the previous `> 0` failed CI on Windows runners. - `telemetry.py`: drop `add_output_event` (only ever called by `record_output` — inlined). Trim verbose comments above `_GEN_AI_USAGE_FIELDS`, `_GEN_AI_TOKEN_TYPES`, `_CAPTURE_CONTENT`, and the two messages-event helpers down to one line each. Net: -210 modified LOC, +72 new helper LOC. PR diff drops from +861 to ~+725. Coverage unchanged (630 tests pass). --- src/celeste/telemetry.py | 32 +-- tests/unit_tests/_telemetry_helpers.py | 49 ++++ tests/unit_tests/conftest.py | 23 ++ .../test_telemetry_content_events.py | 207 ++++++--------- tests/unit_tests/test_telemetry_metrics.py | 244 +++++++----------- tests/unit_tests/test_telemetry_streaming.py | 149 +++-------- 6 files changed, 283 insertions(+), 421 deletions(-) create mode 100644 tests/unit_tests/_telemetry_helpers.py create mode 100644 tests/unit_tests/conftest.py diff --git a/src/celeste/telemetry.py b/src/celeste/telemetry.py index 49a8465..d5e7088 100644 --- a/src/celeste/telemetry.py +++ b/src/celeste/telemetry.py @@ -187,8 +187,7 @@ def span_name(modality: Modality, model: Model) -> str: return f"celeste.{modality.value} {model.id}" -# Maps `UsageField` enum members to GenAI semconv span attribute keys. -# Fields not in this map fall through to `celeste.usage.` (off-spec). +# `Usage` field → semconv span attribute key. Off-spec fields use `celeste.usage.`. _GEN_AI_USAGE_FIELDS: dict[str, str] = { UsageField.INPUT_TOKENS: "gen_ai.usage.input_tokens", UsageField.OUTPUT_TOKENS: "gen_ai.usage.output_tokens", @@ -197,9 +196,7 @@ def span_name(modality: Modality, model: Model) -> str: UsageField.CACHED_TOKENS: "gen_ai.usage.cached_input_tokens", } -# Maps `UsageField` enum members to `gen_ai.token.type` dimension values for metrics. -# Only fields representing token *categories* belong here (input/output/cached/reasoning). -# `total_tokens` is omitted to avoid double-counting in histogram aggregations. +# `Usage` field → `gen_ai.token.type` metric dimension. `total_tokens` excluded to avoid double-count. _GEN_AI_TOKEN_TYPES: dict[str, str] = { UsageField.INPUT_TOKENS: "input", UsageField.OUTPUT_TOKENS: "output", @@ -254,8 +251,7 @@ def record_operation_duration( _operation_duration_histogram.record(duration_seconds, attributes=attrs) -# Read the semconv-standard env flag once at import. Per the GenAI spec, -# content capture is opt-in because prompts/responses often contain user data. +# Opt-in flag, read once at import (semconv-standard env name). _CAPTURE_CONTENT: bool = ( os.environ.get("OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT", "") .strip() @@ -321,10 +317,7 @@ def _message_to_dict(message: Message) -> dict[str, Any]: def _input_messages_event(inputs: Input) -> dict[str, Any] | None: - """Build the ``gen_ai.input.messages`` event attributes from inputs. - - Returns ``None`` when content capture is disabled. - """ + """Build the ``gen_ai.input.messages`` event attributes, or None when capture is off.""" if not _CAPTURE_CONTENT: return None messages: list[dict[str, Any]] = [] @@ -345,10 +338,7 @@ def _input_messages_event(inputs: Input) -> dict[str, Any] | None: def _output_messages_event(output: Output[Any]) -> dict[str, Any] | None: - """Build the ``gen_ai.output.messages`` event attributes from an Output. - - Returns ``None`` when content capture is disabled. - """ + """Build the ``gen_ai.output.messages`` event attributes, or None when capture is off.""" if not _CAPTURE_CONTENT: return None parts = _content_to_parts(output.content) @@ -371,13 +361,6 @@ def add_input_event(span: Any, inputs: Input) -> None: span.add_event("gen_ai.input.messages", attributes=event) -def add_output_event(span: Any, output: Output[Any]) -> None: - """Add the ``gen_ai.output.messages`` event to ``span`` when capture is enabled.""" - event = _output_messages_event(output) - if event is not None: - span.add_event("gen_ai.output.messages", attributes=event) - - def record_output( span: Any, output: Output[Any], @@ -387,7 +370,9 @@ def record_output( ) -> None: """Emit span attrs, content event, and metrics for a successful Output.""" span.set_attributes(output_attributes(output)) - add_output_event(span, output) + output_event = _output_messages_event(output) + if output_event is not None: + span.add_event("gen_ai.output.messages", attributes=output_event) record_token_usage(output.usage, metric_attributes) record_operation_duration(duration_seconds, metric_attributes, error=error) @@ -528,7 +513,6 @@ def trace_stream( "Status", "StatusCode", "add_input_event", - "add_output_event", "meter", "output_attributes", "record_operation_duration", diff --git a/tests/unit_tests/_telemetry_helpers.py b/tests/unit_tests/_telemetry_helpers.py new file mode 100644 index 0000000..56f1c6f --- /dev/null +++ b/tests/unit_tests/_telemetry_helpers.py @@ -0,0 +1,49 @@ +"""Shared fixtures and helpers for telemetry-related unit tests.""" + +from collections.abc import AsyncIterator +from typing import Any, ClassVar + +from celeste.io import Chunk, Output, Usage +from celeste.parameters import Parameters +from celeste.streaming import Stream + + +class TelemetryUsage(Usage): + """Usage with the full GenAI semconv field set used across telemetry tests.""" + + input_tokens: int | None = None + output_tokens: int | None = None + total_tokens: int | None = None + reasoning_tokens: int | None = None + cached_tokens: int | None = None + + +class TelemetryOutput(Output[str]): + """String-content Output used across telemetry tests.""" + + +class TelemetryStream(Stream[TelemetryOutput, Parameters, Chunk]): + """Concrete Stream that aggregates content + usage from per-chunk events.""" + + _usage_class: ClassVar[type[Usage]] = TelemetryUsage + _chunk_class: ClassVar[type[Chunk]] = Chunk + _output_class: ClassVar[type[Output]] = TelemetryOutput + _empty_content: ClassVar[str] = "" + + def _aggregate_content(self, chunks: list[Chunk]) -> str: + """Concatenate chunk content.""" + return "".join(chunk.content for chunk in chunks) + + def _parse_chunk(self, event: dict[str, Any]) -> Chunk | None: + """Parse delta events; lifecycle events return None.""" + content = event.get("delta") + if not content and "usage" not in event: + return None + usage = TelemetryUsage(**event["usage"]) if "usage" in event else None + return Chunk(content=content or "", finish_reason=None, usage=usage) + + +async def async_iter(events: list[dict[str, Any]]) -> AsyncIterator[dict[str, Any]]: + """Convert a list of events into an async iterator.""" + for event in events: + yield event diff --git a/tests/unit_tests/conftest.py b/tests/unit_tests/conftest.py new file mode 100644 index 0000000..c7bdc49 --- /dev/null +++ b/tests/unit_tests/conftest.py @@ -0,0 +1,23 @@ +"""Shared pytest fixtures for unit tests.""" + +import pytest +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 Span + + +@pytest.fixture +def exporter() -> tuple[InMemorySpanExporter, TracerProvider]: + """In-memory span exporter wired into a fresh TracerProvider.""" + span_exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(span_exporter)) + return span_exporter, provider + + +def start_test_span(provider: TracerProvider, name: str = "test") -> Span: + """Start a detached span on the given provider's tracer.""" + return provider.get_tracer("celeste-test").start_span(name) diff --git a/tests/unit_tests/test_telemetry_content_events.py b/tests/unit_tests/test_telemetry_content_events.py index f198274..1d7cd6b 100644 --- a/tests/unit_tests/test_telemetry_content_events.py +++ b/tests/unit_tests/test_telemetry_content_events.py @@ -1,26 +1,21 @@ """Tests for opt-in GenAI content events (`gen_ai.input.messages` / `output.messages`).""" import json -from collections.abc import AsyncIterator -from typing import Any, ClassVar import pytest 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 Span from celeste import telemetry from celeste.artifacts import ImageArtifact -from celeste.io import Chunk, Output, Usage from celeste.mime_types import ImageMimeType from celeste.modalities.text.io import TextInput, TextOutput, TextUsage -from celeste.parameters import Parameters -from celeste.streaming import Stream from celeste.tools import ToolCall from celeste.types import Message, Role +from tests.unit_tests._telemetry_helpers import TelemetryStream, async_iter +from tests.unit_tests.conftest import start_test_span @pytest.fixture @@ -35,88 +30,31 @@ def capture_disabled(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(telemetry, "_CAPTURE_CONTENT", False) -@pytest.fixture -def exporter() -> tuple[InMemorySpanExporter, TracerProvider]: - """Wire an in-memory span exporter.""" - span_exporter = InMemorySpanExporter() - provider = TracerProvider() - provider.add_span_processor(SimpleSpanProcessor(span_exporter)) - return span_exporter, provider - - -def _start_span(provider: TracerProvider, name: str = "test") -> Span: - """Start a detached span on the given provider's tracer.""" - return provider.get_tracer("celeste-test").start_span(name) - - -class _TestStreamUsage(Usage): - """Usage with token fields for streaming tests.""" - - input_tokens: int | None = None - output_tokens: int | None = None - - -class _TestStreamOutput(Output[str]): - """Output for streaming tests.""" - - pass - - -class _TestStream(Stream[_TestStreamOutput, Parameters, Chunk]): - """Concrete Stream that aggregates content + usage.""" - - _usage_class: ClassVar[type[Usage]] = _TestStreamUsage - _chunk_class: ClassVar[type[Chunk]] = Chunk - _output_class: ClassVar[type[Output]] = _TestStreamOutput - _empty_content: ClassVar[str] = "" - - def _aggregate_content(self, chunks: list[Chunk]) -> str: - return "".join(chunk.content for chunk in chunks) - - def _parse_chunk(self, event: dict[str, Any]) -> Chunk | None: - content = event.get("delta") - if not content and "usage" not in event: - return None - usage = _TestStreamUsage(**event["usage"]) if "usage" in event else None - return Chunk(content=content or "", finish_reason=None, usage=usage) - - -async def _async_iter(events: list[dict[str, Any]]) -> AsyncIterator[dict[str, Any]]: - for event in events: - yield event - - class TestInputMessagesEvent: - """`_input_messages_event` builds the input messages payload from typed inputs.""" - def test_capture_disabled_returns_none(self, capture_disabled: None) -> None: """With the env flag off, the helper returns None — no event emitted.""" - inputs = TextInput(prompt="hello") - assert telemetry._input_messages_event(inputs) is None + assert telemetry._input_messages_event(TextInput(prompt="hello")) is None def test_text_prompt_only(self, capture_enabled: None) -> None: """Plain text prompt produces a single user message with one text part.""" - inputs = TextInput(prompt="hello") - - result = telemetry._input_messages_event(inputs) + result = telemetry._input_messages_event(TextInput(prompt="hello")) assert result is not None - messages = json.loads(result["messages"]) - assert messages == [ + assert json.loads(result["messages"]) == [ {"role": "user", "parts": [{"type": "text", "content": "hello"}]} ] def test_messages_field_emitted(self, capture_enabled: None) -> None: """When `inputs.messages` is set, those messages are serialized verbatim.""" - inputs = TextInput( - messages=[ - Message(role=Role.SYSTEM, content="be helpful"), - Message(role=Role.USER, content="hi"), - ] + result = telemetry._input_messages_event( + TextInput( + messages=[ + Message(role=Role.SYSTEM, content="be helpful"), + Message(role=Role.USER, content="hi"), + ] + ) ) - result = telemetry._input_messages_event(inputs) - assert result is not None messages = json.loads(result["messages"]) assert messages[0]["role"] == "system" @@ -125,19 +63,17 @@ def test_messages_field_emitted(self, capture_enabled: None) -> None: def test_image_input_emitted_as_reference(self, capture_enabled: None) -> None: """Images attached to a prompt are emitted as URL references, not bytes.""" - inputs = TextInput( - prompt="describe this", - image=ImageArtifact( - url="https://example.com/img.png", - mime_type=ImageMimeType.PNG, - ), + result = telemetry._input_messages_event( + TextInput( + prompt="describe this", + image=ImageArtifact( + url="https://example.com/img.png", mime_type=ImageMimeType.PNG + ), + ) ) - result = telemetry._input_messages_event(inputs) - assert result is not None - messages = json.loads(result["messages"]) - parts = messages[0]["parts"] + parts = json.loads(result["messages"])[0]["parts"] assert parts[0]["type"] == "text" image_parts = [p for p in parts if p["type"] == "image"] assert len(image_parts) == 1 @@ -146,55 +82,60 @@ def test_image_input_emitted_as_reference(self, capture_enabled: None) -> None: class TestOutputMessagesEvent: - """`_output_messages_event` builds the assistant payload from a typed Output.""" - def test_capture_disabled_returns_none(self, capture_disabled: None) -> None: """With the env flag off, the helper returns None.""" - output = TextOutput(content="reply", usage=TextUsage()) - assert telemetry._output_messages_event(output) is None + assert ( + telemetry._output_messages_event( + TextOutput(content="reply", usage=TextUsage()) + ) + is None + ) def test_text_only(self, capture_enabled: None) -> None: """Plain text output produces a single assistant message with one text part.""" - output = TextOutput(content="reply", usage=TextUsage()) - - result = telemetry._output_messages_event(output) + result = telemetry._output_messages_event( + TextOutput(content="reply", usage=TextUsage()) + ) assert result is not None - messages = json.loads(result["messages"]) - assert messages == [ + assert json.loads(result["messages"]) == [ {"role": "assistant", "parts": [{"type": "text", "content": "reply"}]} ] def test_reasoning_appended(self, capture_enabled: None) -> None: """Reasoning content is appended as a `reasoning` part.""" - output = TextOutput( - content="answer", usage=TextUsage(), reasoning="thinking..." + result = telemetry._output_messages_event( + TextOutput(content="answer", usage=TextUsage(), reasoning="thinking...") ) - result = telemetry._output_messages_event(output) - assert result is not None parts = json.loads(result["messages"])[0]["parts"] - types = [p["type"] for p in parts] - assert "text" in types and "reasoning" in types - reasoning_part = next(p for p in parts if p["type"] == "reasoning") - assert reasoning_part["content"] == "thinking..." + assert {p["type"] for p in parts} == {"text", "reasoning"} + assert ( + next(p for p in parts if p["type"] == "reasoning")["content"] + == "thinking..." + ) def test_tool_calls_emitted_as_tool_call_parts(self, capture_enabled: None) -> None: """Tool calls become semconv `tool_call` parts with id/name/arguments.""" - output = TextOutput( - content="", - usage=TextUsage(), - tool_calls=[ - ToolCall(id="call_1", name="get_weather", arguments={"city": "Paris"}) - ], + result = telemetry._output_messages_event( + TextOutput( + content="", + usage=TextUsage(), + tool_calls=[ + ToolCall( + id="call_1", name="get_weather", arguments={"city": "Paris"} + ) + ], + ) ) - result = telemetry._output_messages_event(output) - assert result is not None - parts = json.loads(result["messages"])[0]["parts"] - tool_parts = [p for p in parts if p["type"] == "tool_call"] + tool_parts = [ + p + for p in json.loads(result["messages"])[0]["parts"] + if p["type"] == "tool_call" + ] assert len(tool_parts) == 1 assert tool_parts[0]["id"] == "call_1" assert tool_parts[0]["name"] == "get_weather" @@ -202,25 +143,25 @@ def test_tool_calls_emitted_as_tool_call_parts(self, capture_enabled: None) -> N def test_image_artifact_output_by_reference(self, capture_enabled: None) -> None: """Image generation outputs emit URL references, never inline bytes.""" - output = TextOutput( - content=ImageArtifact( - url="https://example.com/out.png", mime_type=ImageMimeType.PNG - ), - usage=TextUsage(), + result = telemetry._output_messages_event( + TextOutput( + content=ImageArtifact( + url="https://example.com/out.png", mime_type=ImageMimeType.PNG + ), + usage=TextUsage(), + ) ) - result = telemetry._output_messages_event(output) - assert result is not None parts = json.loads(result["messages"])[0]["parts"] - assert parts[0]["type"] == "image" - assert parts[0]["uri"] == "https://example.com/out.png" - assert parts[0]["mime_type"] == "image/png" + assert parts[0] == { + "type": "image", + "uri": "https://example.com/out.png", + "mime_type": "image/png", + } class TestStreamingContentEvents: - """Streaming flows emit content events at finalize when output is built.""" - async def test_natural_exhaustion_emits_output_event( self, capture_enabled: None, @@ -232,17 +173,16 @@ async def test_natural_exhaustion_emits_output_event( {"delta": "Hello"}, {"delta": " world", "usage": {"input_tokens": 1, "output_tokens": 2}}, ] - stream = _TestStream(_async_iter(events)) - span = _start_span(provider) - wrapped = telemetry.trace_stream(stream, span) + wrapped = telemetry.trace_stream( + TelemetryStream(async_iter(events)), start_test_span(provider) + ) async for _ in wrapped: pass finished = in_memory.get_finished_spans() assert len(finished) == 1 - event_names = [e.name for e in finished[0].events] - assert "gen_ai.output.messages" in event_names + assert "gen_ai.output.messages" in [e.name for e in finished[0].events] async def test_aclose_before_output_skips_output_event( self, @@ -251,15 +191,14 @@ async def test_aclose_before_output_skips_output_event( ) -> None: """When _output is None at finalize, the output event is skipped.""" in_memory, provider = exporter - events = [{"delta": "Partial"}, {"delta": "..."}] - stream = _TestStream(_async_iter(events)) - span = _start_span(provider) - wrapped = telemetry.trace_stream(stream, span) + wrapped = telemetry.trace_stream( + TelemetryStream(async_iter([{"delta": "Partial"}, {"delta": "..."}])), + start_test_span(provider), + ) await wrapped.__anext__() await wrapped.aclose() finished = in_memory.get_finished_spans() assert len(finished) == 1 - event_names = [e.name for e in finished[0].events] - assert "gen_ai.output.messages" not in event_names + assert "gen_ai.output.messages" not in [e.name for e in finished[0].events] diff --git a/tests/unit_tests/test_telemetry_metrics.py b/tests/unit_tests/test_telemetry_metrics.py index e2d7f22..2b407e2 100644 --- a/tests/unit_tests/test_telemetry_metrics.py +++ b/tests/unit_tests/test_telemetry_metrics.py @@ -1,78 +1,33 @@ """Tests for GenAI metric histogram emission.""" -from collections.abc import AsyncIterator -from typing import Any, ClassVar +from typing import Any import pytest from opentelemetry.sdk.metrics import MeterProvider from opentelemetry.sdk.metrics.export import InMemoryMetricReader from celeste import telemetry -from celeste.io import Chunk, Output, Usage -from celeste.parameters import Parameters -from celeste.streaming import Stream - - -class _TestUsage(Usage): - """Usage with input/output/cached/reasoning fields for metric tests.""" - - input_tokens: int | None = None - output_tokens: int | None = None - cached_tokens: int | None = None - reasoning_tokens: int | None = None - - -class _TestOutput(Output[str]): - """Output for tests.""" - - pass - - -class _TestStream(Stream[_TestOutput, Parameters, Chunk]): - """Concrete Stream that aggregates usage from per-chunk metadata.""" - - _usage_class: ClassVar[type[Usage]] = _TestUsage - _chunk_class: ClassVar[type[Chunk]] = Chunk - _output_class: ClassVar[type[Output]] = _TestOutput - _empty_content: ClassVar[str] = "" - - def _aggregate_content(self, chunks: list[Chunk]) -> str: - return "".join(chunk.content for chunk in chunks) - - def _parse_chunk(self, event: dict[str, Any]) -> Chunk | None: - content = event.get("delta") - if not content and "usage" not in event: - return None - usage = _TestUsage(**event["usage"]) if "usage" in event else None - return Chunk(content=content or "", finish_reason=None, usage=usage) - - -async def _async_iter(events: list[dict[str, Any]]) -> AsyncIterator[dict[str, Any]]: - for event in events: - yield event +from tests.unit_tests._telemetry_helpers import ( + TelemetryStream, + TelemetryUsage, + async_iter, +) @pytest.fixture def meter_setup(monkeypatch: pytest.MonkeyPatch) -> InMemoryMetricReader: """Wire fresh histograms backed by InMemoryMetricReader for assertions.""" reader = InMemoryMetricReader() - provider = MeterProvider(metric_readers=[reader]) - test_meter = provider.get_meter("celeste-test") + test_meter = MeterProvider(metric_readers=[reader]).get_meter("celeste-test") monkeypatch.setattr( telemetry, "_token_usage_histogram", - test_meter.create_histogram( - name="gen_ai.client.token.usage", - unit="{token}", - ), + test_meter.create_histogram(name="gen_ai.client.token.usage", unit="{token}"), ) monkeypatch.setattr( telemetry, "_operation_duration_histogram", - test_meter.create_histogram( - name="gen_ai.client.operation.duration", - unit="s", - ), + test_meter.create_histogram(name="gen_ai.client.operation.duration", unit="s"), ) return reader @@ -91,102 +46,85 @@ def _collect_metric(reader: InMemoryMetricReader, name: str) -> list[Any]: return points -class TestRecordTokenUsage: - """`record_token_usage` emits one record per token category.""" - - def test_input_and_output_recorded_separately( - self, meter_setup: InMemoryMetricReader - ) -> None: - """input + output usage produce two records with different `gen_ai.token.type`.""" - usage = _TestUsage(input_tokens=100, output_tokens=50) - attrs = {"gen_ai.request.model": "test-model"} - - telemetry.record_token_usage(usage, attrs) - - points = _collect_metric(meter_setup, "gen_ai.client.token.usage") - token_types = sorted(p.attributes["gen_ai.token.type"] for p in points) - assert token_types == ["input", "output"] - by_type = {p.attributes["gen_ai.token.type"]: p.sum for p in points} - assert by_type["input"] == 100 - assert by_type["output"] == 50 - - def test_cached_and_reasoning_recorded_with_dedicated_token_type( - self, meter_setup: InMemoryMetricReader - ) -> None: - """cached_tokens and reasoning_tokens get distinct token.type values.""" - usage = _TestUsage( - input_tokens=10, output_tokens=5, cached_tokens=8, reasoning_tokens=3 - ) - telemetry.record_token_usage(usage, {}) - - points = _collect_metric(meter_setup, "gen_ai.client.token.usage") - token_types = sorted(p.attributes["gen_ai.token.type"] for p in points) - assert token_types == ["cached_input", "input", "output", "reasoning"] - - def test_none_fields_are_skipped(self, meter_setup: InMemoryMetricReader) -> None: - """Fields set to None do not produce records.""" - usage = _TestUsage(input_tokens=10) - telemetry.record_token_usage(usage, {}) - - points = _collect_metric(meter_setup, "gen_ai.client.token.usage") - assert len(points) == 1 - assert points[0].attributes["gen_ai.token.type"] == "input" - - -class TestRecordOperationDuration: - """`record_operation_duration` records latency with optional error.type.""" - - def test_success_records_duration_without_error_type( - self, meter_setup: InMemoryMetricReader - ) -> None: - """Successful call: duration recorded, no error.type attribute.""" - telemetry.record_operation_duration(1.5, {"gen_ai.request.model": "test-model"}) - - points = _collect_metric(meter_setup, "gen_ai.client.operation.duration") - assert len(points) == 1 - assert "error.type" not in points[0].attributes - assert points[0].sum == pytest.approx(1.5) - - def test_failure_records_duration_with_error_type( - self, meter_setup: InMemoryMetricReader - ) -> None: - """Failed call: duration recorded with error.type set to exception class name.""" - telemetry.record_operation_duration(0.2, {}, error=ValueError("bad input")) - - points = _collect_metric(meter_setup, "gen_ai.client.operation.duration") - assert len(points) == 1 - assert points[0].attributes["error.type"] == "ValueError" - - -class TestStreamingMetrics: - """Streaming flows record duration covering full consumption.""" - - async def test_streaming_records_duration_and_token_usage( - self, meter_setup: InMemoryMetricReader - ) -> None: - """Natural exhaustion records both metrics with correct values.""" - events = [ - {"delta": "Hello"}, - {"delta": " world", "usage": {"input_tokens": 12, "output_tokens": 34}}, - ] - stream = _TestStream(_async_iter(events)) - provider = ( - MeterProvider() - ) # noop meter for span; the spans go to default tracer - del provider - span = telemetry.tracer.start_span("test") - wrapped = telemetry.trace_stream( - stream, span, metric_attributes={"gen_ai.request.model": "test-model"} - ) - - async for _ in wrapped: - pass - - token_points = _collect_metric(meter_setup, "gen_ai.client.token.usage") - duration_points = _collect_metric( - meter_setup, "gen_ai.client.operation.duration" - ) - assert len(token_points) == 2 - assert len(duration_points) == 1 - assert duration_points[0].sum > 0 - assert "error.type" not in duration_points[0].attributes +@pytest.mark.parametrize( + ("usage_kwargs", "expected_token_types", "expected_sums"), + [ + ( + { + "input_tokens": 10, + "output_tokens": 5, + "cached_tokens": 8, + "reasoning_tokens": 3, + }, + ["cached_input", "input", "output", "reasoning"], + {"input": 10, "output": 5, "cached_input": 8, "reasoning": 3}, + ), + ( + {"input_tokens": 100, "output_tokens": 50}, + ["input", "output"], + {"input": 100, "output": 50}, + ), + ({"input_tokens": 10}, ["input"], {"input": 10}), + ], + ids=["all_token_types", "input_output_only", "input_only_others_none"], +) +def test_record_token_usage_emits_one_record_per_category( + meter_setup: InMemoryMetricReader, + usage_kwargs: dict[str, int], + expected_token_types: list[str], + expected_sums: dict[str, int], +) -> None: + """Each populated token field produces one record with a distinct token.type.""" + telemetry.record_token_usage(TelemetryUsage(**usage_kwargs), {}) + + points = _collect_metric(meter_setup, "gen_ai.client.token.usage") + assert ( + sorted(p.attributes["gen_ai.token.type"] for p in points) + == expected_token_types + ) + assert {p.attributes["gen_ai.token.type"]: p.sum for p in points} == expected_sums + + +@pytest.mark.parametrize( + ("error", "expected_error_type"), + [(None, None), (ValueError("bad input"), "ValueError")], + ids=["success", "failure"], +) +def test_record_operation_duration_with_optional_error_type( + meter_setup: InMemoryMetricReader, + error: BaseException | None, + expected_error_type: str | None, +) -> None: + """Duration is recorded; error.type appears on failure only.""" + telemetry.record_operation_duration(1.5, {"gen_ai.request.model": "m"}, error=error) + + points = _collect_metric(meter_setup, "gen_ai.client.operation.duration") + assert len(points) == 1 + assert points[0].sum == pytest.approx(1.5) + assert points[0].attributes.get("error.type") == expected_error_type + + +async def test_streaming_records_duration_and_token_usage( + meter_setup: InMemoryMetricReader, +) -> None: + """Natural exhaustion records both metrics with correct values.""" + events = [ + {"delta": "Hello"}, + {"delta": " world", "usage": {"input_tokens": 12, "output_tokens": 34}}, + ] + span = telemetry.tracer.start_span("test") + wrapped = telemetry.trace_stream( + TelemetryStream(async_iter(events)), + span, + metric_attributes={"gen_ai.request.model": "test-model"}, + ) + + async for _ in wrapped: + pass + + token_points = _collect_metric(meter_setup, "gen_ai.client.token.usage") + duration_points = _collect_metric(meter_setup, "gen_ai.client.operation.duration") + assert len(token_points) == 2 + assert len(duration_points) == 1 + assert duration_points[0].sum >= 0 + assert "error.type" not in duration_points[0].attributes diff --git a/tests/unit_tests/test_telemetry_streaming.py b/tests/unit_tests/test_telemetry_streaming.py index d67b29f..fa91081 100644 --- a/tests/unit_tests/test_telemetry_streaming.py +++ b/tests/unit_tests/test_telemetry_streaming.py @@ -1,31 +1,24 @@ """Tests for `_TracedStream` — span lifecycle and GenAI attribute emission.""" from collections.abc import AsyncIterator -from typing import Any, ClassVar +from typing import Any import pytest 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 Span, StatusCode +from opentelemetry.trace import StatusCode from celeste import telemetry from celeste.exceptions import StreamNotExhaustedError -from celeste.io import Chunk, Output, Usage -from celeste.parameters import Parameters -from celeste.streaming import Stream - - -class _TestUsage(Usage): - """Usage with the full GenAI semconv field set for telemetry assertions.""" - - input_tokens: int | None = None - output_tokens: int | None = None - total_tokens: int | None = None - reasoning_tokens: int | None = None - cached_tokens: int | None = None +from celeste.io import Output, Usage +from tests.unit_tests._telemetry_helpers import ( + TelemetryOutput, + TelemetryStream, + async_iter, +) +from tests.unit_tests.conftest import start_test_span class _OffSpecUsage(Usage): @@ -35,62 +28,13 @@ class _OffSpecUsage(Usage): audio_seconds: float | None = None -class _TestOutput(Output[str]): - """Output for tests.""" - - pass - - -class _TestStream(Stream[_TestOutput, Parameters, Chunk]): - """Stream that aggregates usage from per-chunk metadata for tests.""" - - _usage_class: ClassVar[type[Usage]] = _TestUsage - _chunk_class: ClassVar[type[Chunk]] = Chunk - _output_class: ClassVar[type[Output]] = _TestOutput - _empty_content: ClassVar[str] = "" - - def _aggregate_content(self, chunks: list[Chunk]) -> str: - """Concatenate chunk content.""" - return "".join(chunk.content for chunk in chunks) - - def _parse_chunk(self, event: dict[str, Any]) -> Chunk | None: - """Parse delta events; lifecycle events return None.""" - content = event.get("delta") - if not content and "usage" not in event: - return None - usage = _TestUsage(**event["usage"]) if "usage" in event else None - return Chunk(content=content or "", finish_reason=None, usage=usage) - - -async def _async_iter(events: list[dict[str, Any]]) -> AsyncIterator[dict[str, Any]]: - """Convert a list of events into an async iterator.""" - for event in events: - yield event - - async def _failing_iter() -> AsyncIterator[dict[str, Any]]: """Async iterator that raises before yielding any event.""" raise RuntimeError("boom") yield # pragma: no cover — needed to mark this as an async generator -@pytest.fixture -def exporter() -> tuple[InMemorySpanExporter, TracerProvider]: - """In-memory span exporter wired into a fresh TracerProvider.""" - span_exporter = InMemorySpanExporter() - provider = TracerProvider() - provider.add_span_processor(SimpleSpanProcessor(span_exporter)) - return span_exporter, provider - - -def _start_span(provider: TracerProvider, name: str = "test") -> Span: - """Start a detached span on the given provider's tracer.""" - return provider.get_tracer("celeste-test").start_span(name) - - class TestNaturalExhaustion: - """Span emits usage attrs from inner._output after natural exhaustion.""" - async def test_emits_input_output_tokens_on_natural_exhaustion( self, exporter: tuple[InMemorySpanExporter, TracerProvider] ) -> None: @@ -100,9 +44,9 @@ async def test_emits_input_output_tokens_on_natural_exhaustion( {"delta": "Hello"}, {"delta": " world", "usage": {"input_tokens": 12, "output_tokens": 34}}, ] - stream = _TestStream(_async_iter(events)) - span = _start_span(provider) - wrapped = telemetry.trace_stream(stream, span) + wrapped = telemetry.trace_stream( + TelemetryStream(async_iter(events)), start_test_span(provider) + ) chunks = [chunk async for chunk in wrapped] @@ -119,12 +63,10 @@ async def test_async_context_manager_emits_usage( ) -> None: """`async with wrapped: async for ...` emits usage and ends the span once.""" in_memory, provider = exporter - events = [ - {"delta": "Hi", "usage": {"input_tokens": 1, "output_tokens": 2}}, - ] - stream = _TestStream(_async_iter(events)) - span = _start_span(provider) - wrapped = telemetry.trace_stream(stream, span) + events = [{"delta": "Hi", "usage": {"input_tokens": 1, "output_tokens": 2}}] + wrapped = telemetry.trace_stream( + TelemetryStream(async_iter(events)), start_test_span(provider) + ) async with wrapped as ctx: async for _ in ctx: @@ -138,20 +80,15 @@ async def test_async_context_manager_emits_usage( class TestEarlyAbandonment: - """`aclose()` before exhaustion ends the span without usage attrs.""" - async def test_aclose_before_output_built_emits_no_usage( self, exporter: tuple[InMemorySpanExporter, TracerProvider] ) -> None: """When _output is None at aclose time, span ends without usage attrs.""" in_memory, provider = exporter - events = [ - {"delta": "Partial"}, - {"delta": " result"}, - ] - stream = _TestStream(_async_iter(events)) - span = _start_span(provider) - wrapped = telemetry.trace_stream(stream, span) + events = [{"delta": "Partial"}, {"delta": " result"}] + wrapped = telemetry.trace_stream( + TelemetryStream(async_iter(events)), start_test_span(provider) + ) chunk = await wrapped.__anext__() assert chunk.content == "Partial" @@ -166,16 +103,14 @@ async def test_aclose_before_output_built_emits_no_usage( class TestExceptionPath: - """Exceptions during iteration are recorded with ERROR status.""" - async def test_exception_records_and_sets_error_status( self, exporter: tuple[InMemorySpanExporter, TracerProvider] ) -> None: """Exception during iteration sets ERROR status, ends span, propagates.""" in_memory, provider = exporter - stream = _TestStream(_failing_iter()) - span = _start_span(provider) - wrapped = telemetry.trace_stream(stream, span) + wrapped = telemetry.trace_stream( + TelemetryStream(_failing_iter()), start_test_span(provider) + ) with pytest.raises(RuntimeError, match="boom"): async for _ in wrapped: @@ -184,21 +119,18 @@ async def test_exception_records_and_sets_error_status( finished = in_memory.get_finished_spans() assert len(finished) == 1 assert finished[0].status.status_code == StatusCode.ERROR - events = list(finished[0].events) - assert any(e.name == "exception" for e in events) + assert any(e.name == "exception" for e in finished[0].events) class TestPublicSurface: - """Wrapper preserves the public Stream surface used by consumers.""" - async def test_output_raises_before_exhaustion( self, exporter: tuple[InMemorySpanExporter, TracerProvider] ) -> None: """Accessing .output mid-stream raises StreamNotExhaustedError.""" _, provider = exporter - stream = _TestStream(_async_iter([{"delta": "x"}])) - span = _start_span(provider) - wrapped = telemetry.trace_stream(stream, span) + wrapped = telemetry.trace_stream( + TelemetryStream(async_iter([{"delta": "x"}])), start_test_span(provider) + ) with pytest.raises(StreamNotExhaustedError): _ = wrapped.output @@ -209,31 +141,29 @@ async def test_output_returns_typed_output_after_exhaustion( """After exhaustion, .output returns the typed Output from inner.""" _, provider = exporter events = [{"delta": "abc", "usage": {"input_tokens": 1, "output_tokens": 1}}] - stream = _TestStream(_async_iter(events)) - span = _start_span(provider) - wrapped = telemetry.trace_stream(stream, span) + wrapped = telemetry.trace_stream( + TelemetryStream(async_iter(events)), start_test_span(provider) + ) async for _ in wrapped: pass output = wrapped.output - assert isinstance(output, _TestOutput) + assert isinstance(output, TelemetryOutput) assert output.content == "abc" assert output.usage.input_tokens == 1 class TestIdempotentFinalize: - """Span ends exactly once across natural exhaustion and explicit aclose.""" - async def test_natural_exhaustion_then_aclose_ends_span_once( self, exporter: tuple[InMemorySpanExporter, TracerProvider] ) -> None: """`async for` to completion then `aclose()` produces a single span.""" in_memory, provider = exporter events = [{"delta": "x", "usage": {"input_tokens": 2, "output_tokens": 3}}] - stream = _TestStream(_async_iter(events)) - span = _start_span(provider) - wrapped = telemetry.trace_stream(stream, span) + wrapped = telemetry.trace_stream( + TelemetryStream(async_iter(events)), start_test_span(provider) + ) async for _ in wrapped: pass @@ -266,15 +196,14 @@ async def test_total_reasoning_cached_tokens_emitted( }, } ] - stream = _TestStream(_async_iter(events)) - span = _start_span(provider) - wrapped = telemetry.trace_stream(stream, span) + wrapped = telemetry.trace_stream( + TelemetryStream(async_iter(events)), start_test_span(provider) + ) async for _ in wrapped: pass - finished = in_memory.get_finished_spans() - attrs = finished[0].attributes or {} + attrs = in_memory.get_finished_spans()[0].attributes or {} assert attrs["gen_ai.usage.input_tokens"] == 10 assert attrs["gen_ai.usage.output_tokens"] == 5 assert attrs["gen_ai.usage.total_tokens"] == 15 @@ -284,7 +213,7 @@ async def test_total_reasoning_cached_tokens_emitted( def test_off_spec_usage_emitted_under_celeste_namespace(self) -> None: """Modality-specific Usage fields fall through to `celeste.usage.`.""" usage = _OffSpecUsage(images_generated=4, audio_seconds=1.5) - output = _TestOutput(content="", usage=usage) + output: Output[Any] = TelemetryOutput(content="", usage=usage) attrs = telemetry.output_attributes(output) From b4cf340adb9e7c324669c8b150c57eca2d5b8511 Mon Sep 17 00:00:00 2001 From: kamilbenkirane Date: Tue, 5 May 2026 15:14:27 +0200 Subject: [PATCH 5/5] =?UTF-8?q?refactor(telemetry):=20extract=20gen=5Fai?= =?UTF-8?q?=5Fspan=20context=20manager=20=E2=80=94=20drop=20try/except=20+?= =?UTF-8?q?=20import=20time=20from=20client?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit User pushed back on `_predict` wrapping its entire body in try/except just to record `gen_ai.client.operation.duration` on the error path. The right home for that is a context manager inside `telemetry.py`. `gen_ai_span(model=, provider=, protocol=, modality=)` opens the span via `tracer.start_as_current_span`, captures the start time on enter, and in `finally` records the operation duration with `error.type` populated when the body raised. It yields `(span, request_attrs)` so `_predict` can call `add_input_event` and `record_output` against them. Result: - `_predict` is straight-line code: no try/except, no manual time math. - `client.py` no longer imports `time`. - `record_output` no longer takes `duration_seconds` / `error` params — duration is the span's lifecycle concern, not the output recorder's. - Streaming side (`_TracedStream._finalize`) updates symmetrically: records duration directly, then calls the simplified `record_output`. Real-call validated against gemini-3.1-flash-lite-preview text non-streaming + async streaming: same span attributes, same metric histograms, same content events as before. --- src/celeste/client.py | 78 +++++++++++++++++----------------------- src/celeste/telemetry.py | 56 ++++++++++++++++++++++------- 2 files changed, 76 insertions(+), 58 deletions(-) diff --git a/src/celeste/client.py b/src/celeste/client.py index c3c061c..fa26508 100644 --- a/src/celeste/client.py +++ b/src/celeste/client.py @@ -1,6 +1,5 @@ """Base client for modality-specific AI operations.""" -import time import warnings from abc import ABC, abstractmethod from collections.abc import AsyncIterator @@ -214,55 +213,42 @@ async def _predict( Returns: Output of the parameterized type. """ - request_attrs = telemetry.request_attributes( + with telemetry.gen_ai_span( model=self.model, provider=self.provider, protocol=self.protocol, modality=self.modality, - ) - started = time.monotonic() - with telemetry.tracer.start_as_current_span( - telemetry.span_name(self.modality, self.model), - attributes=request_attrs, - ) as span: - try: - inputs, parameters = self._validate_artifacts(inputs, **parameters) - telemetry.add_input_event(span, inputs) - request_body = self._build_request( - inputs, extra_body=extra_body, **parameters - ) - response_data = await self._make_request( - request_body, - endpoint=endpoint, - extra_headers=extra_headers, - **parameters, - ) - content = self._parse_content(response_data) - content = self._transform_output(content, **parameters) - tool_calls = self._parse_tool_calls(response_data) - reasoning, signature = self._parse_reasoning(response_data) - kwargs: dict[str, Any] = {} - if reasoning is not None: - kwargs["reasoning"] = reasoning - if signature: - kwargs["signature"] = signature - output = self._output_class()( - content=content, - usage=self._get_usage(response_data), - finish_reason=self._get_finish_reason(response_data), - metadata=self._build_metadata(response_data), - tool_calls=tool_calls, - **kwargs, - ) - telemetry.record_output( - span, output, request_attrs, time.monotonic() - started - ) - return output - except BaseException as exc: - telemetry.record_operation_duration( - time.monotonic() - started, request_attrs, error=exc - ) - raise + ) as (span, request_attrs): + inputs, parameters = self._validate_artifacts(inputs, **parameters) + telemetry.add_input_event(span, inputs) + request_body = self._build_request( + inputs, extra_body=extra_body, **parameters + ) + response_data = await self._make_request( + request_body, + endpoint=endpoint, + extra_headers=extra_headers, + **parameters, + ) + content = self._parse_content(response_data) + content = self._transform_output(content, **parameters) + tool_calls = self._parse_tool_calls(response_data) + reasoning, signature = self._parse_reasoning(response_data) + kwargs: dict[str, Any] = {} + if reasoning is not None: + kwargs["reasoning"] = reasoning + if signature: + kwargs["signature"] = signature + output = self._output_class()( + content=content, + usage=self._get_usage(response_data), + finish_reason=self._get_finish_reason(response_data), + metadata=self._build_metadata(response_data), + tool_calls=tool_calls, + **kwargs, + ) + telemetry.record_output(span, output, request_attrs) + return output def _parse_tool_calls(self, response_data: dict[str, Any]) -> list[ToolCall]: """Parse tool calls from response. Override in providers that support tools.""" diff --git a/src/celeste/telemetry.py b/src/celeste/telemetry.py index d5e7088..c6ce838 100644 --- a/src/celeste/telemetry.py +++ b/src/celeste/telemetry.py @@ -365,16 +365,49 @@ def record_output( span: Any, output: Output[Any], metric_attributes: dict[str, Any], - duration_seconds: float, - error: BaseException | None = None, ) -> None: - """Emit span attrs, content event, and metrics for a successful Output.""" + """Emit span attrs, content event, and token usage for a successful Output.""" span.set_attributes(output_attributes(output)) output_event = _output_messages_event(output) if output_event is not None: span.add_event("gen_ai.output.messages", attributes=output_event) record_token_usage(output.usage, metric_attributes) - record_operation_duration(duration_seconds, metric_attributes, error=error) + + +@contextmanager +def gen_ai_span( + *, + model: Model, + provider: Provider | None, + protocol: Protocol | None, + modality: Modality, + extra_attributes: dict[str, Any] | None = None, +) -> Iterator[tuple[Any, dict[str, Any]]]: + """Open a GenAI span and record operation duration on exit. + + Yields ``(span, request_attrs)``. On any exception, the duration is recorded + with ``error.type`` set; on success it's recorded plain. + """ + request_attrs = request_attributes( + model=model, provider=provider, protocol=protocol, modality=modality + ) + span_attrs = ( + {**request_attrs, **extra_attributes} if extra_attributes else request_attrs + ) + started = time.monotonic() + error: BaseException | None = None + with tracer.start_as_current_span( + span_name(modality, model), attributes=span_attrs + ) as span: + try: + yield span, request_attrs + except BaseException as exc: + error = exc + raise + finally: + record_operation_duration( + time.monotonic() - started, request_attrs, error=error + ) class _TracedStream: @@ -484,19 +517,17 @@ def _finalize(self) -> None: if self._ended: return self._ended = True - duration = time.monotonic() - self._started + record_operation_duration( + time.monotonic() - self._started, + self._metric_attributes, + error=self._error, + ) try: output = self._inner.output except StreamNotExhaustedError: output = None if output is not None: - record_output( - self._span, output, self._metric_attributes, duration, error=self._error - ) - else: - record_operation_duration( - duration, self._metric_attributes, error=self._error - ) + record_output(self._span, output, self._metric_attributes) self._span.end() @@ -513,6 +544,7 @@ def trace_stream( "Status", "StatusCode", "add_input_event", + "gen_ai_span", "meter", "output_attributes", "record_operation_duration",