Skip to content

Commit e8b0192

Browse files
google-genai-botcopybara-github
authored andcommitted
refactor(live): rename AudioCacheManager to RealtimeCacheManager
- Move the live audio cache manager to `google.adk.live._realtime_cache_manager` as `RealtimeCacheManager` and `RealtimeCacheConfig`. `google.adk.live._audio_cache_manager` keeps `AudioCacheManager` and `AudioCacheConfig` as aliases. - `flow.audio_cache_manager` still points to the flow's cache manager. Assigning a different manager to it no longer changes the one the flow uses. - No other behavior changes: artifact filenames, flush semantics, and error messages are unchanged. PiperOrigin-RevId: 998293453
1 parent 9c83a65 commit e8b0192

10 files changed

Lines changed: 1012 additions & 883 deletions

‎src/google/adk/agents/invocation_context.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@
3535
from ..events._node_path_builder import _NodePathBuilder
3636
from ..events.event import Event
3737
from ..live._active_streaming_tool import ActiveStreamingTool
38-
from ..live._audio_cache_manager import RealtimeCacheEntry as RealtimeCacheEntry
38+
from ..live._realtime_cache_manager import RealtimeCacheEntry as RealtimeCacheEntry
3939
from ..live._transcription_entry import TranscriptionEntry
4040
from ..live.live_request_queue import LiveRequestQueue
4141
from ..memory.base_memory_service import BaseMemoryService

‎src/google/adk/flows/llm_flows/base_llm_flow.py‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -30,11 +30,11 @@
3030
from ...agents.base_agent import BaseAgent
3131
from ...agents.invocation_context import InvocationContext
3232
from ...events.event import Event
33-
from ...live._audio_cache_manager import AudioCacheManager
3433
from ...live._flow_utils import DEFAULT_ENABLE_CACHE_STATISTICS as DEFAULT_ENABLE_CACHE_STATISTICS
3534
from ...live._flow_utils import DEFAULT_MAX_RECONNECT_ATTEMPTS as DEFAULT_MAX_RECONNECT_ATTEMPTS
3635
from ...live._flow_utils import DEFAULT_TASK_COMPLETION_DELAY as DEFAULT_TASK_COMPLETION_DELAY
3736
from ...live._flow_utils import DEFAULT_TRANSFER_AGENT_DELAY as DEFAULT_TRANSFER_AGENT_DELAY
37+
from ...live._realtime_cache_manager import RealtimeCacheManager as _RealtimeCacheManager
3838
from ...models.base_llm_connection import BaseLlmConnection
3939
from ...models.llm_request import LlmRequest
4040
from ...models.llm_response import LlmResponse
@@ -125,7 +125,10 @@ def __init__(self) -> None:
125125
self.response_processors: list[BaseLlmResponseProcessor] = []
126126

127127
# Initialize configuration and managers
128-
self.audio_cache_manager = AudioCacheManager()
128+
self._realtime_cache_manager = _RealtimeCacheManager()
129+
# Kept for backward compatibility. The flow itself reads
130+
# `_realtime_cache_manager`, so reassigning this attribute has no effect.
131+
self.audio_cache_manager = self._realtime_cache_manager
129132

