mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
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:
committed by
Copybara-Service
parent
b53bc555cc
commit
5269a6b1d6
@@ -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
@@ -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'
|
||||
|
||||
Reference in New Issue
Block a user