mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
Hangfei LinandCopybara-Service
5e8aa15a50
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: 786549455
589 lines
18 KiB
Python
589 lines
18 KiB
Python
# 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.'
|
|
assert len(mock_model.requests) == 1
|
|
# Get the request that was captured
|
|
llm_request_sent_to_mock = mock_model.requests[0]
|
|
|
|
# Assert that the request contained the correct configuration
|
|
assert llm_request_sent_to_mock.live_connect_config is not None
|
|
assert (
|
|
llm_request_sent_to_mock.live_connect_config.output_audio_transcription
|
|
is None
|
|
)
|
|
|
|
|
|
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.'
|
|
assert len(mock_model.requests) == 1
|
|
# Get the request that was captured
|
|
llm_request_sent_to_mock = mock_model.requests[0]
|
|
|
|
# Assert that the request contained the correct configuration
|
|
assert llm_request_sent_to_mock.live_connect_config is not None
|
|
assert (
|
|
llm_request_sent_to_mock.live_connect_config.output_audio_transcription
|
|
is not None
|
|
)
|
|
|
|
|
|
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.'
|
|
assert len(mock_model.requests) == 1
|
|
# Get the request that was captured
|
|
llm_request_sent_to_mock = mock_model.requests[0]
|
|
|
|
# Assert that the request contained the correct configuration
|
|
assert llm_request_sent_to_mock.live_connect_config is not None
|
|
assert (
|
|
llm_request_sent_to_mock.live_connect_config.input_audio_transcription
|
|
is not None
|
|
)
|
|
|
|
|
|
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.'
|
|
assert len(mock_model.requests) == 1
|
|
# Get the request that was captured
|
|
llm_request_sent_to_mock = mock_model.requests[0]
|
|
|
|
# Assert that the request contained the correct configuration
|
|
assert llm_request_sent_to_mock.live_connect_config is not None
|
|
assert (
|
|
llm_request_sent_to_mock.live_connect_config.realtime_input_config.automatic_activity_detection.disabled
|
|
is True
|
|
)
|
|
|
|
|
|
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.'
|
|
assert len(mock_model.requests) == 1
|
|
# Get the request that was captured
|
|
llm_request_sent_to_mock = mock_model.requests[0]
|
|
|
|
# Assert that the request contained the correct configuration
|
|
assert llm_request_sent_to_mock.live_connect_config is not None
|
|
assert (
|
|
llm_request_sent_to_mock.live_connect_config.realtime_input_config.automatic_activity_detection.disabled
|
|
is False
|
|
)
|
|
|
|
|
|
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.'
|
|
assert len(mock_model.requests) == 1
|
|
# Get the request that was captured
|
|
llm_request_sent_to_mock = mock_model.requests[0]
|
|
|
|
# Assert that the request contained the correct configuration
|
|
assert llm_request_sent_to_mock.live_connect_config is not None
|
|
assert (
|
|
llm_request_sent_to_mock.live_connect_config.enable_affective_dialog
|
|
is True
|
|
)
|
|
|
|
|
|
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.'
|
|
assert len(mock_model.requests) == 1
|
|
# Get the request that was captured
|
|
llm_request_sent_to_mock = mock_model.requests[0]
|
|
|
|
# Assert that the request contained the correct configuration
|
|
assert llm_request_sent_to_mock.live_connect_config is not None
|
|
assert (
|
|
llm_request_sent_to_mock.live_connect_config.enable_affective_dialog
|
|
is False
|
|
)
|
|
|
|
|
|
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.'
|
|
assert len(mock_model.requests) == 1
|
|
# Get the request that was captured
|
|
llm_request_sent_to_mock = mock_model.requests[0]
|
|
|
|
# Assert that the request contained the correct configuration
|
|
assert llm_request_sent_to_mock.live_connect_config is not None
|
|
assert llm_request_sent_to_mock.live_connect_config.proactivity is not None
|
|
|
|
|
|
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.'
|
|
assert len(mock_model.requests) == 1
|
|
# Get the request that was captured
|
|
llm_request_sent_to_mock = mock_model.requests[0]
|
|
|
|
# Assert that the request contained the correct configuration
|
|
assert llm_request_sent_to_mock.live_connect_config is not None
|
|
assert (
|
|
llm_request_sent_to_mock.live_connect_config.input_audio_transcription
|
|
is not None
|
|
)
|
|
assert (
|
|
llm_request_sent_to_mock.live_connect_config.output_audio_transcription
|
|
is not None
|
|
)
|
|
|
|
|
|
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.'
|
|
assert len(mock_model.requests) == 1
|
|
# Get the request that was captured
|
|
llm_request_sent_to_mock = mock_model.requests[0]
|
|
|
|
# Assert that the request contained the correct configuration
|
|
assert llm_request_sent_to_mock.live_connect_config is not None
|
|
assert (
|
|
llm_request_sent_to_mock.live_connect_config.realtime_input_config
|
|
is not None
|
|
)
|
|
assert llm_request_sent_to_mock.live_connect_config.proactivity is not None
|
|
assert (
|
|
llm_request_sent_to_mock.live_connect_config.enable_affective_dialog
|
|
is True
|
|
)
|
|
|
|
|
|
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.'
|
|
assert len(mock_model.requests) == 1
|
|
# Get the request that was captured
|
|
llm_request_sent_to_mock = mock_model.requests[0]
|
|
|
|
# Assert that the request contained the correct configuration
|
|
assert llm_request_sent_to_mock.live_connect_config is not None
|
|
assert (
|
|
llm_request_sent_to_mock.live_connect_config.input_audio_transcription
|
|
is not None
|
|
)
|
|
assert (
|
|
llm_request_sent_to_mock.live_connect_config.output_audio_transcription
|
|
is not None
|
|
)
|
|
assert (
|
|
llm_request_sent_to_mock.live_connect_config.enable_affective_dialog
|
|
is True
|
|
)
|
|
|
|
|
|
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.'
|
|
assert len(mock_model.requests) == 1
|
|
# Get the request that was captured
|
|
llm_request_sent_to_mock = mock_model.requests[0]
|
|
|
|
# Assert that the request contained the correct configuration
|
|
assert llm_request_sent_to_mock.live_connect_config is not None
|
|
assert (
|
|
llm_request_sent_to_mock.live_connect_config.session_resumption
|
|
is not None
|
|
)
|
|
assert (
|
|
llm_request_sent_to_mock.live_connect_config.session_resumption.transparent
|
|
is True
|
|
)
|