diff --git a/src/celeste/telemetry.py b/src/celeste/telemetry.py index c6ce838..a4fb07e 100644 --- a/src/celeste/telemetry.py +++ b/src/celeste/telemetry.py @@ -224,6 +224,7 @@ def output_attributes(output: Output[Any]) -> dict[str, Any]: 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,) + attrs["gen_ai.response.model"] = output.metadata["model"] return attrs diff --git a/tests/unit_tests/_telemetry_helpers.py b/tests/unit_tests/_telemetry_helpers.py index 56f1c6f..451c88d 100644 --- a/tests/unit_tests/_telemetry_helpers.py +++ b/tests/unit_tests/_telemetry_helpers.py @@ -1,7 +1,7 @@ """Shared fixtures and helpers for telemetry-related unit tests.""" from collections.abc import AsyncIterator -from typing import Any, ClassVar +from typing import Any, ClassVar, Unpack from celeste.io import Chunk, Output, Usage from celeste.parameters import Parameters @@ -30,6 +30,19 @@ class TelemetryStream(Stream[TelemetryOutput, Parameters, Chunk]): _output_class: ClassVar[type[Output]] = TelemetryOutput _empty_content: ClassVar[str] = "" + def __init__( + self, + sse_iterator: AsyncIterator[dict[str, Any]], + stream_metadata: dict[str, Any] | None = None, + **parameters: Unpack[Parameters], + ) -> None: + """Default `stream_metadata={"model": "test-model"}` to mirror production wiring.""" + super().__init__( + sse_iterator, + stream_metadata=stream_metadata or {"model": "test-model"}, + **parameters, + ) + def _aggregate_content(self, chunks: list[Chunk]) -> str: """Concatenate chunk content.""" return "".join(chunk.content for chunk in chunks) diff --git a/tests/unit_tests/test_telemetry_streaming.py b/tests/unit_tests/test_telemetry_streaming.py index fa91081..ee95a01 100644 --- a/tests/unit_tests/test_telemetry_streaming.py +++ b/tests/unit_tests/test_telemetry_streaming.py @@ -16,6 +16,7 @@ from tests.unit_tests._telemetry_helpers import ( TelemetryOutput, TelemetryStream, + TelemetryUsage, async_iter, ) from tests.unit_tests.conftest import start_test_span @@ -213,10 +214,24 @@ 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: Output[Any] = TelemetryOutput(content="", usage=usage) + output: Output[Any] = TelemetryOutput( + content="", usage=usage, metadata={"model": "test-model"} + ) 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 + + def test_response_model_emitted_from_metadata(self) -> None: + """`gen_ai.response.model` reads from `output.metadata["model"]`.""" + output: Output[Any] = TelemetryOutput( + content="", + usage=TelemetryUsage(), + metadata={"model": "claude-opus-4-1-20250805"}, + ) + + attrs = telemetry.output_attributes(output) + + assert attrs["gen_ai.response.model"] == "claude-opus-4-1-20250805"