encounter/adk-python
0
0
Fork 0
mirror of https://github.com/encounter/adk-python.git synced 2026-07-09 18:19:28 -07:00
Code Issues Packages Projects Releases Wiki Activity

feat: Move livebidi agents esp multi-agent to use session/events

Browse Source 
The old live/bidi agents are using a cache to store context/history during agent transfer etc. As we have added support for session for live/bidi, we are now migrating the context/history cache to it. This improves scalability, efficiency and maintainability.

It introduces several changes:
* AudioTranscriber support is removed as now we are using native transcription from models.
* Transcription is returned as input_transcription/output_transcription fields and no longer as contents.
* We will return a new event with artifact references of file type of audio/pcm.(in addition to existing audio response event. So the users of this api need to do proper filtering here.)

PiperOrigin-RevId: 805997675
This commit is contained in:
Hangfei Lin
2025-09-11 14:58:33 -07:00
committed by Copybara-Service
parent 873551d7b9
commit ab69ef8de8
12 changed files with 435 additions and 417 deletions
Download Patch File Download Diff File Expand all files Collapse all files
contributing/samples/live_bidi_streaming_multi_agent/agent.py
+2 -2
View File
@@ -100,8 +100,8 @@ def get_current_weather(location: str):
root_agent = Agent(
# find supported models here: https://google.github.io/adk-docs/get-started/streaming/quickstart-streaming/
model="gemini-2.0-flash-live-preview-04-09", # for Vertex project
# model="gemini-live-2.5-flash-preview", # for AI studio key
# model="gemini-2.0-flash-live-preview-04-09", # for Vertex project
model="gemini-live-2.5-flash-preview", # for AI studio key
name="root_agent",
instruction="""
You are a helpful assistant that can check time, roll dice and check if numbers are prime.
src/google/adk/agents/run_config.py
+13 -2
View File
@@ -22,6 +22,7 @@ from typing import Optional
from google.genai import types
from pydantic import BaseModel
from pydantic import ConfigDict
from pydantic import Field
from pydantic import field_validator
logger = logging.getLogger('google_adk.' + __name__)
@@ -64,10 +65,14 @@ class RunConfig(BaseModel):
streaming_mode: StreamingMode = StreamingMode.NONE
"""Streaming mode, None or StreamingMode.SSE or StreamingMode.BIDI."""
output_audio_transcription: Optional[types.AudioTranscriptionConfig] = None
output_audio_transcription: Optional[types.AudioTranscriptionConfig] = Field(
default_factory=types.AudioTranscriptionConfig
)
"""Output transcription for live agents with audio response."""
input_audio_transcription: Optional[types.AudioTranscriptionConfig] = None
input_audio_transcription: Optional[types.AudioTranscriptionConfig] = Field(
default_factory=types.AudioTranscriptionConfig
)
"""Input transcription for live agents with audio input from user."""
realtime_input_config: Optional[types.RealtimeInputConfig] = None
@@ -82,6 +87,12 @@ class RunConfig(BaseModel):
session_resumption: Optional[types.SessionResumptionConfig] = None
"""Configures session resumption mechanism. Only support transparent session resumption mode now."""
save_live_audio: bool = False
"""Saves live video and audio data to session and artifact service.
Right now, only audio is supported.
"""
max_llm_calls: int = 500
"""
A limit on the total number of llm calls for a given run.
src/google/adk/flows/llm_flows/audio_cache_manager.py
+8 -14
View File
@@ -87,7 +87,7 @@ class AudioCacheManager:
flush_user_audio: bool = True,
flush_model_audio: bool = True,
) -> None:
"""Flush audio caches to session and artifact services.
"""Flush audio caches to artifact services.
The multimodality data is saved in artifact service in the format of
audio file. The file data reference is added to the session as an event.
@@ -103,24 +103,23 @@ class AudioCacheManager:
flush_model_audio: Whether to flush the output (model) audio cache.
"""
if flush_user_audio and invocation_context.input_realtime_cache:
success = await self._flush_cache_to_services(
flush_success = await self._flush_cache_to_services(
invocation_context,
invocation_context.input_realtime_cache,
'input_audio',
)
if success:
if flush_success:
invocation_context.input_realtime_cache = []
logger.debug('Flushed input audio cache')
if flush_model_audio and invocation_context.output_realtime_cache:
success = await self._flush_cache_to_services(
logger.debug('Flushed output audio cache')
flush_success = await self._flush_cache_to_services(
invocation_context,
invocation_context.output_realtime_cache,
'output_audio',
)
if success:
if flush_success:
invocation_context.output_realtime_cache = []
logger.debug('Flushed output audio cache')
async def _flush_cache_to_services(
self,
@@ -128,7 +127,7 @@ class AudioCacheManager:
audio_cache: list[RealtimeCacheEntry],
cache_type: str,
) -> bool:
"""Flush a list of audio cache entries to session and artifact services.
"""Flush a list of audio cache entries to artifact services.
The artifact service stores the actual blob. The session stores the
reference to the stored blob.
@@ -191,11 +190,6 @@ class AudioCacheManager:
timestamp=audio_cache[0].timestamp,
)
# Add to session
await invocation_context.session_service.append_event(
invocation_context.session, audio_event
)
logger.debug(
'Successfully flushed %s cache: %d chunks, %d bytes, saved as %s',
cache_type,
@@ -203,7 +197,7 @@ class AudioCacheManager:
len(combined_audio_data),
filename,
)
return True
return audio_event
except Exception as e:
logger.error('Failed to flush %s cache: %s', cache_type, e)
src/google/adk/flows/llm_flows/base_llm_flow.py
+79 -58
View File
@@ -69,6 +69,42 @@ DEFAULT_TASK_COMPLETION_DELAY = 1.0
DEFAULT_ENABLE_CACHE_STATISTICS = False
def _get_audio_transcription_from_session(
invocation_context: InvocationContext,
) -> list[types.Content]:
"""Get audio and transcription content from session events.
Collects audio file references and transcription text from session events
to reconstruct the conversation history including multimodal content.
Args:
invocation_context: The invocation context containing session data.
Returns:
A list of Content objects containing audio files and transcriptions.
"""
contents = []
for event in invocation_context.session.events:
# Collect transcription text events
if hasattr(event, 'input_transcription') and event.input_transcription:
contents.append(
types.Content(
role='user',
parts=[types.Part.from_text(text=event.input_transcription.text)],
)
)
if hasattr(event, 'output_transcription') and event.output_transcription:
contents.append(
types.Content(
role='model',
parts=[
types.Part.from_text(text=event.output_transcription.text)
],
)
)
return contents
class BaseLlmFlow(ABC):
"""A basic flow that calls the LLM in a loop until a final response is generated.
@@ -129,21 +165,8 @@ class BaseLlmFlow(ABC):
if llm_request.contents:
# Sends the conversation history to the model.
with tracer.start_as_current_span('send_data'):
if invocation_context.transcription_cache:
from . import audio_transcriber
audio_transcriber = audio_transcriber.AudioTranscriber(
init_client=True
if invocation_context.run_config.input_audio_transcription
is None
else False
)
contents = audio_transcriber.transcribe_file(invocation_context)
logger.debug('Sending history to model: %s', contents)
await llm_connection.send_history(contents)
invocation_context.transcription_cache = None
trace_send_data(invocation_context, event_id, contents)
else:
# Combine regular contents with audio/transcription from session
logger.debug('Sending history to model: %s', llm_request.contents)
await llm_connection.send_history(llm_request.contents)
trace_send_data(
invocation_context, event_id, llm_request.contents
@@ -324,22 +347,6 @@ class BaseLlmFlow(ABC):
author=get_author_for_event(llm_response),
)
# Handle transcription events ONCE per llm_response, outside the event loop
if llm_response.input_transcription:
await self.transcription_manager.handle_input_transcription(
invocation_context, llm_response.input_transcription
)
if llm_response.output_transcription:
await self.transcription_manager.handle_output_transcription(
invocation_context, llm_response.output_transcription
)
# Flush audio caches based on control events using configurable settings
await self._handle_control_event_flush(
invocation_context, llm_response
)
async with Aclosing(
self._postprocess_live(
invocation_context,
@@ -349,28 +356,11 @@ class BaseLlmFlow(ABC):
)
) as agen:
async for event in agen:
if (
event.content
and event.content.parts
and event.content.parts[0].inline_data is None
and not event.partial
):
# This can be either user data or transcription data.
# when output transcription enabled, it will contain model's
# transcription.
# when input transcription enabled, it will contain user
# transcription.
if not invocation_context.transcription_cache:
invocation_context.transcription_cache = []
invocation_context.transcription_cache.append(
TranscriptionEntry(
role=event.content.role, data=event.content
)
)
# Cache output audio chunks from model responses
# TODO: support video data
if (
event.content
invocation_context.run_config.save_live_audio
and event.content
and event.content.parts
and event.content.parts[0].inline_data
and event.content.parts[0].inline_data.mime_type.startswith(
@@ -578,6 +568,36 @@ class BaseLlmFlow(ABC):
):
return
# Handle transcription events ONCE per llm_response, outside the event loop
if llm_response.input_transcription:
input_transcription_event = (
await self.transcription_manager.handle_input_transcription(
invocation_context, llm_response.input_transcription
)
)
yield input_transcription_event
return
if llm_response.output_transcription:
output_transcription_event = (
await self.transcription_manager.handle_output_transcription(
invocation_context, llm_response.output_transcription
)
)
yield output_transcription_event
return
# Flush audio caches based on control events using configurable settings
if invocation_context.run_config.save_live_audio:
_handle_control_event_flush_event = (
await self._handle_control_event_flush(
invocation_context, llm_response
)
)
if _handle_control_event_flush_event:
yield _handle_control_event_flush_event
return
# Builds the event.
model_response_event = self._finalize_model_response_event(
llm_request, llm_response, model_response_event
@@ -877,33 +897,34 @@ class BaseLlmFlow(ABC):
invocation_context: The invocation context containing audio caches.
llm_response: The LLM response containing control event information.
"""
# Log cache statistics if enabled
if DEFAULT_ENABLE_CACHE_STATISTICS:
stats = self.audio_cache_manager.get_cache_stats(invocation_context)
logger.debug('Audio cache stats: %s', stats)
if llm_response.interrupted:
# user interrupts so the model will stop. we can flush model audio here
await self.audio_cache_manager.flush_caches(
return await self.audio_cache_manager.flush_caches(
invocation_context,
flush_user_audio=False,
flush_model_audio=True,
)
elif llm_response.turn_complete:
# turn completes so we can flush both user and model
await self.audio_cache_manager.flush_caches(
return await self.audio_cache_manager.flush_caches(
invocation_context,
flush_user_audio=True,
flush_model_audio=True,
)
elif getattr(llm_response, 'generation_complete', False):
# model generation complete so we can flush model audio
await self.audio_cache_manager.flush_caches(
return await self.audio_cache_manager.flush_caches(
invocation_context,
flush_user_audio=False,
flush_model_audio=True,
)
# Log cache statistics if enabled
if DEFAULT_ENABLE_CACHE_STATISTICS:
stats = self.audio_cache_manager.get_cache_stats(invocation_context)
logger.debug('Audio cache stats: %s', stats)
async def _run_and_handle_error(
self,
response_generator: AsyncGenerator[LlmResponse, None],
src/google/adk/flows/llm_flows/contents.py
+85 -2
View File
@@ -206,6 +206,7 @@ def _contains_empty_content(event: Event) -> bool:
"""Check if an event should be skipped due to missing or empty content.
This can happen to the evnets that only changed session state.
When both content and transcriptions are empty, the event will be considered as empty.
Args:
event: The event to check.
@@ -218,7 +219,7 @@ def _contains_empty_content(event: Event) -> bool:
or not event.content.role
or not event.content.parts
or event.content.parts[0].text == ''
)
) and (not event.output_transcription and not event.input_transcription)
def _get_contents(
@@ -236,9 +237,12 @@ def _get_contents(
Returns:
A list of processed contents.
"""
filtered_events = []
accumulated_input_transcription = ''
accumulated_output_transcription = ''
# Parse the events, leaving the contents and the function calls and
# responses from the current agent.
raw_filtered_events = []
for event in events:
if _contains_empty_content(event):
continue
@@ -252,6 +256,45 @@ def _get_contents(
# Skip request confirmation events.
continue
raw_filtered_events.append(event)
filtered_events = []
# aggregate transcription events
for i in range(len(raw_filtered_events)):
event = raw_filtered_events[i]
if not event.content:
# Convert transcription into normal event
if event.input_transcription and event.input_transcription.text:
accumulated_input_transcription += event.input_transcription.text
if (
i != len(raw_filtered_events) - 1
and raw_filtered_events[i + 1].input_transcription
and raw_filtered_events[i + 1].input_transcription.text
):
continue
event = event.model_copy(deep=True)
event.input_transcription = None
event.content = types.Content(
role='user',
parts=[types.Part(text=accumulated_input_transcription)],
)
accumulated_input_transcription = ''
elif event.output_transcription and event.output_transcription.text:
accumulated_output_transcription += event.output_transcription.text
if (
i != len(raw_filtered_events) - 1
and raw_filtered_events[i + 1].output_transcription
and raw_filtered_events[i + 1].output_transcription.text
):
continue
event = event.model_copy(deep=True)
event.output_transcription = None
event.content = types.Content(
role='model',
parts=[types.Part(text=accumulated_output_transcription)],
)
accumulated_output_transcription = ''
if _is_other_agent_reply(agent_name, event):
if converted_event := _present_other_agent_message(event):
filtered_events.append(converted_event)
@@ -474,3 +517,43 @@ def _is_auth_event(event: Event) -> bool:
def _is_request_confirmation_event(event: Event) -> bool:
"""Checks if the event is a request confirmation event."""
return _is_function_call_event(event, REQUEST_CONFIRMATION_FUNCTION_CALL_NAME)
def _is_live_model_audio_event(event: Event) -> bool:
"""Check if the event is an audio event produced by live/bidi models
There are two possible cases:
content=Content(
parts=[
Part(
file_data=FileData(
file_uri='artifact://live_bidi_streaming_multi_agent/user/cccf0b8b-4a30-449a-890e-e8b8deb661a1/_adk_live/adk_live_audio_storage_input_audio_1756092402277.pcm#1',
mime_type='audio/pcm'
)
),
],
role='user'
)
content=Content(
parts=[
Part(
inline_data=Blob(
data=b'\x01\x00\x00...',
mime_type='audio/pcm'
)
),
],
role='model'
) grounding_metadata=None partial=None turn_complete=None finish_reason=None error_code=None error_message=None ...
"""
if not event.content:
return False
if not event.content.parts:
return False
# If it's audio data, then one event only has one part of audio.
for part in event.content.parts:
if part.inline_data and part.inline_data.mime_type == 'audio/pcm':
return True
if part.file_data and part.file_data.mime_type == 'audio/pcm':
return True
return False
src/google/adk/flows/llm_flows/transcription_manager.py
+3 -5
View File
@@ -42,7 +42,7 @@ class TranscriptionManager:
invocation_context: The current invocation context.
transcription: The transcription data from user input.
"""
await self._create_and_save_transcription_event(
return await self._create_and_save_transcription_event(
invocation_context=invocation_context,
transcription=transcription,
author='user',
@@ -60,7 +60,7 @@ class TranscriptionManager:
invocation_context: The current invocation context.
transcription: The transcription data from model output.
"""
await self._create_and_save_transcription_event(
return await self._create_and_save_transcription_event(
invocation_context=invocation_context,
transcription=transcription,
author=invocation_context.agent.name,
@@ -93,9 +93,6 @@ class TranscriptionManager:
)
# Save transcription event to session
await invocation_context.session_service.append_event(
invocation_context.session, transcription_event
)
logger.debug(
'Saved %s transcription event for %s: %s',
@@ -106,6 +103,7 @@ class TranscriptionManager:
else 'audio transcription',
)
return transcription_event
except Exception as e:
logger.error(
'Failed to save %s transcription event: %s',
src/google/adk/models/gemini_llm_connection.py
+2 -22
View File
@@ -166,38 +166,18 @@ class GeminiLlmConnection(BaseLlmConnection):
message.server_content.input_transcription
and message.server_content.input_transcription.text
):
user_text = message.server_content.input_transcription.text
parts = [
types.Part.from_text(
text=user_text,
)
]
llm_response = LlmResponse(
content=types.Content(role='user', parts=parts)
input_transcription=message.server_content.input_transcription,
)
yield llm_response
if (
message.server_content.output_transcription
and message.server_content.output_transcription.text
):
# TODO: Right now, we just support output_transcription without
# changing interface and data protocol. Later, we can consider to
# support output_transcription as a separate field in LlmResponse.
# Transcription is always considered as partial event
# We rely on other control signals to determine when to yield the
# full text response(turn_complete, interrupted, or tool_call).
text += message.server_content.output_transcription.text
parts = [
types.Part.from_text(
text=message.server_content.output_transcription.text
)
]
llm_response = LlmResponse(
content=types.Content(role='model', parts=parts), partial=True
output_transcription=message.server_content.output_transcription
)
yield llm_response
if message.server_content.turn_complete:
if text:
yield self.__build_full_text_response(text)
src/google/adk/runners.py
+25 -2
View File
@@ -41,6 +41,7 @@ from .auth.credential_service.base_credential_service import BaseCredentialServi
from .code_executors.built_in_code_executor import BuiltInCodeExecutor
from .events.event import Event
from .events.event import EventActions
from .flows.llm_flows import contents
from .flows.llm_flows.functions import find_matching_function_call
from .memory.base_memory_service import BaseMemoryService
from .memory.in_memory_memory_service import InMemoryMemoryService
@@ -305,7 +306,12 @@ class Runner:
yield event
async with Aclosing(
self._exec_with_plugin(invocation_context, session, execute)
self._exec_with_plugin(
invocation_context=invocation_context,
session=session,
execute_fn=execute,
is_live_call=False,
)
) as agen:
async for event in agen:
yield event
@@ -314,11 +320,21 @@ class Runner:
async for event in agen:
yield event
def _should_append_event(self, event: Event, is_live_call: bool) -> bool:
"""Checks if an event should be appended to the session."""
# Don't append audio response from model in live mode to session.
# The data is appended to artifacts with a reference in file_data in the
# event.
if is_live_call and contents._is_live_model_audio_event(event):
return False
return True
async def _exec_with_plugin(
self,
invocation_context: InvocationContext,
session: Session,
execute_fn: Callable[[InvocationContext], AsyncGenerator[Event, None]],
is_live_call: bool = False,
) -> AsyncGenerator[Event, None]:
"""Wraps execution with plugin callbacks.
@@ -343,6 +359,7 @@ class Runner:
author='model',
content=early_exit_result,
)
if self._should_append_event(early_exit_event, is_live_call):
await self.session_service.append_event(
session=session,
event=early_exit_event,
@@ -353,6 +370,7 @@ class Runner:
async with Aclosing(execute_fn(invocation_context)) as agen:
async for event in agen:
if not event.partial:
if self._should_append_event(event, is_live_call):
await self.session_service.append_event(
session=session, event=event
)
@@ -526,7 +544,12 @@ class Runner:
yield event
async with Aclosing(
self._exec_with_plugin(invocation_context, session, execute)
self._exec_with_plugin(
invocation_context=invocation_context,
session=session,
execute_fn=execute,
is_live_call=True,
)
) as agen:
async for event in agen:
yield event
tests/unittests/flows/llm_flows/test_audio_cache_manager.py
+1 -1
View File
@@ -251,7 +251,7 @@ class TestAudioCacheManager:
assert saved_artifact.inline_data.mime_type == 'audio/pcm'
# Verify session event was created
mock_session_service.append_event.assert_called_once()
mock_session_service.append_event.assert_not_called()
def test_get_cache_stats_empty(self):
"""Test getting statistics for empty caches."""
tests/unittests/flows/llm_flows/test_transcription_manager.py
+3 -95
View File
@@ -49,17 +49,7 @@ class TestTranscriptionManager:
)
# Verify session service was called
mock_session_service.append_event.assert_called_once()
# Check the event that was created
call_args = mock_session_service.append_event.call_args
event = call_args[0][1] # Second argument is the event
assert event.author == 'user'
assert event.input_transcription == transcription
assert event.output_transcription is None
assert event.invocation_id == invocation_context.invocation_id
assert isinstance(event.timestamp, float)
mock_session_service.append_event.assert_not_called()
@pytest.mark.asyncio
async def test_handle_output_transcription(self):
@@ -80,17 +70,7 @@ class TestTranscriptionManager:
)
# Verify session service was called
mock_session_service.append_event.assert_called_once()
# Check the event that was created
call_args = mock_session_service.append_event.call_args
event = call_args[0][1] # Second argument is the event
assert event.author == agent.name
assert event.input_transcription is None
assert event.output_transcription == transcription
assert event.invocation_id == invocation_context.invocation_id
assert isinstance(event.timestamp, float)
mock_session_service.append_event.assert_not_called()
@pytest.mark.asyncio
async def test_handle_multiple_transcriptions(self):
@@ -118,53 +98,7 @@ class TestTranscriptionManager:
)
# Verify session service was called for each transcription
assert mock_session_service.append_event.call_count == 5
@pytest.mark.asyncio
async def test_error_handling_input_transcription(self):
"""Test error handling during input transcription processing."""
invocation_context = await testing_utils.create_invocation_context(
testing_utils.create_test_agent()
)
# Set up mock session service that raises an error
mock_session_service = AsyncMock()
mock_session_service.append_event.side_effect = Exception(
'Session service error'
)
invocation_context.session_service = mock_session_service
# Create test transcription
transcription = types.Transcription(text='Test transcription')
# Handle transcription should raise the exception
with pytest.raises(Exception, match='Session service error'):
await self.manager.handle_input_transcription(
invocation_context, transcription
)
@pytest.mark.asyncio
async def test_error_handling_output_transcription(self):
"""Test error handling during output transcription processing."""
invocation_context = await testing_utils.create_invocation_context(
testing_utils.create_test_agent()
)
# Set up mock session service that raises an error
mock_session_service = AsyncMock()
mock_session_service.append_event.side_effect = Exception(
'Session service error'
)
invocation_context.session_service = mock_session_service
# Create test transcription
transcription = types.Transcription(text='Test transcription')
# Handle transcription should raise the exception
with pytest.raises(Exception, match='Session service error'):
await self.manager.handle_output_transcription(
invocation_context, transcription
)
assert mock_session_service.append_event.call_count == 0
def test_get_transcription_stats_empty_session(self):
"""Test getting transcription statistics for empty session."""
@@ -264,25 +198,6 @@ class TestTranscriptionManager:
invocation_context, transcription
)
# Verify the event structure
call_args = mock_session_service.append_event.call_args
event = call_args[0][1]
# Check all required fields are present
assert hasattr(event, 'id')
assert hasattr(event, 'invocation_id')
assert hasattr(event, 'author')
assert hasattr(event, 'input_transcription')
assert hasattr(event, 'output_transcription')
assert hasattr(event, 'timestamp')
# Check values
assert event.id is not None
assert event.invocation_id == invocation_context.invocation_id
assert event.author == 'user'
assert event.input_transcription == transcription
assert event.output_transcription is None
@pytest.mark.asyncio
async def test_transcription_with_different_data_types(self):
"""Test handling transcriptions with different data types."""
@@ -303,10 +218,3 @@ class TestTranscriptionManager:
await self.manager.handle_input_transcription(
invocation_context, transcription
)
# Verify the transcription object is preserved as-is
call_args = mock_session_service.append_event.call_args
event = call_args[0][1]
assert event.input_transcription == transcription
assert event.input_transcription.text == 'Advanced transcription'
tests/unittests/streaming/test_live_streaming_configs.py
+1 -1
View File
@@ -56,7 +56,7 @@ def test_streaming():
assert llm_request_sent_to_mock.live_connect_config is not None
assert (
llm_request_sent_to_mock.live_connect_config.output_audio_transcription
is None
is not None
)
tests/unittests/streaming/test_streaming_audio_storage.py
+202 -202
View File
@@ -1,241 +1,241 @@
# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# # Copyright 2025 Google LLC
# #
# # Licensed under the Apache License, Version 2.0 (the "License");
# # you may not use this file except in compliance with the License.
# # You may obtain a copy of the License at
# #
# # http://www.apache.org/licenses/LICENSE-2.0
# #
# # Unless required by applicable law or agreed to in writing, software
# # distributed under the License is distributed on an "AS IS" BASIS,
# # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# # See the License for the specific language governing permissions and
# # limitations under the License.
import asyncio
import time
# import asyncio
# import time
from google.adk.agents import Agent
from google.adk.agents import LiveRequestQueue
from google.adk.agents.invocation_context import RealtimeCacheEntry
from google.adk.agents.run_config import RunConfig
from google.adk.events.event import Event
from google.adk.models import LlmResponse
from google.genai import types
import pytest
# from google.adk.agents import Agent
# from google.adk.agents import LiveRequestQueue
# from google.adk.agents.invocation_context import RealtimeCacheEntry
# from google.adk.agents.run_config import RunConfig
# from google.adk.events.event import Event
# from google.adk.models import LlmResponse
# from google.genai import types
# import pytest
from .. import testing_utils
# from .. import testing_utils
def test_audio_caching_direct():
"""Test audio caching logic directly without full live streaming."""
# This test directly verifies that our audio caching logic works
audio_data = b'\x00\xFF\x01\x02\x03\x04\x05\x06'
audio_mime_type = 'audio/pcm'
# def test_audio_caching_direct():
# """Test audio caching logic directly without full live streaming."""
# # This test directly verifies that our audio caching logic works
# audio_data = b'\x00\xFF\x01\x02\x03\x04\x05\x06'
# audio_mime_type = 'audio/pcm'
# Create mock responses for successful completion
responses = [
LlmResponse(
content=types.Content(
role='model',
parts=[types.Part.from_text(text='Processing audio...')],
),
turn_complete=False,
),
LlmResponse(turn_complete=True), # This should trigger flush
]
# # Create mock responses for successful completion
# responses = [
# LlmResponse(
# content=types.Content(
# role='model',
# parts=[types.Part.from_text(text='Processing audio...')],
# ),
# turn_complete=False,
# ),
# LlmResponse(turn_complete=True), # This should trigger flush
# ]
mock_model = testing_utils.MockModel.create(responses)
mock_model.model = 'gemini-2.0-flash-exp' # For CFC support
# mock_model = testing_utils.MockModel.create(responses)
# mock_model.model = 'gemini-2.0-flash-exp' # For CFC support
root_agent = Agent(
name='test_agent',
model=mock_model,
tools=[],
)
# root_agent = Agent(
# name='test_agent',
# model=mock_model,
# tools=[],
# )
# Test our implementation by directly calling it
async def test_caching():
# Create context similar to what would be created in real scenario
invocation_context = await testing_utils.create_invocation_context(
root_agent, run_config=RunConfig(support_cfc=True)
)
# # Test our implementation by directly calling it
# async def test_caching():
# # Create context similar to what would be created in real scenario
# invocation_context = await testing_utils.create_invocation_context(
# root_agent, run_config=RunConfig(support_cfc=True)
# )
# Import our caching classes
from google.adk.agents.invocation_context import RealtimeCacheEntry
from google.adk.flows.llm_flows.base_llm_flow import BaseLlmFlow
# # Import our caching classes
# from google.adk.agents.invocation_context import RealtimeCacheEntry
# from google.adk.flows.llm_flows.base_llm_flow import BaseLlmFlow
# Create a mock flow to test our methods
flow = BaseLlmFlow()
# # Create a mock flow to test our methods
# flow = BaseLlmFlow()
# Test adding audio to cache
invocation_context.input_realtime_cache = []
audio_entry = RealtimeCacheEntry(
role='user',
data=types.Blob(data=audio_data, mime_type=audio_mime_type),
timestamp=1234567890.0,
)
invocation_context.input_realtime_cache.append(audio_entry)
# # Test adding audio to cache
# invocation_context.input_realtime_cache = []
# audio_entry = RealtimeCacheEntry(
# role='user',
# data=types.Blob(data=audio_data, mime_type=audio_mime_type),
# timestamp=1234567890.0,
# )
# invocation_context.input_realtime_cache.append(audio_entry)
# Verify cache has data
assert len(invocation_context.input_realtime_cache) == 1
assert invocation_context.input_realtime_cache[0].data.data == audio_data
# # Verify cache has data
# assert len(invocation_context.input_realtime_cache) == 1
# assert invocation_context.input_realtime_cache[0].data.data == audio_data
# Test flushing cache
await flow._handle_control_event_flush(invocation_context, responses[-1])
# # Test flushing cache
# await flow._handle_control_event_flush(invocation_context, responses[-1])
# Verify cache was cleared
assert len(invocation_context.input_realtime_cache) == 0
# # Verify cache was cleared
# assert len(invocation_context.input_realtime_cache) == 0
# Check if artifacts were created
artifact_keys = (
await invocation_context.artifact_service.list_artifact_keys(
app_name=invocation_context.app_name,
user_id=invocation_context.user_id,
session_id=invocation_context.session.id,
)
)
# # Check if artifacts were created
# artifact_keys = (
# await invocation_context.artifact_service.list_artifact_keys(
# app_name=invocation_context.app_name,
# user_id=invocation_context.user_id,
# session_id=invocation_context.session.id,
# )
# )
# Should have at least one audio artifact
audio_artifacts = [key for key in artifact_keys if 'audio' in key.lower()]
assert (
len(audio_artifacts) > 0
), f'Expected audio artifacts, found: {artifact_keys}'
# # Should have at least one audio artifact
# audio_artifacts = [key for key in artifact_keys if 'audio' in key.lower()]
# assert (
# len(audio_artifacts) > 0
# ), f'Expected audio artifacts, found: {artifact_keys}'
# Verify artifact content
if audio_artifacts:
artifact = await invocation_context.artifact_service.load_artifact(
app_name=invocation_context.app_name,
user_id=invocation_context.user_id,
session_id=invocation_context.session.id,
filename=audio_artifacts[0],
)
assert artifact.inline_data.data == audio_data
# # Verify artifact content
# if audio_artifacts:
# artifact = await invocation_context.artifact_service.load_artifact(
# app_name=invocation_context.app_name,
# user_id=invocation_context.user_id,
# session_id=invocation_context.session.id,
# filename=audio_artifacts[0],
# )
# assert artifact.inline_data.data == audio_data
return True
# return True
# Run the async test
result = asyncio.run(test_caching())
assert result is True
# # Run the async test
# result = asyncio.run(test_caching())
# assert result is True
def test_transcription_handling():
"""Test that transcriptions are properly handled and saved to session service."""
# def test_transcription_handling():
# """Test that transcriptions are properly handled and saved to session service."""
# Create mock responses with transcriptions
input_transcription = types.Transcription(
text='Hello, this is transcribed input', finished=True
)
output_transcription = types.Transcription(
text='This is transcribed output', finished=True
)
# # Create mock responses with transcriptions
# input_transcription = types.Transcription(
# text='Hello, this is transcribed input', finished=True
# )
# output_transcription = types.Transcription(
# text='This is transcribed output', finished=True
# )
responses = [
LlmResponse(
content=types.Content(
role='model', parts=[types.Part.from_text(text='Processing...')]
),
turn_complete=False,
),
LlmResponse(input_transcription=input_transcription, turn_complete=False),
LlmResponse(
output_transcription=output_transcription, turn_complete=False
),
LlmResponse(turn_complete=True),
]
# responses = [
# LlmResponse(
# content=types.Content(
# role='model', parts=[types.Part.from_text(text='Processing...')]
# ),
# turn_complete=False,
# ),
# LlmResponse(input_transcription=input_transcription, turn_complete=False),
# LlmResponse(
# output_transcription=output_transcription, turn_complete=False
# ),
# LlmResponse(turn_complete=True),
# ]
mock_model = testing_utils.MockModel.create(responses)
mock_model.model = 'gemini-2.0-flash-exp'
# mock_model = testing_utils.MockModel.create(responses)
# mock_model.model = 'gemini-2.0-flash-exp'
root_agent = Agent(
name='test_agent',
model=mock_model,
tools=[],
)
# root_agent = Agent(
# name='test_agent',
# model=mock_model,
# tools=[],
# )
async def test_transcription():
# Create context
invocation_context = await testing_utils.create_invocation_context(
root_agent, run_config=RunConfig(support_cfc=True)
)
# async def test_transcription():
# # Create context
# invocation_context = await testing_utils.create_invocation_context(
# root_agent, run_config=RunConfig(support_cfc=True)
# )
from google.adk.events.event import Event
from google.adk.flows.llm_flows.base_llm_flow import BaseLlmFlow
# from google.adk.events.event import Event
# from google.adk.flows.llm_flows.base_llm_flow import BaseLlmFlow
flow = BaseLlmFlow()
# flow = BaseLlmFlow()
# Test processing transcription events
session_events_before = len(invocation_context.session.events)
# # Test processing transcription events
# session_events_before = len(invocation_context.session.events)
# Simulate input transcription event
input_event = Event(
id=Event.new_id(),
invocation_id=invocation_context.invocation_id,
author='user',
input_transcription=input_transcription,
)
# # Simulate input transcription event
# input_event = Event(
# id=Event.new_id(),
# invocation_id=invocation_context.invocation_id,
# author='user',
# input_transcription=input_transcription,
# )
# Simulate output transcription event
output_event = Event(
id=Event.new_id(),
invocation_id=invocation_context.invocation_id,
author=invocation_context.agent.name,
output_transcription=output_transcription,
)
# # Simulate output transcription event
# output_event = Event(
# id=Event.new_id(),
# invocation_id=invocation_context.invocation_id,
# author=invocation_context.agent.name,
# output_transcription=output_transcription,
# )
# Save transcription events to session
await invocation_context.session_service.append_event(
invocation_context.session, input_event
)
await invocation_context.session_service.append_event(
invocation_context.session, output_event
)
# # Save transcription events to session
# await invocation_context.session_service.append_event(
# invocation_context.session, input_event
# )
# await invocation_context.session_service.append_event(
# invocation_context.session, output_event
# )
# Verify transcriptions were saved to session
session_events_after = len(invocation_context.session.events)
assert session_events_after == session_events_before + 2
# # Verify transcriptions were saved to session
# session_events_after = len(invocation_context.session.events)
# assert session_events_after == session_events_before + 2
# Check that transcription events were saved
transcription_events = [
event
for event in invocation_context.session.events
if hasattr(event, 'input_transcription')
and event.input_transcription
or hasattr(event, 'output_transcription')
and event.output_transcription
]
assert len(transcription_events) >= 2
# # Check that transcription events were saved
# transcription_events = [
# event
# for event in invocation_context.session.events
# if hasattr(event, 'input_transcription')
# and event.input_transcription
# or hasattr(event, 'output_transcription')
# and event.output_transcription
# ]
# assert len(transcription_events) >= 2
# Verify input transcription
input_transcription_events = [
event
for event in invocation_context.session.events
if hasattr(event, 'input_transcription') and event.input_transcription
]
assert len(input_transcription_events) >= 1
assert (
input_transcription_events[0].input_transcription.text
== 'Hello, this is transcribed input'
)
assert input_transcription_events[0].author == 'user'
# # Verify input transcription
# input_transcription_events = [
# event
# for event in invocation_context.session.events
# if hasattr(event, 'input_transcription') and event.input_transcription
# ]
# assert len(input_transcription_events) >= 1
# assert (
# input_transcription_events[0].input_transcription.text
# == 'Hello, this is transcribed input'
# )
# assert input_transcription_events[0].author == 'user'
# Verify output transcription
output_transcription_events = [
event
for event in invocation_context.session.events
if hasattr(event, 'output_transcription') and event.output_transcription
]
assert len(output_transcription_events) >= 1
assert (
output_transcription_events[0].output_transcription.text
== 'This is transcribed output'
)
assert (
output_transcription_events[0].author == invocation_context.agent.name
)
# # Verify output transcription
# output_transcription_events = [
# event
# for event in invocation_context.session.events
# if hasattr(event, 'output_transcription') and event.output_transcription
# ]
# assert len(output_transcription_events) >= 1
# assert (
# output_transcription_events[0].output_transcription.text
# == 'This is transcribed output'
# )
# assert (
# output_transcription_events[0].author == invocation_context.agent.name
# )
return True
# return True
# Run the async test
result = asyncio.run(test_transcription())
assert result is True
# # Run the async test
# result = asyncio.run(test_transcription())
# assert result is True