feat: Add support for store audio and transcription into events/sessions and artifacts

PiperOrigin-RevId: 798011660

tests/unittests/flows/llm_flows/test_audio_cache_manager.py

@@ -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}'
+    )

tests/unittests/flows/llm_flows/test_transcription_manager.py

@@ -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'

tests/unittests/streaming/test_streaming_audio_storage.py

@@ -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

tests/unittests/testing_utils.py

@@ -41,6 +41,18 @@ from google.genai.types import Part
 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):
 
   def __init__(self, text_or_part: str):