mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
feat: Add support for store audio and transcription into events/sessions and artifacts
PiperOrigin-RevId: 798011660
This commit is contained in:
committed by
Copybara-Service
parent
f660180854
commit
ddb070ff07
@@ -65,8 +65,8 @@ async def check_prime(nums: list[int]) -> str:
|
|||||||
|
|
||||||
|
|
||||||
root_agent = Agent(
|
root_agent = Agent(
|
||||||
model='gemini-2.0-flash-live-preview-04-09', # for Vertex project
|
# model='gemini-2.0-flash-live-preview-04-09', # for Vertex project
|
||||||
# model='gemini-2.0-flash-live-001', # for AI studio key
|
model='gemini-2.0-flash-live-001', # for AI studio key
|
||||||
name='hello_world_agent',
|
name='hello_world_agent',
|
||||||
description=(
|
description=(
|
||||||
'hello world agent that can roll a dice of 8 sides and check prime'
|
'hello world agent that can roll a dice of 8 sides and check prime'
|
||||||
|
|||||||
@@ -40,6 +40,25 @@ class LlmCallsLimitExceededError(Exception):
|
|||||||
"""Error thrown when the number of LLM calls exceed the limit."""
|
"""Error thrown when the number of LLM calls exceed the limit."""
|
||||||
|
|
||||||
|
|
||||||
|
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."""
|
||||||
|
|
||||||
|
|
||||||
class _InvocationCostManager(BaseModel):
|
class _InvocationCostManager(BaseModel):
|
||||||
"""A container to keep track of the cost of invocation.
|
"""A container to keep track of the cost of invocation.
|
||||||
|
|
||||||
@@ -156,6 +175,12 @@ class InvocationContext(BaseModel):
|
|||||||
live_session_resumption_handle: Optional[str] = None
|
live_session_resumption_handle: Optional[str] = None
|
||||||
"""The handle for live session resumption."""
|
"""The handle for live session resumption."""
|
||||||
|
|
||||||
|
input_realtime_cache: Optional[list[RealtimeCacheEntry]] = None
|
||||||
|
"""Caches input audio chunks before flushing to session and artifact services."""
|
||||||
|
|
||||||
|
output_realtime_cache: Optional[list[RealtimeCacheEntry]] = None
|
||||||
|
"""Caches output audio chunks before flushing to session and artifact services."""
|
||||||
|
|
||||||
run_config: Optional[RunConfig] = None
|
run_config: Optional[RunConfig] = None
|
||||||
"""Configurations for live agents under this invocation."""
|
"""Configurations for live agents under this invocation."""
|
||||||
|
|
||||||
|
|||||||
@@ -0,0 +1,264 @@
|
|||||||
|
# Copyright 2025 Google LLC
|
||||||
|
#
|
||||||
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
# you may not use this file except in compliance with the License.
|
||||||
|
# You may obtain a copy of the License at
|
||||||
|
#
|
||||||
|
# http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
#
|
||||||
|
# Unless required by applicable law or agreed to in writing, software
|
||||||
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
# See the License for the specific language governing permissions and
|
||||||
|
# limitations under the License.
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import logging
|
||||||
|
import time
|
||||||
|
from typing import TYPE_CHECKING
|
||||||
|
|
||||||
|
from google.genai import types
|
||||||
|
|
||||||
|
from ...agents.invocation_context import RealtimeCacheEntry
|
||||||
|
from ...events.event import Event
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from ...agents.invocation_context import InvocationContext
|
||||||
|
|
||||||
|
logger = logging.getLogger('google_adk.' + __name__)
|
||||||
|
|
||||||
|
|
||||||
|
class AudioCacheManager:
|
||||||
|
"""Manages audio caching and flushing for live streaming flows."""
|
||||||
|
|
||||||
|
def __init__(self, config: AudioCacheConfig | 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'.
|
||||||
|
"""
|
||||||
|
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=time.time()
|
||||||
|
)
|
||||||
|
cache.append(audio_entry)
|
||||||
|
|
||||||
|
logger.debug(
|
||||||
|
'Cached %s audio chunk: %d bytes, cache size: %d',
|
||||||
|
cache_type,
|
||||||
|
len(audio_blob.data),
|
||||||
|
len(cache),
|
||||||
|
)
|
||||||
|
|
||||||
|
async def flush_caches(
|
||||||
|
self,
|
||||||
|
invocation_context: InvocationContext,
|
||||||
|
flush_user_audio: bool = True,
|
||||||
|
flush_model_audio: bool = True,
|
||||||
|
) -> None:
|
||||||
|
"""Flush audio caches to session and 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.
|
||||||
|
|
||||||
|
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.
|
||||||
|
"""
|
||||||
|
if flush_user_audio and invocation_context.input_realtime_cache:
|
||||||
|
success = await self._flush_cache_to_services(
|
||||||
|
invocation_context,
|
||||||
|
invocation_context.input_realtime_cache,
|
||||||
|
'input_audio',
|
||||||
|
)
|
||||||
|
if success:
|
||||||
|
invocation_context.input_realtime_cache = []
|
||||||
|
logger.debug('Flushed input audio cache')
|
||||||
|
|
||||||
|
if flush_model_audio and invocation_context.output_realtime_cache:
|
||||||
|
success = await self._flush_cache_to_services(
|
||||||
|
invocation_context,
|
||||||
|
invocation_context.output_realtime_cache,
|
||||||
|
'output_audio',
|
||||||
|
)
|
||||||
|
if success:
|
||||||
|
invocation_context.output_realtime_cache = []
|
||||||
|
logger.debug('Flushed output audio cache')
|
||||||
|
|
||||||
|
async def _flush_cache_to_services(
|
||||||
|
self,
|
||||||
|
invocation_context: InvocationContext,
|
||||||
|
audio_cache: list[RealtimeCacheEntry],
|
||||||
|
cache_type: str,
|
||||||
|
) -> bool:
|
||||||
|
"""Flush a list of audio cache entries to session and 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:
|
||||||
|
True if the cache was successfully flushed, False otherwise.
|
||||||
|
"""
|
||||||
|
print('flush cache')
|
||||||
|
if not invocation_context.artifact_service or not audio_cache:
|
||||||
|
logger.debug('Skipping cache flush: no artifact service or empty cache')
|
||||||
|
return False
|
||||||
|
|
||||||
|
try:
|
||||||
|
# Combine audio chunks into a single file
|
||||||
|
combined_audio_data = b''
|
||||||
|
mime_type = audio_cache[0].data.mime_type if audio_cache else 'audio/pcm'
|
||||||
|
|
||||||
|
for entry in audio_cache:
|
||||||
|
combined_audio_data += entry.data.data
|
||||||
|
|
||||||
|
# Generate filename with timestamp from first audio chunk (when recording started)
|
||||||
|
timestamp = int(audio_cache[0].timestamp * 1000) # milliseconds
|
||||||
|
filename = f"adk_live_audio_storage_{cache_type}_{timestamp}.{mime_type.split('/')[-1]}"
|
||||||
|
|
||||||
|
# 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
|
||||||
|
audio_event = Event(
|
||||||
|
id=Event.new_id(),
|
||||||
|
invocation_id=invocation_context.invocation_id,
|
||||||
|
author=audio_cache[0].role,
|
||||||
|
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,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Add to session
|
||||||
|
await invocation_context.session_service.append_event(
|
||||||
|
invocation_context.session, audio_event
|
||||||
|
)
|
||||||
|
|
||||||
|
logger.debug(
|
||||||
|
'Successfully flushed %s cache: %d chunks, %d bytes, saved as %s',
|
||||||
|
cache_type,
|
||||||
|
len(audio_cache),
|
||||||
|
len(combined_audio_data),
|
||||||
|
filename,
|
||||||
|
)
|
||||||
|
return True
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
logger.error('Failed to flush %s cache: %s', cache_type, e)
|
||||||
|
return False
|
||||||
|
|
||||||
|
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(entry.data.data)
|
||||||
|
for entry in (invocation_context.input_realtime_cache or [])
|
||||||
|
)
|
||||||
|
output_bytes = sum(
|
||||||
|
len(entry.data.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 AudioCacheConfig:
|
||||||
|
"""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
|
||||||
|
):
|
||||||
|
"""Initialize audio 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
|
||||||
@@ -47,6 +47,8 @@ from ...telemetry import tracer
|
|||||||
from ...tools.base_toolset import BaseToolset
|
from ...tools.base_toolset import BaseToolset
|
||||||
from ...tools.tool_context import ToolContext
|
from ...tools.tool_context import ToolContext
|
||||||
from ...utils.context_utils import Aclosing
|
from ...utils.context_utils import Aclosing
|
||||||
|
from .audio_cache_manager import AudioCacheManager
|
||||||
|
from .transcription_manager import TranscriptionManager
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from ...agents.llm_agent import LlmAgent
|
from ...agents.llm_agent import LlmAgent
|
||||||
@@ -58,6 +60,14 @@ logger = logging.getLogger('google_adk.' + __name__)
|
|||||||
|
|
||||||
_ADK_AGENT_NAME_LABEL_KEY = 'adk_agent_name'
|
_ADK_AGENT_NAME_LABEL_KEY = 'adk_agent_name'
|
||||||
|
|
||||||
|
# Timing configuration
|
||||||
|
DEFAULT_REQUEST_QUEUE_TIMEOUT = 0.25
|
||||||
|
DEFAULT_TRANSFER_AGENT_DELAY = 1.0
|
||||||
|
DEFAULT_TASK_COMPLETION_DELAY = 1.0
|
||||||
|
|
||||||
|
# Statistics configuration
|
||||||
|
DEFAULT_ENABLE_CACHE_STATISTICS = False
|
||||||
|
|
||||||
|
|
||||||
class BaseLlmFlow(ABC):
|
class BaseLlmFlow(ABC):
|
||||||
"""A basic flow that calls the LLM in a loop until a final response is generated.
|
"""A basic flow that calls the LLM in a loop until a final response is generated.
|
||||||
@@ -69,6 +79,10 @@ class BaseLlmFlow(ABC):
|
|||||||
self.request_processors: list[BaseLlmRequestProcessor] = []
|
self.request_processors: list[BaseLlmRequestProcessor] = []
|
||||||
self.response_processors: list[BaseLlmResponseProcessor] = []
|
self.response_processors: list[BaseLlmResponseProcessor] = []
|
||||||
|
|
||||||
|
# Initialize configuration and managers
|
||||||
|
self.audio_cache_manager = AudioCacheManager()
|
||||||
|
self.transcription_manager = TranscriptionManager()
|
||||||
|
|
||||||
async def run_live(
|
async def run_live(
|
||||||
self,
|
self,
|
||||||
invocation_context: InvocationContext,
|
invocation_context: InvocationContext,
|
||||||
@@ -169,7 +183,7 @@ class BaseLlmFlow(ABC):
|
|||||||
and event.content.parts[0].function_response.name
|
and event.content.parts[0].function_response.name
|
||||||
== 'transfer_to_agent'
|
== 'transfer_to_agent'
|
||||||
):
|
):
|
||||||
await asyncio.sleep(1)
|
await asyncio.sleep(DEFAULT_TRANSFER_AGENT_DELAY)
|
||||||
# cancel the tasks that belongs to the closed connection.
|
# cancel the tasks that belongs to the closed connection.
|
||||||
send_task.cancel()
|
send_task.cancel()
|
||||||
await llm_connection.close()
|
await llm_connection.close()
|
||||||
@@ -181,7 +195,7 @@ class BaseLlmFlow(ABC):
|
|||||||
== 'task_completed'
|
== 'task_completed'
|
||||||
):
|
):
|
||||||
# this is used for sequential agent to signal the end of the agent.
|
# this is used for sequential agent to signal the end of the agent.
|
||||||
await asyncio.sleep(1)
|
await asyncio.sleep(DEFAULT_TASK_COMPLETION_DELAY)
|
||||||
# cancel the tasks that belongs to the closed connection.
|
# cancel the tasks that belongs to the closed connection.
|
||||||
send_task.cancel()
|
send_task.cancel()
|
||||||
return
|
return
|
||||||
@@ -218,7 +232,7 @@ class BaseLlmFlow(ABC):
|
|||||||
# event loop to process events.
|
# event loop to process events.
|
||||||
# TODO: revert back(remove timeout) once we move off streamlit.
|
# TODO: revert back(remove timeout) once we move off streamlit.
|
||||||
live_request = await asyncio.wait_for(
|
live_request = await asyncio.wait_for(
|
||||||
live_request_queue.get(), timeout=0.25
|
live_request_queue.get(), timeout=DEFAULT_REQUEST_QUEUE_TIMEOUT
|
||||||
)
|
)
|
||||||
# duplicate the live_request to all the active streams
|
# duplicate the live_request to all the active streams
|
||||||
logger.debug(
|
logger.debug(
|
||||||
@@ -253,6 +267,12 @@ class BaseLlmFlow(ABC):
|
|||||||
invocation_context.transcription_cache.append(
|
invocation_context.transcription_cache.append(
|
||||||
TranscriptionEntry(role='user', data=live_request.blob)
|
TranscriptionEntry(role='user', data=live_request.blob)
|
||||||
)
|
)
|
||||||
|
|
||||||
|
# Cache input audio chunks before flushing
|
||||||
|
self.audio_cache_manager.cache_audio(
|
||||||
|
invocation_context, live_request.blob, cache_type='input'
|
||||||
|
)
|
||||||
|
|
||||||
await llm_connection.send_realtime(live_request.blob)
|
await llm_connection.send_realtime(live_request.blob)
|
||||||
|
|
||||||
if live_request.content:
|
if live_request.content:
|
||||||
@@ -303,6 +323,23 @@ class BaseLlmFlow(ABC):
|
|||||||
invocation_id=invocation_context.invocation_id,
|
invocation_id=invocation_context.invocation_id,
|
||||||
author=get_author_for_event(llm_response),
|
author=get_author_for_event(llm_response),
|
||||||
)
|
)
|
||||||
|
|
||||||
|
# Handle transcription events ONCE per llm_response, outside the event loop
|
||||||
|
if llm_response.input_transcription:
|
||||||
|
await self.transcription_manager.handle_input_transcription(
|
||||||
|
invocation_context, llm_response.input_transcription
|
||||||
|
)
|
||||||
|
|
||||||
|
if llm_response.output_transcription:
|
||||||
|
await self.transcription_manager.handle_output_transcription(
|
||||||
|
invocation_context, llm_response.output_transcription
|
||||||
|
)
|
||||||
|
|
||||||
|
# Flush audio caches based on control events using configurable settings
|
||||||
|
await self._handle_control_event_flush(
|
||||||
|
invocation_context, llm_response
|
||||||
|
)
|
||||||
|
|
||||||
async with Aclosing(
|
async with Aclosing(
|
||||||
self._postprocess_live(
|
self._postprocess_live(
|
||||||
invocation_context,
|
invocation_context,
|
||||||
@@ -330,6 +367,24 @@ class BaseLlmFlow(ABC):
|
|||||||
role=event.content.role, data=event.content
|
role=event.content.role, data=event.content
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
# Cache output audio chunks from model responses
|
||||||
|
# TODO: support video data
|
||||||
|
if (
|
||||||
|
event.content
|
||||||
|
and event.content.parts
|
||||||
|
and event.content.parts[0].inline_data
|
||||||
|
and event.content.parts[0].inline_data.mime_type.startswith(
|
||||||
|
'audio/'
|
||||||
|
)
|
||||||
|
):
|
||||||
|
audio_blob = types.Blob(
|
||||||
|
data=event.content.parts[0].inline_data.data,
|
||||||
|
mime_type=event.content.parts[0].inline_data.mime_type,
|
||||||
|
)
|
||||||
|
self.audio_cache_manager.cache_audio(
|
||||||
|
invocation_context, audio_blob, cache_type='output'
|
||||||
|
)
|
||||||
|
|
||||||
yield event
|
yield event
|
||||||
# Give opportunity for other tasks to run.
|
# Give opportunity for other tasks to run.
|
||||||
await asyncio.sleep(0)
|
await asyncio.sleep(0)
|
||||||
@@ -512,12 +567,14 @@ class BaseLlmFlow(ABC):
|
|||||||
|
|
||||||
# Skip the model response event if there is no content and no error code.
|
# Skip the model response event if there is no content and no error code.
|
||||||
# This is needed for the code executor to trigger another loop.
|
# This is needed for the code executor to trigger another loop.
|
||||||
# But don't skip control events like turn_complete.
|
# But don't skip control events like turn_complete or transcription events.
|
||||||
if (
|
if (
|
||||||
not llm_response.content
|
not llm_response.content
|
||||||
and not llm_response.error_code
|
and not llm_response.error_code
|
||||||
and not llm_response.interrupted
|
and not llm_response.interrupted
|
||||||
and not llm_response.turn_complete
|
and not llm_response.turn_complete
|
||||||
|
and not llm_response.input_transcription
|
||||||
|
and not llm_response.output_transcription
|
||||||
):
|
):
|
||||||
return
|
return
|
||||||
|
|
||||||
@@ -805,6 +862,42 @@ class BaseLlmFlow(ABC):
|
|||||||
|
|
||||||
return model_response_event
|
return model_response_event
|
||||||
|
|
||||||
|
async def _handle_control_event_flush(
|
||||||
|
self, invocation_context: InvocationContext, llm_response: LlmResponse
|
||||||
|
) -> None:
|
||||||
|
"""Handle audio cache flushing based on control events.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
invocation_context: The invocation context containing audio caches.
|
||||||
|
llm_response: The LLM response containing control event information.
|
||||||
|
"""
|
||||||
|
if llm_response.interrupted:
|
||||||
|
# user interrupts so the model will stop. we can flush model audio here
|
||||||
|
await self.audio_cache_manager.flush_caches(
|
||||||
|
invocation_context,
|
||||||
|
flush_user_audio=False,
|
||||||
|
flush_model_audio=True,
|
||||||
|
)
|
||||||
|
elif llm_response.turn_complete:
|
||||||
|
# turn completes so we can flush both user and model
|
||||||
|
await self.audio_cache_manager.flush_caches(
|
||||||
|
invocation_context,
|
||||||
|
flush_user_audio=True,
|
||||||
|
flush_model_audio=True,
|
||||||
|
)
|
||||||
|
elif getattr(llm_response, 'generation_complete', False):
|
||||||
|
# model generation complete so we can flush model audio
|
||||||
|
await self.audio_cache_manager.flush_caches(
|
||||||
|
invocation_context,
|
||||||
|
flush_user_audio=False,
|
||||||
|
flush_model_audio=True,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Log cache statistics if enabled
|
||||||
|
if DEFAULT_ENABLE_CACHE_STATISTICS:
|
||||||
|
stats = self.audio_cache_manager.get_cache_stats(invocation_context)
|
||||||
|
logger.debug('Audio cache stats: %s', stats)
|
||||||
|
|
||||||
async def _run_and_handle_error(
|
async def _run_and_handle_error(
|
||||||
self,
|
self,
|
||||||
response_generator: AsyncGenerator[LlmResponse, None],
|
response_generator: AsyncGenerator[LlmResponse, None],
|
||||||
|
|||||||
@@ -0,0 +1,141 @@
|
|||||||
|
# Copyright 2025 Google LLC
|
||||||
|
#
|
||||||
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
# you may not use this file except in compliance with the License.
|
||||||
|
# You may obtain a copy of the License at
|
||||||
|
#
|
||||||
|
# http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
#
|
||||||
|
# Unless required by applicable law or agreed to in writing, software
|
||||||
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
# See the License for the specific language governing permissions and
|
||||||
|
# limitations under the License.
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import logging
|
||||||
|
import time
|
||||||
|
from typing import TYPE_CHECKING
|
||||||
|
|
||||||
|
from google.genai import types
|
||||||
|
|
||||||
|
from ...events.event import Event
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from ...agents.invocation_context import InvocationContext
|
||||||
|
|
||||||
|
logger = logging.getLogger('google_adk.' + __name__)
|
||||||
|
|
||||||
|
|
||||||
|
class TranscriptionManager:
|
||||||
|
"""Manages transcription events for live streaming flows."""
|
||||||
|
|
||||||
|
async def handle_input_transcription(
|
||||||
|
self,
|
||||||
|
invocation_context: InvocationContext,
|
||||||
|
transcription: types.Transcription,
|
||||||
|
) -> None:
|
||||||
|
"""Handle user input transcription events.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
invocation_context: The current invocation context.
|
||||||
|
transcription: The transcription data from user input.
|
||||||
|
"""
|
||||||
|
await self._create_and_save_transcription_event(
|
||||||
|
invocation_context=invocation_context,
|
||||||
|
transcription=transcription,
|
||||||
|
author='user',
|
||||||
|
is_input=True,
|
||||||
|
)
|
||||||
|
|
||||||
|
async def handle_output_transcription(
|
||||||
|
self,
|
||||||
|
invocation_context: InvocationContext,
|
||||||
|
transcription: types.Transcription,
|
||||||
|
) -> None:
|
||||||
|
"""Handle model output transcription events.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
invocation_context: The current invocation context.
|
||||||
|
transcription: The transcription data from model output.
|
||||||
|
"""
|
||||||
|
await self._create_and_save_transcription_event(
|
||||||
|
invocation_context=invocation_context,
|
||||||
|
transcription=transcription,
|
||||||
|
author=invocation_context.agent.name,
|
||||||
|
is_input=False,
|
||||||
|
)
|
||||||
|
|
||||||
|
async def _create_and_save_transcription_event(
|
||||||
|
self,
|
||||||
|
invocation_context: InvocationContext,
|
||||||
|
transcription: types.Transcription,
|
||||||
|
author: str,
|
||||||
|
is_input: bool,
|
||||||
|
) -> None:
|
||||||
|
"""Create and save a transcription event to session service.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
invocation_context: The current invocation context.
|
||||||
|
transcription: The transcription data.
|
||||||
|
author: The author of the transcription event.
|
||||||
|
is_input: Whether this is an input (user) or output (model) transcription.
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
transcription_event = Event(
|
||||||
|
id=Event.new_id(),
|
||||||
|
invocation_id=invocation_context.invocation_id,
|
||||||
|
author=author,
|
||||||
|
input_transcription=transcription if is_input else None,
|
||||||
|
output_transcription=transcription if not is_input else None,
|
||||||
|
timestamp=time.time(),
|
||||||
|
)
|
||||||
|
|
||||||
|
# Save transcription event to session
|
||||||
|
await invocation_context.session_service.append_event(
|
||||||
|
invocation_context.session, transcription_event
|
||||||
|
)
|
||||||
|
|
||||||
|
logger.debug(
|
||||||
|
'Saved %s transcription event for %s: %s',
|
||||||
|
'input' if is_input else 'output',
|
||||||
|
author,
|
||||||
|
transcription.text
|
||||||
|
if hasattr(transcription, 'text')
|
||||||
|
else 'audio transcription',
|
||||||
|
)
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(
|
||||||
|
'Failed to save %s transcription event: %s',
|
||||||
|
'input' if is_input else 'output',
|
||||||
|
e,
|
||||||
|
)
|
||||||
|
raise
|
||||||
|
|
||||||
|
def get_transcription_stats(
|
||||||
|
self, invocation_context: InvocationContext
|
||||||
|
) -> dict[str, int]:
|
||||||
|
"""Get statistics about transcription events in the session.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
invocation_context: The current invocation context.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Dictionary containing transcription statistics.
|
||||||
|
"""
|
||||||
|
input_count = 0
|
||||||
|
output_count = 0
|
||||||
|
|
||||||
|
for event in invocation_context.session.events:
|
||||||
|
if hasattr(event, 'input_transcription') and event.input_transcription:
|
||||||
|
input_count += 1
|
||||||
|
if hasattr(event, 'output_transcription') and event.output_transcription:
|
||||||
|
output_count += 1
|
||||||
|
|
||||||
|
return {
|
||||||
|
'input_transcriptions': input_count,
|
||||||
|
'output_transcriptions': output_count,
|
||||||
|
'total_transcriptions': input_count + output_count,
|
||||||
|
}
|
||||||
@@ -164,14 +164,8 @@ class GeminiLlmConnection(BaseLlmConnection):
|
|||||||
message.server_content.input_transcription
|
message.server_content.input_transcription
|
||||||
and message.server_content.input_transcription.text
|
and message.server_content.input_transcription.text
|
||||||
):
|
):
|
||||||
user_text = message.server_content.input_transcription.text
|
|
||||||
parts = [
|
|
||||||
types.Part.from_text(
|
|
||||||
text=user_text,
|
|
||||||
)
|
|
||||||
]
|
|
||||||
llm_response = LlmResponse(
|
llm_response = LlmResponse(
|
||||||
content=types.Content(role='user', parts=parts)
|
input_transcription=message.server_content.input_transcription,
|
||||||
)
|
)
|
||||||
yield llm_response
|
yield llm_response
|
||||||
if (
|
if (
|
||||||
@@ -186,13 +180,8 @@ class GeminiLlmConnection(BaseLlmConnection):
|
|||||||
# We rely on other control signals to determine when to yield the
|
# We rely on other control signals to determine when to yield the
|
||||||
# full text response(turn_complete, interrupted, or tool_call).
|
# full text response(turn_complete, interrupted, or tool_call).
|
||||||
text += message.server_content.output_transcription.text
|
text += message.server_content.output_transcription.text
|
||||||
parts = [
|
|
||||||
types.Part.from_text(
|
|
||||||
text=message.server_content.output_transcription.text
|
|
||||||
)
|
|
||||||
]
|
|
||||||
llm_response = LlmResponse(
|
llm_response = LlmResponse(
|
||||||
content=types.Content(role='model', parts=parts), partial=True
|
output_transcription=message.server_content.output_transcription
|
||||||
)
|
)
|
||||||
yield llm_response
|
yield llm_response
|
||||||
|
|
||||||
|
|||||||
@@ -40,6 +40,8 @@ class LlmResponse(BaseModel):
|
|||||||
interrupted: Flag indicating that LLM was interrupted when generating the
|
interrupted: Flag indicating that LLM was interrupted when generating the
|
||||||
content. Usually it's due to user interruption during a bidi streaming.
|
content. Usually it's due to user interruption during a bidi streaming.
|
||||||
custom_metadata: The custom metadata of the LlmResponse.
|
custom_metadata: The custom metadata of the LlmResponse.
|
||||||
|
input_transcription: Audio transcription of user input.
|
||||||
|
output_transcription: Audio transcription of model output.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
model_config = ConfigDict(
|
model_config = ConfigDict(
|
||||||
@@ -97,6 +99,12 @@ class LlmResponse(BaseModel):
|
|||||||
] = None
|
] = None
|
||||||
"""The session resumption update of the LlmResponse"""
|
"""The session resumption update of the LlmResponse"""
|
||||||
|
|
||||||
|
input_transcription: Optional[types.Transcription] = None
|
||||||
|
"""Audio transcription of user input."""
|
||||||
|
|
||||||
|
output_transcription: Optional[types.Transcription] = None
|
||||||
|
"""Audio transcription of model output."""
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def create(
|
def create(
|
||||||
generate_content_response: types.GenerateContentResponse,
|
generate_content_response: types.GenerateContentResponse,
|
||||||
|
|||||||
@@ -0,0 +1,389 @@
|
|||||||
|
# Copyright 2025 Google LLC
|
||||||
|
#
|
||||||
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
# you may not use this file except in compliance with the License.
|
||||||
|
# You may obtain a copy of the License at
|
||||||
|
#
|
||||||
|
# http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
#
|
||||||
|
# Unless required by applicable law or agreed to in writing, software
|
||||||
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
# See the License for the specific language governing permissions and
|
||||||
|
# limitations under the License.
|
||||||
|
|
||||||
|
import time
|
||||||
|
from unittest.mock import AsyncMock
|
||||||
|
from unittest.mock import Mock
|
||||||
|
|
||||||
|
from google.adk.flows.llm_flows.audio_cache_manager import AudioCacheConfig
|
||||||
|
from google.adk.flows.llm_flows.audio_cache_manager import AudioCacheManager
|
||||||
|
from google.genai import types
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from ... import testing_utils
|
||||||
|
|
||||||
|
|
||||||
|
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_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(self):
|
||||||
|
"""Test flushing when no artifact service is available."""
|
||||||
|
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')
|
||||||
|
self.manager.cache_audio(invocation_context, input_blob, 'input')
|
||||||
|
|
||||||
|
# Flush should not error but should not clear cache either
|
||||||
|
await self.manager.flush_caches(invocation_context)
|
||||||
|
|
||||||
|
# Cache should remain (no actual flushing happened)
|
||||||
|
assert len(invocation_context.input_realtime_cache) == 1
|
||||||
|
|
||||||
|
@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_called_once()
|
||||||
|
|
||||||
|
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 not raise exception but should log error and retain cache
|
||||||
|
await self.manager.flush_caches(invocation_context)
|
||||||
|
|
||||||
|
# Cache should remain since flush failed
|
||||||
|
assert len(invocation_context.input_realtime_cache) == 1
|
||||||
|
|
||||||
|
@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}'
|
||||||
|
)
|
||||||
@@ -0,0 +1,312 @@
|
|||||||
|
# Copyright 2025 Google LLC
|
||||||
|
#
|
||||||
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
# you may not use this file except in compliance with the License.
|
||||||
|
# You may obtain a copy of the License at
|
||||||
|
#
|
||||||
|
# http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
#
|
||||||
|
# Unless required by applicable law or agreed to in writing, software
|
||||||
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
# See the License for the specific language governing permissions and
|
||||||
|
# limitations under the License.
|
||||||
|
|
||||||
|
from unittest.mock import AsyncMock
|
||||||
|
from unittest.mock import Mock
|
||||||
|
|
||||||
|
from google.adk.flows.llm_flows.transcription_manager import TranscriptionManager
|
||||||
|
from google.genai import types
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from ... import testing_utils
|
||||||
|
|
||||||
|
|
||||||
|
class TestTranscriptionManager:
|
||||||
|
"""Test the TranscriptionManager class."""
|
||||||
|
|
||||||
|
def setup_method(self):
|
||||||
|
"""Set up test fixtures."""
|
||||||
|
self.manager = TranscriptionManager()
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_handle_input_transcription(self):
|
||||||
|
"""Test handling user input transcription events."""
|
||||||
|
invocation_context = await testing_utils.create_invocation_context(
|
||||||
|
testing_utils.create_test_agent()
|
||||||
|
)
|
||||||
|
|
||||||
|
# Set up mock session service
|
||||||
|
mock_session_service = AsyncMock()
|
||||||
|
invocation_context.session_service = mock_session_service
|
||||||
|
|
||||||
|
# Create test transcription
|
||||||
|
transcription = types.Transcription(text='Hello from user')
|
||||||
|
|
||||||
|
# Handle transcription
|
||||||
|
await self.manager.handle_input_transcription(
|
||||||
|
invocation_context, transcription
|
||||||
|
)
|
||||||
|
|
||||||
|
# Verify session service was called
|
||||||
|
mock_session_service.append_event.assert_called_once()
|
||||||
|
|
||||||
|
# Check the event that was created
|
||||||
|
call_args = mock_session_service.append_event.call_args
|
||||||
|
event = call_args[0][1] # Second argument is the event
|
||||||
|
|
||||||
|
assert event.author == 'user'
|
||||||
|
assert event.input_transcription == transcription
|
||||||
|
assert event.output_transcription is None
|
||||||
|
assert event.invocation_id == invocation_context.invocation_id
|
||||||
|
assert isinstance(event.timestamp, float)
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_handle_output_transcription(self):
|
||||||
|
"""Test handling model output transcription events."""
|
||||||
|
agent = testing_utils.create_test_agent()
|
||||||
|
invocation_context = await testing_utils.create_invocation_context(agent)
|
||||||
|
|
||||||
|
# Set up mock session service
|
||||||
|
mock_session_service = AsyncMock()
|
||||||
|
invocation_context.session_service = mock_session_service
|
||||||
|
|
||||||
|
# Create test transcription
|
||||||
|
transcription = types.Transcription(text='Hello from model')
|
||||||
|
|
||||||
|
# Handle transcription
|
||||||
|
await self.manager.handle_output_transcription(
|
||||||
|
invocation_context, transcription
|
||||||
|
)
|
||||||
|
|
||||||
|
# Verify session service was called
|
||||||
|
mock_session_service.append_event.assert_called_once()
|
||||||
|
|
||||||
|
# Check the event that was created
|
||||||
|
call_args = mock_session_service.append_event.call_args
|
||||||
|
event = call_args[0][1] # Second argument is the event
|
||||||
|
|
||||||
|
assert event.author == agent.name
|
||||||
|
assert event.input_transcription is None
|
||||||
|
assert event.output_transcription == transcription
|
||||||
|
assert event.invocation_id == invocation_context.invocation_id
|
||||||
|
assert isinstance(event.timestamp, float)
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_handle_multiple_transcriptions(self):
|
||||||
|
"""Test handling multiple transcription events."""
|
||||||
|
invocation_context = await testing_utils.create_invocation_context(
|
||||||
|
testing_utils.create_test_agent()
|
||||||
|
)
|
||||||
|
|
||||||
|
# Set up mock session service
|
||||||
|
mock_session_service = AsyncMock()
|
||||||
|
invocation_context.session_service = mock_session_service
|
||||||
|
|
||||||
|
# Handle multiple input transcriptions
|
||||||
|
for i in range(3):
|
||||||
|
transcription = types.Transcription(text=f'User message {i}')
|
||||||
|
await self.manager.handle_input_transcription(
|
||||||
|
invocation_context, transcription
|
||||||
|
)
|
||||||
|
|
||||||
|
# Handle multiple output transcriptions
|
||||||
|
for i in range(2):
|
||||||
|
transcription = types.Transcription(text=f'Model response {i}')
|
||||||
|
await self.manager.handle_output_transcription(
|
||||||
|
invocation_context, transcription
|
||||||
|
)
|
||||||
|
|
||||||
|
# Verify session service was called for each transcription
|
||||||
|
assert mock_session_service.append_event.call_count == 5
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_error_handling_input_transcription(self):
|
||||||
|
"""Test error handling during input transcription processing."""
|
||||||
|
invocation_context = await testing_utils.create_invocation_context(
|
||||||
|
testing_utils.create_test_agent()
|
||||||
|
)
|
||||||
|
|
||||||
|
# Set up mock session service that raises an error
|
||||||
|
mock_session_service = AsyncMock()
|
||||||
|
mock_session_service.append_event.side_effect = Exception(
|
||||||
|
'Session service error'
|
||||||
|
)
|
||||||
|
invocation_context.session_service = mock_session_service
|
||||||
|
|
||||||
|
# Create test transcription
|
||||||
|
transcription = types.Transcription(text='Test transcription')
|
||||||
|
|
||||||
|
# Handle transcription should raise the exception
|
||||||
|
with pytest.raises(Exception, match='Session service error'):
|
||||||
|
await self.manager.handle_input_transcription(
|
||||||
|
invocation_context, transcription
|
||||||
|
)
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_error_handling_output_transcription(self):
|
||||||
|
"""Test error handling during output transcription processing."""
|
||||||
|
invocation_context = await testing_utils.create_invocation_context(
|
||||||
|
testing_utils.create_test_agent()
|
||||||
|
)
|
||||||
|
|
||||||
|
# Set up mock session service that raises an error
|
||||||
|
mock_session_service = AsyncMock()
|
||||||
|
mock_session_service.append_event.side_effect = Exception(
|
||||||
|
'Session service error'
|
||||||
|
)
|
||||||
|
invocation_context.session_service = mock_session_service
|
||||||
|
|
||||||
|
# Create test transcription
|
||||||
|
transcription = types.Transcription(text='Test transcription')
|
||||||
|
|
||||||
|
# Handle transcription should raise the exception
|
||||||
|
with pytest.raises(Exception, match='Session service error'):
|
||||||
|
await self.manager.handle_output_transcription(
|
||||||
|
invocation_context, transcription
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_get_transcription_stats_empty_session(self):
|
||||||
|
"""Test getting transcription statistics for empty session."""
|
||||||
|
invocation_context = Mock()
|
||||||
|
invocation_context.session.events = []
|
||||||
|
|
||||||
|
stats = self.manager.get_transcription_stats(invocation_context)
|
||||||
|
|
||||||
|
expected = {
|
||||||
|
'input_transcriptions': 0,
|
||||||
|
'output_transcriptions': 0,
|
||||||
|
'total_transcriptions': 0,
|
||||||
|
}
|
||||||
|
assert stats == expected
|
||||||
|
|
||||||
|
def test_get_transcription_stats_with_events(self):
|
||||||
|
"""Test getting transcription statistics for session with events."""
|
||||||
|
invocation_context = Mock()
|
||||||
|
|
||||||
|
# Create mock events
|
||||||
|
input_event1 = Mock()
|
||||||
|
input_event1.input_transcription = types.Transcription(text='User 1')
|
||||||
|
input_event1.output_transcription = None
|
||||||
|
|
||||||
|
input_event2 = Mock()
|
||||||
|
input_event2.input_transcription = types.Transcription(text='User 2')
|
||||||
|
input_event2.output_transcription = None
|
||||||
|
|
||||||
|
output_event = Mock()
|
||||||
|
output_event.input_transcription = None
|
||||||
|
output_event.output_transcription = types.Transcription(
|
||||||
|
text='Model response'
|
||||||
|
)
|
||||||
|
|
||||||
|
regular_event = Mock()
|
||||||
|
regular_event.input_transcription = None
|
||||||
|
regular_event.output_transcription = None
|
||||||
|
|
||||||
|
invocation_context.session.events = [
|
||||||
|
input_event1,
|
||||||
|
output_event,
|
||||||
|
input_event2,
|
||||||
|
regular_event,
|
||||||
|
]
|
||||||
|
|
||||||
|
stats = self.manager.get_transcription_stats(invocation_context)
|
||||||
|
|
||||||
|
expected = {
|
||||||
|
'input_transcriptions': 2,
|
||||||
|
'output_transcriptions': 1,
|
||||||
|
'total_transcriptions': 3,
|
||||||
|
}
|
||||||
|
assert stats == expected
|
||||||
|
|
||||||
|
def test_get_transcription_stats_missing_attributes(self):
|
||||||
|
"""Test getting transcription statistics when events don't have transcription attributes."""
|
||||||
|
invocation_context = Mock()
|
||||||
|
|
||||||
|
# Create mock events and explicitly set transcription attributes to None
|
||||||
|
event1 = Mock()
|
||||||
|
event1.input_transcription = None
|
||||||
|
event1.output_transcription = None
|
||||||
|
|
||||||
|
event2 = Mock()
|
||||||
|
event2.input_transcription = None
|
||||||
|
event2.output_transcription = None
|
||||||
|
|
||||||
|
invocation_context.session.events = [event1, event2]
|
||||||
|
|
||||||
|
stats = self.manager.get_transcription_stats(invocation_context)
|
||||||
|
|
||||||
|
expected = {
|
||||||
|
'input_transcriptions': 0,
|
||||||
|
'output_transcriptions': 0,
|
||||||
|
'total_transcriptions': 0,
|
||||||
|
}
|
||||||
|
assert stats == expected
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_transcription_event_fields(self):
|
||||||
|
"""Test that transcription events have correct field values."""
|
||||||
|
invocation_context = await testing_utils.create_invocation_context(
|
||||||
|
testing_utils.create_test_agent()
|
||||||
|
)
|
||||||
|
|
||||||
|
# Set up mock session service
|
||||||
|
mock_session_service = AsyncMock()
|
||||||
|
invocation_context.session_service = mock_session_service
|
||||||
|
|
||||||
|
# Create test transcription with specific content
|
||||||
|
transcription = types.Transcription(
|
||||||
|
text='Test transcription content', finished=True
|
||||||
|
)
|
||||||
|
|
||||||
|
# Handle input transcription
|
||||||
|
await self.manager.handle_input_transcription(
|
||||||
|
invocation_context, transcription
|
||||||
|
)
|
||||||
|
|
||||||
|
# Verify the event structure
|
||||||
|
call_args = mock_session_service.append_event.call_args
|
||||||
|
event = call_args[0][1]
|
||||||
|
|
||||||
|
# Check all required fields are present
|
||||||
|
assert hasattr(event, 'id')
|
||||||
|
assert hasattr(event, 'invocation_id')
|
||||||
|
assert hasattr(event, 'author')
|
||||||
|
assert hasattr(event, 'input_transcription')
|
||||||
|
assert hasattr(event, 'output_transcription')
|
||||||
|
assert hasattr(event, 'timestamp')
|
||||||
|
|
||||||
|
# Check values
|
||||||
|
assert event.id is not None
|
||||||
|
assert event.invocation_id == invocation_context.invocation_id
|
||||||
|
assert event.author == 'user'
|
||||||
|
assert event.input_transcription == transcription
|
||||||
|
assert event.output_transcription is None
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_transcription_with_different_data_types(self):
|
||||||
|
"""Test handling transcriptions with different data types."""
|
||||||
|
invocation_context = await testing_utils.create_invocation_context(
|
||||||
|
testing_utils.create_test_agent()
|
||||||
|
)
|
||||||
|
|
||||||
|
# Set up mock session service
|
||||||
|
mock_session_service = AsyncMock()
|
||||||
|
invocation_context.session_service = mock_session_service
|
||||||
|
|
||||||
|
# Test with transcription that has basic fields only
|
||||||
|
transcription = types.Transcription(
|
||||||
|
text='Advanced transcription', finished=True
|
||||||
|
)
|
||||||
|
|
||||||
|
# Handle transcription
|
||||||
|
await self.manager.handle_input_transcription(
|
||||||
|
invocation_context, transcription
|
||||||
|
)
|
||||||
|
|
||||||
|
# Verify the transcription object is preserved as-is
|
||||||
|
call_args = mock_session_service.append_event.call_args
|
||||||
|
event = call_args[0][1]
|
||||||
|
|
||||||
|
assert event.input_transcription == transcription
|
||||||
|
assert event.input_transcription.text == 'Advanced transcription'
|
||||||
@@ -0,0 +1,241 @@
|
|||||||
|
# Copyright 2025 Google LLC
|
||||||
|
#
|
||||||
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
# you may not use this file except in compliance with the License.
|
||||||
|
# You may obtain a copy of the License at
|
||||||
|
#
|
||||||
|
# http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
#
|
||||||
|
# Unless required by applicable law or agreed to in writing, software
|
||||||
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
# See the License for the specific language governing permissions and
|
||||||
|
# limitations under the License.
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import time
|
||||||
|
|
||||||
|
from google.adk.agents import Agent
|
||||||
|
from google.adk.agents import LiveRequestQueue
|
||||||
|
from google.adk.agents.invocation_context import RealtimeCacheEntry
|
||||||
|
from google.adk.agents.run_config import RunConfig
|
||||||
|
from google.adk.events.event import Event
|
||||||
|
from google.adk.models import LlmResponse
|
||||||
|
from google.genai import types
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from .. import testing_utils
|
||||||
|
|
||||||
|
|
||||||
|
def test_audio_caching_direct():
|
||||||
|
"""Test audio caching logic directly without full live streaming."""
|
||||||
|
# This test directly verifies that our audio caching logic works
|
||||||
|
audio_data = b'\x00\xFF\x01\x02\x03\x04\x05\x06'
|
||||||
|
audio_mime_type = 'audio/pcm'
|
||||||
|
|
||||||
|
# Create mock responses for successful completion
|
||||||
|
responses = [
|
||||||
|
LlmResponse(
|
||||||
|
content=types.Content(
|
||||||
|
role='model',
|
||||||
|
parts=[types.Part.from_text(text='Processing audio...')],
|
||||||
|
),
|
||||||
|
turn_complete=False,
|
||||||
|
),
|
||||||
|
LlmResponse(turn_complete=True), # This should trigger flush
|
||||||
|
]
|
||||||
|
|
||||||
|
mock_model = testing_utils.MockModel.create(responses)
|
||||||
|
mock_model.model = 'gemini-2.0-flash-exp' # For CFC support
|
||||||
|
|
||||||
|
root_agent = Agent(
|
||||||
|
name='test_agent',
|
||||||
|
model=mock_model,
|
||||||
|
tools=[],
|
||||||
|
)
|
||||||
|
|
||||||
|
# Test our implementation by directly calling it
|
||||||
|
async def test_caching():
|
||||||
|
# Create context similar to what would be created in real scenario
|
||||||
|
invocation_context = await testing_utils.create_invocation_context(
|
||||||
|
root_agent, run_config=RunConfig(support_cfc=True)
|
||||||
|
)
|
||||||
|
|
||||||
|
# Import our caching classes
|
||||||
|
from google.adk.agents.invocation_context import RealtimeCacheEntry
|
||||||
|
from google.adk.flows.llm_flows.base_llm_flow import BaseLlmFlow
|
||||||
|
|
||||||
|
# Create a mock flow to test our methods
|
||||||
|
flow = BaseLlmFlow()
|
||||||
|
|
||||||
|
# Test adding audio to cache
|
||||||
|
invocation_context.input_realtime_cache = []
|
||||||
|
audio_entry = RealtimeCacheEntry(
|
||||||
|
role='user',
|
||||||
|
data=types.Blob(data=audio_data, mime_type=audio_mime_type),
|
||||||
|
timestamp=1234567890.0,
|
||||||
|
)
|
||||||
|
invocation_context.input_realtime_cache.append(audio_entry)
|
||||||
|
|
||||||
|
# Verify cache has data
|
||||||
|
assert len(invocation_context.input_realtime_cache) == 1
|
||||||
|
assert invocation_context.input_realtime_cache[0].data.data == audio_data
|
||||||
|
|
||||||
|
# Test flushing cache
|
||||||
|
await flow._handle_control_event_flush(invocation_context, responses[-1])
|
||||||
|
|
||||||
|
# Verify cache was cleared
|
||||||
|
assert len(invocation_context.input_realtime_cache) == 0
|
||||||
|
|
||||||
|
# Check if artifacts were created
|
||||||
|
artifact_keys = (
|
||||||
|
await invocation_context.artifact_service.list_artifact_keys(
|
||||||
|
app_name=invocation_context.app_name,
|
||||||
|
user_id=invocation_context.user_id,
|
||||||
|
session_id=invocation_context.session.id,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
# Should have at least one audio artifact
|
||||||
|
audio_artifacts = [key for key in artifact_keys if 'audio' in key.lower()]
|
||||||
|
assert (
|
||||||
|
len(audio_artifacts) > 0
|
||||||
|
), f'Expected audio artifacts, found: {artifact_keys}'
|
||||||
|
|
||||||
|
# Verify artifact content
|
||||||
|
if audio_artifacts:
|
||||||
|
artifact = await invocation_context.artifact_service.load_artifact(
|
||||||
|
app_name=invocation_context.app_name,
|
||||||
|
user_id=invocation_context.user_id,
|
||||||
|
session_id=invocation_context.session.id,
|
||||||
|
filename=audio_artifacts[0],
|
||||||
|
)
|
||||||
|
assert artifact.inline_data.data == audio_data
|
||||||
|
|
||||||
|
return True
|
||||||
|
|
||||||
|
# Run the async test
|
||||||
|
result = asyncio.run(test_caching())
|
||||||
|
assert result is True
|
||||||
|
|
||||||
|
|
||||||
|
def test_transcription_handling():
|
||||||
|
"""Test that transcriptions are properly handled and saved to session service."""
|
||||||
|
|
||||||
|
# Create mock responses with transcriptions
|
||||||
|
input_transcription = types.Transcription(
|
||||||
|
text='Hello, this is transcribed input', finished=True
|
||||||
|
)
|
||||||
|
output_transcription = types.Transcription(
|
||||||
|
text='This is transcribed output', finished=True
|
||||||
|
)
|
||||||
|
|
||||||
|
responses = [
|
||||||
|
LlmResponse(
|
||||||
|
content=types.Content(
|
||||||
|
role='model', parts=[types.Part.from_text(text='Processing...')]
|
||||||
|
),
|
||||||
|
turn_complete=False,
|
||||||
|
),
|
||||||
|
LlmResponse(input_transcription=input_transcription, turn_complete=False),
|
||||||
|
LlmResponse(
|
||||||
|
output_transcription=output_transcription, turn_complete=False
|
||||||
|
),
|
||||||
|
LlmResponse(turn_complete=True),
|
||||||
|
]
|
||||||
|
|
||||||
|
mock_model = testing_utils.MockModel.create(responses)
|
||||||
|
mock_model.model = 'gemini-2.0-flash-exp'
|
||||||
|
|
||||||
|
root_agent = Agent(
|
||||||
|
name='test_agent',
|
||||||
|
model=mock_model,
|
||||||
|
tools=[],
|
||||||
|
)
|
||||||
|
|
||||||
|
async def test_transcription():
|
||||||
|
# Create context
|
||||||
|
invocation_context = await testing_utils.create_invocation_context(
|
||||||
|
root_agent, run_config=RunConfig(support_cfc=True)
|
||||||
|
)
|
||||||
|
|
||||||
|
from google.adk.events.event import Event
|
||||||
|
from google.adk.flows.llm_flows.base_llm_flow import BaseLlmFlow
|
||||||
|
|
||||||
|
flow = BaseLlmFlow()
|
||||||
|
|
||||||
|
# Test processing transcription events
|
||||||
|
session_events_before = len(invocation_context.session.events)
|
||||||
|
|
||||||
|
# Simulate input transcription event
|
||||||
|
input_event = Event(
|
||||||
|
id=Event.new_id(),
|
||||||
|
invocation_id=invocation_context.invocation_id,
|
||||||
|
author='user',
|
||||||
|
input_transcription=input_transcription,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Simulate output transcription event
|
||||||
|
output_event = Event(
|
||||||
|
id=Event.new_id(),
|
||||||
|
invocation_id=invocation_context.invocation_id,
|
||||||
|
author=invocation_context.agent.name,
|
||||||
|
output_transcription=output_transcription,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Save transcription events to session
|
||||||
|
await invocation_context.session_service.append_event(
|
||||||
|
invocation_context.session, input_event
|
||||||
|
)
|
||||||
|
await invocation_context.session_service.append_event(
|
||||||
|
invocation_context.session, output_event
|
||||||
|
)
|
||||||
|
|
||||||
|
# Verify transcriptions were saved to session
|
||||||
|
session_events_after = len(invocation_context.session.events)
|
||||||
|
assert session_events_after == session_events_before + 2
|
||||||
|
|
||||||
|
# Check that transcription events were saved
|
||||||
|
transcription_events = [
|
||||||
|
event
|
||||||
|
for event in invocation_context.session.events
|
||||||
|
if hasattr(event, 'input_transcription')
|
||||||
|
and event.input_transcription
|
||||||
|
or hasattr(event, 'output_transcription')
|
||||||
|
and event.output_transcription
|
||||||
|
]
|
||||||
|
assert len(transcription_events) >= 2
|
||||||
|
|
||||||
|
# Verify input transcription
|
||||||
|
input_transcription_events = [
|
||||||
|
event
|
||||||
|
for event in invocation_context.session.events
|
||||||
|
if hasattr(event, 'input_transcription') and event.input_transcription
|
||||||
|
]
|
||||||
|
assert len(input_transcription_events) >= 1
|
||||||
|
assert (
|
||||||
|
input_transcription_events[0].input_transcription.text
|
||||||
|
== 'Hello, this is transcribed input'
|
||||||
|
)
|
||||||
|
assert input_transcription_events[0].author == 'user'
|
||||||
|
|
||||||
|
# Verify output transcription
|
||||||
|
output_transcription_events = [
|
||||||
|
event
|
||||||
|
for event in invocation_context.session.events
|
||||||
|
if hasattr(event, 'output_transcription') and event.output_transcription
|
||||||
|
]
|
||||||
|
assert len(output_transcription_events) >= 1
|
||||||
|
assert (
|
||||||
|
output_transcription_events[0].output_transcription.text
|
||||||
|
== 'This is transcribed output'
|
||||||
|
)
|
||||||
|
assert (
|
||||||
|
output_transcription_events[0].author == invocation_context.agent.name
|
||||||
|
)
|
||||||
|
|
||||||
|
return True
|
||||||
|
|
||||||
|
# Run the async test
|
||||||
|
result = asyncio.run(test_transcription())
|
||||||
|
assert result is True
|
||||||
@@ -41,6 +41,18 @@ from google.genai.types import Part
|
|||||||
from typing_extensions import override
|
from typing_extensions import override
|
||||||
|
|
||||||
|
|
||||||
|
def create_test_agent(name: str = 'test_agent') -> LlmAgent:
|
||||||
|
"""Create a simple test agent for use in unit tests.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
name: The name of the test agent.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
A configured LlmAgent instance suitable for testing.
|
||||||
|
"""
|
||||||
|
return LlmAgent(name=name)
|
||||||
|
|
||||||
|
|
||||||
class UserContent(types.Content):
|
class UserContent(types.Content):
|
||||||
|
|
||||||
def __init__(self, text_or_part: str):
|
def __init__(self, text_or_part: str):
|
||||||
|
|||||||
Reference in New Issue
Block a user