From 6d2d2d0bdb0a2e47223a7ce410893462823af46e Mon Sep 17 00:00:00 2001 From: Hangfei Lin Date: Wed, 16 Jul 2025 13:33:46 -0700 Subject: [PATCH] feat: add support for session resumption(only transparent mode) config to run_config This commit adds support for the session resumption configuration in the run_config. The SessionResumptionConfig is added to RunConfig to allow the user to set up a configuration for session resumption(only transparent mode for now). There are two modes of session resumption: manual and transparent. In manual mode, you have to manually bookkeeping the session information and restarts the session which is tricky to do right now. In transparent mode, the server does the bookkeeping for you and no hassle on ADK side. For now, the transparent mode should be enough. Also, added the relevant unit tests to check that every possible configuration is set properly and the run_config is correctly populated. This is needed for supporting the new session resumption feature. PiperOrigin-RevId: 783876503 --- src/google/adk/agents/run_config.py | 3 + src/google/adk/flows/llm_flows/basic.py | 3 + .../streaming/test_live_streaming_configs.py | 477 ++++++++++++++++++ 3 files changed, 483 insertions(+) create mode 100644 tests/unittests/streaming/test_live_streaming_configs.py diff --git a/src/google/adk/agents/run_config.py b/src/google/adk/agents/run_config.py index c9a50a0ae..52d8a9f57 100644 --- a/src/google/adk/agents/run_config.py +++ b/src/google/adk/agents/run_config.py @@ -79,6 +79,9 @@ class RunConfig(BaseModel): proactivity: Optional[types.ProactivityConfig] = None """Configures the proactivity of the model. This allows the model to respond proactively to the input and to ignore irrelevant input.""" + session_resumption: Optional[types.SessionResumptionConfig] = None + """Configures session resumption mechanism. Only support transparent session resumption mode now.""" + max_llm_calls: int = 500 """ A limit on the total number of llm calls for a given run. diff --git a/src/google/adk/flows/llm_flows/basic.py b/src/google/adk/flows/llm_flows/basic.py index ee5c83da1..2557088b1 100644 --- a/src/google/adk/flows/llm_flows/basic.py +++ b/src/google/adk/flows/llm_flows/basic.py @@ -74,6 +74,9 @@ async def run_async( llm_request.live_connect_config.proactivity = ( invocation_context.run_config.proactivity ) + llm_request.live_connect_config.proactivity = ( + invocation_context.run_config.session_resumption + ) # TODO: handle tool append here, instead of in BaseTool.process_llm_request. diff --git a/tests/unittests/streaming/test_live_streaming_configs.py b/tests/unittests/streaming/test_live_streaming_configs.py new file mode 100644 index 000000000..eb5a49402 --- /dev/null +++ b/tests/unittests/streaming/test_live_streaming_configs.py @@ -0,0 +1,477 @@ +# 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 google.adk.agents import Agent +from google.adk.agents import LiveRequestQueue +from google.adk.agents.run_config import RunConfig +from google.adk.models import LlmResponse +from google.genai import types +import pytest + +from .. import testing_utils + + +def test_streaming(): + response1 = LlmResponse( + turn_complete=True, + ) + + mock_model = testing_utils.MockModel.create([response1]) + + root_agent = Agent( + name='root_agent', + model=mock_model, + tools=[], + ) + + runner = testing_utils.InMemoryRunner( + root_agent=root_agent, response_modalities=['AUDIO'] + ) + live_request_queue = LiveRequestQueue() + live_request_queue.send_realtime( + blob=types.Blob(data=b'\x00\xFF', mime_type='audio/pcm') + ) + res_events = runner.run_live(live_request_queue) + + assert res_events is not None, 'Expected a list of events, got None.' + assert ( + len(res_events) > 0 + ), 'Expected at least one response, but got an empty list.' + + +def test_streaming_with_output_audio_transcription(): + """Test streaming with output audio transcription configuration.""" + response1 = LlmResponse( + turn_complete=True, + ) + + mock_model = testing_utils.MockModel.create([response1]) + + root_agent = Agent( + name='root_agent', + model=mock_model, + tools=[], + ) + + runner = testing_utils.InMemoryRunner( + root_agent=root_agent, response_modalities=['AUDIO'] + ) + + # Create run config with output audio transcription + run_config = RunConfig( + output_audio_transcription=types.AudioTranscriptionConfig() + ) + + live_request_queue = LiveRequestQueue() + live_request_queue.send_realtime( + blob=types.Blob(data=b'\x00\xFF', mime_type='audio/pcm') + ) + res_events = runner.run_live(live_request_queue, run_config) + + assert res_events is not None, 'Expected a list of events, got None.' + assert ( + len(res_events) > 0 + ), 'Expected at least one response, but got an empty list.' + + +def test_streaming_with_input_audio_transcription(): + """Test streaming with input audio transcription configuration.""" + response1 = LlmResponse( + turn_complete=True, + ) + + mock_model = testing_utils.MockModel.create([response1]) + + root_agent = Agent( + name='root_agent', + model=mock_model, + tools=[], + ) + + runner = testing_utils.InMemoryRunner( + root_agent=root_agent, response_modalities=['AUDIO'] + ) + + # Create run config with input audio transcription + run_config = RunConfig( + input_audio_transcription=types.AudioTranscriptionConfig() + ) + + live_request_queue = LiveRequestQueue() + live_request_queue.send_realtime( + blob=types.Blob(data=b'\x00\xFF', mime_type='audio/pcm') + ) + res_events = runner.run_live(live_request_queue, run_config) + + assert res_events is not None, 'Expected a list of events, got None.' + assert ( + len(res_events) > 0 + ), 'Expected at least one response, but got an empty list.' + + +def test_streaming_with_realtime_input_config(): + """Test streaming with realtime input configuration.""" + response1 = LlmResponse( + turn_complete=True, + ) + + mock_model = testing_utils.MockModel.create([response1]) + + root_agent = Agent( + name='root_agent', + model=mock_model, + tools=[], + ) + + runner = testing_utils.InMemoryRunner( + root_agent=root_agent, response_modalities=['AUDIO'] + ) + + # Create run config with realtime input config + run_config = RunConfig( + realtime_input_config=types.RealtimeInputConfig( + automatic_activity_detection=types.AutomaticActivityDetection( + disabled=True + ) + ) + ) + + live_request_queue = LiveRequestQueue() + live_request_queue.send_realtime( + blob=types.Blob(data=b'\x00\xFF', mime_type='audio/pcm') + ) + res_events = runner.run_live(live_request_queue, run_config) + + assert res_events is not None, 'Expected a list of events, got None.' + assert ( + len(res_events) > 0 + ), 'Expected at least one response, but got an empty list.' + + +def test_streaming_with_realtime_input_config_vad_enabled(): + """Test streaming with realtime input configuration with VAD enabled.""" + response1 = LlmResponse( + turn_complete=True, + ) + + mock_model = testing_utils.MockModel.create([response1]) + + root_agent = Agent( + name='root_agent', + model=mock_model, + tools=[], + ) + + runner = testing_utils.InMemoryRunner( + root_agent=root_agent, response_modalities=['AUDIO'] + ) + + # Create run config with realtime input config with VAD enabled + run_config = RunConfig( + realtime_input_config=types.RealtimeInputConfig( + automatic_activity_detection=types.AutomaticActivityDetection( + disabled=False + ) + ) + ) + + live_request_queue = LiveRequestQueue() + live_request_queue.send_realtime( + blob=types.Blob(data=b'\x00\xFF', mime_type='audio/pcm') + ) + res_events = runner.run_live(live_request_queue, run_config) + + assert res_events is not None, 'Expected a list of events, got None.' + assert ( + len(res_events) > 0 + ), 'Expected at least one response, but got an empty list.' + + +def test_streaming_with_enable_affective_dialog_true(): + """Test streaming with affective dialog enabled.""" + response1 = LlmResponse( + turn_complete=True, + ) + + mock_model = testing_utils.MockModel.create([response1]) + + root_agent = Agent( + name='root_agent', + model=mock_model, + tools=[], + ) + + runner = testing_utils.InMemoryRunner( + root_agent=root_agent, response_modalities=['AUDIO'] + ) + + # Create run config with affective dialog enabled + run_config = RunConfig(enable_affective_dialog=True) + + live_request_queue = LiveRequestQueue() + live_request_queue.send_realtime( + blob=types.Blob(data=b'\x00\xFF', mime_type='audio/pcm') + ) + res_events = runner.run_live(live_request_queue, run_config) + + assert res_events is not None, 'Expected a list of events, got None.' + assert ( + len(res_events) > 0 + ), 'Expected at least one response, but got an empty list.' + + +def test_streaming_with_enable_affective_dialog_false(): + """Test streaming with affective dialog disabled.""" + response1 = LlmResponse( + turn_complete=True, + ) + + mock_model = testing_utils.MockModel.create([response1]) + + root_agent = Agent( + name='root_agent', + model=mock_model, + tools=[], + ) + + runner = testing_utils.InMemoryRunner( + root_agent=root_agent, response_modalities=['AUDIO'] + ) + + # Create run config with affective dialog disabled + run_config = RunConfig(enable_affective_dialog=False) + + live_request_queue = LiveRequestQueue() + live_request_queue.send_realtime( + blob=types.Blob(data=b'\x00\xFF', mime_type='audio/pcm') + ) + res_events = runner.run_live(live_request_queue, run_config) + + assert res_events is not None, 'Expected a list of events, got None.' + assert ( + len(res_events) > 0 + ), 'Expected at least one response, but got an empty list.' + + +def test_streaming_with_proactivity_config(): + """Test streaming with proactivity configuration.""" + response1 = LlmResponse( + turn_complete=True, + ) + + mock_model = testing_utils.MockModel.create([response1]) + + root_agent = Agent( + name='root_agent', + model=mock_model, + tools=[], + ) + + runner = testing_utils.InMemoryRunner( + root_agent=root_agent, response_modalities=['AUDIO'] + ) + + # Create run config with proactivity config + run_config = RunConfig(proactivity=types.ProactivityConfig()) + + live_request_queue = LiveRequestQueue() + live_request_queue.send_realtime( + blob=types.Blob(data=b'\x00\xFF', mime_type='audio/pcm') + ) + res_events = runner.run_live(live_request_queue, run_config) + + assert res_events is not None, 'Expected a list of events, got None.' + assert ( + len(res_events) > 0 + ), 'Expected at least one response, but got an empty list.' + + +def test_streaming_with_combined_audio_transcription_configs(): + """Test streaming with both input and output audio transcription configurations.""" + response1 = LlmResponse( + turn_complete=True, + ) + + mock_model = testing_utils.MockModel.create([response1]) + + root_agent = Agent( + name='root_agent', + model=mock_model, + tools=[], + ) + + runner = testing_utils.InMemoryRunner( + root_agent=root_agent, response_modalities=['AUDIO'] + ) + + # Create run config with both input and output audio transcription + run_config = RunConfig( + input_audio_transcription=types.AudioTranscriptionConfig(), + output_audio_transcription=types.AudioTranscriptionConfig(), + ) + + live_request_queue = LiveRequestQueue() + live_request_queue.send_realtime( + blob=types.Blob(data=b'\x00\xFF', mime_type='audio/pcm') + ) + res_events = runner.run_live(live_request_queue, run_config) + + assert res_events is not None, 'Expected a list of events, got None.' + assert ( + len(res_events) > 0 + ), 'Expected at least one response, but got an empty list.' + + +def test_streaming_with_all_configs_combined(): + """Test streaming with all the new configurations combined.""" + response1 = LlmResponse( + turn_complete=True, + ) + + mock_model = testing_utils.MockModel.create([response1]) + + root_agent = Agent( + name='root_agent', + model=mock_model, + tools=[], + ) + + runner = testing_utils.InMemoryRunner( + root_agent=root_agent, response_modalities=['AUDIO'] + ) + + # Create run config with all configurations + run_config = RunConfig( + output_audio_transcription=types.AudioTranscriptionConfig(), + input_audio_transcription=types.AudioTranscriptionConfig(), + realtime_input_config=types.RealtimeInputConfig( + automatic_activity_detection=types.AutomaticActivityDetection( + disabled=True + ) + ), + enable_affective_dialog=True, + proactivity=types.ProactivityConfig(), + ) + + live_request_queue = LiveRequestQueue() + live_request_queue.send_realtime( + blob=types.Blob(data=b'\x00\xFF', mime_type='audio/pcm') + ) + res_events = runner.run_live(live_request_queue, run_config) + + assert res_events is not None, 'Expected a list of events, got None.' + assert ( + len(res_events) > 0 + ), 'Expected at least one response, but got an empty list.' + + +def test_streaming_config_validation(): + """Test that run_config values are properly set and accessible.""" + # Test that RunConfig properly validates and stores the configurations + run_config = RunConfig( + output_audio_transcription=types.AudioTranscriptionConfig(), + input_audio_transcription=types.AudioTranscriptionConfig(), + realtime_input_config=types.RealtimeInputConfig( + automatic_activity_detection=types.AutomaticActivityDetection( + disabled=False + ) + ), + enable_affective_dialog=True, + proactivity=types.ProactivityConfig(), + ) + + # Verify configurations are properly set + assert run_config.output_audio_transcription is not None + assert run_config.input_audio_transcription is not None + assert run_config.realtime_input_config is not None + assert ( + run_config.realtime_input_config.automatic_activity_detection.disabled + == False + ) + assert run_config.enable_affective_dialog == True + assert run_config.proactivity is not None + + +def test_streaming_with_multiple_audio_configs(): + """Test streaming with multiple audio transcription configurations.""" + response1 = LlmResponse( + turn_complete=True, + ) + + mock_model = testing_utils.MockModel.create([response1]) + + root_agent = Agent( + name='root_agent', + model=mock_model, + tools=[], + ) + + runner = testing_utils.InMemoryRunner( + root_agent=root_agent, response_modalities=['AUDIO'] + ) + + # Create run config with multiple audio transcription configs + run_config = RunConfig( + input_audio_transcription=types.AudioTranscriptionConfig(), + output_audio_transcription=types.AudioTranscriptionConfig(), + enable_affective_dialog=True, + ) + + live_request_queue = LiveRequestQueue() + live_request_queue.send_realtime( + blob=types.Blob(data=b'\x00\xFF', mime_type='audio/pcm') + ) + + res_events = runner.run_live(live_request_queue, run_config) + + assert res_events is not None, 'Expected a list of events, got None.' + assert ( + len(res_events) > 0 + ), 'Expected at least one response, but got an empty list.' + + +def test_streaming_with_session_resumption_config(): + """Test streaming with multiple audio transcription configurations.""" + response1 = LlmResponse( + turn_complete=True, + ) + + mock_model = testing_utils.MockModel.create([response1]) + + root_agent = Agent( + name='root_agent', + model=mock_model, + tools=[], + ) + + runner = testing_utils.InMemoryRunner( + root_agent=root_agent, response_modalities=['AUDIO'] + ) + + # Create run config with multiple audio transcription configs + run_config = RunConfig( + session_resumption=types.SessionResumptionConfig(transparent=True), + ) + + live_request_queue = LiveRequestQueue() + live_request_queue.send_realtime( + blob=types.Blob(data=b'\x00\xFF', mime_type='audio/pcm') + ) + + res_events = runner.run_live(live_request_queue, run_config) + + assert res_events is not None, 'Expected a list of events, got None.' + assert ( + len(res_events) > 0 + ), 'Expected at least one response, but got an empty list.'