mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
Xiang (Sean) ZhouandCopybara-Service
5269a6b1d6
previously we only register streaming tool that accept stream input at runner, now uniformly register all streaming tool at runner. Co-authored-by: Xiang (Sean) Zhou <seanzhougoogle@google.com> PiperOrigin-RevId: 869447996
1548 lines
48 KiB
Python
1548 lines
48 KiB
Python
# Copyright 2026 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
|
|
from typing import Any
|
|
from typing import AsyncGenerator
|
|
from typing import Awaitable
|
|
|
|
from google.adk.agents.live_request_queue import LiveRequestQueue
|
|
from google.adk.agents.llm_agent import Agent
|
|
from google.adk.models.llm_response 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_live_streaming_function_call_single():
|
|
"""Test live streaming with a single function call response."""
|
|
# Create a function call response
|
|
function_call = types.Part.from_function_call(
|
|
name='get_weather', args={'location': 'San Francisco', 'unit': 'celsius'}
|
|
)
|
|
|
|
# Create LLM responses: function call followed by turn completion
|
|
response1 = LlmResponse(
|
|
content=types.Content(role='model', parts=[function_call]),
|
|
turn_complete=False,
|
|
)
|
|
response2 = LlmResponse(
|
|
turn_complete=True,
|
|
)
|
|
|
|
mock_model = testing_utils.MockModel.create([response1, response2])
|
|
|
|
# Mock function that would be called
|
|
def get_weather(location: str, unit: str = 'celsius') -> dict:
|
|
return {
|
|
'temperature': 22,
|
|
'condition': 'sunny',
|
|
'location': location,
|
|
'unit': unit,
|
|
}
|
|
|
|
root_agent = Agent(
|
|
name='root_agent',
|
|
model=mock_model,
|
|
tools=[get_weather],
|
|
)
|
|
|
|
# Create a custom runner class that collects all events
|
|
class CustomTestRunner(testing_utils.InMemoryRunner):
|
|
|
|
def _run_with_loop(self, coro):
|
|
try:
|
|
old_loop = asyncio.get_event_loop()
|
|
except RuntimeError:
|
|
old_loop = None
|
|
|
|
loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(loop)
|
|
try:
|
|
loop.run_until_complete(coro)
|
|
except (asyncio.TimeoutError, asyncio.CancelledError):
|
|
pass
|
|
finally:
|
|
loop.close()
|
|
asyncio.set_event_loop(old_loop)
|
|
|
|
def run_live(
|
|
self,
|
|
live_request_queue: LiveRequestQueue,
|
|
run_config: testing_utils.RunConfig = None,
|
|
) -> list[testing_utils.Event]:
|
|
collected_responses = []
|
|
|
|
async def consume_responses(session: testing_utils.Session):
|
|
run_res = self.runner.run_live(
|
|
session=session,
|
|
live_request_queue=live_request_queue,
|
|
run_config=run_config or testing_utils.RunConfig(),
|
|
)
|
|
|
|
async for response in run_res:
|
|
collected_responses.append(response)
|
|
# Collect a reasonable number of events, don't wait for too many
|
|
if len(collected_responses) >= 3:
|
|
return
|
|
|
|
self._run_with_loop(
|
|
asyncio.wait_for(consume_responses(self.session), timeout=5.0)
|
|
)
|
|
|
|
return collected_responses
|
|
|
|
runner = CustomTestRunner(root_agent=root_agent)
|
|
live_request_queue = LiveRequestQueue()
|
|
live_request_queue.send_realtime(
|
|
blob=types.Blob(
|
|
data=b'What is the weather in San Francisco?', 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) >= 1, 'Expected at least one event.'
|
|
|
|
# Check that we got a function call event
|
|
function_call_found = False
|
|
function_response_found = False
|
|
|
|
for event in res_events:
|
|
if event.content and event.content.parts:
|
|
for part in event.content.parts:
|
|
if part.function_call and part.function_call.name == 'get_weather':
|
|
function_call_found = True
|
|
assert part.function_call.args['location'] == 'San Francisco'
|
|
assert part.function_call.args['unit'] == 'celsius'
|
|
elif (
|
|
part.function_response
|
|
and part.function_response.name == 'get_weather'
|
|
):
|
|
function_response_found = True
|
|
assert part.function_response.response['temperature'] == 22
|
|
assert part.function_response.response['condition'] == 'sunny'
|
|
|
|
assert function_call_found, 'Expected a function call event.'
|
|
# Note: In live streaming, function responses might be handled differently,
|
|
# so we check for the function call which is the primary indicator of function calling working
|
|
|
|
|
|
def test_live_streaming_function_call_multiple():
|
|
"""Test live streaming with multiple function calls in sequence."""
|
|
# Create multiple function call responses
|
|
function_call1 = types.Part.from_function_call(
|
|
name='get_weather', args={'location': 'San Francisco'}
|
|
)
|
|
function_call2 = types.Part.from_function_call(
|
|
name='get_time', args={'timezone': 'PST'}
|
|
)
|
|
|
|
# Create LLM responses: two function calls followed by turn completion
|
|
response1 = LlmResponse(
|
|
content=types.Content(role='model', parts=[function_call1]),
|
|
turn_complete=False,
|
|
)
|
|
response2 = LlmResponse(
|
|
content=types.Content(role='model', parts=[function_call2]),
|
|
turn_complete=False,
|
|
)
|
|
response3 = LlmResponse(
|
|
turn_complete=True,
|
|
)
|
|
|
|
mock_model = testing_utils.MockModel.create([response1, response2, response3])
|
|
|
|
# Mock functions
|
|
def get_weather(location: str) -> dict:
|
|
return {'temperature': 22, 'condition': 'sunny', 'location': location}
|
|
|
|
def get_time(timezone: str) -> dict:
|
|
return {'time': '14:30', 'timezone': timezone}
|
|
|
|
root_agent = Agent(
|
|
name='root_agent',
|
|
model=mock_model,
|
|
tools=[get_weather, get_time],
|
|
)
|
|
|
|
# Use the custom runner
|
|
class CustomTestRunner(testing_utils.InMemoryRunner):
|
|
|
|
def _run_with_loop(self, coro):
|
|
try:
|
|
old_loop = asyncio.get_event_loop()
|
|
except RuntimeError:
|
|
old_loop = None
|
|
|
|
loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(loop)
|
|
try:
|
|
loop.run_until_complete(coro)
|
|
except (asyncio.TimeoutError, asyncio.CancelledError):
|
|
pass
|
|
finally:
|
|
loop.close()
|
|
asyncio.set_event_loop(old_loop)
|
|
|
|
def run_live(
|
|
self,
|
|
live_request_queue: LiveRequestQueue,
|
|
run_config: testing_utils.RunConfig = None,
|
|
) -> list[testing_utils.Event]:
|
|
collected_responses = []
|
|
|
|
async def consume_responses(session: testing_utils.Session):
|
|
run_res = self.runner.run_live(
|
|
session=session,
|
|
live_request_queue=live_request_queue,
|
|
run_config=run_config or testing_utils.RunConfig(),
|
|
)
|
|
|
|
async for response in run_res:
|
|
collected_responses.append(response)
|
|
if len(collected_responses) >= 3:
|
|
return
|
|
|
|
self._run_with_loop(
|
|
asyncio.wait_for(consume_responses(self.session), timeout=5.0)
|
|
)
|
|
|
|
return collected_responses
|
|
|
|
runner = CustomTestRunner(root_agent=root_agent)
|
|
live_request_queue = LiveRequestQueue()
|
|
live_request_queue.send_realtime(
|
|
blob=types.Blob(
|
|
data=b'What is the weather and current time?', 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) >= 1, 'Expected at least one event.'
|
|
|
|
# Check function calls
|
|
weather_call_found = False
|
|
time_call_found = False
|
|
|
|
for event in res_events:
|
|
if event.content and event.content.parts:
|
|
for part in event.content.parts:
|
|
if part.function_call:
|
|
if part.function_call.name == 'get_weather':
|
|
weather_call_found = True
|
|
assert part.function_call.args['location'] == 'San Francisco'
|
|
elif part.function_call.name == 'get_time':
|
|
time_call_found = True
|
|
assert part.function_call.args['timezone'] == 'PST'
|
|
|
|
# In live streaming, we primarily check that function calls are generated correctly
|
|
assert (
|
|
weather_call_found or time_call_found
|
|
), 'Expected at least one function call.'
|
|
|
|
|
|
def test_live_streaming_function_call_parallel():
|
|
"""Test live streaming with parallel function calls."""
|
|
# Create parallel function calls in the same response
|
|
function_call1 = types.Part.from_function_call(
|
|
name='get_weather', args={'location': 'San Francisco'}
|
|
)
|
|
function_call2 = types.Part.from_function_call(
|
|
name='get_weather', args={'location': 'New York'}
|
|
)
|
|
|
|
# Create LLM response with parallel function calls
|
|
response1 = LlmResponse(
|
|
content=types.Content(
|
|
role='model', parts=[function_call1, function_call2]
|
|
),
|
|
turn_complete=False,
|
|
)
|
|
response2 = LlmResponse(
|
|
turn_complete=True,
|
|
)
|
|
|
|
mock_model = testing_utils.MockModel.create([response1, response2])
|
|
|
|
# Mock function
|
|
def get_weather(location: str) -> dict:
|
|
temperatures = {'San Francisco': 22, 'New York': 15}
|
|
return {'temperature': temperatures.get(location, 20), 'location': location}
|
|
|
|
root_agent = Agent(
|
|
name='root_agent',
|
|
model=mock_model,
|
|
tools=[get_weather],
|
|
)
|
|
|
|
# Use the custom runner
|
|
class CustomTestRunner(testing_utils.InMemoryRunner):
|
|
|
|
def _run_with_loop(self, coro):
|
|
try:
|
|
old_loop = asyncio.get_event_loop()
|
|
except RuntimeError:
|
|
old_loop = None
|
|
|
|
loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(loop)
|
|
try:
|
|
loop.run_until_complete(coro)
|
|
except (asyncio.TimeoutError, asyncio.CancelledError):
|
|
pass
|
|
finally:
|
|
loop.close()
|
|
asyncio.set_event_loop(old_loop)
|
|
|
|
def run_live(
|
|
self,
|
|
live_request_queue: LiveRequestQueue,
|
|
run_config: testing_utils.RunConfig = None,
|
|
) -> list[testing_utils.Event]:
|
|
collected_responses = []
|
|
|
|
async def consume_responses(session: testing_utils.Session):
|
|
run_res = self.runner.run_live(
|
|
session=session,
|
|
live_request_queue=live_request_queue,
|
|
run_config=run_config or testing_utils.RunConfig(),
|
|
)
|
|
|
|
async for response in run_res:
|
|
collected_responses.append(response)
|
|
if len(collected_responses) >= 3:
|
|
return
|
|
|
|
self._run_with_loop(
|
|
asyncio.wait_for(consume_responses(self.session), timeout=5.0)
|
|
)
|
|
|
|
return collected_responses
|
|
|
|
runner = CustomTestRunner(root_agent=root_agent)
|
|
live_request_queue = LiveRequestQueue()
|
|
live_request_queue.send_realtime(
|
|
blob=types.Blob(
|
|
data=b'Compare weather in SF and NYC', 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) >= 1, 'Expected at least one event.'
|
|
|
|
# Check parallel function calls
|
|
sf_call_found = False
|
|
nyc_call_found = False
|
|
|
|
for event in res_events:
|
|
if event.content and event.content.parts:
|
|
for part in event.content.parts:
|
|
if part.function_call and part.function_call.name == 'get_weather':
|
|
location = part.function_call.args['location']
|
|
if location == 'San Francisco':
|
|
sf_call_found = True
|
|
elif location == 'New York':
|
|
nyc_call_found = True
|
|
|
|
assert (
|
|
sf_call_found and nyc_call_found
|
|
), 'Expected both location function calls.'
|
|
|
|
|
|
def test_live_streaming_function_call_with_error():
|
|
"""Test live streaming with function call that returns an error."""
|
|
# Create a function call response
|
|
function_call = types.Part.from_function_call(
|
|
name='get_weather', args={'location': 'Invalid Location'}
|
|
)
|
|
|
|
# Create LLM responses
|
|
response1 = LlmResponse(
|
|
content=types.Content(role='model', parts=[function_call]),
|
|
turn_complete=False,
|
|
)
|
|
response2 = LlmResponse(
|
|
turn_complete=True,
|
|
)
|
|
|
|
mock_model = testing_utils.MockModel.create([response1, response2])
|
|
|
|
# Mock function that returns an error for invalid locations
|
|
def get_weather(location: str) -> dict:
|
|
if location == 'Invalid Location':
|
|
return {'error': 'Location not found'}
|
|
return {'temperature': 22, 'condition': 'sunny', 'location': location}
|
|
|
|
root_agent = Agent(
|
|
name='root_agent',
|
|
model=mock_model,
|
|
tools=[get_weather],
|
|
)
|
|
|
|
# Use the custom runner
|
|
class CustomTestRunner(testing_utils.InMemoryRunner):
|
|
|
|
def _run_with_loop(self, coro):
|
|
try:
|
|
old_loop = asyncio.get_event_loop()
|
|
except RuntimeError:
|
|
old_loop = None
|
|
|
|
loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(loop)
|
|
try:
|
|
loop.run_until_complete(coro)
|
|
except (asyncio.TimeoutError, asyncio.CancelledError):
|
|
pass
|
|
finally:
|
|
loop.close()
|
|
asyncio.set_event_loop(old_loop)
|
|
|
|
def run_live(
|
|
self,
|
|
live_request_queue: LiveRequestQueue,
|
|
run_config: testing_utils.RunConfig = None,
|
|
) -> list[testing_utils.Event]:
|
|
collected_responses = []
|
|
|
|
async def consume_responses(session: testing_utils.Session):
|
|
run_res = self.runner.run_live(
|
|
session=session,
|
|
live_request_queue=live_request_queue,
|
|
run_config=run_config or testing_utils.RunConfig(),
|
|
)
|
|
|
|
async for response in run_res:
|
|
collected_responses.append(response)
|
|
if len(collected_responses) >= 3:
|
|
return
|
|
|
|
self._run_with_loop(
|
|
asyncio.wait_for(consume_responses(self.session), timeout=5.0)
|
|
)
|
|
|
|
return collected_responses
|
|
|
|
runner = CustomTestRunner(root_agent=root_agent)
|
|
live_request_queue = LiveRequestQueue()
|
|
live_request_queue.send_realtime(
|
|
blob=types.Blob(
|
|
data=b'What is weather in Invalid Location?', 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) >= 1, 'Expected at least one event.'
|
|
|
|
# Check that we got the function call (error handling happens at execution time)
|
|
function_call_found = False
|
|
for event in res_events:
|
|
if event.content and event.content.parts:
|
|
for part in event.content.parts:
|
|
if part.function_call and part.function_call.name == 'get_weather':
|
|
function_call_found = True
|
|
assert part.function_call.args['location'] == 'Invalid Location'
|
|
|
|
assert function_call_found, 'Expected function call event with error case.'
|
|
|
|
|
|
def test_live_streaming_function_call_sync_tool():
|
|
"""Test live streaming with synchronous function call."""
|
|
# Create a function call response
|
|
function_call = types.Part.from_function_call(
|
|
name='calculate', args={'x': 5, 'y': 3}
|
|
)
|
|
|
|
# Create LLM responses
|
|
response1 = LlmResponse(
|
|
content=types.Content(role='model', parts=[function_call]),
|
|
turn_complete=False,
|
|
)
|
|
response2 = LlmResponse(
|
|
turn_complete=True,
|
|
)
|
|
|
|
mock_model = testing_utils.MockModel.create([response1, response2])
|
|
|
|
# Mock sync function
|
|
def calculate(x: int, y: int) -> dict:
|
|
return {'result': x + y, 'operation': 'addition'}
|
|
|
|
root_agent = Agent(
|
|
name='root_agent',
|
|
model=mock_model,
|
|
tools=[calculate],
|
|
)
|
|
|
|
# Use the custom runner
|
|
class CustomTestRunner(testing_utils.InMemoryRunner):
|
|
|
|
def _run_with_loop(self, coro):
|
|
try:
|
|
old_loop = asyncio.get_event_loop()
|
|
except RuntimeError:
|
|
old_loop = None
|
|
|
|
loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(loop)
|
|
try:
|
|
loop.run_until_complete(coro)
|
|
except (asyncio.TimeoutError, asyncio.CancelledError):
|
|
pass
|
|
finally:
|
|
loop.close()
|
|
asyncio.set_event_loop(old_loop)
|
|
|
|
def run_live(
|
|
self,
|
|
live_request_queue: LiveRequestQueue,
|
|
run_config: testing_utils.RunConfig = None,
|
|
) -> list[testing_utils.Event]:
|
|
collected_responses = []
|
|
|
|
async def consume_responses(session: testing_utils.Session):
|
|
run_res = self.runner.run_live(
|
|
session=session,
|
|
live_request_queue=live_request_queue,
|
|
run_config=run_config or testing_utils.RunConfig(),
|
|
)
|
|
|
|
async for response in run_res:
|
|
collected_responses.append(response)
|
|
if len(collected_responses) >= 3:
|
|
return
|
|
|
|
self._run_with_loop(
|
|
asyncio.wait_for(consume_responses(self.session), timeout=5.0)
|
|
)
|
|
|
|
return collected_responses
|
|
|
|
runner = CustomTestRunner(root_agent=root_agent)
|
|
live_request_queue = LiveRequestQueue()
|
|
live_request_queue.send_realtime(
|
|
blob=types.Blob(data=b'Calculate 5 plus 3', 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) >= 1, 'Expected at least one event.'
|
|
|
|
# Check function call
|
|
function_call_found = False
|
|
for event in res_events:
|
|
if event.content and event.content.parts:
|
|
for part in event.content.parts:
|
|
if part.function_call and part.function_call.name == 'calculate':
|
|
function_call_found = True
|
|
assert part.function_call.args['x'] == 5
|
|
assert part.function_call.args['y'] == 3
|
|
|
|
assert function_call_found, 'Expected calculate function call event.'
|
|
|
|
|
|
def test_live_streaming_simple_streaming_tool():
|
|
"""Test live streaming with a simple streaming tool (non-video)."""
|
|
# Create a function call response for the streaming tool
|
|
function_call = types.Part.from_function_call(
|
|
name='monitor_stock_price', args={'stock_symbol': 'AAPL'}
|
|
)
|
|
|
|
# Create LLM responses
|
|
response1 = LlmResponse(
|
|
content=types.Content(role='model', parts=[function_call]),
|
|
turn_complete=False,
|
|
)
|
|
response2 = LlmResponse(
|
|
turn_complete=True,
|
|
)
|
|
|
|
mock_model = testing_utils.MockModel.create([response1, response2])
|
|
|
|
# Mock simple streaming tool (without return type annotation to avoid parsing issues)
|
|
async def monitor_stock_price(stock_symbol: str):
|
|
"""Mock streaming tool that monitors stock prices."""
|
|
# Simulate some streaming updates
|
|
yield f'Stock {stock_symbol} price: $150'
|
|
await asyncio.sleep(0.1)
|
|
yield f'Stock {stock_symbol} price: $155'
|
|
await asyncio.sleep(0.1)
|
|
yield f'Stock {stock_symbol} price: $160'
|
|
|
|
def stop_streaming(function_name: str):
|
|
"""Stop the streaming tool."""
|
|
pass
|
|
|
|
root_agent = Agent(
|
|
name='root_agent',
|
|
model=mock_model,
|
|
tools=[monitor_stock_price, stop_streaming],
|
|
)
|
|
|
|
# Use the custom runner
|
|
class CustomTestRunner(testing_utils.InMemoryRunner):
|
|
|
|
def _run_with_loop(self, coro):
|
|
try:
|
|
old_loop = asyncio.get_event_loop()
|
|
except RuntimeError:
|
|
old_loop = None
|
|
|
|
loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(loop)
|
|
try:
|
|
loop.run_until_complete(coro)
|
|
except (asyncio.TimeoutError, asyncio.CancelledError):
|
|
pass
|
|
finally:
|
|
loop.close()
|
|
asyncio.set_event_loop(old_loop)
|
|
|
|
def run_live(
|
|
self,
|
|
live_request_queue: LiveRequestQueue,
|
|
run_config: testing_utils.RunConfig = None,
|
|
) -> list[testing_utils.Event]:
|
|
collected_responses = []
|
|
|
|
async def consume_responses(session: testing_utils.Session):
|
|
run_res = self.runner.run_live(
|
|
session=session,
|
|
live_request_queue=live_request_queue,
|
|
run_config=run_config or testing_utils.RunConfig(),
|
|
)
|
|
|
|
async for response in run_res:
|
|
collected_responses.append(response)
|
|
if len(collected_responses) >= 3:
|
|
return
|
|
|
|
self._run_with_loop(
|
|
asyncio.wait_for(consume_responses(self.session), timeout=5.0)
|
|
)
|
|
|
|
return collected_responses
|
|
|
|
runner = CustomTestRunner(root_agent=root_agent)
|
|
live_request_queue = LiveRequestQueue()
|
|
live_request_queue.send_realtime(
|
|
blob=types.Blob(data=b'Monitor AAPL stock price', 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) >= 1, 'Expected at least one event.'
|
|
|
|
# Check that we got the streaming tool function call
|
|
function_call_found = False
|
|
for event in res_events:
|
|
if event.content and event.content.parts:
|
|
for part in event.content.parts:
|
|
if (
|
|
part.function_call
|
|
and part.function_call.name == 'monitor_stock_price'
|
|
):
|
|
function_call_found = True
|
|
assert part.function_call.args['stock_symbol'] == 'AAPL'
|
|
|
|
assert (
|
|
function_call_found
|
|
), 'Expected monitor_stock_price function call event.'
|
|
|
|
|
|
def test_live_streaming_video_streaming_tool():
|
|
"""Test live streaming with a video streaming tool."""
|
|
# Create a function call response for the video streaming tool
|
|
function_call = types.Part.from_function_call(
|
|
name='monitor_video_stream', args={}
|
|
)
|
|
|
|
# Create LLM responses
|
|
response1 = LlmResponse(
|
|
content=types.Content(role='model', parts=[function_call]),
|
|
turn_complete=False,
|
|
)
|
|
response2 = LlmResponse(
|
|
turn_complete=True,
|
|
)
|
|
|
|
mock_model = testing_utils.MockModel.create([response1, response2])
|
|
|
|
# Mock video streaming tool (without return type annotation to avoid parsing issues)
|
|
async def monitor_video_stream(input_stream: LiveRequestQueue):
|
|
"""Mock video streaming tool that processes video frames."""
|
|
# Simulate processing a few frames from the input stream
|
|
frame_count = 0
|
|
while frame_count < 3: # Process a few frames
|
|
try:
|
|
# Try to get a frame from the queue with timeout
|
|
live_req = await asyncio.wait_for(input_stream.get(), timeout=0.1)
|
|
if live_req.blob and live_req.blob.mime_type == 'image/jpeg':
|
|
frame_count += 1
|
|
yield f'Processed frame {frame_count}: detected 2 people'
|
|
except asyncio.TimeoutError:
|
|
# No more frames, simulate detection anyway for testing
|
|
frame_count += 1
|
|
yield f'Simulated frame {frame_count}: detected 1 person'
|
|
await asyncio.sleep(0.1)
|
|
|
|
def stop_streaming(function_name: str):
|
|
"""Stop the streaming tool."""
|
|
pass
|
|
|
|
root_agent = Agent(
|
|
name='root_agent',
|
|
model=mock_model,
|
|
tools=[monitor_video_stream, stop_streaming],
|
|
)
|
|
|
|
# Use the custom runner
|
|
class CustomTestRunner(testing_utils.InMemoryRunner):
|
|
|
|
def _run_with_loop(self, coro):
|
|
try:
|
|
old_loop = asyncio.get_event_loop()
|
|
except RuntimeError:
|
|
old_loop = None
|
|
|
|
loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(loop)
|
|
try:
|
|
loop.run_until_complete(coro)
|
|
except (asyncio.TimeoutError, asyncio.CancelledError):
|
|
pass
|
|
finally:
|
|
loop.close()
|
|
asyncio.set_event_loop(old_loop)
|
|
|
|
def run_live(
|
|
self,
|
|
live_request_queue: LiveRequestQueue,
|
|
run_config: testing_utils.RunConfig = None,
|
|
) -> list[testing_utils.Event]:
|
|
collected_responses = []
|
|
|
|
async def consume_responses(session: testing_utils.Session):
|
|
run_res = self.runner.run_live(
|
|
session=session,
|
|
live_request_queue=live_request_queue,
|
|
run_config=run_config or testing_utils.RunConfig(),
|
|
)
|
|
|
|
async for response in run_res:
|
|
collected_responses.append(response)
|
|
if len(collected_responses) >= 3:
|
|
return
|
|
|
|
self._run_with_loop(
|
|
asyncio.wait_for(consume_responses(self.session), timeout=5.0)
|
|
)
|
|
|
|
return collected_responses
|
|
|
|
runner = CustomTestRunner(root_agent=root_agent)
|
|
live_request_queue = LiveRequestQueue()
|
|
|
|
# Send some mock video frames
|
|
live_request_queue.send_realtime(
|
|
blob=types.Blob(data=b'fake_jpeg_data_1', mime_type='image/jpeg')
|
|
)
|
|
live_request_queue.send_realtime(
|
|
blob=types.Blob(data=b'fake_jpeg_data_2', mime_type='image/jpeg')
|
|
)
|
|
live_request_queue.send_realtime(
|
|
blob=types.Blob(data=b'Monitor video stream', 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) >= 1, 'Expected at least one event.'
|
|
|
|
# Check that we got the video streaming tool function call
|
|
function_call_found = False
|
|
for event in res_events:
|
|
if event.content and event.content.parts:
|
|
for part in event.content.parts:
|
|
if (
|
|
part.function_call
|
|
and part.function_call.name == 'monitor_video_stream'
|
|
):
|
|
function_call_found = True
|
|
|
|
assert (
|
|
function_call_found
|
|
), 'Expected monitor_video_stream function call event.'
|
|
|
|
|
|
def test_live_streaming_stop_streaming_tool():
|
|
"""Test live streaming with stop_streaming functionality."""
|
|
# Create function calls for starting and stopping a streaming tool
|
|
start_function_call = types.Part.from_function_call(
|
|
name='monitor_stock_price', args={'stock_symbol': 'TSLA'}
|
|
)
|
|
stop_function_call = types.Part.from_function_call(
|
|
name='stop_streaming', args={'function_name': 'monitor_stock_price'}
|
|
)
|
|
|
|
# Create LLM responses: start streaming, then stop streaming
|
|
response1 = LlmResponse(
|
|
content=types.Content(role='model', parts=[start_function_call]),
|
|
turn_complete=False,
|
|
)
|
|
response2 = LlmResponse(
|
|
content=types.Content(role='model', parts=[stop_function_call]),
|
|
turn_complete=False,
|
|
)
|
|
response3 = LlmResponse(
|
|
turn_complete=True,
|
|
)
|
|
|
|
mock_model = testing_utils.MockModel.create([response1, response2, response3])
|
|
|
|
# Mock streaming tool and stop function
|
|
async def monitor_stock_price(stock_symbol: str):
|
|
"""Mock streaming tool that monitors stock prices."""
|
|
yield f'Started monitoring {stock_symbol}'
|
|
while True: # Infinite stream (would be stopped by stop_streaming)
|
|
yield f'Stock {stock_symbol} price update'
|
|
await asyncio.sleep(0.1)
|
|
|
|
def stop_streaming(function_name: str):
|
|
"""Stop the streaming tool."""
|
|
return f'Stopped streaming for {function_name}'
|
|
|
|
root_agent = Agent(
|
|
name='root_agent',
|
|
model=mock_model,
|
|
tools=[monitor_stock_price, stop_streaming],
|
|
)
|
|
|
|
# Use the custom runner
|
|
class CustomTestRunner(testing_utils.InMemoryRunner):
|
|
|
|
def _run_with_loop(self, coro):
|
|
try:
|
|
old_loop = asyncio.get_event_loop()
|
|
except RuntimeError:
|
|
old_loop = None
|
|
|
|
loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(loop)
|
|
try:
|
|
loop.run_until_complete(coro)
|
|
except (asyncio.TimeoutError, asyncio.CancelledError):
|
|
pass
|
|
finally:
|
|
loop.close()
|
|
asyncio.set_event_loop(old_loop)
|
|
|
|
def run_live(
|
|
self,
|
|
live_request_queue: LiveRequestQueue,
|
|
run_config: testing_utils.RunConfig = None,
|
|
) -> list[testing_utils.Event]:
|
|
collected_responses = []
|
|
|
|
async def consume_responses(session: testing_utils.Session):
|
|
run_res = self.runner.run_live(
|
|
session=session,
|
|
live_request_queue=live_request_queue,
|
|
run_config=run_config or testing_utils.RunConfig(),
|
|
)
|
|
|
|
async for response in run_res:
|
|
collected_responses.append(response)
|
|
if len(collected_responses) >= 3:
|
|
return
|
|
|
|
self._run_with_loop(
|
|
asyncio.wait_for(consume_responses(self.session), timeout=5.0)
|
|
)
|
|
|
|
return collected_responses
|
|
|
|
runner = CustomTestRunner(root_agent=root_agent)
|
|
live_request_queue = LiveRequestQueue()
|
|
live_request_queue.send_realtime(
|
|
blob=types.Blob(data=b'Monitor TSLA and then stop', 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) >= 1, 'Expected at least one event.'
|
|
|
|
# Check that we got both function calls
|
|
monitor_call_found = False
|
|
stop_call_found = False
|
|
|
|
for event in res_events:
|
|
if event.content and event.content.parts:
|
|
for part in event.content.parts:
|
|
if part.function_call:
|
|
if part.function_call.name == 'monitor_stock_price':
|
|
monitor_call_found = True
|
|
assert part.function_call.args['stock_symbol'] == 'TSLA'
|
|
elif part.function_call.name == 'stop_streaming':
|
|
stop_call_found = True
|
|
assert (
|
|
part.function_call.args['function_name']
|
|
== 'monitor_stock_price'
|
|
)
|
|
|
|
assert monitor_call_found, 'Expected monitor_stock_price function call event.'
|
|
assert stop_call_found, 'Expected stop_streaming function call event.'
|
|
|
|
|
|
def test_live_streaming_multiple_streaming_tools():
|
|
"""Test live streaming with multiple streaming tools running simultaneously."""
|
|
# Create function calls for multiple streaming tools
|
|
stock_function_call = types.Part.from_function_call(
|
|
name='monitor_stock_price', args={'stock_symbol': 'NVDA'}
|
|
)
|
|
video_function_call = types.Part.from_function_call(
|
|
name='monitor_video_stream', args={}
|
|
)
|
|
|
|
# Create LLM responses: start both streaming tools
|
|
response1 = LlmResponse(
|
|
content=types.Content(
|
|
role='model', parts=[stock_function_call, video_function_call]
|
|
),
|
|
turn_complete=False,
|
|
)
|
|
response2 = LlmResponse(
|
|
turn_complete=True,
|
|
)
|
|
|
|
mock_model = testing_utils.MockModel.create([response1, response2])
|
|
|
|
# Mock streaming tools
|
|
async def monitor_stock_price(stock_symbol: str):
|
|
"""Mock streaming tool that monitors stock prices."""
|
|
yield f'Stock {stock_symbol} price: $800'
|
|
await asyncio.sleep(0.1)
|
|
yield f'Stock {stock_symbol} price: $805'
|
|
|
|
async def monitor_video_stream(input_stream: LiveRequestQueue):
|
|
"""Mock video streaming tool."""
|
|
yield 'Video monitoring started'
|
|
await asyncio.sleep(0.1)
|
|
yield 'Detected motion in video stream'
|
|
|
|
def stop_streaming(function_name: str):
|
|
"""Stop the streaming tool."""
|
|
pass
|
|
|
|
root_agent = Agent(
|
|
name='root_agent',
|
|
model=mock_model,
|
|
tools=[monitor_stock_price, monitor_video_stream, stop_streaming],
|
|
)
|
|
|
|
# Use the custom runner
|
|
class CustomTestRunner(testing_utils.InMemoryRunner):
|
|
|
|
def _run_with_loop(self, coro):
|
|
try:
|
|
old_loop = asyncio.get_event_loop()
|
|
except RuntimeError:
|
|
old_loop = None
|
|
|
|
loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(loop)
|
|
try:
|
|
loop.run_until_complete(coro)
|
|
except (asyncio.TimeoutError, asyncio.CancelledError):
|
|
pass
|
|
finally:
|
|
loop.close()
|
|
asyncio.set_event_loop(old_loop)
|
|
|
|
def run_live(
|
|
self,
|
|
live_request_queue: LiveRequestQueue,
|
|
run_config: testing_utils.RunConfig = None,
|
|
) -> list[testing_utils.Event]:
|
|
collected_responses = []
|
|
|
|
async def consume_responses(session: testing_utils.Session):
|
|
run_res = self.runner.run_live(
|
|
session=session,
|
|
live_request_queue=live_request_queue,
|
|
run_config=run_config or testing_utils.RunConfig(),
|
|
)
|
|
|
|
async for response in run_res:
|
|
collected_responses.append(response)
|
|
if len(collected_responses) >= 3:
|
|
return
|
|
|
|
self._run_with_loop(
|
|
asyncio.wait_for(consume_responses(self.session), timeout=5.0)
|
|
)
|
|
|
|
return collected_responses
|
|
|
|
runner = CustomTestRunner(root_agent=root_agent)
|
|
live_request_queue = LiveRequestQueue()
|
|
live_request_queue.send_realtime(
|
|
blob=types.Blob(
|
|
data=b'Monitor both stock and video', 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) >= 1, 'Expected at least one event.'
|
|
|
|
# Check that we got both streaming tool function calls
|
|
stock_call_found = False
|
|
video_call_found = False
|
|
|
|
for event in res_events:
|
|
if event.content and event.content.parts:
|
|
for part in event.content.parts:
|
|
if part.function_call:
|
|
if part.function_call.name == 'monitor_stock_price':
|
|
stock_call_found = True
|
|
assert part.function_call.args['stock_symbol'] == 'NVDA'
|
|
elif part.function_call.name == 'monitor_video_stream':
|
|
video_call_found = True
|
|
|
|
assert stock_call_found, 'Expected monitor_stock_price function call event.'
|
|
assert video_call_found, 'Expected monitor_video_stream function call event.'
|
|
|
|
|
|
def test_live_streaming_buffered_function_call_yielded_during_transcription():
|
|
"""Test that function calls buffered during transcription are yielded.
|
|
|
|
This tests the fix for the bug where function_call and function_response
|
|
events were buffered during active transcription but never yielded to the
|
|
caller. The fix ensures buffered events are yielded after transcription ends.
|
|
"""
|
|
function_call = types.Part.from_function_call(
|
|
name='get_weather', args={'location': 'San Francisco'}
|
|
)
|
|
|
|
response1 = LlmResponse(
|
|
input_transcription=types.Transcription(text='Show'),
|
|
partial=True, # ← Triggers is_transcribing = True
|
|
)
|
|
response2 = LlmResponse(
|
|
content=types.Content(
|
|
role='model', parts=[function_call]
|
|
), # ← Gets buffered
|
|
turn_complete=False,
|
|
)
|
|
response3 = LlmResponse(
|
|
input_transcription=types.Transcription(text='Show me the weather'),
|
|
partial=False, # ← Transcription ends, buffered events yielded
|
|
)
|
|
response4 = LlmResponse(
|
|
turn_complete=True,
|
|
)
|
|
|
|
mock_model = testing_utils.MockModel.create(
|
|
[response1, response2, response3, response4]
|
|
)
|
|
|
|
def get_weather(location: str) -> dict:
|
|
return {'temperature': 22, 'location': location}
|
|
|
|
root_agent = Agent(
|
|
name='root_agent',
|
|
model=mock_model,
|
|
tools=[get_weather],
|
|
)
|
|
|
|
class CustomTestRunner(testing_utils.InMemoryRunner):
|
|
|
|
def _run_with_loop(self, coro):
|
|
try:
|
|
old_loop = asyncio.get_event_loop()
|
|
except RuntimeError:
|
|
old_loop = None
|
|
|
|
loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(loop)
|
|
try:
|
|
loop.run_until_complete(coro)
|
|
except (asyncio.TimeoutError, asyncio.CancelledError):
|
|
pass
|
|
finally:
|
|
loop.close()
|
|
asyncio.set_event_loop(old_loop)
|
|
|
|
def run_live(
|
|
self,
|
|
live_request_queue: LiveRequestQueue,
|
|
run_config: testing_utils.RunConfig = None,
|
|
) -> list[testing_utils.Event]:
|
|
collected_responses = []
|
|
|
|
async def consume_responses(session: testing_utils.Session):
|
|
run_res = self.runner.run_live(
|
|
session=session,
|
|
live_request_queue=live_request_queue,
|
|
run_config=run_config or testing_utils.RunConfig(),
|
|
)
|
|
|
|
async for response in run_res:
|
|
collected_responses.append(response)
|
|
if len(collected_responses) >= 5:
|
|
return
|
|
|
|
self._run_with_loop(
|
|
asyncio.wait_for(consume_responses(self.session), timeout=5.0)
|
|
)
|
|
|
|
return collected_responses
|
|
|
|
runner = CustomTestRunner(root_agent=root_agent)
|
|
live_request_queue = LiveRequestQueue()
|
|
live_request_queue.send_realtime(
|
|
blob=types.Blob(data=b'Show me the weather', 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) >= 1, 'Expected at least one event.'
|
|
|
|
function_call_found = False
|
|
function_response_found = False
|
|
|
|
for event in res_events:
|
|
if event.content and event.content.parts:
|
|
for part in event.content.parts:
|
|
if part.function_call and part.function_call.name == 'get_weather':
|
|
function_call_found = True
|
|
assert part.function_call.args['location'] == 'San Francisco'
|
|
if (
|
|
part.function_response
|
|
and part.function_response.name == 'get_weather'
|
|
):
|
|
function_response_found = True
|
|
assert part.function_response.response['temperature'] == 22
|
|
|
|
assert function_call_found, 'Buffered function_call event was not yielded.'
|
|
assert (
|
|
function_response_found
|
|
), 'Buffered function_response event was not yielded.'
|
|
|
|
|
|
def test_live_streaming_text_content_persisted_in_session():
|
|
"""Test that user text content sent via send_content is persisted in session."""
|
|
response1 = LlmResponse(
|
|
content=types.Content(
|
|
role='model', parts=[types.Part(text='Hello! How can I help you?')]
|
|
),
|
|
turn_complete=True,
|
|
)
|
|
|
|
mock_model = testing_utils.MockModel.create([response1])
|
|
|
|
root_agent = Agent(
|
|
name='root_agent',
|
|
model=mock_model,
|
|
tools=[],
|
|
)
|
|
|
|
class CustomTestRunner(testing_utils.InMemoryRunner):
|
|
|
|
def _run_with_loop(self, coro):
|
|
try:
|
|
old_loop = asyncio.get_event_loop()
|
|
except RuntimeError:
|
|
old_loop = None
|
|
|
|
loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(loop)
|
|
try:
|
|
loop.run_until_complete(coro)
|
|
except (asyncio.TimeoutError, asyncio.CancelledError):
|
|
pass
|
|
finally:
|
|
loop.close()
|
|
asyncio.set_event_loop(old_loop)
|
|
|
|
def run_live_and_get_session(
|
|
self,
|
|
live_request_queue: LiveRequestQueue,
|
|
run_config: testing_utils.RunConfig = None,
|
|
) -> tuple[list[testing_utils.Event], testing_utils.Session]:
|
|
collected_responses = []
|
|
|
|
async def consume_responses(session: testing_utils.Session):
|
|
run_res = self.runner.run_live(
|
|
session=session,
|
|
live_request_queue=live_request_queue,
|
|
run_config=run_config or testing_utils.RunConfig(),
|
|
)
|
|
async for response in run_res:
|
|
collected_responses.append(response)
|
|
if len(collected_responses) >= 1:
|
|
return
|
|
|
|
self._run_with_loop(
|
|
asyncio.wait_for(consume_responses(self.session), timeout=5.0)
|
|
)
|
|
|
|
# Get the updated session
|
|
updated_session = self.runner.session_service.get_session_sync(
|
|
app_name=self.app_name,
|
|
user_id=self.session.user_id,
|
|
session_id=self.session.id,
|
|
)
|
|
return collected_responses, updated_session
|
|
|
|
runner = CustomTestRunner(root_agent=root_agent)
|
|
live_request_queue = LiveRequestQueue()
|
|
|
|
# Send text content (not audio blob)
|
|
user_text = 'Hello, this is a test message'
|
|
live_request_queue.send_content(
|
|
types.Content(role='user', parts=[types.Part(text=user_text)])
|
|
)
|
|
|
|
res_events, session = runner.run_live_and_get_session(live_request_queue)
|
|
|
|
assert res_events is not None, 'Expected a list of events, got None.'
|
|
|
|
# Check that user text content was persisted in the session
|
|
user_content_found = False
|
|
for event in session.events:
|
|
if event.author == 'user' and event.content:
|
|
for part in event.content.parts:
|
|
if part.text and user_text in part.text:
|
|
user_content_found = True
|
|
break
|
|
|
|
assert user_content_found, (
|
|
f'Expected user text content "{user_text}" to be persisted in session. '
|
|
f'Session events: {[e.content for e in session.events]}'
|
|
)
|
|
|
|
|
|
def _collect_function_call_names(events):
|
|
"""Extract the set of function call names from a list of events."""
|
|
return {fc.name for event in events for fc in event.get_function_calls()}
|
|
|
|
|
|
class _LiveTestRunner(testing_utils.InMemoryRunner):
|
|
"""Test runner with custom event loop management for live streaming tests."""
|
|
|
|
def _run_with_loop(self, coro: Awaitable[Any]) -> None:
|
|
"""Run a coroutine in a new event loop, suppressing timeouts."""
|
|
try:
|
|
old_loop = asyncio.get_event_loop()
|
|
except RuntimeError:
|
|
old_loop = None
|
|
loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(loop)
|
|
try:
|
|
loop.run_until_complete(coro)
|
|
except (asyncio.TimeoutError, asyncio.CancelledError):
|
|
pass
|
|
finally:
|
|
loop.close()
|
|
asyncio.set_event_loop(old_loop)
|
|
|
|
def run_live(
|
|
self,
|
|
live_request_queue: LiveRequestQueue,
|
|
max_responses: int = 3,
|
|
) -> list[testing_utils.Event]:
|
|
"""Run live and collect up to max_responses events."""
|
|
collected = []
|
|
|
|
async def consume(session: testing_utils.Session):
|
|
async for response in self.runner.run_live(
|
|
session=session,
|
|
live_request_queue=live_request_queue,
|
|
):
|
|
collected.append(response)
|
|
if len(collected) >= max_responses:
|
|
return
|
|
|
|
self._run_with_loop(asyncio.wait_for(consume(self.session), timeout=5.0))
|
|
return collected
|
|
|
|
|
|
def test_input_streaming_tool_stream_is_none_before_model_calls():
|
|
"""Test that input-streaming tools have stream=None until the model calls them."""
|
|
# Add a text response before the function call so we can observe stream
|
|
# state between registration and tool invocation.
|
|
text_response = LlmResponse(
|
|
content=types.Content(
|
|
role='model',
|
|
parts=[types.Part(text='Processing...')],
|
|
),
|
|
turn_complete=False,
|
|
)
|
|
function_call = types.Part.from_function_call(
|
|
name='monitor_video_stream', args={}
|
|
)
|
|
call_response = LlmResponse(
|
|
content=types.Content(role='model', parts=[function_call]),
|
|
turn_complete=False,
|
|
)
|
|
done_response = LlmResponse(turn_complete=True)
|
|
|
|
mock_model = testing_utils.MockModel.create(
|
|
[text_response, call_response, done_response]
|
|
)
|
|
|
|
stream_state_during_call = None
|
|
|
|
async def monitor_video_stream(input_stream: LiveRequestQueue):
|
|
"""Record whether input_stream was provided."""
|
|
nonlocal stream_state_during_call
|
|
stream_state_during_call = input_stream is not None
|
|
yield 'monitoring started'
|
|
|
|
root_agent = Agent(
|
|
name='root_agent',
|
|
model=mock_model,
|
|
tools=[monitor_video_stream],
|
|
)
|
|
|
|
runner = _LiveTestRunner(root_agent=root_agent)
|
|
|
|
# Capture the invocation context to inspect stream state.
|
|
captured_context = None
|
|
original_method = runner.runner._new_invocation_context_for_live
|
|
|
|
def capturing_method(*args, **kwargs):
|
|
nonlocal captured_context
|
|
ctx = original_method(*args, **kwargs)
|
|
captured_context = ctx
|
|
return ctx
|
|
|
|
runner.runner._new_invocation_context_for_live = capturing_method
|
|
|
|
live_request_queue = LiveRequestQueue()
|
|
live_request_queue.send_realtime(
|
|
blob=types.Blob(data=b'test_data', mime_type='audio/pcm')
|
|
)
|
|
|
|
# Collect events and capture stream state before the tool is called.
|
|
collected = []
|
|
stream_states_before_call = []
|
|
|
|
async def consume(session: testing_utils.Session):
|
|
async for response in runner.runner.run_live(
|
|
session=session,
|
|
live_request_queue=live_request_queue,
|
|
):
|
|
collected.append(response)
|
|
# On a non-function-call event, the tool is registered but not
|
|
# yet invoked — capture the stream value at that point.
|
|
active = (
|
|
captured_context.active_streaming_tools if captured_context else {}
|
|
)
|
|
if (
|
|
not stream_states_before_call
|
|
and not response.get_function_calls()
|
|
and 'monitor_video_stream' in active
|
|
):
|
|
stream_states_before_call.append(active['monitor_video_stream'].stream)
|
|
if len(collected) >= 4:
|
|
return
|
|
|
|
runner._run_with_loop(asyncio.wait_for(consume(runner.session), timeout=5.0))
|
|
|
|
# Before the model calls the tool, stream should be None.
|
|
assert (
|
|
stream_states_before_call
|
|
), 'Stream state was never observed before the tool call'
|
|
assert (
|
|
stream_states_before_call[0] is None
|
|
), 'Expected stream to be None before the model calls the tool'
|
|
# When the model calls the tool, input_stream should be provided.
|
|
assert (
|
|
stream_state_during_call is True
|
|
), 'Expected input_stream to be provided to the streaming tool when called'
|
|
|
|
|
|
def test_stop_streaming_resets_stream_to_none():
|
|
"""Test that stop_streaming sets stream back to None."""
|
|
start_call = types.Part.from_function_call(
|
|
name='monitor_stock_price', args={'stock_symbol': 'GOOG'}
|
|
)
|
|
stop_call = types.Part.from_function_call(
|
|
name='stop_streaming', args={'function_name': 'monitor_stock_price'}
|
|
)
|
|
|
|
response1 = LlmResponse(
|
|
content=types.Content(role='model', parts=[start_call]),
|
|
turn_complete=False,
|
|
)
|
|
response2 = LlmResponse(
|
|
content=types.Content(role='model', parts=[stop_call]),
|
|
turn_complete=False,
|
|
)
|
|
response3 = LlmResponse(turn_complete=True)
|
|
|
|
mock_model = testing_utils.MockModel.create([response1, response2, response3])
|
|
|
|
async def monitor_stock_price(stock_symbol: str):
|
|
"""Yield periodic price updates for the given stock symbol."""
|
|
yield f'Monitoring {stock_symbol}'
|
|
while True:
|
|
await asyncio.sleep(0.1)
|
|
yield f'{stock_symbol} price update'
|
|
|
|
def stop_streaming(function_name: str):
|
|
"""Stop a running streaming tool by name."""
|
|
pass
|
|
|
|
root_agent = Agent(
|
|
name='root_agent',
|
|
model=mock_model,
|
|
tools=[monitor_stock_price, stop_streaming],
|
|
)
|
|
|
|
runner = _LiveTestRunner(root_agent=root_agent)
|
|
|
|
# Capture invocation context to verify stream is reset.
|
|
captured_context = None
|
|
original_method = runner.runner._new_invocation_context_for_live
|
|
|
|
def capturing_method(*args, **kwargs):
|
|
nonlocal captured_context
|
|
ctx = original_method(*args, **kwargs)
|
|
captured_context = ctx
|
|
return ctx
|
|
|
|
runner.runner._new_invocation_context_for_live = capturing_method
|
|
|
|
live_request_queue = LiveRequestQueue()
|
|
live_request_queue.send_realtime(
|
|
blob=types.Blob(data=b'Monitor GOOG then stop', mime_type='audio/pcm')
|
|
)
|
|
|
|
res_events = runner.run_live(live_request_queue, max_responses=4)
|
|
|
|
# Verify both function calls were processed.
|
|
call_names = _collect_function_call_names(res_events)
|
|
assert (
|
|
'monitor_stock_price' in call_names
|
|
), 'Expected monitor_stock_price function call.'
|
|
assert (
|
|
'stop_streaming' in call_names
|
|
), 'Expected stop_streaming function call.'
|
|
|
|
# Verify that stop_streaming reset the stream to None.
|
|
assert (
|
|
captured_context is not None
|
|
), 'Expected invocation context to be captured'
|
|
active_tools = captured_context.active_streaming_tools or {}
|
|
assert (
|
|
'monitor_stock_price' in active_tools
|
|
), 'Expected monitor_stock_price in active_streaming_tools'
|
|
assert (
|
|
active_tools['monitor_stock_price'].stream is None
|
|
), 'Expected stream to be reset to None after stop_streaming'
|
|
|
|
|
|
def test_output_streaming_tool_registered_at_startup():
|
|
"""Test that output-streaming tools (async generators without LiveRequestQueue) are registered at startup."""
|
|
response1 = LlmResponse(turn_complete=True)
|
|
|
|
mock_model = testing_utils.MockModel.create([response1])
|
|
|
|
async def monitor_stock_price(stock_symbol: str):
|
|
"""Yield periodic price updates."""
|
|
yield f'price for {stock_symbol}'
|
|
|
|
root_agent = Agent(
|
|
name='root_agent',
|
|
model=mock_model,
|
|
tools=[monitor_stock_price],
|
|
)
|
|
|
|
runner = _LiveTestRunner(root_agent=root_agent)
|
|
|
|
# Capture invocation context to verify registration.
|
|
captured_context = None
|
|
original_method = runner.runner._new_invocation_context_for_live
|
|
|
|
def capturing_method(*args, **kwargs):
|
|
nonlocal captured_context
|
|
ctx = original_method(*args, **kwargs)
|
|
captured_context = ctx
|
|
return ctx
|
|
|
|
runner.runner._new_invocation_context_for_live = capturing_method
|
|
|
|
live_request_queue = LiveRequestQueue()
|
|
live_request_queue.send_realtime(
|
|
blob=types.Blob(data=b'test', mime_type='audio/pcm')
|
|
)
|
|
|
|
runner.run_live(live_request_queue, max_responses=1)
|
|
|
|
# Output-streaming tool should be registered with stream=None.
|
|
assert captured_context is not None
|
|
active_tools = captured_context.active_streaming_tools or {}
|
|
assert (
|
|
'monitor_stock_price' in active_tools
|
|
), 'Expected output-streaming tool to be registered at startup'
|
|
assert (
|
|
active_tools['monitor_stock_price'].stream is None
|
|
), 'Expected stream to be None for output-streaming tool'
|