chore: Register all streaming tool at runner

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
This commit is contained in:
Xiang (Sean) Zhou
2026-02-12 16:59:36 -08:00
committed by Copybara-Service
parent b53bc555cc
commit 5269a6b1d6
3 changed files with 72 additions and 38 deletions
+4 -10
View File
@@ -825,17 +825,11 @@ async def _process_function_live_helper(
run_tool_and_update_queue(tool, function_args, tool_context)
)
# Register streaming tool using original logic
# The tool is already registered in active_streaming_tools by
# runners.py at startup (all async-generator tools are registered
# there). Just attach the background task.
async with streaming_lock:
if invocation_context.active_streaming_tools is None:
invocation_context.active_streaming_tools = {}
if tool.name in invocation_context.active_streaming_tools:
invocation_context.active_streaming_tools[tool.name].task = task
else:
invocation_context.active_streaming_tools[tool.name] = (
ActiveStreamingTool(task=task)
)
invocation_context.active_streaming_tools[tool.name].task = task
# Immediately return a pending response.
# This is required by current live model.
+20 -28
View File
@@ -1023,38 +1023,30 @@ class Runner:
canonical_tools = await invocation_context.agent.canonical_tools(
invocation_context
)
# Register all async-generator tools as streaming tools.
# A streaming tool is any tool whose underlying function is an
# async generator (i.e. uses `yield`). There are two sub-types:
# 1. Input-streaming tools: accept a `input_stream:
# LiveRequestQueue` parameter to consume the live audio/video
# stream. The stream is created lazily in `_call_live` when
# the model actually calls the tool.
# 2. Output-streaming tools: async generators that yield results
# over time but don't consume the live stream. They are run
# as background tasks when called.
# Both types are registered here with `stream=None`. The
# distinction between them is made at call time.
for tool in canonical_tools:
# We use `inspect.signature()` to examine the tool's underlying function (`tool.func`).
# This approach is deliberately chosen over `typing.get_type_hints()` for robustness.
#
# The Problem with `get_type_hints()`:
# `get_type_hints()` attempts to resolve forward-referenced (string-based) type
# annotations. This resolution can easily fail with a `NameError` (e.g., "Union not found")
# if the type isn't available in the scope where `get_type_hints()` is called.
# This is a common and brittle issue in framework code that inspects functions
# defined in separate user modules.
#
# Why `inspect.signature()` is Better Here:
# `inspect.signature()` does NOT resolve the annotations; it retrieves the raw
# annotation object as it was defined on the function. This allows us to
# perform a direct and reliable identity check (`param.annotation is LiveRequestQueue`)
# without risking a `NameError`.
callable_to_inspect = tool.func if hasattr(tool, 'func') else tool
# Ensure the target is actually callable before inspecting to avoid errors.
if not callable(callable_to_inspect):
continue
for param in inspect.signature(callable_to_inspect).parameters.values():
if param.annotation is LiveRequestQueue:
if not invocation_context.active_streaming_tools:
invocation_context.active_streaming_tools = {}
logger.debug(
'Register streaming tool with input stream: %s', tool.name
)
active_streaming_tool = ActiveStreamingTool()
invocation_context.active_streaming_tools[tool.name] = (
active_streaming_tool
)
if inspect.isasyncgenfunction(callable_to_inspect):
if not invocation_context.active_streaming_tools:
invocation_context.active_streaming_tools = {}
logger.debug('Register streaming tool: %s', tool.name)
active_streaming_tool = ActiveStreamingTool()
invocation_context.active_streaming_tools[tool.name] = (
active_streaming_tool
)
async def execute(ctx: InvocationContext) -> AsyncGenerator[Event]:
async with Aclosing(ctx.agent.run_live(ctx)) as agen:
@@ -1497,3 +1497,51 @@ def test_stop_streaming_resets_stream_to_none():
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'