fix: Ignore AsyncGenerator return types in function declarations

For Vertex model backend, we send response back. This doesn't work for streaming tools that the return type is AsyncGenerator. So the fix here is to ignore the return type when it's AsyncGenerator.

We can't distinguish streaming vs non-streaming tool with AsyncGenerator though as LiveRequestQueue is optional in streaming tool.

Adds an `ignore_response` option to `build_function_declaration` to skip including the return type in the function declaration. This is enabled for tools that return `AsyncGenerator`, as the model does not yet support understanding these return types, while streaming tools can still handle them. Also, removes redundant return statements in `_get_mandatory_params`.

PiperOrigin-RevId: 794392846
This commit is contained in:
Hangfei Lin
2025-08-12 21:45:52 -07:00
committed by Copybara-Service
parent 8c65967cdc
commit e2518dc371
22 changed files with 134 additions and 54 deletions
+24 -4
View File
@@ -14,6 +14,9 @@
from __future__ import annotations
from collections.abc import AsyncGenerator as ABCAsyncGenerator
import inspect
from typing import get_origin
from typing import Optional
from google.genai import types
@@ -22,6 +25,7 @@ from pydantic import ConfigDict
from pydantic import Field
from ..tools.base_tool import BaseTool
from ..tools.function_tool import FunctionTool
def _find_tool_with_function_declarations(
@@ -66,13 +70,13 @@ class LlmRequest(BaseModel):
config: types.GenerateContentConfig = Field(
default_factory=types.GenerateContentConfig
)
live_connect_config: types.LiveConnectConfig = Field(
default_factory=types.LiveConnectConfig
)
"""Additional config for the generate content request.
tools in generate_content_config should not be set.
"""
live_connect_config: Optional[types.LiveConnectConfig] = None
"""Live connection config.
"""
tools_dict: dict[str, BaseTool] = Field(default_factory=dict, exclude=True)
"""The tools dictionary."""
@@ -99,7 +103,23 @@ class LlmRequest(BaseModel):
return
declarations = []
for tool in tools:
declaration = tool._get_declaration()
if self.live_connect_config is not None:
# ignore response for tools that returns AsyncGenerator that the model
# can't understand yet even though the model can't handle it, streaming
# tools can handle it.
# to check type, use typing.collections.abc.AsyncGenerator and not
# typing.AsyncGenerator
is_async_generator_return = False
if isinstance(tool, FunctionTool):
signature = inspect.signature(tool.func)
is_async_generator_return = (
get_origin(signature.return_annotation) is ABCAsyncGenerator
)
declaration = tool._get_declaration(
ignore_return_declaration=is_async_generator_return
)
else:
declaration = tool._get_declaration()
if declaration:
declarations.append(declaration)
self.tools_dict[tool.name] = tool