Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion src/google/adk/agents/invocation_context.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
7 changes: 5 additions & 2 deletions src/google/adk/flows/llm_flows/base_llm_flow.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down
320 changes: 18 additions & 302 deletions src/google/adk/live/_audio_cache_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -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',
]
8 changes: 4 additions & 4 deletions src/google/adk/live/_flow_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
Loading
Loading