From 03fe909ada7f61659c0ca4f6ae4fa70bd18ab512 Mon Sep 17 00:00:00 2001
From: Stephen Smith <stephen.smith@newfront.com>
Date: Tue, 27 May 2025 13:21:10 -0700
Subject: [PATCH] Copybara import of the project:

--
b781880d9bfb9786bd5e50314eaedc441fc2a93e by Stephen Smith <stephen.smith@newfront.com>:

Telemetry unit test for non-serializable data.

--
179da9db997bb3f992e126c9c64193ff7df67b3d by Stephen Smith <stephen.smith@newfront.com>:

When converting the llm_request to JSON, skip non-serializable data.

--
5dc68f4f5a6d12b753fdb81d1449716d13490afb by Stephen Smith <stephen.smith@newfront.com>:

Update _create_invocation_context() return type to InvocationContext.

--
23a33f754409fcd2a7641098d68cef7e4f1c72c6 by Stephen Smith <stephen.smith@newfront.com>:

Oops, remove unnecessary import.

COPYBARA_INTEGRATE_REVIEW=https://github.com/google/adk-python/pull/879 from stephensmithnewfront:main f71e195a9ed157e4c0b3abfa74ae078b0c1a920a
PiperOrigin-RevId: 763922003
---
 src/google/adk/telemetry.py       |  5 ++-
 tests/unittests/test_telemetry.py | 62 +++++++++++++++++++++++++++++++
 2 files changed, 66 insertions(+), 1 deletion(-)
 create mode 100644 tests/unittests/test_telemetry.py

diff --git a/src/google/adk/telemetry.py b/src/google/adk/telemetry.py
index a899ea6..2744ea5 100644
--- a/src/google/adk/telemetry.py
+++ b/src/google/adk/telemetry.py
@@ -117,7 +117,10 @@ def trace_call_llm(
   # Consider removing once GenAI SDK provides a way to record this info.
   span.set_attribute(
       'gcp.vertex.agent.llm_request',
-      json.dumps(_build_llm_request_for_trace(llm_request)),
+      json.dumps(
+          _build_llm_request_for_trace(llm_request),
+          default=lambda o: '<not serializable>',
+      ),
   )
   # Consider removing once GenAI SDK provides a way to record this info.
   span.set_attribute(
diff --git a/tests/unittests/test_telemetry.py b/tests/unittests/test_telemetry.py
new file mode 100644
index 0000000..5a77584
--- /dev/null
+++ b/tests/unittests/test_telemetry.py
@@ -0,0 +1,62 @@
+from typing import Any
+from typing import Optional
+
+from google.adk.sessions import InMemorySessionService
+from google.adk.agents.invocation_context import InvocationContext
+from google.adk.agents.llm_agent import LlmAgent
+from google.adk.models.llm_request import LlmRequest
+from google.adk.models.llm_response import LlmResponse
+from google.adk.telemetry import trace_call_llm
+from google.genai import types
+import pytest
+
+
+async def _create_invocation_context(
+    agent: LlmAgent, state: Optional[dict[str, Any]] = None
+) -> InvocationContext:
+  session_service = InMemorySessionService()
+  session = await session_service.create_session(
+      app_name='test_app', user_id='test_user', state=state
+  )
+  invocation_context = InvocationContext(
+      invocation_id='test_id',
+      agent=agent,
+      session=session,
+      session_service=session_service,
+  )
+  return invocation_context
+
+
+@pytest.mark.asyncio
+async def test_trace_call_llm_function_response_includes_part_from_bytes():
+  agent = LlmAgent(name='test_agent')
+  invocation_context = await _create_invocation_context(agent)
+  llm_request = LlmRequest(
+    contents=[
+      types.Content(
+        role="user",
+        parts=[
+          types.Part.from_function_response(
+            name="test_function_1",
+            response={
+              "result": b"test_data",
+            },
+          ),
+        ],
+      ),
+      types.Content(
+        role="user",
+        parts=[
+          types.Part.from_function_response(
+            name="test_function_2",
+            response={
+              "result": types.Part.from_bytes(data=b"test_data", mime_type="application/octet-stream"),
+            },
+          ),
+        ],
+      ),
+    ],
+    config=types.GenerateContentConfig(system_instruction=""),
+  )
+  llm_response = LlmResponse(turn_complete=True)
+  trace_call_llm(invocation_context, 'test_event_id', llm_request, llm_response)