adk-python/tests/unittests/test_telemetry.py
Stephen Smith 03fe909ada 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
2025-05-27 13:21:44 -07:00

63 lines
1.9 KiB
Python

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)