diff --git a/src/google/adk/agents/invocation_context.py b/src/google/adk/agents/invocation_context.py index 3984d51f0b..e810c2f07d 100644 --- a/src/google/adk/agents/invocation_context.py +++ b/src/google/adk/agents/invocation_context.py @@ -35,7 +35,7 @@ from ..events._node_path_builder import _NodePathBuilder from ..events.event import Event from ..live._active_streaming_tool import ActiveStreamingTool -from ..live._audio_cache_manager import RealtimeCacheEntry as RealtimeCacheEntry +from ..live._realtime_cache_manager import RealtimeCacheEntry as RealtimeCacheEntry from ..live._transcription_entry import TranscriptionEntry from ..live.live_request_queue import LiveRequestQueue from ..memory.base_memory_service import BaseMemoryService diff --git a/src/google/adk/flows/llm_flows/base_llm_flow.py b/src/google/adk/flows/llm_flows/base_llm_flow.py index 6ce845fb8d..2f1cb2f6f9 100644 --- a/src/google/adk/flows/llm_flows/base_llm_flow.py +++ b/src/google/adk/flows/llm_flows/base_llm_flow.py @@ -30,11 +30,11 @@ from ...agents.base_agent import BaseAgent from ...agents.invocation_context import InvocationContext from ...events.event import Event -from ...live._audio_cache_manager import AudioCacheManager from ...live._flow_utils import DEFAULT_ENABLE_CACHE_STATISTICS as DEFAULT_ENABLE_CACHE_STATISTICS from ...live._flow_utils import DEFAULT_MAX_RECONNECT_ATTEMPTS as DEFAULT_MAX_RECONNECT_ATTEMPTS from ...live._flow_utils import DEFAULT_TASK_COMPLETION_DELAY as DEFAULT_TASK_COMPLETION_DELAY from ...live._flow_utils import DEFAULT_TRANSFER_AGENT_DELAY as DEFAULT_TRANSFER_AGENT_DELAY +from ...live._realtime_cache_manager import RealtimeCacheManager as _RealtimeCacheManager from ...models.base_llm_connection import BaseLlmConnection from ...models.llm_request import LlmRequest from ...models.llm_response import LlmResponse @@ -125,7 +125,10 @@ def __init__(self) -> None: self.response_processors: list[BaseLlmResponseProcessor] = [] # Initialize configuration and managers - self.audio_cache_manager = AudioCacheManager() + self._realtime_cache_manager = _RealtimeCacheManager() + # Kept for backward compatibility. The flow itself reads + # `_realtime_cache_manager`, so reassigning this attribute has no effect. + self.audio_cache_manager = self._realtime_cache_manager def _request_processor_lists( self, diff --git a/src/google/adk/live/_audio_cache_manager.py b/src/google/adk/live/_audio_cache_manager.py index 025e67ff35..be035731e4 100644 --- a/src/google/adk/live/_audio_cache_manager.py +++ b/src/google/adk/live/_audio_cache_manager.py @@ -12,310 +12,26 @@ # See the License for the specific language governing permissions and # limitations under the License. -from __future__ import annotations - -import logging -from typing import TYPE_CHECKING - -from google.adk.platform import time as platform_time -from google.genai import types -from pydantic import BaseModel -from pydantic import ConfigDict - -from ..events.event import Event - -if TYPE_CHECKING: - from ..agents.invocation_context import InvocationContext - -logger = logging.getLogger('google_adk.' + __name__) - - -class RealtimeCacheEntry(BaseModel): - """Store audio data chunks for caching before flushing.""" - - model_config = ConfigDict( - arbitrary_types_allowed=True, - extra='forbid', - ) - """The pydantic model config.""" - - role: str - """The role that created this audio data, typically "user" or "model".""" - - data: types.Blob - """The audio data chunk.""" - - timestamp: float - """Timestamp when the audio chunk was received.""" - - -def _require_audio_data(blob: types.Blob) -> bytes: - data = blob.data - if not isinstance(data, bytes): - raise ValueError('Audio blobs must contain byte data.') - return data - - -# Deliberately duplicated from `flows.llm_flows.core._utils` -# rather than imported: `live` sits below `flows.llm_flows` in the -# layering, and importing upward would put a cycle back in. -def _require_agent_name(invocation_context: InvocationContext) -> str: - agent = invocation_context.agent - if agent is None: - raise TypeError('Live audio requires an agent in InvocationContext.') - return agent.name - - -class AudioCacheManager: - """Manages audio caching and flushing for live streaming flows.""" - - def __init__(self, config: AudioCacheConfig | None = None) -> None: - """Initialize the audio cache manager. - - Args: - config: Configuration for audio caching behavior. - """ - self.config = config or AudioCacheConfig() - - def cache_audio( - self, - invocation_context: InvocationContext, - audio_blob: types.Blob, - cache_type: str, - ) -> None: - """Cache incoming user or outgoing model audio data. - - Args: - invocation_context: The current invocation context. - audio_blob: The audio data to cache. - cache_type: Type of audio to cache, either 'input' or 'output'. - - Raises: - ValueError: If cache_type is not 'input' or 'output'. - """ - audio_data = _require_audio_data(audio_blob) - if cache_type == 'input': - if not invocation_context.input_realtime_cache: - invocation_context.input_realtime_cache = [] - cache = invocation_context.input_realtime_cache - role = 'user' - elif cache_type == 'output': - if not invocation_context.output_realtime_cache: - invocation_context.output_realtime_cache = [] - cache = invocation_context.output_realtime_cache - role = 'model' - else: - raise ValueError("cache_type must be either 'input' or 'output'") - - audio_entry = RealtimeCacheEntry( - role=role, data=audio_blob, timestamp=platform_time.get_time() - ) - cache.append(audio_entry) - - logger.debug( - 'Cached %s audio chunk: %d bytes, cache size: %d', - cache_type, - len(audio_data), - len(cache), - ) +"""Backward compatibility module for the audio-only cache manager names. - async def flush_caches( - self, - invocation_context: InvocationContext, - flush_user_audio: bool = True, - flush_model_audio: bool = True, - ) -> list[Event]: - """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. - The audio file follows the naming convention: artifact_ref = - f"artifact://{invocation_context.app_name}/{invocation_context.user_id}/ - {invocation_context.session.id}/_adk_live/{filename}#{revision_id}" - - Note: video data is not supported yet. - - A flushed cache is emptied whether or not an artifact was written, so a - cache the artifact service could not accept is dropped rather than kept - for a later attempt. - - Args: - invocation_context: The invocation context containing audio caches. - flush_user_audio: Whether to flush the input (user) audio cache. - flush_model_audio: Whether to flush the output (model) audio cache. - - Returns: - A list of Event objects created from the flushed caches. - """ - flushed_events: list[Event] = [] - if flush_user_audio and invocation_context.input_realtime_cache: - # Swap before the await so audio cached during it is kept for the next - # flush rather than discarded by the unconditional drop. - input_cache = invocation_context.input_realtime_cache - invocation_context.input_realtime_cache = [] - audio_event = await self._flush_cache_to_services( - invocation_context, - input_cache, - 'input_audio', - ) - if audio_event: - flushed_events.append(audio_event) - - if flush_model_audio and invocation_context.output_realtime_cache: - logger.debug('Flushed output audio cache') - output_cache = invocation_context.output_realtime_cache - invocation_context.output_realtime_cache = [] - audio_event = await self._flush_cache_to_services( - invocation_context, - output_cache, - 'output_audio', - ) - if audio_event: - flushed_events.append(audio_event) - - return flushed_events - - async def _flush_cache_to_services( - self, - invocation_context: InvocationContext, - audio_cache: list[RealtimeCacheEntry], - cache_type: str, - ) -> Event | None: - """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. - - Args: - invocation_context: The invocation context. - audio_cache: The audio cache to flush. - cache_type: Type identifier for the cache ('input_audio' or - 'output_audio'). - - Returns: - The created Event if the cache was successfully flushed, None otherwise. - """ - if not invocation_context.artifact_service or not audio_cache: - logger.debug('Skipping cache flush: no artifact service or empty cache') - return None - - try: - # Combine audio chunks into a single file. Use join rather than repeated - # `+=`, which is O(n^2) over the total audio size. - mime_type = audio_cache[0].data.mime_type or 'audio/pcm' - combined_audio_data = b''.join( - entry.data.data or b'' for entry in audio_cache - ) - - # Generate filename with timestamp from first audio chunk (when recording started) - timestamp = int(audio_cache[0].timestamp * 1000) # milliseconds - # Drop MIME parameters such as `;rate=24000` so they don't end up in the - # file extension. The full MIME type is kept on the saved artifact. - extension = mime_type.split(';', 1)[0].strip().split('/')[-1] - filename = f'adk_live_audio_storage_{cache_type}_{timestamp}.{extension}' - - # Save to artifact service - combined_audio_part = types.Part( - inline_data=types.Blob(data=combined_audio_data, mime_type=mime_type) - ) - - revision_id = await invocation_context.artifact_service.save_artifact( - app_name=invocation_context.app_name, - user_id=invocation_context.user_id, - session_id=invocation_context.session.id, - filename=filename, - artifact=combined_audio_part, - ) - - # Create artifact reference for session service - artifact_ref = f'artifact://{invocation_context.app_name}/{invocation_context.user_id}/{invocation_context.session.id}/_adk_live/{filename}#{revision_id}' - - # Create event with file data reference to add to session - # For model events, author should be the agent name, not the role - author = ( - _require_agent_name(invocation_context) - if audio_cache[0].role == 'model' - else audio_cache[0].role - ) - audio_event = Event( - id=Event.new_id(), - invocation_id=invocation_context.invocation_id, - author=author, - content=types.Content( - role=audio_cache[0].role, - parts=[ - types.Part( - file_data=types.FileData( - file_uri=artifact_ref, mime_type=mime_type - ) - ) - ], - ), - timestamp=audio_cache[0].timestamp, - ) - - logger.debug( - 'Successfully flushed %s cache: %d chunks, %d bytes, saved as %s', - cache_type, - len(audio_cache), - len(combined_audio_data), - filename, - ) - return audio_event - - except Exception as e: - logger.error('Failed to flush %s cache: %s', cache_type, e) - return None - - def get_cache_stats( - self, invocation_context: InvocationContext - ) -> dict[str, int]: - """Get statistics about current cache state. - - Args: - invocation_context: The invocation context. - - Returns: - Dictionary containing cache statistics. - """ - input_count = len(invocation_context.input_realtime_cache or []) - output_count = len(invocation_context.output_realtime_cache or []) - - input_bytes = sum( - len(_require_audio_data(entry.data)) - for entry in invocation_context.input_realtime_cache or [] - ) - output_bytes = sum( - len(_require_audio_data(entry.data)) - for entry in invocation_context.output_realtime_cache or [] - ) - - return { - 'input_chunks': input_count, - 'output_chunks': output_count, - 'input_bytes': input_bytes, - 'output_bytes': output_bytes, - 'total_chunks': input_count + output_count, - 'total_bytes': input_bytes + output_bytes, - } +The cache manager now lives in ``_realtime_cache_manager`` as +``RealtimeCacheManager`` and ``RealtimeCacheConfig``. This module keeps the +former names importable; it holds no implementation of its own. +""" +from __future__ import annotations -class AudioCacheConfig: - """Configuration for audio caching behavior.""" +from ._realtime_cache_manager import logger +from ._realtime_cache_manager import RealtimeCacheConfig +from ._realtime_cache_manager import RealtimeCacheEntry +from ._realtime_cache_manager import RealtimeCacheManager - def __init__( - self, - max_cache_size_bytes: int = 10 * 1024 * 1024, # 10MB - max_cache_duration_seconds: float = 300.0, # 5 minutes - auto_flush_threshold: int = 100, # Number of chunks - ) -> None: - """Initialize audio cache configuration. +AudioCacheConfig = RealtimeCacheConfig +AudioCacheManager = RealtimeCacheManager - Args: - max_cache_size_bytes: Maximum cache size in bytes before auto-flush. - max_cache_duration_seconds: Maximum duration to keep data in cache. - auto_flush_threshold: Number of chunks that triggers auto-flush. - """ - self.max_cache_size_bytes = max_cache_size_bytes - self.max_cache_duration_seconds = max_cache_duration_seconds - self.auto_flush_threshold = auto_flush_threshold +__all__ = [ + 'AudioCacheConfig', + 'AudioCacheManager', + 'RealtimeCacheEntry', + 'logger', +] diff --git a/src/google/adk/live/_flow_utils.py b/src/google/adk/live/_flow_utils.py index 2b994d4a50..aa6d5b3c91 100644 --- a/src/google/adk/live/_flow_utils.py +++ b/src/google/adk/live/_flow_utils.py @@ -175,23 +175,23 @@ async def handle_control_event_flush( Returns: A list of Event objects created from the flushed caches. """ - audio_cache_manager = flow.audio_cache_manager + realtime_cache_manager = flow._realtime_cache_manager # Log cache statistics if enabled if DEFAULT_ENABLE_CACHE_STATISTICS: - stats = audio_cache_manager.get_cache_stats(invocation_context) + stats = realtime_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 - return await audio_cache_manager.flush_caches( + return await realtime_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 - return await audio_cache_manager.flush_caches( + return await realtime_cache_manager.flush_caches( invocation_context, flush_user_audio=True, flush_model_audio=True, diff --git a/src/google/adk/live/_live_llm_flow.py b/src/google/adk/live/_live_llm_flow.py index 6aa973bffc..97f6f83984 100644 --- a/src/google/adk/live/_live_llm_flow.py +++ b/src/google/adk/live/_live_llm_flow.py @@ -96,7 +96,7 @@ async def send_to_model( ) -> None: """Sends data to model.""" run_config = _require_run_config(invocation_context) - audio_cache_manager = flow.audio_cache_manager + realtime_cache_manager = flow._realtime_cache_manager while True: live_request_queue = invocation_context.live_request_queue assert live_request_queue is not None @@ -160,7 +160,7 @@ async def send_to_model( and live_request.blob.mime_type and live_request.blob.mime_type.startswith('audio/') ): - audio_cache_manager.cache_audio( + realtime_cache_manager.cache_audio( invocation_context, live_request.blob, cache_type='input' ) @@ -220,7 +220,7 @@ async def receive_from_model( ) -> AsyncGenerator[Event, None]: """Receive data from model and process events using BaseLlmConnection.""" run_config = _require_run_config(invocation_context) - audio_cache_manager = flow.audio_cache_manager + realtime_cache_manager = flow._realtime_cache_manager def get_author_for_event(llm_response: LlmResponse) -> str: """Get the author of the event. @@ -329,7 +329,7 @@ def get_author_for_event(llm_response: LlmResponse) -> str: data=part.inline_data.data, mime_type=part.inline_data.mime_type, ) - audio_cache_manager.cache_audio( + realtime_cache_manager.cache_audio( invocation_context, audio_blob, cache_type='output' ) diff --git a/src/google/adk/live/_realtime_cache_manager.py b/src/google/adk/live/_realtime_cache_manager.py new file mode 100644 index 0000000000..cebee1b7e0 --- /dev/null +++ b/src/google/adk/live/_realtime_cache_manager.py @@ -0,0 +1,321 @@ +# Copyright 2026 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. + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING + +from google.adk.platform import time as platform_time +from google.genai import types +from pydantic import BaseModel +from pydantic import ConfigDict + +from ..events.event import Event + +if TYPE_CHECKING: + from ..agents.invocation_context import InvocationContext + +logger = logging.getLogger('google_adk.' + __name__) + + +class RealtimeCacheEntry(BaseModel): + """Store audio data chunks for caching before flushing.""" + + model_config = ConfigDict( + arbitrary_types_allowed=True, + extra='forbid', + ) + """The pydantic model config.""" + + role: str + """The role that created this audio data, typically "user" or "model".""" + + data: types.Blob + """The audio data chunk.""" + + timestamp: float + """Timestamp when the audio chunk was received.""" + + +def _require_audio_data(blob: types.Blob) -> bytes: + data = blob.data + if not isinstance(data, bytes): + raise ValueError('Audio blobs must contain byte data.') + return data + + +# Deliberately duplicated from `flows.llm_flows.core._utils` +# rather than imported: `live` sits below `flows.llm_flows` in the +# layering, and importing upward would put a cycle back in. +def _require_agent_name(invocation_context: InvocationContext) -> str: + agent = invocation_context.agent + if agent is None: + raise TypeError('Live audio requires an agent in InvocationContext.') + return agent.name + + +class RealtimeCacheManager: + """Manages audio caching and flushing for live streaming flows.""" + + def __init__(self, config: RealtimeCacheConfig | None = None) -> None: + """Initialize the realtime cache manager. + + Args: + config: Configuration for audio caching behavior. + """ + self.config = config or RealtimeCacheConfig() + + def cache_audio( + self, + invocation_context: InvocationContext, + audio_blob: types.Blob, + cache_type: str, + ) -> None: + """Cache incoming user or outgoing model audio data. + + Args: + invocation_context: The current invocation context. + audio_blob: The audio data to cache. + cache_type: Type of audio to cache, either 'input' or 'output'. + + Raises: + ValueError: If cache_type is not 'input' or 'output'. + """ + audio_data = _require_audio_data(audio_blob) + if cache_type == 'input': + if not invocation_context.input_realtime_cache: + invocation_context.input_realtime_cache = [] + cache = invocation_context.input_realtime_cache + role = 'user' + elif cache_type == 'output': + if not invocation_context.output_realtime_cache: + invocation_context.output_realtime_cache = [] + cache = invocation_context.output_realtime_cache + role = 'model' + else: + raise ValueError("cache_type must be either 'input' or 'output'") + + audio_entry = RealtimeCacheEntry( + role=role, data=audio_blob, timestamp=platform_time.get_time() + ) + cache.append(audio_entry) + + logger.debug( + 'Cached %s audio chunk: %d bytes, cache size: %d', + cache_type, + len(audio_data), + len(cache), + ) + + async def flush_caches( + self, + invocation_context: InvocationContext, + flush_user_audio: bool = True, + flush_model_audio: bool = True, + ) -> list[Event]: + """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. + The audio file follows the naming convention: artifact_ref = + f"artifact://{invocation_context.app_name}/{invocation_context.user_id}/ + {invocation_context.session.id}/_adk_live/{filename}#{revision_id}" + + Note: video data is not supported yet. + + A flushed cache is emptied whether or not an artifact was written, so a + cache the artifact service could not accept is dropped rather than kept + for a later attempt. + + Args: + invocation_context: The invocation context containing audio caches. + flush_user_audio: Whether to flush the input (user) audio cache. + flush_model_audio: Whether to flush the output (model) audio cache. + + Returns: + A list of Event objects created from the flushed caches. + """ + flushed_events: list[Event] = [] + if flush_user_audio and invocation_context.input_realtime_cache: + # Swap before the await so audio cached during it is kept for the next + # flush rather than discarded by the unconditional drop. + input_cache = invocation_context.input_realtime_cache + invocation_context.input_realtime_cache = [] + audio_event = await self._flush_cache_to_services( + invocation_context, + input_cache, + 'input_audio', + ) + if audio_event: + flushed_events.append(audio_event) + + if flush_model_audio and invocation_context.output_realtime_cache: + logger.debug('Flushed output audio cache') + output_cache = invocation_context.output_realtime_cache + invocation_context.output_realtime_cache = [] + audio_event = await self._flush_cache_to_services( + invocation_context, + output_cache, + 'output_audio', + ) + if audio_event: + flushed_events.append(audio_event) + + return flushed_events + + async def _flush_cache_to_services( + self, + invocation_context: InvocationContext, + audio_cache: list[RealtimeCacheEntry], + cache_type: str, + ) -> Event | None: + """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. + + Args: + invocation_context: The invocation context. + audio_cache: The audio cache to flush. + cache_type: Type identifier for the cache ('input_audio' or + 'output_audio'). + + Returns: + The created Event if the cache was successfully flushed, None otherwise. + """ + if not invocation_context.artifact_service or not audio_cache: + logger.debug('Skipping cache flush: no artifact service or empty cache') + return None + + try: + # Combine audio chunks into a single file. Use join rather than repeated + # `+=`, which is O(n^2) over the total audio size. + mime_type = audio_cache[0].data.mime_type or 'audio/pcm' + combined_audio_data = b''.join( + entry.data.data or b'' for entry in audio_cache + ) + + # Generate filename with timestamp from first audio chunk (when recording started) + timestamp = int(audio_cache[0].timestamp * 1000) # milliseconds + # Drop MIME parameters such as `;rate=24000` so they don't end up in the + # file extension. The full MIME type is kept on the saved artifact. + extension = mime_type.split(';', 1)[0].strip().split('/')[-1] + filename = f'adk_live_audio_storage_{cache_type}_{timestamp}.{extension}' + + # Save to artifact service + combined_audio_part = types.Part( + inline_data=types.Blob(data=combined_audio_data, mime_type=mime_type) + ) + + revision_id = await invocation_context.artifact_service.save_artifact( + app_name=invocation_context.app_name, + user_id=invocation_context.user_id, + session_id=invocation_context.session.id, + filename=filename, + artifact=combined_audio_part, + ) + + # Create artifact reference for session service + artifact_ref = f'artifact://{invocation_context.app_name}/{invocation_context.user_id}/{invocation_context.session.id}/_adk_live/{filename}#{revision_id}' + + # Create event with file data reference to add to session + # For model events, author should be the agent name, not the role + author = ( + _require_agent_name(invocation_context) + if audio_cache[0].role == 'model' + else audio_cache[0].role + ) + audio_event = Event( + id=Event.new_id(), + invocation_id=invocation_context.invocation_id, + author=author, + content=types.Content( + role=audio_cache[0].role, + parts=[ + types.Part( + file_data=types.FileData( + file_uri=artifact_ref, mime_type=mime_type + ) + ) + ], + ), + timestamp=audio_cache[0].timestamp, + ) + + logger.debug( + 'Successfully flushed %s cache: %d chunks, %d bytes, saved as %s', + cache_type, + len(audio_cache), + len(combined_audio_data), + filename, + ) + return audio_event + + except Exception as e: + logger.error('Failed to flush %s cache: %s', cache_type, e) + return None + + def get_cache_stats( + self, invocation_context: InvocationContext + ) -> dict[str, int]: + """Get statistics about current cache state. + + Args: + invocation_context: The invocation context. + + Returns: + Dictionary containing cache statistics. + """ + input_count = len(invocation_context.input_realtime_cache or []) + output_count = len(invocation_context.output_realtime_cache or []) + + input_bytes = sum( + len(_require_audio_data(entry.data)) + for entry in invocation_context.input_realtime_cache or [] + ) + output_bytes = sum( + len(_require_audio_data(entry.data)) + for entry in invocation_context.output_realtime_cache or [] + ) + + return { + 'input_chunks': input_count, + 'output_chunks': output_count, + 'input_bytes': input_bytes, + 'output_bytes': output_bytes, + 'total_chunks': input_count + output_count, + 'total_bytes': input_bytes + output_bytes, + } + + +class RealtimeCacheConfig: + """Configuration for audio caching behavior.""" + + def __init__( + self, + max_cache_size_bytes: int = 10 * 1024 * 1024, # 10MB + max_cache_duration_seconds: float = 300.0, # 5 minutes + auto_flush_threshold: int = 100, # Number of chunks + ) -> None: + """Initialize realtime cache configuration. + + Args: + max_cache_size_bytes: Maximum cache size in bytes before auto-flush. + max_cache_duration_seconds: Maximum duration to keep data in cache. + auto_flush_threshold: Number of chunks that triggers auto-flush. + """ + self.max_cache_size_bytes = max_cache_size_bytes + self.max_cache_duration_seconds = max_cache_duration_seconds + self.auto_flush_threshold = auto_flush_threshold diff --git a/tests/unittests/live/test_audio_cache_manager.py b/tests/unittests/live/test_audio_cache_manager.py index 8875997e89..8f7e1117e5 100644 --- a/tests/unittests/live/test_audio_cache_manager.py +++ b/tests/unittests/live/test_audio_cache_manager.py @@ -12,563 +12,29 @@ # See the License for the specific language governing permissions and # limitations under the License. -from __future__ import annotations +"""Backward-compatibility tests for the `_audio_cache_manager` aliases.""" -import time -from unittest.mock import AsyncMock -from unittest.mock import Mock +from __future__ import annotations +from google.adk.live import _audio_cache_manager as compat_module from google.adk.live._audio_cache_manager import AudioCacheConfig from google.adk.live._audio_cache_manager import AudioCacheManager from google.adk.live._audio_cache_manager import RealtimeCacheEntry -from google.genai import types -import pydantic -import pytest - -from .. import testing_utils - - -class TestRealtimeCacheEntry: - """RealtimeCacheEntry moved here verbatim; its validation must not drift.""" - - def test_unknown_fields_are_rejected(self): - # `extra="forbid"` came with the model from agents.invocation_context. - # Losing it in the move would silently accept malformed cache entries. - with pytest.raises(pydantic.ValidationError): - RealtimeCacheEntry( - role='user', - data=types.Blob(data=b'x', mime_type='audio/pcm'), - timestamp=0.0, - not_a_field=1, - ) - - def test_accepts_the_declared_fields(self): - entry = RealtimeCacheEntry( - role='user', - data=types.Blob(data=b'x', mime_type='audio/pcm'), - timestamp=1.5, - ) - assert entry.role == 'user' - assert entry.timestamp == 1.5 - - -class TestAudioCacheConfig: - """Test the AudioCacheConfig class.""" - - def test_default_values(self): - """Test that default configuration values are set correctly.""" - config = AudioCacheConfig() - assert config.max_cache_size_bytes == 10 * 1024 * 1024 # 10MB - assert config.max_cache_duration_seconds == 300.0 # 5 minutes - assert config.auto_flush_threshold == 100 - - def test_custom_values(self): - """Test that custom configuration values are set correctly.""" - config = AudioCacheConfig( - max_cache_size_bytes=5 * 1024 * 1024, - max_cache_duration_seconds=120.0, - auto_flush_threshold=50, - ) - assert config.max_cache_size_bytes == 5 * 1024 * 1024 - assert config.max_cache_duration_seconds == 120.0 - assert config.auto_flush_threshold == 50 - - -class TestAudioCacheManager: - """Test the AudioCacheManager class.""" - - def setup_method(self): - """Set up test fixtures.""" - self.config = AudioCacheConfig() - self.manager = AudioCacheManager(self.config) - - @pytest.mark.asyncio - async def test_cache_input_audio(self): - """Test caching input audio data.""" - invocation_context = await testing_utils.create_invocation_context( - testing_utils.create_test_agent() - ) - - audio_blob = types.Blob(data=b'test_audio_data', mime_type='audio/pcm') - - # Initially no cache - assert invocation_context.input_realtime_cache is None - - # Cache audio - self.manager.cache_audio(invocation_context, audio_blob, 'input') - - # Verify cache is created and populated - assert invocation_context.input_realtime_cache is not None - assert len(invocation_context.input_realtime_cache) == 1 - - entry = invocation_context.input_realtime_cache[0] - assert entry.role == 'user' - assert entry.data == audio_blob - assert isinstance(entry.timestamp, float) - - @pytest.mark.asyncio - async def test_cache_audio_rejects_missing_byte_data(self): - invocation_context = await testing_utils.create_invocation_context( - testing_utils.create_test_agent() - ) - - with pytest.raises(ValueError, match='must contain byte data'): - self.manager.cache_audio( - invocation_context, - types.Blob(data=None, mime_type='audio/pcm'), - 'input', - ) - - assert invocation_context.input_realtime_cache is None - - @pytest.mark.asyncio - async def test_cache_output_audio(self): - """Test caching output audio data.""" - invocation_context = await testing_utils.create_invocation_context( - testing_utils.create_test_agent() - ) - - audio_blob = types.Blob(data=b'test_model_audio', mime_type='audio/wav') - - # Initially no cache - assert invocation_context.output_realtime_cache is None - - # Cache audio - self.manager.cache_audio(invocation_context, audio_blob, 'output') - - # Verify cache is created and populated - assert invocation_context.output_realtime_cache is not None - assert len(invocation_context.output_realtime_cache) == 1 - - entry = invocation_context.output_realtime_cache[0] - assert entry.role == 'model' - assert entry.data == audio_blob - assert isinstance(entry.timestamp, float) - - @pytest.mark.asyncio - async def test_multiple_audio_caching(self): - """Test caching multiple audio chunks.""" - invocation_context = await testing_utils.create_invocation_context( - testing_utils.create_test_agent() - ) - - # Cache multiple input audio chunks - for i in range(3): - audio_blob = types.Blob(data=f'input_{i}'.encode(), mime_type='audio/pcm') - self.manager.cache_audio(invocation_context, audio_blob, 'input') - - # Cache multiple output audio chunks - for i in range(2): - audio_blob = types.Blob( - data=f'output_{i}'.encode(), mime_type='audio/wav' - ) - self.manager.cache_audio(invocation_context, audio_blob, 'output') - - # Verify all chunks are cached - assert len(invocation_context.input_realtime_cache) == 3 - assert len(invocation_context.output_realtime_cache) == 2 - - @pytest.mark.asyncio - async def test_flush_caches_both(self): - """Test flushing both input and output caches.""" - invocation_context = await testing_utils.create_invocation_context( - testing_utils.create_test_agent() - ) - - # Set up mock artifact service - mock_artifact_service = AsyncMock() - mock_artifact_service.save_artifact.return_value = 123 - invocation_context.artifact_service = mock_artifact_service - - # Cache some audio - input_blob = types.Blob(data=b'input_data', mime_type='audio/pcm') - output_blob = types.Blob(data=b'output_data', mime_type='audio/wav') - self.manager.cache_audio(invocation_context, input_blob, 'input') - self.manager.cache_audio(invocation_context, output_blob, 'output') - - # Flush caches - await self.manager.flush_caches(invocation_context) - - # Verify caches are cleared - assert invocation_context.input_realtime_cache == [] - assert invocation_context.output_realtime_cache == [] - - # Verify artifact service was called twice (once for each cache) - assert mock_artifact_service.save_artifact.call_count == 2 - - @pytest.mark.asyncio - async def test_flush_caches_selective(self): - """Test selectively flushing only one cache.""" - invocation_context = await testing_utils.create_invocation_context( - testing_utils.create_test_agent() - ) - - # Set up mock artifact service - mock_artifact_service = AsyncMock() - mock_artifact_service.save_artifact.return_value = 123 - invocation_context.artifact_service = mock_artifact_service - - # Cache some audio - input_blob = types.Blob(data=b'input_data', mime_type='audio/pcm') - output_blob = types.Blob(data=b'output_data', mime_type='audio/wav') - self.manager.cache_audio(invocation_context, input_blob, 'input') - self.manager.cache_audio(invocation_context, output_blob, 'output') - - # Flush only input cache - await self.manager.flush_caches( - invocation_context, flush_user_audio=True, flush_model_audio=False - ) - - # Verify only input cache is cleared - assert invocation_context.input_realtime_cache == [] - assert len(invocation_context.output_realtime_cache) == 1 - - # Verify artifact service was called once - assert mock_artifact_service.save_artifact.call_count == 1 - - @pytest.mark.asyncio - async def test_flush_empty_caches(self): - """Test flushing when caches are empty.""" - invocation_context = await testing_utils.create_invocation_context( - testing_utils.create_test_agent() - ) - - # Set up mock artifact service - mock_artifact_service = AsyncMock() - invocation_context.artifact_service = mock_artifact_service - - # Flush empty caches (should not error) - await self.manager.flush_caches(invocation_context) - - # Verify artifact service was not called - mock_artifact_service.save_artifact.assert_not_called() - - @pytest.mark.asyncio - async def test_flush_without_artifact_service_still_clears_cache(self): - """Audio that cannot be saved is dropped rather than accumulated.""" - invocation_context = await testing_utils.create_invocation_context( - testing_utils.create_test_agent() - ) - - # No artifact service - invocation_context.artifact_service = None - - # Cache some audio - input_blob = types.Blob(data=b'input_data', mime_type='audio/pcm') - output_blob = types.Blob(data=b'output_data', mime_type='audio/wav') - self.manager.cache_audio(invocation_context, input_blob, 'input') - self.manager.cache_audio(invocation_context, output_blob, 'output') - - # Flush should not error. - events = await self.manager.flush_caches(invocation_context) - - # There is nowhere to write the audio, so no event is produced, but both - # caches are emptied. Retaining them would grow for the whole session. - assert not events - assert not invocation_context.input_realtime_cache - assert not invocation_context.output_realtime_cache - - @pytest.mark.asyncio - async def test_flush_after_failure_does_not_repeat_earlier_audio(self): - """A failed flush must not leave its audio in the next artifact.""" - invocation_context = await testing_utils.create_invocation_context( - testing_utils.create_test_agent() - ) - - mock_artifact_service = AsyncMock() - mock_artifact_service.save_artifact.side_effect = [ - Exception('Artifact service error'), - 123, - ] - invocation_context.artifact_service = mock_artifact_service - - self.manager.cache_audio( - invocation_context, - types.Blob(data=b'first_turn', mime_type='audio/pcm'), - 'input', - ) - await self.manager.flush_caches(invocation_context) - - self.manager.cache_audio( - invocation_context, - types.Blob(data=b'second_turn', mime_type='audio/pcm'), - 'input', - ) - await self.manager.flush_caches(invocation_context) - - saved = mock_artifact_service.save_artifact.call_args.kwargs['artifact'] - assert saved.inline_data.data == b'second_turn' - - @pytest.mark.asyncio - async def test_flush_keeps_audio_cached_while_saving(self): - """Audio arriving during the save belongs to the next flush.""" - invocation_context = await testing_utils.create_invocation_context( - testing_utils.create_test_agent() - ) - - def cache_while_saving(**kwargs): - self.manager.cache_audio( - invocation_context, - types.Blob(data=b'mid_flush', mime_type='audio/pcm'), - 'input', - ) - return 123 - - mock_artifact_service = AsyncMock() - mock_artifact_service.save_artifact.side_effect = cache_while_saving - invocation_context.artifact_service = mock_artifact_service - - self.manager.cache_audio( - invocation_context, - types.Blob(data=b'first_turn', mime_type='audio/pcm'), - 'input', - ) - await self.manager.flush_caches(invocation_context) - - saved = mock_artifact_service.save_artifact.call_args.kwargs['artifact'] - assert saved.inline_data.data == b'first_turn' - assert [ - entry.data.data for entry in invocation_context.input_realtime_cache - ] == [b'mid_flush'] - - @pytest.mark.asyncio - async def test_flush_artifact_creation(self): - """Test that artifacts are created correctly during flush.""" - invocation_context = await testing_utils.create_invocation_context( - testing_utils.create_test_agent() - ) - - # Set up mock services - mock_artifact_service = AsyncMock() - mock_artifact_service.save_artifact.return_value = 456 - mock_session_service = AsyncMock() - - invocation_context.artifact_service = mock_artifact_service - invocation_context.session_service = mock_session_service - - # Cache audio with specific data - test_data = b'specific_test_audio_data' - audio_blob = types.Blob(data=test_data, mime_type='audio/pcm') - self.manager.cache_audio(invocation_context, audio_blob, 'input') - - # Flush cache - await self.manager.flush_caches(invocation_context) - - # Verify artifact was saved with correct data - mock_artifact_service.save_artifact.assert_called_once() - call_args = mock_artifact_service.save_artifact.call_args - saved_artifact = call_args.kwargs['artifact'] - assert saved_artifact.inline_data.data == test_data - assert saved_artifact.inline_data.mime_type == 'audio/pcm' - - # Verify session event was created - mock_session_service.append_event.assert_not_called() - - def test_get_cache_stats_empty(self): - """Test getting statistics for empty caches.""" - invocation_context = Mock() - invocation_context.input_realtime_cache = None - invocation_context.output_realtime_cache = None - - stats = self.manager.get_cache_stats(invocation_context) - - expected = { - 'input_chunks': 0, - 'output_chunks': 0, - 'input_bytes': 0, - 'output_bytes': 0, - 'total_chunks': 0, - 'total_bytes': 0, - } - assert stats == expected - - @pytest.mark.asyncio - async def test_get_cache_stats_with_data(self): - """Test getting statistics for caches with data.""" - invocation_context = await testing_utils.create_invocation_context( - testing_utils.create_test_agent() - ) - - # Cache some audio data of different sizes - input_blob1 = types.Blob(data=b'12345', mime_type='audio/pcm') # 5 bytes - input_blob2 = types.Blob( - data=b'1234567890', mime_type='audio/pcm' - ) # 10 bytes - output_blob = types.Blob(data=b'abc', mime_type='audio/wav') # 3 bytes - - self.manager.cache_audio(invocation_context, input_blob1, 'input') - self.manager.cache_audio(invocation_context, input_blob2, 'input') - self.manager.cache_audio(invocation_context, output_blob, 'output') - - stats = self.manager.get_cache_stats(invocation_context) - - expected = { - 'input_chunks': 2, - 'output_chunks': 1, - 'input_bytes': 15, # 5 + 10 - 'output_bytes': 3, - 'total_chunks': 3, - 'total_bytes': 18, # 15 + 3 - } - assert stats == expected - - @pytest.mark.asyncio - async def test_error_handling_in_flush(self): - """Test error handling during cache flush operations.""" - invocation_context = await testing_utils.create_invocation_context( - testing_utils.create_test_agent() - ) - - # Set up mock artifact service that raises an error - mock_artifact_service = AsyncMock() - mock_artifact_service.save_artifact.side_effect = Exception( - 'Artifact service error' - ) - invocation_context.artifact_service = mock_artifact_service - - # Cache some audio - audio_blob = types.Blob(data=b'test_data', mime_type='audio/pcm') - self.manager.cache_audio(invocation_context, audio_blob, 'input') - - # Flush should log the error rather than raise. - events = await self.manager.flush_caches(invocation_context) - - # The save failed, so no event is produced, but the cache is still - # emptied. A retry would fail the same way and the audio would pile up. - assert not events - assert not invocation_context.input_realtime_cache - - @pytest.mark.asyncio - async def test_filename_drops_mime_type_parameters(self): - """MIME parameters like `;rate=24000` stay out of the file extension.""" - invocation_context = await testing_utils.create_invocation_context( - testing_utils.create_test_agent() - ) - mock_artifact_service = AsyncMock() - mock_artifact_service.save_artifact.return_value = 0 - invocation_context.artifact_service = mock_artifact_service - audio_blob = types.Blob(data=b'audio', mime_type='audio/pcm;rate=24000') - self.manager.cache_audio(invocation_context, audio_blob, 'output') - - await self.manager.flush_caches(invocation_context) - - call_args = mock_artifact_service.save_artifact.call_args - assert call_args.kwargs['filename'].endswith('.pcm') - assert ( - call_args.kwargs['artifact'].inline_data.mime_type - == 'audio/pcm;rate=24000' - ) - - @pytest.mark.asyncio - async def test_filename_uses_first_chunk_timestamp(self): - """Test that the filename timestamp comes from the first audio chunk, not flush time.""" - invocation_context = await testing_utils.create_invocation_context( - testing_utils.create_test_agent() - ) - - # Set up mock services - mock_artifact_service = AsyncMock() - mock_artifact_service.save_artifact.return_value = 789 - mock_session_service = AsyncMock() - - invocation_context.artifact_service = mock_artifact_service - invocation_context.session_service = mock_session_service - - # Cache multiple audio chunks with specific timestamps - first_timestamp = 1234567890.123 # First chunk timestamp - second_timestamp = 1234567891.456 # Second chunk timestamp (later) - - # Manually create audio cache entries with specific timestamps - invocation_context.input_realtime_cache = [] - - from google.adk.agents.invocation_context import RealtimeCacheEntry - - first_entry = RealtimeCacheEntry( - role='user', - data=types.Blob(data=b'first_chunk', mime_type='audio/pcm'), - timestamp=first_timestamp, - ) - - second_entry = RealtimeCacheEntry( - role='user', - data=types.Blob(data=b'second_chunk', mime_type='audio/pcm'), - timestamp=second_timestamp, - ) - - invocation_context.input_realtime_cache.extend([first_entry, second_entry]) - - # Sleep briefly to ensure current time is different from first timestamp - time.sleep(0.01) - - # Flush cache - await self.manager.flush_caches(invocation_context) - - # Verify artifact was saved - mock_artifact_service.save_artifact.assert_called_once() - call_args = mock_artifact_service.save_artifact.call_args - filename = call_args.kwargs['filename'] - - # Extract timestamp from filename (format: input_audio_{timestamp}.pcm) - expected_timestamp_ms = int(first_timestamp * 1000) - assert ( - f'adk_live_audio_storage_input_audio_{expected_timestamp_ms}.pcm' - == filename - ) - - # Verify the timestamp in filename matches first chunk, not current time - current_timestamp_ms = int(time.time() * 1000) - assert expected_timestamp_ms != current_timestamp_ms # Should be different - assert filename.startswith( - f'adk_live_audio_storage_input_audio_{expected_timestamp_ms}' - ) - - @pytest.mark.asyncio - async def test_flush_event_author_for_user_audio(self): - """Test that flushed user audio events have 'user' as author.""" - invocation_context = await testing_utils.create_invocation_context( - testing_utils.create_test_agent() - ) - - # Set up mock artifact service - mock_artifact_service = AsyncMock() - mock_artifact_service.save_artifact.return_value = 123 - invocation_context.artifact_service = mock_artifact_service - - # Cache user input audio - input_blob = types.Blob(data=b'user_audio_data', mime_type='audio/pcm') - self.manager.cache_audio(invocation_context, input_blob, 'input') - - # Flush cache and get events - events = await self.manager.flush_caches( - invocation_context, flush_user_audio=True, flush_model_audio=False - ) - - # Verify event author is 'user' for user audio - assert len(events) == 1 - assert events[0].author == 'user' - assert events[0].content.role == 'user' - - @pytest.mark.asyncio - async def test_flush_event_author_for_model_audio(self): - """Test that flushed model audio events have agent name as author, not 'model'.""" - agent = testing_utils.create_test_agent(name='my_test_agent') - invocation_context = await testing_utils.create_invocation_context(agent) - - # Set up mock artifact service - mock_artifact_service = AsyncMock() - mock_artifact_service.save_artifact.return_value = 123 - invocation_context.artifact_service = mock_artifact_service - - # Cache model output audio - output_blob = types.Blob(data=b'model_audio_data', mime_type='audio/wav') - self.manager.cache_audio(invocation_context, output_blob, 'output') - - # Flush cache and get events - events = await self.manager.flush_caches( - invocation_context, flush_user_audio=False, flush_model_audio=True - ) - - # Verify event author is agent name (not 'model') for model audio - assert len(events) == 1 - assert events[0].author == 'my_test_agent' # Agent name, not 'model' - assert events[0].content.role == 'model' # Role is still 'model' +from google.adk.live._realtime_cache_manager import logger as canonical_logger +from google.adk.live._realtime_cache_manager import RealtimeCacheConfig +from google.adk.live._realtime_cache_manager import RealtimeCacheEntry as CanonicalRealtimeCacheEntry +from google.adk.live._realtime_cache_manager import RealtimeCacheManager + + +def test_audio_cache_manager_reexports_realtime_cache_manager_symbols(): + """The legacy `_audio_cache_manager` module aliases `_realtime_cache_manager`.""" + assert AudioCacheManager is RealtimeCacheManager + assert AudioCacheConfig is RealtimeCacheConfig + assert RealtimeCacheEntry is CanonicalRealtimeCacheEntry + assert compat_module.logger is canonical_logger + assert set(compat_module.__all__) == { + 'AudioCacheConfig', + 'AudioCacheManager', + 'RealtimeCacheEntry', + 'logger', + } diff --git a/tests/unittests/live/test_flow_utils.py b/tests/unittests/live/test_flow_utils.py index 13957da6d4..e6a126f834 100644 --- a/tests/unittests/live/test_flow_utils.py +++ b/tests/unittests/live/test_flow_utils.py @@ -47,11 +47,11 @@ def test_require_live_request_queue_raises_when_missing(): async def test_handle_control_event_flush_flushes_model_audio_on_interrupt(): """Flushes only model audio when the response indicates an interruption.""" flushed_event = Event(author='model') - audio_cache_manager = mock.Mock() - audio_cache_manager.flush_caches = mock.AsyncMock( + realtime_cache_manager = mock.Mock() + realtime_cache_manager.flush_caches = mock.AsyncMock( return_value=[flushed_event] ) - flow = mock.Mock(audio_cache_manager=audio_cache_manager) + flow = mock.Mock(_realtime_cache_manager=realtime_cache_manager) invocation_context = mock.Mock() llm_response = LlmResponse(interrupted=True) @@ -60,7 +60,7 @@ async def test_handle_control_event_flush_flushes_model_audio_on_interrupt(): ) assert events == [flushed_event] - audio_cache_manager.flush_caches.assert_awaited_once_with( + realtime_cache_manager.flush_caches.assert_awaited_once_with( invocation_context, flush_user_audio=False, flush_model_audio=True, @@ -70,11 +70,11 @@ async def test_handle_control_event_flush_flushes_model_audio_on_interrupt(): async def test_handle_control_event_flush_flushes_both_on_turn_complete(): """Flushes both user and model audio when the turn completes.""" flushed_event = Event(author='model') - audio_cache_manager = mock.Mock() - audio_cache_manager.flush_caches = mock.AsyncMock( + realtime_cache_manager = mock.Mock() + realtime_cache_manager.flush_caches = mock.AsyncMock( return_value=[flushed_event] ) - flow = mock.Mock(audio_cache_manager=audio_cache_manager) + flow = mock.Mock(_realtime_cache_manager=realtime_cache_manager) invocation_context = mock.Mock() llm_response = LlmResponse(turn_complete=True) @@ -83,7 +83,7 @@ async def test_handle_control_event_flush_flushes_both_on_turn_complete(): ) assert events == [flushed_event] - audio_cache_manager.flush_caches.assert_awaited_once_with( + realtime_cache_manager.flush_caches.assert_awaited_once_with( invocation_context, flush_user_audio=True, flush_model_audio=True, diff --git a/tests/unittests/live/test_live_llm_flow.py b/tests/unittests/live/test_live_llm_flow.py index 6acfce02bb..6563f0cc76 100644 --- a/tests/unittests/live/test_live_llm_flow.py +++ b/tests/unittests/live/test_live_llm_flow.py @@ -181,7 +181,7 @@ async def test_handle_control_event_flush_on_interrupted(): response = LlmResponse(interrupted=True) with mock.patch.object( - flow.audio_cache_manager, 'flush_caches', new_callable=mock.AsyncMock + flow._realtime_cache_manager, 'flush_caches', new_callable=mock.AsyncMock ) as mock_flush: mock_flush.return_value = [Event(id='flushed-event')] events = await _live_llm_flow.handle_control_event_flush( @@ -201,7 +201,7 @@ async def test_handle_control_event_flush_on_turn_complete(): response = LlmResponse(turn_complete=True) with mock.patch.object( - flow.audio_cache_manager, 'flush_caches', new_callable=mock.AsyncMock + flow._realtime_cache_manager, 'flush_caches', new_callable=mock.AsyncMock ) as mock_flush: mock_flush.return_value = [Event(id='flushed-event')] events = await _live_llm_flow.handle_control_event_flush( @@ -308,10 +308,10 @@ async def test_handle_control_event_flush_logs_stats_when_enabled(): with ( mock.patch.object(_flow_utils, 'DEFAULT_ENABLE_CACHE_STATISTICS', True), mock.patch.object( - flow.audio_cache_manager, 'get_cache_stats' + flow._realtime_cache_manager, 'get_cache_stats' ) as mock_get_stats, mock.patch.object( - flow.audio_cache_manager, 'flush_caches', return_value=[] + flow._realtime_cache_manager, 'flush_caches', return_value=[] ), ): await _live_llm_flow.handle_control_event_flush(flow, context, response) @@ -319,8 +319,8 @@ async def test_handle_control_event_flush_logs_stats_when_enabled(): mock_get_stats.assert_called_once_with(context) -async def test_send_to_model_uses_flow_audio_cache_manager(): - """send_to_model accesses the audio cache manager directly from the flow instance.""" +async def test_send_to_model_uses_flow_realtime_cache_manager(): + """send_to_model accesses the realtime cache manager directly from the flow instance.""" flow = _TestBaseLlmFlow() queue = LiveRequestQueue() queue.send_realtime(types.Blob(mime_type='audio/pcm', data=b'audio_bytes')) @@ -330,7 +330,7 @@ async def test_send_to_model_uses_flow_audio_cache_manager(): mock_connection = mock.AsyncMock() with mock.patch.object( - flow.audio_cache_manager, 'cache_audio' + flow._realtime_cache_manager, 'cache_audio' ) as mock_cache_audio: # Run send_to_model briefly and cancel it after processing the queued item send_task = asyncio.create_task( @@ -398,3 +398,52 @@ async def test_send_to_model_does_not_mutate_caller_content_role(): assert content.role is None sent_content = mock_connection._send_content.await_args.args[0] assert sent_content.role == 'user' + + +async def test_receive_from_model_caches_only_audio_blobs(): + """Model audio blobs are cached; other inline blobs are not.""" + flow = _TestBaseLlmFlow() + context = _create_test_context(run_config=RunConfig(save_live_blob=True)) + calls = 0 + + async def _receive(): + nonlocal calls + calls += 1 + if calls == 1: + yield LlmResponse( + content=types.Content( + role='model', + parts=[ + types.Part( + inline_data=types.Blob( + mime_type='image/png', data=b'model_image' + ) + ), + types.Part( + inline_data=types.Blob( + mime_type='audio/pcm', data=b'model_audio' + ) + ), + ], + ) + ) + + connection = mock.MagicMock(receive=_receive) + events = [ + event + async for event in _live_llm_flow.receive_from_model( + flow, connection, context, LlmRequest() + ) + ] + + assert len(events) == 1 + assert [entry.data.data for entry in context.output_realtime_cache] == [ + b'model_audio' + ] + + +def test_audio_cache_manager_alias_points_to_realtime_cache_manager(): + """`flow.audio_cache_manager` is kept as an alias of the flow's manager.""" + flow = _TestBaseLlmFlow() + + assert flow.audio_cache_manager is flow._realtime_cache_manager diff --git a/tests/unittests/live/test_realtime_cache_manager.py b/tests/unittests/live/test_realtime_cache_manager.py new file mode 100644 index 0000000000..b5ffa33118 --- /dev/null +++ b/tests/unittests/live/test_realtime_cache_manager.py @@ -0,0 +1,574 @@ +# Copyright 2026 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. + +from __future__ import annotations + +import time +from unittest.mock import AsyncMock +from unittest.mock import Mock + +from google.adk.live._realtime_cache_manager import RealtimeCacheConfig +from google.adk.live._realtime_cache_manager import RealtimeCacheEntry +from google.adk.live._realtime_cache_manager import RealtimeCacheManager +from google.genai import types +import pydantic +import pytest + +from .. import testing_utils + + +class TestRealtimeCacheEntry: + """RealtimeCacheEntry moved here verbatim; its validation must not drift.""" + + def test_unknown_fields_are_rejected(self): + # `extra="forbid"` came with the model from agents.invocation_context. + # Losing it in the move would silently accept malformed cache entries. + with pytest.raises(pydantic.ValidationError): + RealtimeCacheEntry( + role='user', + data=types.Blob(data=b'x', mime_type='audio/pcm'), + timestamp=0.0, + not_a_field=1, + ) + + def test_accepts_the_declared_fields(self): + entry = RealtimeCacheEntry( + role='user', + data=types.Blob(data=b'x', mime_type='audio/pcm'), + timestamp=1.5, + ) + assert entry.role == 'user' + assert entry.timestamp == 1.5 + + +class TestRealtimeCacheConfig: + """Test the RealtimeCacheConfig class.""" + + def test_default_values(self): + """Test that default configuration values are set correctly.""" + config = RealtimeCacheConfig() + assert config.max_cache_size_bytes == 10 * 1024 * 1024 # 10MB + assert config.max_cache_duration_seconds == 300.0 # 5 minutes + assert config.auto_flush_threshold == 100 + + def test_custom_values(self): + """Test that custom configuration values are set correctly.""" + config = RealtimeCacheConfig( + max_cache_size_bytes=5 * 1024 * 1024, + max_cache_duration_seconds=120.0, + auto_flush_threshold=50, + ) + assert config.max_cache_size_bytes == 5 * 1024 * 1024 + assert config.max_cache_duration_seconds == 120.0 + assert config.auto_flush_threshold == 50 + + +class TestRealtimeCacheManager: + """Test the RealtimeCacheManager class.""" + + def setup_method(self): + """Set up test fixtures.""" + self.config = RealtimeCacheConfig() + self.manager = RealtimeCacheManager(self.config) + + @pytest.mark.asyncio + async def test_cache_input_audio(self): + """Test caching input audio data.""" + invocation_context = await testing_utils.create_invocation_context( + testing_utils.create_test_agent() + ) + + audio_blob = types.Blob(data=b'test_audio_data', mime_type='audio/pcm') + + # Initially no cache + assert invocation_context.input_realtime_cache is None + + # Cache audio + self.manager.cache_audio(invocation_context, audio_blob, 'input') + + # Verify cache is created and populated + assert invocation_context.input_realtime_cache is not None + assert len(invocation_context.input_realtime_cache) == 1 + + entry = invocation_context.input_realtime_cache[0] + assert entry.role == 'user' + assert entry.data == audio_blob + assert isinstance(entry.timestamp, float) + + @pytest.mark.asyncio + async def test_cache_audio_rejects_missing_byte_data(self): + invocation_context = await testing_utils.create_invocation_context( + testing_utils.create_test_agent() + ) + + with pytest.raises(ValueError, match='must contain byte data'): + self.manager.cache_audio( + invocation_context, + types.Blob(data=None, mime_type='audio/pcm'), + 'input', + ) + + assert invocation_context.input_realtime_cache is None + + @pytest.mark.asyncio + async def test_cache_output_audio(self): + """Test caching output audio data.""" + invocation_context = await testing_utils.create_invocation_context( + testing_utils.create_test_agent() + ) + + audio_blob = types.Blob(data=b'test_model_audio', mime_type='audio/wav') + + # Initially no cache + assert invocation_context.output_realtime_cache is None + + # Cache audio + self.manager.cache_audio(invocation_context, audio_blob, 'output') + + # Verify cache is created and populated + assert invocation_context.output_realtime_cache is not None + assert len(invocation_context.output_realtime_cache) == 1 + + entry = invocation_context.output_realtime_cache[0] + assert entry.role == 'model' + assert entry.data == audio_blob + assert isinstance(entry.timestamp, float) + + @pytest.mark.asyncio + async def test_multiple_audio_caching(self): + """Test caching multiple audio chunks.""" + invocation_context = await testing_utils.create_invocation_context( + testing_utils.create_test_agent() + ) + + # Cache multiple input audio chunks + for i in range(3): + audio_blob = types.Blob(data=f'input_{i}'.encode(), mime_type='audio/pcm') + self.manager.cache_audio(invocation_context, audio_blob, 'input') + + # Cache multiple output audio chunks + for i in range(2): + audio_blob = types.Blob( + data=f'output_{i}'.encode(), mime_type='audio/wav' + ) + self.manager.cache_audio(invocation_context, audio_blob, 'output') + + # Verify all chunks are cached + assert len(invocation_context.input_realtime_cache) == 3 + assert len(invocation_context.output_realtime_cache) == 2 + + @pytest.mark.asyncio + async def test_flush_caches_both(self): + """Test flushing both input and output caches.""" + invocation_context = await testing_utils.create_invocation_context( + testing_utils.create_test_agent() + ) + + # Set up mock artifact service + mock_artifact_service = AsyncMock() + mock_artifact_service.save_artifact.return_value = 123 + invocation_context.artifact_service = mock_artifact_service + + # Cache some audio + input_blob = types.Blob(data=b'input_data', mime_type='audio/pcm') + output_blob = types.Blob(data=b'output_data', mime_type='audio/wav') + self.manager.cache_audio(invocation_context, input_blob, 'input') + self.manager.cache_audio(invocation_context, output_blob, 'output') + + # Flush caches + await self.manager.flush_caches(invocation_context) + + # Verify caches are cleared + assert invocation_context.input_realtime_cache == [] + assert invocation_context.output_realtime_cache == [] + + # Verify artifact service was called twice (once for each cache) + assert mock_artifact_service.save_artifact.call_count == 2 + + @pytest.mark.asyncio + async def test_flush_caches_selective(self): + """Test selectively flushing only one cache.""" + invocation_context = await testing_utils.create_invocation_context( + testing_utils.create_test_agent() + ) + + # Set up mock artifact service + mock_artifact_service = AsyncMock() + mock_artifact_service.save_artifact.return_value = 123 + invocation_context.artifact_service = mock_artifact_service + + # Cache some audio + input_blob = types.Blob(data=b'input_data', mime_type='audio/pcm') + output_blob = types.Blob(data=b'output_data', mime_type='audio/wav') + self.manager.cache_audio(invocation_context, input_blob, 'input') + self.manager.cache_audio(invocation_context, output_blob, 'output') + + # Flush only input cache + await self.manager.flush_caches( + invocation_context, flush_user_audio=True, flush_model_audio=False + ) + + # Verify only input cache is cleared + assert invocation_context.input_realtime_cache == [] + assert len(invocation_context.output_realtime_cache) == 1 + + # Verify artifact service was called once + assert mock_artifact_service.save_artifact.call_count == 1 + + @pytest.mark.asyncio + async def test_flush_empty_caches(self): + """Test flushing when caches are empty.""" + invocation_context = await testing_utils.create_invocation_context( + testing_utils.create_test_agent() + ) + + # Set up mock artifact service + mock_artifact_service = AsyncMock() + invocation_context.artifact_service = mock_artifact_service + + # Flush empty caches (should not error) + await self.manager.flush_caches(invocation_context) + + # Verify artifact service was not called + mock_artifact_service.save_artifact.assert_not_called() + + @pytest.mark.asyncio + async def test_flush_without_artifact_service_still_clears_cache(self): + """Audio that cannot be saved is dropped rather than accumulated.""" + invocation_context = await testing_utils.create_invocation_context( + testing_utils.create_test_agent() + ) + + # No artifact service + invocation_context.artifact_service = None + + # Cache some audio + input_blob = types.Blob(data=b'input_data', mime_type='audio/pcm') + output_blob = types.Blob(data=b'output_data', mime_type='audio/wav') + self.manager.cache_audio(invocation_context, input_blob, 'input') + self.manager.cache_audio(invocation_context, output_blob, 'output') + + # Flush should not error. + events = await self.manager.flush_caches(invocation_context) + + # There is nowhere to write the audio, so no event is produced, but both + # caches are emptied. Retaining them would grow for the whole session. + assert not events + assert not invocation_context.input_realtime_cache + assert not invocation_context.output_realtime_cache + + @pytest.mark.asyncio + async def test_flush_after_failure_does_not_repeat_earlier_audio(self): + """A failed flush must not leave its audio in the next artifact.""" + invocation_context = await testing_utils.create_invocation_context( + testing_utils.create_test_agent() + ) + + mock_artifact_service = AsyncMock() + mock_artifact_service.save_artifact.side_effect = [ + Exception('Artifact service error'), + 123, + ] + invocation_context.artifact_service = mock_artifact_service + + self.manager.cache_audio( + invocation_context, + types.Blob(data=b'first_turn', mime_type='audio/pcm'), + 'input', + ) + await self.manager.flush_caches(invocation_context) + + self.manager.cache_audio( + invocation_context, + types.Blob(data=b'second_turn', mime_type='audio/pcm'), + 'input', + ) + await self.manager.flush_caches(invocation_context) + + saved = mock_artifact_service.save_artifact.call_args.kwargs['artifact'] + assert saved.inline_data.data == b'second_turn' + + @pytest.mark.asyncio + async def test_flush_keeps_audio_cached_while_saving(self): + """Audio arriving during the save belongs to the next flush.""" + invocation_context = await testing_utils.create_invocation_context( + testing_utils.create_test_agent() + ) + + def cache_while_saving(**kwargs): + self.manager.cache_audio( + invocation_context, + types.Blob(data=b'mid_flush', mime_type='audio/pcm'), + 'input', + ) + return 123 + + mock_artifact_service = AsyncMock() + mock_artifact_service.save_artifact.side_effect = cache_while_saving + invocation_context.artifact_service = mock_artifact_service + + self.manager.cache_audio( + invocation_context, + types.Blob(data=b'first_turn', mime_type='audio/pcm'), + 'input', + ) + await self.manager.flush_caches(invocation_context) + + saved = mock_artifact_service.save_artifact.call_args.kwargs['artifact'] + assert saved.inline_data.data == b'first_turn' + assert [ + entry.data.data for entry in invocation_context.input_realtime_cache + ] == [b'mid_flush'] + + @pytest.mark.asyncio + async def test_flush_artifact_creation(self): + """Test that artifacts are created correctly during flush.""" + invocation_context = await testing_utils.create_invocation_context( + testing_utils.create_test_agent() + ) + + # Set up mock services + mock_artifact_service = AsyncMock() + mock_artifact_service.save_artifact.return_value = 456 + mock_session_service = AsyncMock() + + invocation_context.artifact_service = mock_artifact_service + invocation_context.session_service = mock_session_service + + # Cache audio with specific data + test_data = b'specific_test_audio_data' + audio_blob = types.Blob(data=test_data, mime_type='audio/pcm') + self.manager.cache_audio(invocation_context, audio_blob, 'input') + + # Flush cache + await self.manager.flush_caches(invocation_context) + + # Verify artifact was saved with correct data + mock_artifact_service.save_artifact.assert_called_once() + call_args = mock_artifact_service.save_artifact.call_args + saved_artifact = call_args.kwargs['artifact'] + assert saved_artifact.inline_data.data == test_data + assert saved_artifact.inline_data.mime_type == 'audio/pcm' + + # Verify session event was created + mock_session_service.append_event.assert_not_called() + + def test_get_cache_stats_empty(self): + """Test getting statistics for empty caches.""" + invocation_context = Mock() + invocation_context.input_realtime_cache = None + invocation_context.output_realtime_cache = None + + stats = self.manager.get_cache_stats(invocation_context) + + expected = { + 'input_chunks': 0, + 'output_chunks': 0, + 'input_bytes': 0, + 'output_bytes': 0, + 'total_chunks': 0, + 'total_bytes': 0, + } + assert stats == expected + + @pytest.mark.asyncio + async def test_get_cache_stats_with_data(self): + """Test getting statistics for caches with data.""" + invocation_context = await testing_utils.create_invocation_context( + testing_utils.create_test_agent() + ) + + # Cache some audio data of different sizes + input_blob1 = types.Blob(data=b'12345', mime_type='audio/pcm') # 5 bytes + input_blob2 = types.Blob( + data=b'1234567890', mime_type='audio/pcm' + ) # 10 bytes + output_blob = types.Blob(data=b'abc', mime_type='audio/wav') # 3 bytes + + self.manager.cache_audio(invocation_context, input_blob1, 'input') + self.manager.cache_audio(invocation_context, input_blob2, 'input') + self.manager.cache_audio(invocation_context, output_blob, 'output') + + stats = self.manager.get_cache_stats(invocation_context) + + expected = { + 'input_chunks': 2, + 'output_chunks': 1, + 'input_bytes': 15, # 5 + 10 + 'output_bytes': 3, + 'total_chunks': 3, + 'total_bytes': 18, # 15 + 3 + } + assert stats == expected + + @pytest.mark.asyncio + async def test_error_handling_in_flush(self): + """Test error handling during cache flush operations.""" + invocation_context = await testing_utils.create_invocation_context( + testing_utils.create_test_agent() + ) + + # Set up mock artifact service that raises an error + mock_artifact_service = AsyncMock() + mock_artifact_service.save_artifact.side_effect = Exception( + 'Artifact service error' + ) + invocation_context.artifact_service = mock_artifact_service + + # Cache some audio + audio_blob = types.Blob(data=b'test_data', mime_type='audio/pcm') + self.manager.cache_audio(invocation_context, audio_blob, 'input') + + # Flush should log the error rather than raise. + events = await self.manager.flush_caches(invocation_context) + + # The save failed, so no event is produced, but the cache is still + # emptied. A retry would fail the same way and the audio would pile up. + assert not events + assert not invocation_context.input_realtime_cache + + @pytest.mark.asyncio + async def test_filename_drops_mime_type_parameters(self): + """MIME parameters like `;rate=24000` stay out of the file extension.""" + invocation_context = await testing_utils.create_invocation_context( + testing_utils.create_test_agent() + ) + mock_artifact_service = AsyncMock() + mock_artifact_service.save_artifact.return_value = 0 + invocation_context.artifact_service = mock_artifact_service + audio_blob = types.Blob(data=b'audio', mime_type='audio/pcm;rate=24000') + self.manager.cache_audio(invocation_context, audio_blob, 'output') + + await self.manager.flush_caches(invocation_context) + + call_args = mock_artifact_service.save_artifact.call_args + assert call_args.kwargs['filename'].endswith('.pcm') + assert ( + call_args.kwargs['artifact'].inline_data.mime_type + == 'audio/pcm;rate=24000' + ) + + @pytest.mark.asyncio + async def test_filename_uses_first_chunk_timestamp(self): + """Test that the filename timestamp comes from the first audio chunk, not flush time.""" + invocation_context = await testing_utils.create_invocation_context( + testing_utils.create_test_agent() + ) + + # Set up mock services + mock_artifact_service = AsyncMock() + mock_artifact_service.save_artifact.return_value = 789 + mock_session_service = AsyncMock() + + invocation_context.artifact_service = mock_artifact_service + invocation_context.session_service = mock_session_service + + # Cache multiple audio chunks with specific timestamps + first_timestamp = 1234567890.123 # First chunk timestamp + second_timestamp = 1234567891.456 # Second chunk timestamp (later) + + # Manually create audio cache entries with specific timestamps + invocation_context.input_realtime_cache = [] + + from google.adk.agents.invocation_context import RealtimeCacheEntry + + first_entry = RealtimeCacheEntry( + role='user', + data=types.Blob(data=b'first_chunk', mime_type='audio/pcm'), + timestamp=first_timestamp, + ) + + second_entry = RealtimeCacheEntry( + role='user', + data=types.Blob(data=b'second_chunk', mime_type='audio/pcm'), + timestamp=second_timestamp, + ) + + invocation_context.input_realtime_cache.extend([first_entry, second_entry]) + + # Sleep briefly to ensure current time is different from first timestamp + time.sleep(0.01) + + # Flush cache + await self.manager.flush_caches(invocation_context) + + # Verify artifact was saved + mock_artifact_service.save_artifact.assert_called_once() + call_args = mock_artifact_service.save_artifact.call_args + filename = call_args.kwargs['filename'] + + # Extract timestamp from filename (format: input_audio_{timestamp}.pcm) + expected_timestamp_ms = int(first_timestamp * 1000) + assert ( + f'adk_live_audio_storage_input_audio_{expected_timestamp_ms}.pcm' + == filename + ) + + # Verify the timestamp in filename matches first chunk, not current time + current_timestamp_ms = int(time.time() * 1000) + assert expected_timestamp_ms != current_timestamp_ms # Should be different + assert filename.startswith( + f'adk_live_audio_storage_input_audio_{expected_timestamp_ms}' + ) + + @pytest.mark.asyncio + async def test_flush_event_author_for_user_audio(self): + """Test that flushed user audio events have 'user' as author.""" + invocation_context = await testing_utils.create_invocation_context( + testing_utils.create_test_agent() + ) + + # Set up mock artifact service + mock_artifact_service = AsyncMock() + mock_artifact_service.save_artifact.return_value = 123 + invocation_context.artifact_service = mock_artifact_service + + # Cache user input audio + input_blob = types.Blob(data=b'user_audio_data', mime_type='audio/pcm') + self.manager.cache_audio(invocation_context, input_blob, 'input') + + # Flush cache and get events + events = await self.manager.flush_caches( + invocation_context, flush_user_audio=True, flush_model_audio=False + ) + + # Verify event author is 'user' for user audio + assert len(events) == 1 + assert events[0].author == 'user' + assert events[0].content.role == 'user' + + @pytest.mark.asyncio + async def test_flush_event_author_for_model_audio(self): + """Test that flushed model audio events have agent name as author, not 'model'.""" + agent = testing_utils.create_test_agent(name='my_test_agent') + invocation_context = await testing_utils.create_invocation_context(agent) + + # Set up mock artifact service + mock_artifact_service = AsyncMock() + mock_artifact_service.save_artifact.return_value = 123 + invocation_context.artifact_service = mock_artifact_service + + # Cache model output audio + output_blob = types.Blob(data=b'model_audio_data', mime_type='audio/wav') + self.manager.cache_audio(invocation_context, output_blob, 'output') + + # Flush cache and get events + events = await self.manager.flush_caches( + invocation_context, flush_user_audio=False, flush_model_audio=True + ) + + # Verify event author is agent name (not 'model') for model audio + assert len(events) == 1 + assert events[0].author == 'my_test_agent' # Agent name, not 'model' + assert events[0].content.role == 'model' # Role is still 'model'