From 2c5773f0fa32f35559a95efd0954bebf415fa673 Mon Sep 17 00:00:00 2001 From: Ishan Katoch <98383932+ishan-1010@users.noreply.github.com> Date: Thu, 22 Oct 2026 20:56:53 +0530 Subject: [PATCH] fix(telemetry): record a LiteLlm streamed answer once, not chunks plus aggregate LiteLlm streams chunks without a finish reason and then yields their aggregate with one. record_llm_response kept recording until that finish reason, so gen_ai.output.messages held every chunk and then the whole answer again (a 5-chunk answer was 6 messages, a streamed tool call 4 half-built calls plus the real one). When a complete response with a finish reason follows only partial chunks without one, drop the chunks' output messages and keep the aggregate. Streams whose last chunk carries the finish reason (Gemini) and streams that fail before their aggregate keep their chunks. Fixes #7630 --- src/google/adk/telemetry/_instrumentation.py | 19 +++ .../telemetry/test_instrumentation.py | 157 ++++++++++++++++++ 2 files changed, 176 insertions(+) diff --git a/src/google/adk/telemetry/_instrumentation.py b/src/google/adk/telemetry/_instrumentation.py index 5adb3099917..b07a5242e33 100644 --- a/src/google/adk/telemetry/_instrumentation.py +++ b/src/google/adk/telemetry/_instrumentation.py @@ -29,6 +29,7 @@ from opentelemetry import trace import opentelemetry.context as context_api +from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import GEN_AI_OUTPUT_MESSAGES from typing_extensions import assert_never from . import _adk_attributes @@ -574,6 +575,24 @@ def record_llm_response( if self._inference_span_ended: return + # A model that streams chunks without a finish reason and then yields the + # aggregate of those chunks with one (LiteLlm does) would report the + # answer twice. Keep only the aggregate. A stream whose last chunk carries + # the finish reason never reaches its aggregate, and one cut off before + # its aggregate keeps its chunks. + if ( + isinstance(self.span, tracing.GenerateContentSpan) + and not response.partial + and is_reported_finish_reason(response.finish_reason) + and len(self._llm_responses) > 1 + and all( + earlier.partial + and not is_reported_finish_reason(earlier.finish_reason) + for earlier in self._llm_responses[:-1] + ) + ): + self.span.operation_details_attributes.pop(GEN_AI_OUTPUT_MESSAGES, None) + tracing.trace_inference_result(invocation_context, self.span, response) # The inference ends at the response that says why generation stopped. End # the span there, before the response is handed back to the caller, so what diff --git a/tests/unittests/telemetry/test_instrumentation.py b/tests/unittests/telemetry/test_instrumentation.py index fccacb568da..2561279e288 100644 --- a/tests/unittests/telemetry/test_instrumentation.py +++ b/tests/unittests/telemetry/test_instrumentation.py @@ -2673,3 +2673,160 @@ async def test_a_run_nobody_is_collecting_for_still_works(): assert _about(busy.tool_seconds(), _SLICE) assert busy.model_seconds() == 0.0 + + +# --- a stream whose aggregate repeats its chunks (LiteLlm) ------------------ + + +def _text_chunk(text: str) -> LlmResponse: + return LlmResponse( + content=types.Content(role="model", parts=[types.Part(text=text)]), + partial=True, + ) + + +async def _recorded_output_messages( + monkeypatch: pytest.MonkeyPatch, responses, error: Exception | None = None +) -> list: + monkeypatch.setenv( + "OTEL_SEMCONV_STABILITY_OPT_IN", "gen_ai_latest_experimental" + ) + monkeypatch.setenv( + "OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT", "SPAN_ONLY" + ) + agent = _agent() + ctx = await _invocation_context(agent) + llm_request = LlmRequest( + model="some-model", + contents=[types.Content(role="user", parts=[types.Part(text="hi")])], + ) + model_response_event = mock.MagicMock() + model_response_event.id = "event-1" + recorded = [] + try: + async with _instrumentation.record_inference_telemetry( + llm_request, ctx, model_response_event + ) as tel_ctx: + for response in responses: + tel_ctx.record_llm_response(ctx, response) + recorded = tel_ctx.span.operation_details_attributes.get( + "gen_ai.output.messages", [] + ) + if error is not None: + raise error + except type(error) if error is not None else (): + pass + return recorded + + +@pytest.mark.asyncio +async def test_aggregate_after_unfinished_chunks_is_recorded_once( + telemetry: _Telemetry, monkeypatch: pytest.MonkeyPatch +): + """The aggregate repeated every chunk, so the answer appeared twice.""" + del telemetry + messages = await _recorded_output_messages( + monkeypatch, + [ + _text_chunk("Oslo is "), + _text_chunk("4C and cloudy."), + _llm_response( + content=types.Content( + role="model", + parts=[types.Part(text="Oslo is 4C and cloudy.")], + ), + partial=False, + ), + ], + ) + + assert [m["parts"][0]["content"] for m in messages] == [ + "Oslo is 4C and cloudy." + ] + assert messages[0]["finish_reason"] == "stop" + + +@pytest.mark.asyncio +async def test_tool_call_aggregate_replaces_its_partial_call_chunks( + telemetry: _Telemetry, monkeypatch: pytest.MonkeyPatch +): + del telemetry + call = types.FunctionCall( + id="call-1", name="get_weather", args={"city": "Oslo"} + ) + messages = await _recorded_output_messages( + monkeypatch, + [ + LlmResponse( + content=types.Content( + role="model", + parts=[ + types.Part( + function_call=types.FunctionCall( + id="call-1", + name="get_weather", + will_continue=True, + ) + ) + ], + ), + partial=True, + ), + _llm_response( + content=types.Content( + role="model", parts=[types.Part(function_call=call)] + ), + partial=False, + ), + ], + ) + + assert len(messages) == 1 + assert messages[0]["parts"][0]["type"] == "tool_call" + assert messages[0]["parts"][0]["arguments"] == {"city": "Oslo"} + + +@pytest.mark.asyncio +async def test_stream_cut_off_before_its_aggregate_keeps_its_chunks( + telemetry: _Telemetry, monkeypatch: pytest.MonkeyPatch +): + del telemetry + messages = await _recorded_output_messages( + monkeypatch, + [_text_chunk("Oslo is "), _text_chunk("4C")], + error=ConnectionError("stream dropped"), + ) + + assert [m["parts"][0]["content"] for m in messages] == ["Oslo is ", "4C"] + + +@pytest.mark.asyncio +async def test_stream_whose_last_chunk_finishes_keeps_its_chunks( + telemetry: _Telemetry, monkeypatch: pytest.MonkeyPatch +): + """Gemini's last chunk carries the finish reason; its aggregate is unseen.""" + del telemetry + messages = await _recorded_output_messages( + monkeypatch, + [ + _text_chunk("Oslo is "), + _llm_response( + content=types.Content( + role="model", parts=[types.Part(text="4C and cloudy.")] + ), + partial=True, + ), + _llm_response( + content=types.Content( + role="model", + parts=[types.Part(text="Oslo is 4C and cloudy.")], + ), + partial=False, + ), + ], + ) + + assert [m["parts"][0]["content"] for m in messages] == [ + "Oslo is ", + "4C and cloudy.", + ]