From 65232686eee121b37fea0dd93c73bf1620152be4 Mon Sep 17 00:00:00 2001 From: Aishanee Shah Date: Wed, 4 Feb 2026 20:13:57 +0000 Subject: [PATCH] feat: implement model-dependent tool definitions architecture and refactor read_file and shell tools --- package-lock.json | 25 ++++- packages/core/src/core/client.ts | 8 +- .../core/src/tools/definitions/coreTools.ts | 91 +++++++++++++++++++ .../core/src/tools/definitions/resolver.ts | 59 ++++++++++++ packages/core/src/tools/definitions/types.ts | 25 +++++ packages/core/src/tools/read-file.ts | 9 ++ packages/core/src/tools/shell.ts | 17 ++++ packages/core/src/tools/tool-registry.ts | 13 ++- packages/core/src/tools/tools.ts | 7 +- 9 files changed, 244 insertions(+), 10 deletions(-) create mode 100644 packages/core/src/tools/definitions/coreTools.ts create mode 100644 packages/core/src/tools/definitions/resolver.ts create mode 100644 packages/core/src/tools/definitions/types.ts diff --git a/package-lock.json b/package-lock.json index 882e0e55b1..4c27a4fd95 100644 --- a/package-lock.json +++ b/package-lock.json @@ -2254,6 +2254,7 @@ "integrity": "sha512-t54CUOsFMappY1Jbzb7fetWeO0n6K0k/4+/ZpkS+3Joz8I4VcvY9OiEBFRYISqaI2fq5sCiPtAjRDOzVYG8m+Q==", "dev": true, "license": "MIT", + "peer": true, "dependencies": { "@octokit/auth-token": "^6.0.0", "@octokit/graphql": "^9.0.2", @@ -2434,6 +2435,7 @@ "resolved": "https://registry.npmjs.org/@opentelemetry/api/-/api-1.9.0.tgz", "integrity": "sha512-3giAOQvZiH5F9bMlMiv8+GSPMeqg0dbaeo58/0SlA9sxSqZhnUtxzX9/2FzyhS9sWQf5S0GJE0AKBrFqjpeYcg==", "license": "Apache-2.0", + "peer": true, "engines": { "node": ">=8.0.0" } @@ -2467,6 +2469,7 @@ "resolved": "https://registry.npmjs.org/@opentelemetry/core/-/core-2.0.1.tgz", "integrity": "sha512-MaZk9SJIDgo1peKevlbhP6+IwIiNPNmswNL4AF0WaQJLbHXjr9SrZMgS12+iqr9ToV4ZVosCcc0f8Rg67LXjxw==", "license": "Apache-2.0", + "peer": true, "dependencies": { "@opentelemetry/semantic-conventions": "^1.29.0" }, @@ -2835,6 +2838,7 @@ "resolved": "https://registry.npmjs.org/@opentelemetry/resources/-/resources-2.0.1.tgz", "integrity": "sha512-dZOB3R6zvBwDKnHDTB4X1xtMArB/d324VsbiPkX/Yu0Q8T2xceRthoIVFhJdvgVM2QhGVUyX9tzwiNxGtoBJUw==", "license": "Apache-2.0", + "peer": true, "dependencies": { "@opentelemetry/core": "2.0.1", "@opentelemetry/semantic-conventions": "^1.29.0" @@ -2868,6 +2872,7 @@ "resolved": "https://registry.npmjs.org/@opentelemetry/sdk-metrics/-/sdk-metrics-2.0.1.tgz", "integrity": "sha512-wf8OaJoSnujMAHWR3g+/hGvNcsC16rf9s1So4JlMiFaFHiE4HpIA3oUh+uWZQ7CNuK8gVW/pQSkgoa5HkkOl0g==", "license": "Apache-2.0", + "peer": true, "dependencies": { "@opentelemetry/core": "2.0.1", "@opentelemetry/resources": "2.0.1" @@ -2920,6 +2925,7 @@ "resolved": "https://registry.npmjs.org/@opentelemetry/sdk-trace-base/-/sdk-trace-base-2.0.1.tgz", "integrity": "sha512-xYLlvk/xdScGx1aEqvxLwf6sXQLXCjk3/1SQT9X9AoN5rXRhkdvIFShuNNmtTEPRBqcsMbS4p/gJLNI2wXaDuQ==", "license": "Apache-2.0", + "peer": true, "dependencies": { "@opentelemetry/core": "2.0.1", "@opentelemetry/resources": "2.0.1", @@ -4135,6 +4141,7 @@ "integrity": "sha512-6mDvHUFSjyT2B2yeNx2nUgMxh9LtOWvkhIU3uePn2I2oyNymUAX1NIsdgviM4CH+JSrp2D2hsMvJOkxY+0wNRA==", "devOptional": true, "license": "MIT", + "peer": true, "dependencies": { "csstype": "^3.0.2" } @@ -4429,6 +4436,7 @@ "integrity": "sha512-6sMvZePQrnZH2/cJkwRpkT7DxoAWh+g6+GFRK6bV3YQo7ogi3SX5rgF6099r5Q53Ma5qeT7LGmOmuIutF4t3lA==", "dev": true, "license": "MIT", + "peer": true, "dependencies": { "@typescript-eslint/scope-manager": "8.35.0", "@typescript-eslint/types": "8.35.0", @@ -5421,6 +5429,7 @@ "resolved": "https://registry.npmjs.org/acorn/-/acorn-8.15.0.tgz", "integrity": "sha512-NZyJarBfL7nWwIq+FDL6Zp/yHEhePMNnnJ0y3qfieCrmNvYct8uvtiV41UvlSe6apAfk0fY1FbWx+NwfmpvtTg==", "license": "MIT", + "peer": true, "bin": { "acorn": "bin/acorn" }, @@ -8430,6 +8439,7 @@ "integrity": "sha512-GsGizj2Y1rCWDu6XoEekL3RLilp0voSePurjZIkxL3wlm5o5EC9VpgaP7lrCvjnkuLvzFBQWB3vWB3K5KQTveQ==", "dev": true, "license": "MIT", + "peer": true, "dependencies": { "@eslint-community/eslint-utils": "^4.2.0", "@eslint-community/regexpp": "^4.12.1", @@ -8970,6 +8980,7 @@ "resolved": "https://registry.npmjs.org/express/-/express-5.2.1.tgz", "integrity": "sha512-hIS4idWWai69NezIdRt2xFVofaF4j+6INOpJlVOLDO8zXGpUVEVzIYk12UUi2JzjEzWL3IOAxcTubgz9Po0yXw==", "license": "MIT", + "peer": true, "dependencies": { "accepts": "^2.0.0", "body-parser": "^2.2.1", @@ -10571,6 +10582,7 @@ "resolved": "https://registry.npmjs.org/@jrichman/ink/-/ink-6.4.8.tgz", "integrity": "sha512-v0thcXIKl9hqF/1w4HqA6MKxIcMoWSP3YtEZIAA+eeJngXpN5lGnMkb6rllB7FnOdwyEyYaFTcu1ZVr4/JZpWQ==", "license": "MIT", + "peer": true, "dependencies": { "@alcalzone/ansi-tokenize": "^0.2.1", "ansi-escapes": "^7.0.0", @@ -14355,6 +14367,7 @@ "resolved": "https://registry.npmjs.org/react/-/react-19.2.0.tgz", "integrity": "sha512-tmbWg6W31tQLeB5cdIBOicJDJRR2KzXsV7uSK9iNfLWQ5bIZfxuPEHp7M8wiHyHnn0DD1i7w3Zmin0FtkrwoCQ==", "license": "MIT", + "peer": true, "engines": { "node": ">=0.10.0" } @@ -14365,6 +14378,7 @@ "integrity": "sha512-ePrwPfxAnB+7hgnEr8vpKxL9cmnp7F322t8oqcPshbIQQhDKgFDW4tjhF2wjVbdXF9O/nyuy3sQWd9JGpiLPvA==", "devOptional": true, "license": "MIT", + "peer": true, "dependencies": { "shell-quote": "^1.6.1", "ws": "^7" @@ -16601,6 +16615,7 @@ "resolved": "https://registry.npmjs.org/picomatch/-/picomatch-4.0.3.tgz", "integrity": "sha512-5gTmgEY/sqK6gFXLIsQNH19lWb4ebPDLA4SdLP7dsWkIXHWlG66oPuVvXSGFPppYZz8ZDZq0dYYrbHfBCVUb1Q==", "license": "MIT", + "peer": true, "engines": { "node": ">=12" }, @@ -16824,7 +16839,8 @@ "resolved": "https://registry.npmjs.org/tslib/-/tslib-2.8.1.tgz", "integrity": "sha512-oJFu94HQb+KVduSUQL7wnpmqnfmLsOA/nAh6b6EH0wCEoK0/mPeXU6c3wKDV83MkOuHPRHtSXKKU99IBazS/2w==", "dev": true, - "license": "0BSD" + "license": "0BSD", + "peer": true }, "node_modules/tsx": { "version": "4.20.3", @@ -16832,6 +16848,7 @@ "integrity": "sha512-qjbnuR9Tr+FJOMBqJCW5ehvIo/buZq7vH7qD7JziU98h6l3qGy0a/yPFjwO+y0/T7GFpNgNAvEcPPVfyT8rrPQ==", "devOptional": true, "license": "MIT", + "peer": true, "dependencies": { "esbuild": "~0.25.0", "get-tsconfig": "^4.7.5" @@ -17004,6 +17021,7 @@ "integrity": "sha512-p1diW6TqL9L07nNxvRMM7hMMw4c5XOo/1ibL4aAIGmSAt9slTE1Xgw5KWuof2uTOvCg9BY7ZRi+GaF+7sfgPeQ==", "devOptional": true, "license": "Apache-2.0", + "peer": true, "bin": { "tsc": "bin/tsc", "tsserver": "bin/tsserver" @@ -17211,6 +17229,7 @@ "resolved": "https://registry.npmjs.org/vite/-/vite-7.2.2.tgz", "integrity": "sha512-BxAKBWmIbrDgrokdGZH1IgkIk/5mMHDreLDmCJ0qpyJaAteP8NvMhkwr/ZCQNqNH97bw/dANTE9PDzqwJghfMQ==", "license": "MIT", + "peer": true, "dependencies": { "esbuild": "^0.25.0", "fdir": "^6.5.0", @@ -17324,6 +17343,7 @@ "resolved": "https://registry.npmjs.org/picomatch/-/picomatch-4.0.3.tgz", "integrity": "sha512-5gTmgEY/sqK6gFXLIsQNH19lWb4ebPDLA4SdLP7dsWkIXHWlG66oPuVvXSGFPppYZz8ZDZq0dYYrbHfBCVUb1Q==", "license": "MIT", + "peer": true, "engines": { "node": ">=12" }, @@ -17336,6 +17356,7 @@ "resolved": "https://registry.npmjs.org/vitest/-/vitest-3.2.4.tgz", "integrity": "sha512-LUCP5ev3GURDysTWiP47wRRUpLKMOfPh+yKTx3kVIEiu5KOMeqzpnYNsKyOoVrULivR8tLcks4+lga33Whn90A==", "license": "MIT", + "peer": true, "dependencies": { "@types/chai": "^5.2.2", "@vitest/expect": "3.2.4", @@ -18040,6 +18061,7 @@ "resolved": "https://registry.npmjs.org/zod/-/zod-3.25.76.tgz", "integrity": "sha512-gzUt/qt81nXsFGKIFcC3YnfEAx5NkunCfnDlvuBSSFS02bcXu4Lmea0AFIUwbLWxWPx3d9p8S5QoaujKcNQxcQ==", "license": "MIT", + "peer": true, "funding": { "url": "https://github.com/sponsors/colinhacks" } @@ -18338,6 +18360,7 @@ "resolved": "https://registry.npmjs.org/picomatch/-/picomatch-4.0.3.tgz", "integrity": "sha512-5gTmgEY/sqK6gFXLIsQNH19lWb4ebPDLA4SdLP7dsWkIXHWlG66oPuVvXSGFPppYZz8ZDZq0dYYrbHfBCVUb1Q==", "license": "MIT", + "peer": true, "engines": { "node": ">=12" }, diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts index 91434d12b3..c793ed0bd7 100644 --- a/packages/core/src/core/client.ts +++ b/packages/core/src/core/client.ts @@ -256,9 +256,9 @@ export class GeminiClient { this.forceFullIdeContext = true; } - async setTools(): Promise { + async setTools(modelId?: string): Promise { const toolRegistry = this.config.getToolRegistry(); - const toolDeclarations = toolRegistry.getFunctionDeclarations(); + const toolDeclarations = toolRegistry.getFunctionDeclarations(modelId); const tools: Tool[] = [{ functionDeclarations: toolDeclarations }]; this.getChat().setTools(tools); } @@ -653,6 +653,10 @@ export class GeminiClient { yield { type: GeminiEventType.ModelInfo, value: modelToUse }; } this.currentSequenceModel = modelToUse; + + // Update tools with the final modelId to ensure model-dependent descriptions are used. + await this.setTools(modelToUse); + const resultStream = turn.run( modelConfigKey, request, diff --git a/packages/core/src/tools/definitions/coreTools.ts b/packages/core/src/tools/definitions/coreTools.ts new file mode 100644 index 0000000000..828b0c90d8 --- /dev/null +++ b/packages/core/src/tools/definitions/coreTools.ts @@ -0,0 +1,91 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { ToolDefinition } from './types.js'; +import { READ_FILE_TOOL_NAME, SHELL_TOOL_NAME } from '../tool-names.js'; + +export const READ_FILE_DEFINITION: ToolDefinition = { + base: { + name: READ_FILE_TOOL_NAME, + description: `Reads and returns the content of a specified file. If the file is large, the content will be truncated. The tool's response will clearly indicate if truncation has occurred and will provide details on how to read more of the file using the 'offset' and 'limit' parameters. Handles text, images (PNG, JPG, GIF, WEBP, SVG, BMP), audio files (MP3, WAV, AIFF, AAC, OGG, FLAC), and PDF files. For text files, it can read specific line ranges.`, + parameters: { + type: 'object', + properties: { + file_path: { + description: 'The path to the file to read.', + type: 'string', + }, + offset: { + description: + "Optional: For text files, the 0-based line number to start reading from. Requires 'limit' to be set. Use for paginating through large files.", + type: 'number', + }, + limit: { + description: + "Optional: For text files, maximum number of lines to read. Use with 'offset' to paginate through large files. If omitted, reads the entire file (if feasible, up to a default limit).", + type: 'number', + }, + }, + required: ['file_path'], + }, + }, + variants: { + flash: { + description: + 'Reads a file from the local filesystem. Fast and efficient for checking file content.', + }, + pro: { + description: + 'Reads and returns the content of a specified file. Use this for comprehensive analysis of source code, configuration, or documentation.', + }, + }, +}; + +/** + * Note: Shell tool has platform-specific and dynamic parts. + * The base here contains the core schema. + */ +export const SHELL_DEFINITION: ToolDefinition = { + base: { + name: SHELL_TOOL_NAME, + description: 'Executes a shell command.', + parameters: { + type: 'object', + properties: { + command: { + type: 'string', + description: 'The command to execute.', + }, + description: { + type: 'string', + description: + 'Brief description of the command for the user. Be specific and concise. Ideally a single sentence. Can be up to 3 sentences for clarity. No line breaks.', + }, + dir_path: { + type: 'string', + description: + '(OPTIONAL) The path of the directory to run the command in. If not provided, the project root directory is used. Must be a directory within the workspace and must already exist.', + }, + is_background: { + type: 'boolean', + description: + 'Set to true if this command should be run in the background (e.g. for long-running servers or watchers). The command will be started, allowed to run for a brief moment to check for immediate errors, and then moved to the background.', + }, + }, + required: ['command'], + }, + }, + variants: { + flash: { + description: + 'Executes a single shell command. Use for simple operations like listing files or moving data.', + }, + pro: { + description: + 'Executes a shell command. Can be used for complex workflows, multi-step installations, or deep system investigations.', + }, + }, +}; diff --git a/packages/core/src/tools/definitions/resolver.ts b/packages/core/src/tools/definitions/resolver.ts new file mode 100644 index 0000000000..2883b66c58 --- /dev/null +++ b/packages/core/src/tools/definitions/resolver.ts @@ -0,0 +1,59 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { type FunctionDeclaration } from '@google/genai'; +import type { ToolDefinition } from './types.js'; + +/** + * Resolves a model-specific declaration for a tool. + * + * @param definition The tool definition containing base and variants. + * @param modelId The concrete model ID (e.g., 'gemini-1.5-flash'). + * @returns The final FunctionDeclaration to be sent to the API. + */ +export function resolveToolDeclaration( + definition: ToolDefinition, + modelId: string, +): FunctionDeclaration { + const { base, variants } = definition; + + if (!variants) { + return base; + } + + // Simplified mapping logic: check if the modelId contains 'flash' or 'pro'. + // This can be made more robust as needed. + let variantKey: 'flash' | 'pro' | undefined; + if (modelId.toLowerCase().includes('flash')) { + variantKey = 'flash'; + } else if (modelId.toLowerCase().includes('pro')) { + variantKey = 'pro'; + } + + const variant = variantKey ? variants[variantKey] : undefined; + + if (!variant) { + return base; + } + + // Deep merge strategy for the declaration. + return { + ...base, + ...variant, + parameters: + variant.parameters && base.parameters + ? { + ...base.parameters, + ...variant.parameters, + properties: { + ...(base.parameters.properties || {}), + ...(variant.parameters.properties || {}), + }, + required: variant.parameters.required || base.parameters.required, + } + : (variant.parameters ?? base.parameters), + }; +} diff --git a/packages/core/src/tools/definitions/types.ts b/packages/core/src/tools/definitions/types.ts new file mode 100644 index 0000000000..428d224ec9 --- /dev/null +++ b/packages/core/src/tools/definitions/types.ts @@ -0,0 +1,25 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { type FunctionDeclaration } from '@google/genai'; + +/** + * Defines a tool's identity with potential model-specific flavor variants. + */ +export interface ToolDefinition { + /** The base declaration used by default. */ + base: FunctionDeclaration; + + /** + * Model-specific overrides for the tool declaration. + * Can override description, parameters, or any other field. + */ + variants?: { + flash?: Partial; + pro?: Partial; + [modelKey: string]: Partial | undefined; + }; +} diff --git a/packages/core/src/tools/read-file.ts b/packages/core/src/tools/read-file.ts index 2fa5772187..ed9fc93f18 100644 --- a/packages/core/src/tools/read-file.ts +++ b/packages/core/src/tools/read-file.ts @@ -23,6 +23,8 @@ import { logFileOperation } from '../telemetry/loggers.js'; import { FileOperationEvent } from '../telemetry/types.js'; import { READ_FILE_TOOL_NAME } from './tool-names.js'; import { FileDiscoveryService } from '../services/fileDiscoveryService.js'; +import { READ_FILE_DEFINITION } from './definitions/coreTools.js'; +import { resolveToolDeclaration } from './definitions/resolver.js'; /** * Parameters for the ReadFile tool @@ -252,4 +254,11 @@ export class ReadFileTool extends BaseDeclarativeTool< _toolDisplayName, ); } + + override getSchema(modelId?: string) { + if (!modelId) { + return super.getSchema(); + } + return resolveToolDeclaration(READ_FILE_DEFINITION, modelId); + } } diff --git a/packages/core/src/tools/shell.ts b/packages/core/src/tools/shell.ts index e29419913e..994c4a9a90 100644 --- a/packages/core/src/tools/shell.ts +++ b/packages/core/src/tools/shell.ts @@ -43,6 +43,8 @@ import { } from '../utils/shell-utils.js'; import { SHELL_TOOL_NAME } from './tool-names.js'; import type { MessageBus } from '../confirmation-bus/message-bus.js'; +import { SHELL_DEFINITION } from './definitions/coreTools.js'; +import { resolveToolDeclaration } from './definitions/resolver.js'; export const OUTPUT_UPDATE_INTERVAL_MS = 1000; @@ -564,4 +566,19 @@ export class ShellTool extends BaseDeclarativeTool< _toolDisplayName, ); } + + override getSchema(modelId?: string) { + const declaration = modelId + ? resolveToolDeclaration(SHELL_DEFINITION, modelId) + : super.getSchema(); + + // Append platform-specific info which is currently not in the static definition + const platformInfo = getShellToolDescription( + this.config.getEnableInteractiveShell(), + ); + return { + ...declaration, + description: `${declaration.description}\n\n${platformInfo}`, + }; + } } diff --git a/packages/core/src/tools/tool-registry.ts b/packages/core/src/tools/tool-registry.ts index ae4278986b..94082dcb57 100644 --- a/packages/core/src/tools/tool-registry.ts +++ b/packages/core/src/tools/tool-registry.ts @@ -498,12 +498,13 @@ export class ToolRegistry { * Retrieves the list of tool schemas (FunctionDeclaration array). * Extracts the declarations from the ToolListUnion structure. * Includes discovered (vs registered) tools if configured. + * @param modelId Optional model identifier to get model-specific schemas. * @returns An array of FunctionDeclarations. */ - getFunctionDeclarations(): FunctionDeclaration[] { + getFunctionDeclarations(modelId?: string): FunctionDeclaration[] { const declarations: FunctionDeclaration[] = []; this.getActiveTools().forEach((tool) => { - declarations.push(tool.schema); + declarations.push(tool.getSchema(modelId)); }); return declarations; } @@ -511,14 +512,18 @@ export class ToolRegistry { /** * Retrieves a filtered list of tool schemas based on a list of tool names. * @param toolNames - An array of tool names to include. + * @param modelId Optional model identifier to get model-specific schemas. * @returns An array of FunctionDeclarations for the specified tools. */ - getFunctionDeclarationsFiltered(toolNames: string[]): FunctionDeclaration[] { + getFunctionDeclarationsFiltered( + toolNames: string[], + modelId?: string, + ): FunctionDeclaration[] { const declarations: FunctionDeclaration[] = []; for (const name of toolNames) { const tool = this.getTool(name); if (tool) { - declarations.push(tool.schema); + declarations.push(tool.getSchema(modelId)); } } return declarations; diff --git a/packages/core/src/tools/tools.ts b/packages/core/src/tools/tools.ts index 65aeb0884f..ec4e308c9a 100644 --- a/packages/core/src/tools/tools.ts +++ b/packages/core/src/tools/tools.ts @@ -312,8 +312,9 @@ export interface ToolBuilder< /** * Function declaration schema from @google/genai. + * @param modelId Optional model identifier to get a model-specific schema. */ - schema: FunctionDeclaration; + getSchema(modelId?: string): FunctionDeclaration; /** * Whether the tool's output should be rendered as markdown. @@ -355,7 +356,7 @@ export abstract class DeclarativeTool< readonly extensionId?: string, ) {} - get schema(): FunctionDeclaration { + getSchema(_modelId?: string): FunctionDeclaration { return { name: this.name, description: this.description, @@ -486,7 +487,7 @@ export abstract class BaseDeclarativeTool< override validateToolParams(params: TParams): string | null { const errors = SchemaValidator.validate( - this.schema.parametersJsonSchema, + this.getSchema().parametersJsonSchema, params, );