From 96c5db5a07f7f851751ccd68f176dad1634885cb Mon Sep 17 00:00:00 2001 From: Sasha Sobran Date: Wed, 7 Jan 2026 07:37:38 -0800 Subject: [PATCH] fix: use mode='json' in model_dump to serialize bytes correctly when using telemetry fixes #4043 Co-authored-by: Sasha Sobran PiperOrigin-RevId: 853256045 --- src/google/adk/telemetry/tracing.py | 6 +- tests/unittests/telemetry/test_spans.py | 118 ++++++++++++++++++++---- 2 files changed, 101 insertions(+), 23 deletions(-) diff --git a/src/google/adk/telemetry/tracing.py b/src/google/adk/telemetry/tracing.py index f03cdc80..386ae3b4 100644 --- a/src/google/adk/telemetry/tracing.py +++ b/src/google/adk/telemetry/tracing.py @@ -340,7 +340,7 @@ def trace_send_data( 'gcp.vertex.agent.data', _safe_json_serialize([ types.Content(role=content.role, parts=content.parts).model_dump( - exclude_none=True + exclude_none=True, mode='json' ) for content in data ]), @@ -366,7 +366,7 @@ def _build_llm_request_for_trace(llm_request: LlmRequest) -> dict[str, Any]: result = { 'model': llm_request.model, 'config': llm_request.config.model_dump( - exclude_none=True, exclude='response_schema' + exclude_none=True, exclude='response_schema', mode='json' ), 'contents': [], } @@ -375,7 +375,7 @@ def _build_llm_request_for_trace(llm_request: LlmRequest) -> dict[str, Any]: parts = [part for part in content.parts if not part.inline_data] result['contents'].append( types.Content(role=content.role, parts=parts).model_dump( - exclude_none=True + exclude_none=True, mode='json' ) ) return result diff --git a/tests/unittests/telemetry/test_spans.py b/tests/unittests/telemetry/test_spans.py index 38a8358f..c87730a5 100644 --- a/tests/unittests/telemetry/test_spans.py +++ b/tests/unittests/telemetry/test_spans.py @@ -13,7 +13,6 @@ # limitations under the License. import json -import os from typing import Any from typing import Dict from typing import Optional @@ -207,18 +206,87 @@ async def test_trace_call_llm_with_binary_content( assert mock_span_fixture.set_attribute.call_count == 7 mock_span_fixture.set_attribute.assert_has_calls(expected_calls) - # Verify binary content is replaced with '' in JSON + # Verify binary values are properly serialized as base64 llm_request_json_str = None for call_obj in mock_span_fixture.set_attribute.call_args_list: - if call_obj.args[0] == 'gcp.vertex.agent.llm_request': - llm_request_json_str = call_obj.args[1] + arg_name, arg_value = call_obj.args + if arg_name == 'gcp.vertex.agent.llm_request': + llm_request_json_str = arg_value + break + + assert llm_request_json_str is not None + + # Verify bytes are base64 encoded (b'test_data' -> 'dGVzdF9kYXRh') + assert 'dGVzdF9kYXRh' in llm_request_json_str + + # Verify no serialization failures + assert '' not in llm_request_json_str + + +@pytest.mark.asyncio +async def test_trace_call_llm_with_thought_signature( + monkeypatch, mock_span_fixture +): + """Test trace_call_llm handles thought_signature bytes correctly. + + This test verifies that thought_signature bytes from Gemini 3.0 models + are properly serialized as base64 in telemetry traces. + """ + monkeypatch.setattr( + 'opentelemetry.trace.get_current_span', lambda: mock_span_fixture + ) + + agent = LlmAgent(name='test_agent') + invocation_context = await _create_invocation_context(agent) + + # multi-turn conversation where the model's response contains + # thought_signature bytes + thought_signature_bytes = b'thought_signature' + llm_request = LlmRequest( + model='gemini-3-pro-preview', + contents=[ + types.Content( + role='user', + parts=[types.Part(text='Hello')], + ), + types.Content( + role='model', + parts=[ + types.Part( + thought=True, + thought_signature=thought_signature_bytes, + ) + ], + ), + types.Content( + role='user', + parts=[types.Part(text='Follow up question')], + ), + ], + config=types.GenerateContentConfig(), + ) + llm_response = LlmResponse(turn_complete=True) + + # should not raise TypeError for bytes serialization + trace_call_llm(invocation_context, 'test_event_id', llm_request, llm_response) + + llm_request_json_str = None + for call_obj in mock_span_fixture.set_attribute.call_args_list: + arg_name, arg_value = call_obj.args + if arg_name == 'gcp.vertex.agent.llm_request': + llm_request_json_str = arg_value break assert ( llm_request_json_str is not None ), "Attribute 'gcp.vertex.agent.llm_request' was not set on the span." - assert llm_request_json_str.count('') == 2 + # no serialization failures + assert '' not in llm_request_json_str + # llm request is valid JSON + parsed = json.loads(llm_request_json_str) + assert parsed['model'] == 'gemini-3-pro-preview' + assert len(parsed['contents']) == 3 def test_trace_tool_call_with_scalar_response( @@ -407,15 +475,19 @@ async def test_call_llm_disabling_request_response_content( # Assert assert not any( - call_obj.args[0] == 'gcp.vertex.agent.llm_request' - and call_obj.args[1] != {} - for call_obj in mock_span_fixture.set_attribute.call_args_list + arg_name == 'gcp.vertex.agent.llm_request' and arg_value != {} + for arg_name, arg_value in ( + call_obj.args + for call_obj in mock_span_fixture.set_attribute.call_args_list + ) ), "Attribute 'gcp.vertex.agent.llm_request' was incorrectly set on the span." assert not any( - call_obj.args[0] == 'gcp.vertex.agent.llm_response' - and call_obj.args[1] != {} - for call_obj in mock_span_fixture.set_attribute.call_args_list + arg_name == 'gcp.vertex.agent.llm_response' and arg_value != {} + for arg_name, arg_value in ( + call_obj.args + for call_obj in mock_span_fixture.set_attribute.call_args_list + ) ), ( "Attribute 'gcp.vertex.agent.llm_response' was incorrectly set on the" ' span.' @@ -466,18 +538,22 @@ def test_trace_tool_call_disabling_request_response_content( # Assert assert not any( - call_obj.args[0] == 'gcp.vertex.agent.tool_call_args' - and call_obj.args[1] != {} - for call_obj in mock_span_fixture.set_attribute.call_args_list + arg_name == 'gcp.vertex.agent.tool_call_args' and arg_value != {} + for arg_name, arg_value in ( + call_obj.args + for call_obj in mock_span_fixture.set_attribute.call_args_list + ) ), ( "Attribute 'gcp.vertex.agent.tool_call_args' was incorrectly set on the" ' span.' ) assert not any( - call_obj.args[0] == 'gcp.vertex.agent.tool_response' - and call_obj.args[1] != {} - for call_obj in mock_span_fixture.set_attribute.call_args_list + arg_name == 'gcp.vertex.agent.tool_response' and arg_value != {} + for arg_name, arg_value in ( + call_obj.args + for call_obj in mock_span_fixture.set_attribute.call_args_list + ) ), ( "Attribute 'gcp.vertex.agent.tool_response' was incorrectly set on the" ' span.' @@ -510,9 +586,11 @@ def test_trace_merged_tool_disabling_request_response_content( # Assert assert not any( - call_obj.args[0] == 'gcp.vertex.agent.tool_response' - and call_obj.args[1] != {} - for call_obj in mock_span_fixture.set_attribute.call_args_list + arg_name == 'gcp.vertex.agent.tool_response' and arg_value != {} + for arg_name, arg_value in ( + call_obj.args + for call_obj in mock_span_fixture.set_attribute.call_args_list + ) ), ( "Attribute 'gcp.vertex.agent.tool_response' was incorrectly set on the" ' span.'