You've already forked adk-python
mirror of
https://github.com/encounter/adk-python.git
synced 2026-03-30 10:57:20 -07:00
feat: Allow google search tool to set different model
Merge https://github.com/google/adk-python/pull/4136 **Please ensure you have read the [contribution guide](https://github.com/google/adk-python/blob/main/CONTRIBUTING.md) before creating a pull request.** ### Link to Issue or Description of Change **1. Link to an existing issue (if applicable):** Allow google search tool to set different model #4135 **2. Or, if no issue exists, describe the change:** _If applicable, please follow the issue templates to provide as much detail as possible._ **Problem:** Currently, the Google Search tool inherits and uses the same LLM model set from the parent agent for processing and summarizing search results. This creates a limitation for users who wish to decouple the agent's reasoning model from the model used for search summarization (e.g., for cost optimization or using a lightweight model for simpler summarization tasks). **Solution:** I have updated the Google Search tool to accept an optional LLM model parameter. Custom Model: Users can now explicitly specify which model should be used for processing search results. Default Behavior: If no model is specified, the tool defaults to the parent agent's model, ensuring backward compatibility. ``` # If a custom model is specified, use it instead of the original model if self.model is not None: llm_request.model = self.model ``` ### Testing Plan Added a new test case test_process_llm_request_with_custom_model in [test_google_search_tool.py] that verifies: When a custom model parameter is provided to GoogleSearchTool, it overrides the model from the incoming llm_request during process_llm_request The tool correctly uses the custom model for LLM calls while maintaining other request parameters **Unit Tests:** - [X] I have added or updated unit tests for my change. - [X] All unit tests pass locally. (base) wanglu2:adk-python/ (feature/allow-google-search-tool-set-different-llm✗) $ uv run pytest ./tests/unittests/tools/test_google_search_tool.py [22:07:32] ======================================================================== test session starts ======================================================================== platform darwin -- Python 3.13.1, pytest-9.0.2, pluggy-1.6.0 rootdir: /Users/wanglu2/Documents/Git/adk-python configfile: pyproject.toml plugins: mock-3.15.1, anyio-4.12.0, xdist-3.8.0, asyncio-1.3.0, langsmith-0.6.0 asyncio: mode=Mode.AUTO, debug=False, asyncio_default_fixture_loop_scope=function, asyncio_default_test_loop_scope=function collected 21 items tests/unittests/tools/test_google_search_tool.py ..................... [100%] ======================================================================== 21 passed in 7.91s ========================================================================= ### Checklist - [X] I have read the [CONTRIBUTING.md](https://github.com/google/adk-python/blob/main/CONTRIBUTING.md) document. - [X] I have performed a self-review of my own code. - [X] I have commented my code, particularly in hard-to-understand areas. - [X] I have added tests that prove my fix is effective or that my feature works. - [X] New and existing unit tests pass locally with my changes. - [X] I have manually tested my changes end-to-end. - [ ] Any dependent changes have been merged and published in downstream modules. ### Additional context Co-authored-by: Kathy Wu <wukathy@google.com> COPYBARA_INTEGRATE_REVIEW=https://github.com/google/adk-python/pull/4136 from lwangverizon:feature/allow-google-search-tool-set-different-llm 239ea9577bd7befc9a3574aca48dc8243d55192c PiperOrigin-RevId: 859265757
This commit is contained in:
committed by
Copybara-Service
parent
9579bea05d
commit
b57a3d43e4
@@ -35,17 +35,26 @@ class GoogleSearchTool(BaseTool):
|
||||
local code execution.
|
||||
"""
|
||||
|
||||
def __init__(self, *, bypass_multi_tools_limit: bool = False):
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
bypass_multi_tools_limit: bool = False,
|
||||
model: str | None = None,
|
||||
):
|
||||
"""Initializes the Google search tool.
|
||||
|
||||
Args:
|
||||
bypass_multi_tools_limit: Whether to bypass the multi tools limitation,
|
||||
so that the tool can be used with other tools in the same agent.
|
||||
model: Optional model name to use for processing the LLM request. If
|
||||
provided, this model will be used instead of the model from the
|
||||
incoming llm_request.
|
||||
"""
|
||||
|
||||
# Name and description are not used because this is a model built-in tool.
|
||||
super().__init__(name='google_search', description='google_search')
|
||||
self.bypass_multi_tools_limit = bypass_multi_tools_limit
|
||||
self.model = model
|
||||
|
||||
@override
|
||||
async def process_llm_request(
|
||||
@@ -54,6 +63,10 @@ class GoogleSearchTool(BaseTool):
|
||||
tool_context: ToolContext,
|
||||
llm_request: LlmRequest,
|
||||
) -> None:
|
||||
# If a custom model is specified, use it instead of the original model
|
||||
if self.model is not None:
|
||||
llm_request.model = self.model
|
||||
|
||||
llm_request.config = llm_request.config or types.GenerateContentConfig()
|
||||
llm_request.config.tools = llm_request.config.tools or []
|
||||
if is_gemini_1_model(llm_request.model):
|
||||
|
||||
@@ -432,3 +432,46 @@ class TestGoogleSearchTool:
|
||||
assert len(llm_request.config.tools) == 1
|
||||
assert llm_request.config.tools[0].google_search is not None
|
||||
assert llm_request.config.tools[0].google_search_retrieval is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
(
|
||||
'tool_model',
|
||||
'request_model',
|
||||
'expected_model',
|
||||
),
|
||||
[
|
||||
(
|
||||
'gemini-2.5-flash-lite',
|
||||
'gemini-2.5-flash',
|
||||
'gemini-2.5-flash-lite',
|
||||
),
|
||||
(
|
||||
None,
|
||||
'gemini-2.5-flash',
|
||||
'gemini-2.5-flash',
|
||||
),
|
||||
],
|
||||
ids=['with_custom_model', 'without_custom_model'],
|
||||
)
|
||||
async def test_process_llm_request_custom_model_behavior(
|
||||
self,
|
||||
tool_model,
|
||||
request_model,
|
||||
expected_model,
|
||||
):
|
||||
"""Tests custom model parameter behavior in process_llm_request."""
|
||||
tool = GoogleSearchTool(model=tool_model)
|
||||
tool_context = await _create_tool_context()
|
||||
|
||||
llm_request = LlmRequest(
|
||||
model=request_model, config=types.GenerateContentConfig()
|
||||
)
|
||||
|
||||
await tool.process_llm_request(
|
||||
tool_context=tool_context, llm_request=llm_request
|
||||
)
|
||||
|
||||
assert llm_request.model == expected_model
|
||||
assert llm_request.config.tools is not None
|
||||
assert len(llm_request.config.tools) == 1
|
||||
|
||||
Reference in New Issue
Block a user