diff --git a/src/celeste/client.py b/src/celeste/client.py index fbc6621..fa26508 100644 --- a/src/celeste/client.py +++ b/src/celeste/client.py @@ -213,16 +213,14 @@ async def _predict( Returns: Output of the parameterized type. """ - 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, - ), - ) as span: + with telemetry.gen_ai_span( + model=self.model, + provider=self.provider, + protocol=self.protocol, + modality=self.modality, + ) 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 ) @@ -249,7 +247,7 @@ async def _predict( tool_calls=tool_calls, **kwargs, ) - span.set_attributes(telemetry.output_attributes(output)) + telemetry.record_output(span, output, request_attrs) return output def _parse_tool_calls(self, response_data: dict[str, Any]) -> list[ToolCall]: @@ -297,18 +295,17 @@ 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}, ) + telemetry.add_input_event(span, inputs) sse_iterator = self._make_stream_request( request_body, endpoint=endpoint, @@ -326,7 +323,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..c6ce838 100644 --- a/src/celeste/telemetry.py +++ b/src/celeste/telemetry.py @@ -1,17 +1,24 @@ """OpenTelemetry GenAI telemetry for Celeste.""" import asyncio +import json +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 celeste.core import Modality, Protocol, Provider +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 -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 +53,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 +92,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 +119,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 +187,246 @@ def span_name(modality: Modality, model: Model) -> str: return f"celeste.{modality.value} {model.id}" +# `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", + 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", +} + +# `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", + UsageField.REASONING_TOKENS: "reasoning", + UsageField.CACHED_TOKENS: "cached_input", +} + + +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] = {} - 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 + 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: 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.""" + for field_name, value in _iter_usage_numeric(usage): + token_type = _GEN_AI_TOKEN_TYPES.get(field_name) + if token_type is None: + 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) + + +# Opt-in flag, read once at import (semconv-standard env name). +_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, or None when capture is off.""" + 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, or None when capture is off.""" + 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) + } + + +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 record_output( + span: Any, + output: Output[Any], + metric_attributes: dict[str, Any], +) -> None: + """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) + + +@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: """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 +441,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() @@ -210,8 +469,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.""" @@ -242,28 +513,43 @@ 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 + 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: - self._span.set_attributes(output_attributes(output)) + record_output(self._span, output, self._metric_attributes) 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", + "add_input_event", + "gen_ai_span", + "meter", "output_attributes", + "record_operation_duration", + "record_output", + "record_token_usage", "request_attributes", "span_name", "trace_stream", 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 new file mode 100644 index 0000000..1d7cd6b --- /dev/null +++ b/tests/unit_tests/test_telemetry_content_events.py @@ -0,0 +1,204 @@ +"""Tests for opt-in GenAI content events (`gen_ai.input.messages` / `output.messages`).""" + +import json + +import pytest +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( + InMemorySpanExporter, +) + +from celeste import telemetry +from celeste.artifacts import ImageArtifact +from celeste.mime_types import ImageMimeType +from celeste.modalities.text.io import TextInput, TextOutput, TextUsage +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 +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) + + +class TestInputMessagesEvent: + def test_capture_disabled_returns_none(self, capture_disabled: None) -> None: + """With the env flag off, the helper returns None — no event emitted.""" + 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.""" + result = telemetry._input_messages_event(TextInput(prompt="hello")) + + assert result is not None + 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.""" + result = telemetry._input_messages_event( + TextInput( + messages=[ + Message(role=Role.SYSTEM, content="be helpful"), + Message(role=Role.USER, content="hi"), + ] + ) + ) + + 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.""" + result = telemetry._input_messages_event( + TextInput( + prompt="describe this", + image=ImageArtifact( + url="https://example.com/img.png", mime_type=ImageMimeType.PNG + ), + ) + ) + + assert result is not None + 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 + assert image_parts[0]["uri"] == "https://example.com/img.png" + assert image_parts[0]["mime_type"] == "image/png" + + +class TestOutputMessagesEvent: + def test_capture_disabled_returns_none(self, capture_disabled: None) -> None: + """With the env flag off, the helper returns 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.""" + result = telemetry._output_messages_event( + TextOutput(content="reply", usage=TextUsage()) + ) + + assert result is not None + 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.""" + result = telemetry._output_messages_event( + TextOutput(content="answer", usage=TextUsage(), reasoning="thinking...") + ) + + assert result is not None + parts = json.loads(result["messages"])[0]["parts"] + 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.""" + result = telemetry._output_messages_event( + TextOutput( + content="", + usage=TextUsage(), + tool_calls=[ + ToolCall( + id="call_1", name="get_weather", arguments={"city": "Paris"} + ) + ], + ) + ) + + assert result is not None + 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" + 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.""" + result = telemetry._output_messages_event( + TextOutput( + content=ImageArtifact( + url="https://example.com/out.png", mime_type=ImageMimeType.PNG + ), + usage=TextUsage(), + ) + ) + + assert result is not None + parts = json.loads(result["messages"])[0]["parts"] + assert parts[0] == { + "type": "image", + "uri": "https://example.com/out.png", + "mime_type": "image/png", + } + + +class TestStreamingContentEvents: + 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}}, + ] + 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 + assert "gen_ai.output.messages" in [e.name for e in finished[0].events] + + 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 + 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 + 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 new file mode 100644 index 0000000..2b407e2 --- /dev/null +++ b/tests/unit_tests/test_telemetry_metrics.py @@ -0,0 +1,130 @@ +"""Tests for GenAI metric histogram emission.""" + +from typing import Any + +import pytest +from opentelemetry.sdk.metrics import MeterProvider +from opentelemetry.sdk.metrics.export import InMemoryMetricReader + +from celeste import telemetry +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() + 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}"), + ) + 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 + + +@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 2026563..fa91081 100644 --- a/tests/unit_tests/test_telemetry_streaming.py +++ b/tests/unit_tests/test_telemetry_streaming.py @@ -1,61 +1,31 @@ """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 input/output token fields for telemetry assertions.""" - - input_tokens: int | None = None - output_tokens: int | 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) +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 - 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) +class _OffSpecUsage(Usage): + """Modality-specific usage with fields outside the GenAI semconv set.""" -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 + images_generated: int | None = None + audio_seconds: float | None = None async def _failing_iter() -> AsyncIterator[dict[str, Any]]: @@ -64,23 +34,7 @@ async def _failing_iter() -> AsyncIterator[dict[str, Any]]: 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: @@ -90,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] @@ -109,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: @@ -128,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" @@ -156,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: @@ -174,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 @@ -199,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 @@ -234,3 +174,49 @@ 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, + }, + } + ] + wrapped = telemetry.trace_stream( + TelemetryStream(async_iter(events)), start_test_span(provider) + ) + + async for _ in wrapped: + pass + + 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 + 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: Output[Any] = TelemetryOutput(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