130133
def _request_processor_lists(
131134
self,

‎src/google/adk/live/_audio_cache_manager.py‎

Lines changed: 18 additions & 302 deletions
Original file line numberDiff line numberDiff line change
@@ -12,310 +12,26 @@
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
1414

15-
from __future__ import annotations
16-
17-
import logging
18-
from typing import TYPE_CHECKING
19-
20-
from google.adk.platform import time as platform_time
21-
from google.genai import types
22-
from pydantic import BaseModel
23-
from pydantic import ConfigDict
24-
25-
from ..events.event import Event
26-
27-
if TYPE_CHECKING:
28-
from ..agents.invocation_context import InvocationContext
29-
30-
logger = logging.getLogger('google_adk.' + __name__)
31-
32-
33-
class RealtimeCacheEntry(BaseModel):
34-
"""Store audio data chunks for caching before flushing."""
35-
36-
model_config = ConfigDict(
37-
arbitrary_types_allowed=True,
38-
extra='forbid',
39-
)
40-
"""The pydantic model config."""
41-
42-
role: str
43-
"""The role that created this audio data, typically "user" or "model"."""
44-
45-
data: types.Blob
46-
"""The audio data chunk."""
47-
48-
timestamp: float
49-
"""Timestamp when the audio chunk was received."""
50-
51-
52-
def _require_audio_data(blob: types.Blob) -> bytes:
53-
data = blob.data
54-
if not isinstance(data, bytes):
55-
raise ValueError('Audio blobs must contain byte data.')
56-
return data
57-
58-
59-
# Deliberately duplicated from `flows.llm_flows.core._utils`
60-
# rather than imported: `live` sits below `flows.llm_flows` in the
61-
# layering, and importing upward would put a cycle back in.
62-
def _require_agent_name(invocation_context: InvocationContext) -> str:
63-
agent = invocation_context.agent
64-
if agent is None:
65-
raise TypeError('Live audio requires an agent in InvocationContext.')
66-
return agent.name
67-
68-
69-
class AudioCacheManager:
70-
"""Manages audio caching and flushing for live streaming flows."""
71-
72-
def __init__(self, config: AudioCacheConfig | None = None) -> None:
73-
"""Initialize the audio cache manager.
74-
75-
Args:
76-
config: Configuration for audio caching behavior.
77-
"""
78-
self.config = config or AudioCacheConfig()
79-
80-
def cache_audio(
81-
self,
82-
invocation_context: InvocationContext,
83-
audio_blob: types.Blob,
84-
cache_type: str,
85-
) -> None:
86-
"""Cache incoming user or outgoing model audio data.
87-
88-
Args:
89-
invocation_context: The current invocation context.
90-
audio_blob: The audio data to cache.
91-
cache_type: Type of audio to cache, either 'input' or 'output'.
92-
93-
Raises:
94-
ValueError: If cache_type is not 'input' or 'output'.
95-
"""
96-
audio_data = _require_audio_data(audio_blob)
97-
if cache_type == 'input':
98-
if not invocation_context.input_realtime_cache:
99-
invocation_context.input_realtime_cache = []
100-
cache = invocation_context.input_realtime_cache
101-
role = 'user'
102-
elif cache_type == 'output':
103-
if not invocation_context.output_realtime_cache:
104-
invocation_context.output_realtime_cache = []
105-
cache = invocation_context.output_realtime_cache
106-
role = 'model'
107-
else:
108-
raise ValueError("cache_type must be either 'input' or 'output'")
109-
110-
audio_entry = RealtimeCacheEntry(
111-
role=role, data=audio_blob, timestamp=platform_time.get_time()
112-
)
113-
cache.append(audio_entry)
114-
115-
logger.debug(
116-
'Cached %s audio chunk: %d bytes, cache size: %d',
117-
cache_type,
118-
len(audio_data),
119-
len(cache),
120-
)
15+
"""Backward compatibility module for the audio-only cache manager names.
12116
122-
async def flush_caches(
123-
self,
124-
invocation_context: InvocationContext,
125-
flush_user_audio: bool = True,
126-
flush_model_audio: bool = True,
127-
) -> list[Event]:
128-
"""Flush audio caches to artifact services.
129-
130-
The multimodality data is saved in artifact service in the format of
131-
audio file. The file data reference is added to the session as an event.
132-
The audio file follows the naming convention: artifact_ref =
133-
f"artifact://{invocation_context.app_name}/{invocation_context.user_id}/
134-
{invocation_context.session.id}/_adk_live/{filename}#{revision_id}"
135-
136-
Note: video data is not supported yet.
137-
138-
A flushed cache is emptied whether or not an artifact was written, so a
139-
cache the artifact service could not accept is dropped rather than kept
140-
for a later attempt.
141-
142-
Args:
143-
invocation_context: The invocation context containing audio caches.
144-
flush_user_audio: Whether to flush the input (user) audio cache.
145-
flush_model_audio: Whether to flush the output (model) audio cache.
146-
147-
Returns:
148-
A list of Event objects created from the flushed caches.
149-
"""
150-
flushed_events: list[Event] = []
151-
if flush_user_audio and invocation_context.input_realtime_cache:
152-
# Swap before the await so audio cached during it is kept for the next
153-
# flush rather than discarded by the unconditional drop.
154-
input_cache = invocation_context.input_realtime_cache
155-
invocation_context.input_realtime_cache = []
156-
audio_event = await self._flush_cache_to_services(
157-
invocation_context,
158-
input_cache,
159-
'input_audio',
160-
)
161-
if audio_event:
162-
flushed_events.append(audio_event)
163-
164-
if flush_model_audio and invocation_context.output_realtime_cache:
165-
logger.debug('Flushed output audio cache')
166-
output_cache = invocation_context.output_realtime_cache
167-
invocation_context.output_realtime_cache = []
168-
audio_event = await self._flush_cache_to_services(
169-
invocation_context,
170-
output_cache,
171-
'output_audio',
172-
)
173-
if audio_event:
174-
flushed_events.append(audio_event)
175-
176-
return flushed_events
177-
178-
async def _flush_cache_to_services(
179-
self,
180-
invocation_context: InvocationContext,
181-
audio_cache: list[RealtimeCacheEntry],
182-
cache_type: str,
183-
) -> Event | None:
184-
"""Flush a list of audio cache entries to artifact services.
185-
186-
The artifact service stores the actual blob. The session stores the
187-
reference to the stored blob.
188-
189-
Args:
190-
invocation_context: The invocation context.
191-
audio_cache: The audio cache to flush.
192-
cache_type: Type identifier for the cache ('input_audio' or
193-
'output_audio').
194-
195-
Returns:
196-
The created Event if the cache was successfully flushed, None otherwise.
197-
"""
198-
if not invocation_context.artifact_service or not audio_cache:
199-
logger.debug('Skipping cache flush: no artifact service or empty cache')
200-
return None
201-
202-
try:
203-
# Combine audio chunks into a single file. Use join rather than repeated
204-
# `+=`, which is O(n^2) over the total audio size.
205-
mime_type = audio_cache[0].data.mime_type or 'audio/pcm'
206-
combined_audio_data = b''.join(
207-
entry.data.data or b'' for entry in audio_cache
208-
)
209-
210-
# Generate filename with timestamp from first audio chunk (when recording started)
211-
timestamp = int(audio_cache[0].timestamp * 1000) # milliseconds
212-
# Drop MIME parameters such as `;rate=24000` so they don't end up in the
213-
# file extension. The full MIME type is kept on the saved artifact.
214-
extension = mime_type.split(';', 1)[0].strip().split('/')[-1]
215-
filename = f'adk_live_audio_storage_{cache_type}_{timestamp}.{extension}'
216-
217-
# Save to artifact service
218-
combined_audio_part = types.Part(
219-
inline_data=types.Blob(data=combined_audio_data, mime_type=mime_type)
220-
)
221-
222-
revision_id = await invocation_context.artifact_service.save_artifact(
223-
app_name=invocation_context.app_name,
224-
user_id=invocation_context.user_id,
225-
session_id=invocation_context.session.id,
226-
filename=filename,
227-
artifact=combined_audio_part,
228-
)
229-
230-
# Create artifact reference for session service
231-
artifact_ref = f'artifact://{invocation_context.app_name}/{invocation_context.user_id}/{invocation_context.session.id}/_adk_live/{filename}#{revision_id}'
232-
233-
# Create event with file data reference to add to session
234-
# For model events, author should be the agent name, not the role
235-
author = (
236-
_require_agent_name(invocation_context)
237-
if audio_cache[0].role == 'model'
238-
else audio_cache[0].role
239-
)
240-
audio_event = Event(
241-
id=Event.new_id(),
242-
invocation_id=invocation_context.invocation_id,
243-
author=author,
244-
content=types.Content(
245-
role=audio_cache[0].role,
246-
parts=[
247-
types.Part(
248-
file_data=types.FileData(
249-
file_uri=artifact_ref, mime_type=mime_type
250-
)
251-
)
252-
],
253-
),
254-
timestamp=audio_cache[0].timestamp,
255-
)
256-
257-
logger.debug(
258-
'Successfully flushed %s cache: %d chunks, %d bytes, saved as %s',
259-
cache_type,
260-
len(audio_cache),
261-
len(combined_audio_data),
262-
filename,
263-
)
264-
return audio_event
265-
266-
except Exception as e:
267-
logger.error('Failed to flush %s cache: %s', cache_type, e)
268-
return None
269-
270-
def get_cache_stats(
271-
self, invocation_context: InvocationContext
272-
) -> dict[str, int]:
273-
"""Get statistics about current cache state.
274-
275-
Args:
276-
invocation_context: The invocation context.
277-
278-
Returns:
279-
Dictionary containing cache statistics.
280-
"""
281-
input_count = len(invocation_context.input_realtime_cache or [])
282-
output_count = len(invocation_context.output_realtime_cache or [])
283-
284-
input_bytes = sum(
285-
len(_require_audio_data(entry.data))
286-
for entry in invocation_context.input_realtime_cache or []
287-
)
288-
output_bytes = sum(
289-
len(_require_audio_data(entry.data))
290-
for entry in invocation_context.output_realtime_cache or []
291-
)
292-
293-
return {
294-
'input_chunks': input_count,
295-
'output_chunks': output_count,
296-
'input_bytes': input_bytes,
297-
'output_bytes': output_bytes,
298-
'total_chunks': input_count + output_count,
299-
'total_bytes': input_bytes + output_bytes,
300-
}
17+
The cache manager now lives in ``_realtime_cache_manager`` as
18+
``RealtimeCacheManager`` and ``RealtimeCacheConfig``. This module keeps the
19+
former names importable; it holds no implementation of its own.
20+
"""
30121

22+
from __future__ import annotations
30223

303-
class AudioCacheConfig:
304-
"""Configuration for audio caching behavior."""
24+
from ._realtime_cache_manager import logger
25+
from ._realtime_cache_manager import RealtimeCacheConfig
26+
from ._realtime_cache_manager import RealtimeCacheEntry
27+
from ._realtime_cache_manager import RealtimeCacheManager
30528

306-
def __init__(
307-
self,
308-
max_cache_size_bytes: int = 10 * 1024 * 1024, # 10MB
309-
max_cache_duration_seconds: float = 300.0, # 5 minutes
310-
auto_flush_threshold: int = 100, # Number of chunks
311-
) -> None:
312-
"""Initialize audio cache configuration.
29+
AudioCacheConfig = RealtimeCacheConfig
30+
AudioCacheManager = RealtimeCacheManager
31331

314-
Args:
315-
max_cache_size_bytes: Maximum cache size in bytes before auto-flush.
316-
max_cache_duration_seconds: Maximum duration to keep data in cache.
317-
auto_flush_threshold: Number of chunks that triggers auto-flush.
318-
"""
319-
self.max_cache_size_bytes = max_cache_size_bytes
320-
self.max_cache_duration_seconds = max_cache_duration_seconds
321-
self.auto_flush_threshold = auto_flush_threshold
32+
__all__ = [
33+
'AudioCacheConfig',
34+
'AudioCacheManager',
35+
'RealtimeCacheEntry',
36+
'logger',
37+
]

‎src/google/adk/live/_flow_utils.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -175,23 +175,23 @@ async def handle_control_event_flush(
175175
Returns:
176176
A list of Event objects created from the flushed caches.
177177
"""
178-
audio_cache_manager = flow.audio_cache_manager
178+
realtime_cache_manager = flow._realtime_cache_manager
179179

180180
# Log cache statistics if enabled
181181
if DEFAULT_ENABLE_CACHE_STATISTICS:
182-
stats = audio_cache_manager.get_cache_stats(invocation_context)
182+
stats = realtime_cache_manager.get_cache_stats(invocation_context)
183183
logger.debug('Audio cache stats: %s', stats)
184184

185185
if llm_response.interrupted:
186186
# user interrupts so the model will stop. we can flush model audio here
187-
return await audio_cache_manager.flush_caches(
187+
return await realtime_cache_manager.flush_caches(
188188
invocation_context,
189189
flush_user_audio=False,
190190
flush_model_audio=True,
191191
)
192192
elif llm_response.turn_complete:
193193
# turn completes so we can flush both user and model
194-
return await audio_cache_manager.flush_caches(
194+
return await realtime_cache_manager.flush_caches(
195195
invocation_context,
196196
flush_user_audio=True,
197197
flush_model_audio=True,

0 commit comments

Comments
 (0)