diff --git a/src/google/adk/telemetry/_instrumentation.py b/src/google/adk/telemetry/_instrumentation.py index 270cd732cc9..36cb5b70cfa 100644 --- a/src/google/adk/telemetry/_instrumentation.py +++ b/src/google/adk/telemetry/_instrumentation.py @@ -611,6 +611,16 @@ def record_llm_response( exit_stack.close() self._inference_span_ended = True + def _record_unfinished_stream( + self, invocation_context: InvocationContext + ) -> None: + """Records the chunks of a stream cut short before its aggregate.""" + if self._inference_span_ended: + return + for response in self._llm_responses: + if response.partial: + tracing.trace_inference_result(invocation_context, self.span, response) + def _record_chunk_timing(self, response: LlmResponse) -> None: """Records the wait for ``response``, if it is a chunk being timed. @@ -1139,7 +1149,10 @@ async def record_inference_telemetry( else: # A span that is not recorded has no start time. tel_ctx.span_started_ns = time.time_ns() - yield tel_ctx + try: + yield tel_ctx + finally: + tel_ctx._record_unfinished_stream(invocation_context) # pylint: disable=protected-access finally: context_api.detach(token) inference_error = sys.exc_info()[1] diff --git a/tests/unittests/telemetry/test_instrumentation.py b/tests/unittests/telemetry/test_instrumentation.py index 381fae3ec98..8d69aacf73f 100644 --- a/tests/unittests/telemetry/test_instrumentation.py +++ b/tests/unittests/telemetry/test_instrumentation.py @@ -1793,6 +1793,42 @@ async def test_chunk_timing_is_left_to_the_genai_instrumentor( assert not telemetry.points(_TIME_PER_OUTPUT_CHUNK) +@pytest.mark.asyncio +async def test_stream_cut_short_before_its_aggregate_records_its_chunks( + telemetry: _Telemetry, monkeypatch: pytest.MonkeyPatch +): + """Chunks wait for the aggregate; without one they are all there is.""" + monkeypatch.setenv( + "OTEL_SEMCONV_STABILITY_OPT_IN", "gen_ai_latest_experimental" + ) + ctx = await _invocation_context(_agent()) + model_response_event = mock.MagicMock() + model_response_event.id = "event-1" + + with pytest.raises(RuntimeError): + async with _instrumentation.record_inference_telemetry( + LlmRequest(model="some-model"), ctx, model_response_event + ) as tel_ctx: + for text in ("text ", "response"): + tel_ctx.record_llm_response( + ctx, + _llm_response( + content=types.Content( + role="model", parts=[types.Part(text=text)] + ), + partial=True, + finish_reason=None, + ), + ) + raise RuntimeError("stream interrupted") + + messages = tel_ctx.span.operation_details_attributes["gen_ai.output.messages"] + assert [part["content"] for part in messages[0]["parts"]] == [ + "text ", + "response", + ] + + # --- record_invocation -----------------------------------------------------