diff --git a/core/src/tools/google_search_tool.ts b/core/src/tools/google_search_tool.ts index 69408fdf9..3ee46be27 100644 --- a/core/src/tools/google_search_tool.ts +++ b/core/src/tools/google_search_tool.ts @@ -5,7 +5,11 @@ */ import {GenerateContentConfig} from '@google/genai'; -import {isGemini1Model, isGeminiModel} from '../utils/model_name.js'; +import { + isGemini1Model, + isGeminiModel, + isGeminiModelIdCheckDisabled, +} from '../utils/model_name.js'; import {BaseTool, ToolProcessLlmRequest} from './base_tool.js'; @@ -34,6 +38,7 @@ export class GoogleSearchTool extends BaseTool { return; } + const modelCheckDisabled = isGeminiModelIdCheckDisabled(); llmRequest.config = llmRequest.config || ({} as GenerateContentConfig); llmRequest.config.tools = llmRequest.config.tools || []; @@ -51,7 +56,7 @@ export class GoogleSearchTool extends BaseTool { return; } - if (isGeminiModel(llmRequest.model)) { + if (isGeminiModel(llmRequest.model) || modelCheckDisabled) { llmRequest.config.tools.push({ googleSearch: {}, }); diff --git a/core/test/tools/google_search_tool_test.ts b/core/test/tools/google_search_tool_test.ts index f304d94ba..e51044e82 100644 --- a/core/test/tools/google_search_tool_test.ts +++ b/core/test/tools/google_search_tool_test.ts @@ -4,8 +4,19 @@ * SPDX-License-Identifier: Apache-2.0 */ -import {GOOGLE_SEARCH, GoogleSearchTool, LlmRequest} from '@google/adk'; -import {describe, expect, it} from 'vitest'; +import { + Context, + createSession, + GOOGLE_SEARCH, + GoogleSearchTool, + InvocationContext, + LlmAgent, + LlmRequest, + PluginManager, +} from '@google/adk'; +import {afterEach, describe, expect, it, vi} from 'vitest'; + +const MODEL_ID_CHECK_ENV_VAR = 'ADK_DISABLE_GEMINI_MODEL_ID_CHECK'; function makeRequest(model?: string, tools = []): LlmRequest { return { @@ -17,6 +28,18 @@ function makeRequest(model?: string, tools = []): LlmRequest { } as unknown as LlmRequest; } +/** Builds a real `Context` backed by real ADK plumbing, with no stubs. */ +function makeToolContext(): Context { + return new Context({ + invocationContext: new InvocationContext({ + invocationId: 'test-invocation', + agent: new LlmAgent({name: 'test_agent'}), + session: createSession({id: 'test-session', appName: 'test-app'}), + pluginManager: new PluginManager([]), + }), + }); +} + describe('GoogleSearchTool', () => { describe('processLlmRequest', () => { it('returns early when model is not set', async () => { @@ -91,6 +114,53 @@ describe('GoogleSearchTool', () => { expect(req.config!.tools).toEqual([{googleSearch: {}}]); }); + + describe('ADK_DISABLE_GEMINI_MODEL_ID_CHECK', () => { + afterEach(() => { + vi.unstubAllEnvs(); + }); + + it('adds googleSearch for a non-Gemini model when the check is disabled', async () => { + vi.stubEnv(MODEL_ID_CHECK_ENV_VAR, 'true'); + const tool = new GoogleSearchTool(); + const req = makeRequest('internal-model-v1'); + + await tool.processLlmRequest({ + llmRequest: req, + toolContext: makeToolContext(), + }); + + expect(req.config!.tools).toEqual([{googleSearch: {}}]); + }); + + it('keeps Gemini 1.x handling when the check is disabled', async () => { + vi.stubEnv(MODEL_ID_CHECK_ENV_VAR, 'true'); + const tool = new GoogleSearchTool(); + const req = makeRequest('gemini-1.5-pro'); + + await tool.processLlmRequest({ + llmRequest: req, + toolContext: makeToolContext(), + }); + + expect(req.config!.tools).toEqual([{googleSearchRetrieval: {}}]); + }); + + it('still throws for a non-Gemini model when the value is falsy', async () => { + vi.stubEnv(MODEL_ID_CHECK_ENV_VAR, 'false'); + const tool = new GoogleSearchTool(); + const req = makeRequest('internal-model-v1'); + + await expect( + tool.processLlmRequest({ + llmRequest: req, + toolContext: makeToolContext(), + }), + ).rejects.toThrow( + 'Google search tool is not supported for model internal-model-v1', + ); + }); + }); }); it('has a global instance GOOGLE_SEARCH', () => {