1 files changed, 81 insertions, 65 deletions

diff --git a/packages/core/src/tools/web-search.ts b/packages/core/src/tools/web-search.ts
index a2306894..3ecaf21e 100644
--- a/packages/core/src/tools/web-search.ts
+++ b/packages/core/src/tools/web-search.ts
@@ -5,8 +5,13 @@
  */
 
 import { GroundingMetadata } from '@google/genai';
-import { BaseTool, Kind, ToolResult } from './tools.js';
-import { Type } from '@google/genai';
+import {
+  BaseDeclarativeTool,
+  BaseToolInvocation,
+  Kind,
+  ToolInvocation,
+  ToolResult,
+} from './tools.js';
 import { SchemaValidator } from '../utils/schemaValidator.js';
 
 import { getErrorMessage } from '../utils/errors.js';
@@ -55,74 +60,27 @@ export interface WebSearchToolResult extends ToolResult {
     : GroundingChunkItem[];
 }
 
-/**
- * A tool to perform web searches using Google Search via the Gemini API.
- */
-export class WebSearchTool extends BaseTool<
+class WebSearchToolInvocation extends BaseToolInvocation<
   WebSearchToolParams,
   WebSearchToolResult
 > {
-  static readonly Name: string = 'google_web_search';
-
-  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.',
-      Kind.Search,
-      {
-        type: Type.OBJECT,
-        properties: {
-          query: {
-            type: Type.STRING,
-            description: 'The search query to find information on the web.',
-          },
-        },
-        required: ['query'],
-      },
-    );
-  }
-
-  /**
-   * Validates the parameters for the WebSearchTool.
-   * @param params The parameters to validate
-   * @returns An error message string if validation fails, null if valid
-   */
-  validateParams(params: WebSearchToolParams): string | null {
-    const errors = SchemaValidator.validate(
-      this.schema.parametersJsonSchema,
-      params,
-    );
-    if (errors) {
-      return errors;
-    }
-
-    if (!params.query || params.query.trim() === '') {
-      return "The 'query' parameter cannot be empty.";
-    }
-    return null;
+  constructor(
+    private readonly config: Config,
+    params: WebSearchToolParams,
+  ) {
+    super(params);
   }
 
-  override getDescription(params: WebSearchToolParams): string {
-    return `Searching the web for: "${params.query}"`;
+  override getDescription(): string {
+    return `Searching the web for: "${this.params.query}"`;
   }
 
-  async execute(
-    params: WebSearchToolParams,
-    signal: AbortSignal,
-  ): Promise<WebSearchToolResult> {
-    const validationError = this.validateToolParams(params);
-    if (validationError) {
-      return {
-        llmContent: `Error: Invalid parameters provided. Reason: ${validationError}`,
-        returnDisplay: validationError,
-      };
-    }
+  async execute(signal: AbortSignal): Promise<WebSearchToolResult> {
     const geminiClient = this.config.getGeminiClient();
 
     try {
       const response = await geminiClient.generateContent(
-        [{ role: 'user', parts: [{ text: params.query }] }],
+        [{ role: 'user', parts: [{ text: this.params.query }] }],
         { tools: [{ googleSearch: {} }] },
         signal,
       );
@@ -138,7 +96,7 @@ export class WebSearchTool extends BaseTool<
 
       if (!responseText || !responseText.trim()) {
         return {
-          llmContent: `No search results or information found for query: "${params.query}"`,
+          llmContent: `No search results or information found for query: "${this.params.query}"`,
           returnDisplay: 'No information found.',
         };
       }
@@ -172,7 +130,6 @@ export class WebSearchTool extends BaseTool<
 
           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
@@ -180,17 +137,19 @@ export class WebSearchTool extends BaseTool<
 
         if (sourceListFormatted.length > 0) {
           modifiedResponseText +=
-            '\n\nSources:\n' + sourceListFormatted.join('\n'); // Fixed string concatenation
+            '\n\nSources:\n' + sourceListFormatted.join('\n');
         }
       }
 
       return {
-        llmContent: `Web search results for "${params.query}":\n\n${modifiedResponseText}`,
-        returnDisplay: `Search results for "${params.query}" returned.`,
+        llmContent: `Web search results for "${this.params.query}":\n\n${modifiedResponseText}`,
+        returnDisplay: `Search results for "${this.params.query}" returned.`,
         sources,
       };
     } catch (error: unknown) {
-      const errorMessage = `Error during web search for query "${params.query}": ${getErrorMessage(error)}`;
+      const errorMessage = `Error during web search for query "${
+        this.params.query
+      }": ${getErrorMessage(error)}`;
       console.error(errorMessage, error);
       return {
         llmContent: `Error: ${errorMessage}`,
@@ -199,3 +158,60 @@ export class WebSearchTool extends BaseTool<
     }
   }
 }
+
+/**
+ * A tool to perform web searches using Google Search via the Gemini API.
+ */
+export class WebSearchTool extends BaseDeclarativeTool<
+  WebSearchToolParams,
+  WebSearchToolResult
+> {
+  static readonly Name: string = 'google_web_search';
+
+  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.',
+      Kind.Search,
+      {
+        type: 'object',
+        properties: {
+          query: {
+            type: 'string',
+            description: 'The search query to find information on the web.',
+          },
+        },
+        required: ['query'],
+      },
+    );
+  }
+
+  /**
+   * Validates the parameters for the WebSearchTool.
+   * @param params The parameters to validate
+   * @returns An error message string if validation fails, null if valid
+   */
+  protected override validateToolParams(
+    params: WebSearchToolParams,
+  ): string | null {
+    const errors = SchemaValidator.validate(
+      this.schema.parametersJsonSchema,
+      params,
+    );
+    if (errors) {
+      return errors;
+    }
+
+    if (!params.query || params.query.trim() === '') {
+      return "The 'query' parameter cannot be empty.";
+    }
+    return null;
+  }
+
+  protected createInvocation(
+    params: WebSearchToolParams,
+  ): ToolInvocation<WebSearchToolParams, WebSearchToolResult> {
+    return new WebSearchToolInvocation(this.config, params);
+  }
+}