Implementation of web search as a tool (#307)

3 files changed, 212 insertions, 3 deletions

diff --git a/packages/server/src/config/config.ts b/packages/server/src/config/config.ts
index 6ed270cc..0a9bc2bf 100644
--- a/packages/server/src/config/config.ts
+++ b/packages/server/src/config/config.ts
@@ -20,6 +20,7 @@ import { WriteFileTool } from '../tools/write-file.js';
 import { WebFetchTool } from '../tools/web-fetch.js';
 import { ReadManyFilesTool } from '../tools/read-many-files.js';
 import { MemoryTool } from '../tools/memoryTool.js';
+import { WebSearchTool } from '../tools/web-search.js';
 
 export class MCPServerConfig {
   constructor(
@@ -202,7 +203,7 @@ export function createServerConfig(
   );
 }
 
-function createToolRegistry(config: Config): ToolRegistry {
+export function createToolRegistry(config: Config): ToolRegistry {
   const registry = new ToolRegistry(config);
   const targetDir = config.getTargetDir();
   const tools = config.getCoreTools()
@@ -228,7 +229,7 @@ function createToolRegistry(config: Config): ToolRegistry {
   registerCoreTool(ReadManyFilesTool, targetDir);
   registerCoreTool(ShellTool, config);
   registerCoreTool(MemoryTool);
-
+  registerCoreTool(WebSearchTool, config);
   registry.discoverTools();
   return registry;
 }
diff --git a/packages/server/src/tools/tool-registry.ts b/packages/server/src/tools/tool-registry.ts
index a544c8c1..ca59d666 100644
--- a/packages/server/src/tools/tool-registry.ts
+++ b/packages/server/src/tools/tool-registry.ts
@@ -133,8 +133,11 @@ Returns the MCP server response as a json string.
 
 export class ToolRegistry {
   private tools: Map<string, Tool> = new Map();
+  private config: Config;
 
-  constructor(private readonly config: Config) {}
+  constructor(config: Config) {
+    this.config = config;
+  }
 
   /**
    * Registers a tool definition.
diff --git a/packages/server/src/tools/web-search.ts b/packages/server/src/tools/web-search.ts
new file mode 100644
index 00000000..3e32ae30
--- /dev/null
+++ b/packages/server/src/tools/web-search.ts
@@ -0,0 +1,205 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { GoogleGenAI, GroundingMetadata } from '@google/genai';
+import { BaseTool, ToolResult } from './tools.js';
+import { SchemaValidator } from '../utils/schemaValidator.js';
+
+import { getErrorMessage } from '../utils/errors.js';
+import { Config } from '../config/config.js';
+import { getResponseText } from '../utils/generateContentResponseUtilities.js';
+
+interface GroundingChunkWeb {
+  uri?: string;
+  title?: string;
+}
+
+interface GroundingChunkItem {
+  web?: GroundingChunkWeb;
+  // Other properties might exist if needed in the future
+}
+
+interface GroundingSupportSegment {
+  startIndex: number;
+  endIndex: number;
+  text?: string; // text is optional as per the example
+}
+
+interface GroundingSupportItem {
+  segment?: GroundingSupportSegment;
+  groundingChunkIndices?: number[];
+  confidenceScores?: number[]; // Optional as per example
+}
+
+/**
+ * Parameters for the WebSearchTool.
+ */
+export interface WebSearchToolParams {
+  /**
+   * The search query.
+   */
+
+  query: string;
+}
+
+/**
+ * Extends ToolResult to include sources for web search.
+ */
+export interface WebSearchToolResult extends ToolResult {
+  sources?: GroundingMetadata extends { groundingChunks: GroundingChunkItem[] }
+    ? GroundingMetadata['groundingChunks']
+    : GroundingChunkItem[];
+}
+
+/**
+ * A tool to perform web searches using Google Search via the Gemini API.
+ */
+export class WebSearchTool extends BaseTool<
+  WebSearchToolParams,
+  WebSearchToolResult
+> {
+  static readonly Name: string = 'google_web_search';
+
+  private ai: GoogleGenAI;
+  private modelName: string;
+
+  constructor(private readonly config: Config) {
+    super(
+      WebSearchTool.Name,
+      'GoogleSearch',
+      'Performs a web search using Google Search (via the Gemini API) and returns the results. This tool is useful for finding information on the internet based on a query.',
+      {
+        type: 'object',
+        properties: {
+          query: {
+            type: 'string',
+            description: 'The search query to find information on the web.',
+          },
+        },
+        required: ['query'],
+      },
+    );
+
+    const apiKey = this.config.getApiKey();
+    if (!apiKey) {
+      throw new Error(
+        'Google AI API key is not configured. WebSearchTool cannot be initialized.',
+      );
+    }
+    this.ai = new GoogleGenAI({ apiKey });
+    this.modelName = this.config.getModel();
+  }
+
+  validateParams(params: WebSearchToolParams): string | null {
+    if (
+      this.schema.parameters &&
+      !SchemaValidator.validate(
+        this.schema.parameters as Record<string, unknown>,
+        params,
+      )
+    ) {
+      return "Parameters failed schema validation. Ensure 'query' is a string.";
+    }
+    if (!params.query || params.query.trim() === '') {
+      return "The 'query' parameter cannot be empty.";
+    }
+    return null;
+  }
+
+  getDescription(params: WebSearchToolParams): string {
+    return `Searching the web for: "${params.query}"`;
+  }
+
+  async execute(params: WebSearchToolParams): Promise<WebSearchToolResult> {
+    const validationError = this.validateParams(params);
+    if (validationError) {
+      return {
+        llmContent: `Error: Invalid parameters provided. Reason: ${validationError}`,
+        returnDisplay: validationError,
+      };
+    }
+
+    try {
+      const response = await this.ai.models.generateContent({
+        model: this.modelName,
+        contents: [{ role: 'user', parts: [{ text: params.query }] }],
+        config: {
+          tools: [{ googleSearch: {} }],
+        },
+      });
+
+      const responseText = getResponseText(response);
+      const groundingMetadata = response.candidates?.[0]?.groundingMetadata;
+      const sources = groundingMetadata?.groundingChunks as
+        | GroundingChunkItem[]
+        | undefined;
+      const groundingSupports = groundingMetadata?.groundingSupports as
+        | GroundingSupportItem[]
+        | undefined;
+
+      if (!responseText || !responseText.trim()) {
+        return {
+          llmContent: `No search results or information found for query: "${params.query}"`,
+          returnDisplay: 'No information found.',
+        };
+      }
+
+      let modifiedResponseText = responseText;
+      const sourceListFormatted: string[] = [];
+
+      if (sources && sources.length > 0) {
+        sources.forEach((source: GroundingChunkItem, index: number) => {
+          const title = source.web?.title || 'Untitled';
+          const uri = source.web?.uri || 'No URI';
+          sourceListFormatted.push(`[${index + 1}] ${title} (${uri})`);
+        });
+
+        if (groundingSupports && groundingSupports.length > 0) {
+          const insertions: Array<{ index: number; marker: string }> = [];
+          groundingSupports.forEach((support: GroundingSupportItem) => {
+            if (support.segment && support.groundingChunkIndices) {
+              const citationMarker = support.groundingChunkIndices
+                .map((chunkIndex: number) => `[${chunkIndex + 1}]`)
+                .join('');
+              insertions.push({
+                index: support.segment.endIndex,
+                marker: citationMarker,
+              });
+            }
+          });
+
+          // Sort insertions by index in descending order to avoid shifting subsequent indices
+          insertions.sort((a, b) => b.index - a.index);
+
+          const responseChars = modifiedResponseText.split(''); // Use new variable
+          insertions.forEach((insertion) => {
+            // Fixed arrow function syntax
+            responseChars.splice(insertion.index, 0, insertion.marker);
+          });
+          modifiedResponseText = responseChars.join(''); // Assign back to modifiedResponseText
+        }
+
+        if (sourceListFormatted.length > 0) {
+          modifiedResponseText +=
+            '\n\nSources:\n' + sourceListFormatted.join('\n'); // Fixed string concatenation
+        }
+      }
+
+      return {
+        llmContent: `Web search results for "${params.query}":\n\n${modifiedResponseText}`,
+        returnDisplay: `Search results for "${params.query}" returned.`,
+        sources,
+      };
+    } catch (error: unknown) {
+      const errorMessage = `Error during web search for query "${params.query}": ${getErrorMessage(error)}`;
+      console.error(errorMessage, error);
+      return {
+        llmContent: `Error: ${errorMessage}`,
+        returnDisplay: `Error performing web search.`,
+      };
+    }
+  }
+}