|
12 | 12 | # See the License for the specific language governing permissions and |
13 | 13 | # limitations under the License. |
14 | 14 |
|
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. |
121 | 16 |
|
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 | +""" |
301 | 21 |
|
| 22 | +from __future__ import annotations |
302 | 23 |
|
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 |
305 | 28 |
|
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 |
313 | 31 |
|
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 | +] |
0 commit comments