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:
lwangverizon
2026-01-21 14:31:13 -08:00
committed by Copybara-Service
parent 9579bea05d
commit b57a3d43e4
2 changed files with 57 additions and 1 deletions
+14 -1
View File
@@ -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