Add Google Search option for Gemini

Author: netham45

package-lock.json

@@ -157,6 +157,7 @@ const lmStudioSchema = baseProviderSettingsSchema.extend({
 const geminiSchema = apiModelIdProviderModelSchema.extend({
 	geminiApiKey: z.string().optional(),
 	googleGeminiBaseUrl: z.string().optional(),
+	geminiEnableGoogleSearch: z.boolean().optional(),
 })
 
 const geminiCliSchema = apiModelIdProviderModelSchema.extend({

packages/types/src/provider-settings.ts

@@ -66,9 +66,9 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl
 		const { id: model, info, reasoning: thinkingConfig, maxTokens } = this.getModel()
 
 		const contents = messages.map(convertAnthropicMessageToGemini)
-
 		const config: GenerateContentConfig = {
 			systemInstruction,
+			...(this.options.geminiEnableGoogleSearch && { tools: [{ googleSearch: {} }] }),
 			httpOptions: this.options.googleGeminiBaseUrl ? { baseUrl: this.options.googleGeminiBaseUrl } : undefined,
 			thinkingConfig,
 			maxOutputTokens: this.options.modelMaxTokens ?? maxTokens ?? undefined,
@@ -144,20 +144,52 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl
 
 	async completePrompt(prompt: string): Promise<string> {
 		try {
-			const { id: model } = this.getModel()
-
+			const { id: model, reasoning: thinkingConfig } = this.getModel()
 			const result = await this.client.models.generateContent({
 				model,
 				contents: [{ role: "user", parts: [{ text: prompt }] }],
 				config: {
+					thinkingConfig,
+					...(this.options.geminiEnableGoogleSearch && { tools: [{ googleSearch: {} }] }),
 					httpOptions: this.options.googleGeminiBaseUrl
 						? { baseUrl: this.options.googleGeminiBaseUrl }
 						: undefined,
 					temperature: this.options.modelTemperature ?? 0,
 				},
 			})
 
-			return result.text ?? ""
+			let text = result.text ?? ""
+			const candidate = result.candidates?.[0]
+			const supports = candidate?.groundingMetadata?.groundingSupports
+			const chunks = candidate?.groundingMetadata?.groundingChunks
+
+			if (text && supports && chunks) {
+				// Sort supports by end index in descending order to avoid shifting issues when inserting.
+				const sortedSupports = [...supports].sort(
+					(a, b) => (b.segment?.endIndex ?? 0) - (a.segment?.endIndex ?? 0),
+				)
+
+				for (const support of sortedSupports) {
+					const endIndex = support.segment?.endIndex
+					if (endIndex === undefined || !support.groundingChunkIndices?.length) {
+						continue
+					}
+
+					const citationLinks = support.groundingChunkIndices
+						.map((i: number) => {
+							const uri = chunks[i]?.web?.uri
+							return uri ? `[${i + 1}](${uri})` : null
+						})
+						.filter(Boolean)
+
+					if (citationLinks.length > 0) {
+						const citationString = " " + citationLinks.join(" ")
+						text = text.slice(0, endIndex) + citationString + text.slice(endIndex)
+					}
+				}
+			}
+
+			return text
 		} catch (error) {
 			if (error instanceof Error) {
 				throw new Error(`Gemini completion error: ${error.message}`)

src/api/providers/gemini.ts

@@ -58,6 +58,6 @@ export const getGeminiReasoning = ({
 	reasoningBudget,
 	settings,
 }: GetModelReasoningOptions): GeminiReasoningParams | undefined =>
-	shouldUseReasoningBudget({ model, settings })
+	shouldUseReasoningBudget({ model, settings }) && !settings.geminiEnableGoogleSearch
 		? { thinkingBudget: reasoningBudget!, includeThoughts: true }
 		: undefined

src/api/transform/reasoning.ts

@@ -166,11 +166,29 @@ export const SYSTEM_PROMPT = async (
 			{ language: language ?? formatLanguage(vscode.env.language), rooIgnoreInstructions },
 		)
 
-		// For file-based prompts, don't include the tool sections
+		const effectiveDiffStrategy = diffEnabled ? diffStrategy : undefined
+		const codeIndexManager = CodeIndexManager.getInstance(context)
+		const toolDescriptions = getToolDescriptionsForMode(
+			mode,
+			cwd,
+			supportsComputerUse,
+			codeIndexManager,
+			effectiveDiffStrategy,
+			browserViewportSize,
+			mcpHub,
+			customModes,
+			experiments,
+			partialReadsEnabled,
+			settings,
+		)
+
+		// For file-based prompts, include the tool sections
 		return `${roleDefinition}
 
 ${fileCustomSystemPrompt}
 
+${toolDescriptions}
+
 ${customInstructions}`
 	}
 

src/core/prompts/system.ts

@@ -0,0 +1,8 @@
+import { ToolArgs } from "./types"
+
+export function getGoogleSearchDescription(args: ToolArgs): string {
+	return `
+# google_search
+Description: When you need to answer questions about current events or things that have happened since your knowledge cutoff, you can use this tool to get information from the web.
+`
+}

src/core/prompts/tools/google-search.ts

@@ -14,6 +14,7 @@ import { getListFilesDescription } from "./list-files"
 import { getInsertContentDescription } from "./insert-content"
 import { getSearchAndReplaceDescription } from "./search-and-replace"
 import { getListCodeDefinitionNamesDescription } from "./list-code-definition-names"
+import { getGoogleSearchDescription } from "./google-search"
 import { getBrowserActionDescription } from "./browser-action"
 import { getAskFollowupQuestionDescription } from "./ask-followup-question"
 import { getAttemptCompletionDescription } from "./attempt-completion"
@@ -43,6 +44,7 @@ const toolDescriptionMap: Record<string, (args: ToolArgs) => string | undefined>
 	new_task: (args) => getNewTaskDescription(args),
 	insert_content: (args) => getInsertContentDescription(args),
 	search_and_replace: (args) => getSearchAndReplaceDescription(args),
+	google_search: (args) => getGoogleSearchDescription(args),
 	apply_diff: (args) =>
 		args.diffStrategy ? args.diffStrategy.getToolDescription({ cwd: args.cwd, toolOptions: args.toolOptions }) : "",
 }
@@ -107,6 +109,10 @@ export function getToolDescriptionsForMode(
 		tools.delete("codebase_search")
 	}
 
+	if (settings?.geminiEnableGoogleSearch) {
+		tools.add("google_search")
+	}
+
 	// Map tool descriptions for allowed tools
 	const descriptions = Array.from(tools).map((toolName) => {
 		const descriptionFn = toolDescriptionMap[toolName]

src/core/prompts/tools/index.ts

@@ -1636,6 +1636,7 @@ export class Task extends EventEmitter<ClineEvents> {
 				maxReadFileLine !== -1,
 				{
 					maxConcurrentFileReads,
+					...this.apiConfiguration,
 				},
 			)
 		})()
@@ -1716,7 +1717,9 @@ export class Task extends EventEmitter<ClineEvents> {
 
 			const contextWindow = modelInfo.contextWindow
 
-			const currentProfileId = state?.listApiConfigMeta.find((profile) => profile.name === state?.currentApiConfigName)?.id ?? "default";
+			const currentProfileId =
+				state?.listApiConfigMeta.find((profile) => profile.name === state?.currentApiConfigName)?.id ??
+				"default"
 
 			const truncateResult = await truncateConversationIfNeeded({
 				messages: this.apiConversationHistory,

src/core/task/Task.ts

@@ -82,6 +82,7 @@ export const generateSystemPrompt = async (provider: ClineProvider, message: Web
 		maxReadFileLine !== -1,
 		{
 			maxConcurrentFileReads,
+			...apiConfiguration,
 		},
 	)
 

src/core/webview/generateSystemPrompt.ts

@@ -50,6 +50,18 @@ export const Gemini = ({ apiConfiguration, setApiConfigurationField }: GeminiPro
 					{t("settings:providers.getGeminiApiKey")}
 				</VSCodeButtonLink>
 			)}
+			<Checkbox
+				checked={!!apiConfiguration.geminiEnableGoogleSearch}
+				onChange={(checked: boolean) => {
+					setApiConfigurationField("geminiEnableGoogleSearch", checked)
+				}}>
+				<div className="flex flex-col">
+					<span className="font-medium">{t("settings:providers.geminiEnableGoogleSearch.label")}</span>
+					<span className="text-sm text-vscode-descriptionForeground">
+						{t("settings:providers.geminiEnableGoogleSearch.description")}
+					</span>
+				</div>
+			</Checkbox>
 			<div>
 				<Checkbox
 					checked={googleGeminiBaseUrlSelected}