Skip to content

Commit b57a3d4

Browse files
lwangverizonwukath
authored andcommitted
feat: Allow google search tool to set different model
Merge #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=#4136 from lwangverizon:feature/allow-google-search-tool-set-different-llm 239ea95 PiperOrigin-RevId: 859265757
1 parent 9579bea commit b57a3d4

File tree

2 files changed

+57
-1
lines changed

2 files changed

+57
-1
lines changed

src/google/adk/tools/google_search_tool.py

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,17 +35,26 @@ class GoogleSearchTool(BaseTool):
3535
local code execution.
3636
"""
3737

38-
def __init__(self, *, bypass_multi_tools_limit: bool = False):
38+
def __init__(
39+
self,
40+
*,
41+
bypass_multi_tools_limit: bool = False,
42+
model: str | None = None,
43+
):
3944
"""Initializes the Google search tool.
4045
4146
Args:
4247
bypass_multi_tools_limit: Whether to bypass the multi tools limitation,
4348
so that the tool can be used with other tools in the same agent.
49+
model: Optional model name to use for processing the LLM request. If
50+
provided, this model will be used instead of the model from the
51+
incoming llm_request.
4452
"""
4553

4654
# Name and description are not used because this is a model built-in tool.
4755
super().__init__(name='google_search', description='google_search')
4856
self.bypass_multi_tools_limit = bypass_multi_tools_limit
57+
self.model = model
4958

5059
@override
5160
async def process_llm_request(
@@ -54,6 +63,10 @@ async def process_llm_request(
5463
tool_context: ToolContext,
5564
llm_request: LlmRequest,
5665
) -> None:
66+
# If a custom model is specified, use it instead of the original model
67+
if self.model is not None:
68+
llm_request.model = self.model
69+
5770
llm_request.config = llm_request.config or types.GenerateContentConfig()
5871
llm_request.config.tools = llm_request.config.tools or []
5972
if is_gemini_1_model(llm_request.model):

tests/unittests/tools/test_google_search_tool.py

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -432,3 +432,46 @@ async def test_process_llm_request_gemini_version_specifics(self):
432432
assert len(llm_request.config.tools) == 1
433433
assert llm_request.config.tools[0].google_search is not None
434434
assert llm_request.config.tools[0].google_search_retrieval is None
435+
436+
@pytest.mark.asyncio
437+
@pytest.mark.parametrize(
438+
(
439+
'tool_model',
440+
'request_model',
441+
'expected_model',
442+
),
443+
[
444+
(
445+
'gemini-2.5-flash-lite',
446+
'gemini-2.5-flash',
447+
'gemini-2.5-flash-lite',
448+
),
449+
(
450+
None,
451+
'gemini-2.5-flash',
452+
'gemini-2.5-flash',
453+
),
454+
],
455+
ids=['with_custom_model', 'without_custom_model'],
456+
)
457+
async def test_process_llm_request_custom_model_behavior(
458+
self,
459+
tool_model,
460+
request_model,
461+
expected_model,
462+
):
463+
"""Tests custom model parameter behavior in process_llm_request."""
464+
tool = GoogleSearchTool(model=tool_model)
465+
tool_context = await _create_tool_context()
466+
467+
llm_request = LlmRequest(
468+
model=request_model, config=types.GenerateContentConfig()
469+
)
470+
471+
await tool.process_llm_request(
472+
tool_context=tool_context, llm_request=llm_request
473+
)
474+
475+
assert llm_request.model == expected_model
476+
assert llm_request.config.tools is not None
477+
assert len(llm_request.config.tools) == 1

0 commit comments

Comments
 (0)