fix: use mode='json' in model_dump to serialize bytes correctly when using telemetry

fixes #4043

Co-authored-by: Sasha Sobran <asobran@google.com>
PiperOrigin-RevId: 853256045

Author: sasha-gitg

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

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 '<not serializable>' 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 serializable>' 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('<not serializable>') == 2
+  # no serialization failures
+  assert '<not serializable>' 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.'