From 49b18a3a16651538ae97a78b2578958e5b88f692 Mon Sep 17 00:00:00 2001 From: kamilbenkirane Date: Tue, 5 May 2026 11:15:04 +0200 Subject: [PATCH] feat(telemetry): emit gen_ai.usage attrs on streaming spans MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Move telemetry wrap site from the upstream SSE iterator to the constructed Stream. _TracedStream delegates the public Stream surface (__aiter__/__anext__/aclose/output/__aenter__/__aexit__/__iter__/ __enter__/__exit__/__repr__) to the inner stream and owns span lifecycle so output_attributes can read Stream.output after _parse_output runs on natural exhaustion — same call _predict already makes for non-streaming. Streaming spans now carry gen_ai.usage.input_tokens and gen_ai.usage.output_tokens, matching the non-streaming path. Zero changes to streaming.py. Add opentelemetry-sdk as a dev dep for InMemorySpanExporter-based unit tests. --- pyproject.toml | 1 + src/celeste/client.py | 4 +- src/celeste/telemetry.py | 148 ++++++++++-- tests/unit_tests/test_telemetry_streaming.py | 236 +++++++++++++++++++ 4 files changed, 364 insertions(+), 25 deletions(-) create mode 100644 tests/unit_tests/test_telemetry_streaming.py diff --git a/pyproject.toml b/pyproject.toml index f12b87d..177e1aa 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -50,6 +50,7 @@ dev = [ "bandit[toml]>=1.7.5", "pre-commit>=3.5.0", "pytest-rerunfailures>=16.1", + "opentelemetry-sdk>=1.30", ] [build-system] diff --git a/src/celeste/client.py b/src/celeste/client.py index 0f8f5af..fbc6621 100644 --- a/src/celeste/client.py +++ b/src/celeste/client.py @@ -316,8 +316,7 @@ def _stream( **parameters, ) sse_iterator = enrich_stream_errors(sse_iterator, self._handle_error_response) - sse_iterator = telemetry.trace_stream(sse_iterator, span) - return stream_class( + stream = stream_class( sse_iterator, transform_output=self._transform_output, stream_metadata={ @@ -327,6 +326,7 @@ def _stream( }, **parameters, ) + return telemetry.trace_stream(stream, span) # type: ignore[return-value] @classmethod @abstractmethod diff --git a/src/celeste/telemetry.py b/src/celeste/telemetry.py index 8b60abc..7bfd363 100644 --- a/src/celeste/telemetry.py +++ b/src/celeste/telemetry.py @@ -1,13 +1,17 @@ """OpenTelemetry GenAI telemetry for Celeste.""" +import asyncio import time -from collections.abc import AsyncIterator, Iterator +from collections.abc import Iterator from contextlib import contextmanager +from types import TracebackType from typing import Any from celeste.core import Modality, Protocol, Provider +from celeste.exceptions import StreamNotExhaustedError from celeste.io import Output from celeste.models import Model +from celeste.streaming import Stream _PROVIDER_NAME_MAP: dict[Provider, str] = { Provider.OPENAI: "openai", @@ -65,14 +69,33 @@ def _noop_use_span(span: Any, end_on_exit: bool = False) -> Iterator[Any]: yield span +class _NoOpStatus: + """Status stand-in used when ``opentelemetry-api`` is not installed.""" + + def __init__(self, *args: Any, **kwargs: Any) -> None: + """Discard the status arguments.""" + + +class _NoOpStatusCode: + """StatusCode stand-in used when ``opentelemetry-api`` is not installed.""" + + ERROR = "ERROR" + + try: from opentelemetry import trace as _otel_trace + from opentelemetry.trace import Status as _OtelStatus + from opentelemetry.trace import StatusCode as _OtelStatusCode tracer: Any = _otel_trace.get_tracer("celeste") use_span: Any = _otel_trace.use_span + Status: Any = _OtelStatus + StatusCode: Any = _OtelStatusCode except ImportError: tracer = _NoOpTracer() use_span = _noop_use_span + Status = _NoOpStatus + StatusCode = _NoOpStatusCode def request_attributes( @@ -136,31 +159,110 @@ def output_attributes(output: Output[Any]) -> dict[str, Any]: return attrs -async def trace_stream( - sse_iterator: AsyncIterator[dict[str, Any]], - span: Any, -) -> AsyncIterator[dict[str, Any]]: - """Wrap an SSE iterator to record TTFC and end the span at iteration's end. - - OTel's ``use_span(span, end_on_exit=True)`` context manager makes the - span current, auto-records exceptions, sets ERROR status on exception, - and ends the span on exit (including ``GeneratorExit`` from consumer - abandonment). No try/except needed. - """ - with use_span(span, end_on_exit=True): - started = time.monotonic() - seen_first = False - async for event in sse_iterator: - if not seen_first: - seen_first = True - span.set_attribute( - "gen_ai.response.time_to_first_chunk", - time.monotonic() - started, - ) - yield event +class _TracedStream: + """Stream-shaped wrapper that emits GenAI telemetry against an OTel span.""" + + def __init__(self, inner: Stream[Any, Any, Any], span: Any) -> None: + """Initialize wrapper around a constructed Stream and a detached span.""" + self._inner = inner + self._span = span + self._started = time.monotonic() + self._seen_first = False + self._ended = False + + def __aiter__(self) -> "_TracedStream": + """Return self as async iterator.""" + return self + + async def __anext__(self) -> Any: + """Yield next chunk; emit TTFC on first, finalize span on terminal events.""" + with use_span(self._span, end_on_exit=False): + try: + chunk = await self._inner.__anext__() + except (StopAsyncIteration, asyncio.CancelledError): + self._finalize() + raise + except Exception as exc: + self._span.record_exception(exc) + self._span.set_status(Status(StatusCode.ERROR, str(exc))) + self._finalize() + raise + if not self._seen_first: + self._seen_first = True + self._span.set_attribute( + "gen_ai.response.time_to_first_chunk", + time.monotonic() - self._started, + ) + return chunk + + async def __aenter__(self) -> "_TracedStream": + """Enter async context — return self for iteration.""" + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc_val: BaseException | None, + exc_tb: TracebackType | None, + ) -> bool: + """Exit async context — close stream and finalize span.""" + await self.aclose() + return False + + def __iter__(self) -> Iterator[Any]: + """Delegate sync iteration to the inner Stream.""" + return iter(self._inner) + + def __enter__(self) -> "_TracedStream": + """Enter sync context — delegate to inner Stream.""" + self._inner.__enter__() + return self + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc_val: BaseException | None, + exc_tb: TracebackType | None, + ) -> None: + """Exit sync context — delegate to inner Stream.""" + self._inner.__exit__(exc_type, exc_val, exc_tb) + + def __repr__(self) -> str: + """Developer-friendly representation showing inner stream state.""" + return repr(self._inner) + + @property + def output(self) -> Any: + """Access final Output after stream exhaustion.""" + return self._inner.output + + async def aclose(self) -> None: + """Close the inner stream and finalize the span.""" + await self._inner.aclose() + self._finalize() + + def _finalize(self) -> None: + """End the span exactly once; emit usage attrs if Output is populated.""" + if self._ended: + return + self._ended = True + try: + output = self._inner.output + except StreamNotExhaustedError: + output = None + if output is not None: + self._span.set_attributes(output_attributes(output)) + self._span.end() + + +def trace_stream(stream: Stream[Any, Any, Any], span: Any) -> _TracedStream: + """Wrap a Stream to emit GenAI telemetry against ``span``.""" + return _TracedStream(stream, span) __all__ = [ + "Status", + "StatusCode", "output_attributes", "request_attributes", "span_name", diff --git a/tests/unit_tests/test_telemetry_streaming.py b/tests/unit_tests/test_telemetry_streaming.py new file mode 100644 index 0000000..2026563 --- /dev/null +++ b/tests/unit_tests/test_telemetry_streaming.py @@ -0,0 +1,236 @@ +"""Tests for `_TracedStream` — span lifecycle and GenAI attribute emission.""" + +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, 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) + + 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: + """After `async for` to completion, span carries gen_ai.usage.*_tokens.""" + in_memory, provider = exporter + events = [ + {"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) + + chunks = [chunk async for chunk in wrapped] + + assert len(chunks) == 2 + finished = in_memory.get_finished_spans() + assert len(finished) == 1 + attrs = finished[0].attributes or {} + assert attrs["gen_ai.usage.input_tokens"] == 12 + assert attrs["gen_ai.usage.output_tokens"] == 34 + assert "gen_ai.response.time_to_first_chunk" in attrs + + async def test_async_context_manager_emits_usage( + self, exporter: tuple[InMemorySpanExporter, TracerProvider] + ) -> 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) + + async with wrapped as ctx: + async for _ in ctx: + pass + + finished = in_memory.get_finished_spans() + assert len(finished) == 1 + attrs = finished[0].attributes or {} + assert attrs["gen_ai.usage.input_tokens"] == 1 + assert attrs["gen_ai.usage.output_tokens"] == 2 + + +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) + + chunk = await wrapped.__anext__() + assert chunk.content == "Partial" + await wrapped.aclose() + + finished = in_memory.get_finished_spans() + assert len(finished) == 1 + attrs = finished[0].attributes or {} + assert "gen_ai.usage.input_tokens" not in attrs + assert "gen_ai.usage.output_tokens" not in attrs + assert finished[0].status.status_code != StatusCode.ERROR + + +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) + + with pytest.raises(RuntimeError, match="boom"): + async for _ in wrapped: + pass + + 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) + + +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) + + with pytest.raises(StreamNotExhaustedError): + _ = wrapped.output + + async def test_output_returns_typed_output_after_exhaustion( + self, exporter: tuple[InMemorySpanExporter, TracerProvider] + ) -> None: + """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) + + async for _ in wrapped: + pass + + output = wrapped.output + assert isinstance(output, _TestOutput) + 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) + + async for _ in wrapped: + pass + await wrapped.aclose() + + finished = in_memory.get_finished_spans() + assert len(finished) == 1 + attrs = finished[0].attributes or {} + assert attrs["gen_ai.usage.input_tokens"] == 2 + assert attrs["gen_ai.usage.output_tokens"] == 3