mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
fix: Only relay the LiveRequest after tools is invoked
currently we started to relay live request to streaming tool even when the tool was not called yet. Co-authored-by: Xiang (Sean) Zhou <seanzhougoogle@google.com> PiperOrigin-RevId: 869421826
This commit is contained in:
committed by
Copybara-Service
parent
647d3b16a3
commit
b53bc555cc
@@ -784,6 +784,9 @@ async def _process_function_live_helper(
|
||||
and function_name in invocation_context.active_streaming_tools
|
||||
):
|
||||
invocation_context.active_streaming_tools[function_name].task = None
|
||||
invocation_context.active_streaming_tools[function_name].stream = (
|
||||
None
|
||||
)
|
||||
|
||||
function_response = {
|
||||
'status': f'Successfully stopped streaming function {function_name}'
|
||||
|
||||
@@ -1051,9 +1051,7 @@ class Runner:
|
||||
logger.debug(
|
||||
'Register streaming tool with input stream: %s', tool.name
|
||||
)
|
||||
active_streaming_tool = ActiveStreamingTool(
|
||||
stream=LiveRequestQueue()
|
||||
)
|
||||
active_streaming_tool = ActiveStreamingTool()
|
||||
invocation_context.active_streaming_tools[tool.name] = (
|
||||
active_streaming_tool
|
||||
)
|
||||
|
||||
@@ -27,6 +27,7 @@ from google.genai import types
|
||||
import pydantic
|
||||
from typing_extensions import override
|
||||
|
||||
from ..agents.live_request_queue import LiveRequestQueue
|
||||
from ..utils.context_utils import Aclosing
|
||||
from ._automatic_function_calling_util import build_function_declaration
|
||||
from .base_tool import BaseTool
|
||||
@@ -246,8 +247,13 @@ You could retry calling this tool, but it is IMPORTANT for you to provide all th
|
||||
signature = inspect.signature(self.func)
|
||||
if (
|
||||
self.name in invocation_context.active_streaming_tools
|
||||
and invocation_context.active_streaming_tools[self.name].stream
|
||||
and 'input_stream' in signature.parameters
|
||||
):
|
||||
# Create the stream now that the model has called the tool.
|
||||
# _send_to_model will start duplicating LiveRequests to this stream.
|
||||
invocation_context.active_streaming_tools[self.name].stream = (
|
||||
LiveRequestQueue()
|
||||
)
|
||||
args_to_call['input_stream'] = invocation_context.active_streaming_tools[
|
||||
self.name
|
||||
].stream
|
||||
|
||||
@@ -13,7 +13,9 @@
|
||||
# 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
|
||||
@@ -1273,3 +1275,225 @@ def test_live_streaming_text_content_persisted_in_session():
|
||||
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'
|
||||
|
||||
Reference in New Issue
Block a user