refactor(core): model-dependent tool definitions

[aishaneeshah]

Summary

Refactors the tool definition architecture to support model-dependent schemas and centralized configuration.

Details

• Model-Dependent Definitions: Centralized read_file and shell tool definitions into packages/core/src/tools/definitions/coreTools.ts.
• Schema Resolution: Introduced getSchema(modelId) to support model-specific overrides while maintaining backward compatibility via the existing schema getter.
• Late-Binding Updates: Implemented tool updates that refresh based on the final resolved model before API requests.
• Bug Fix: Fixed a bug where lastUsedModelId would become stale after resetChat() or tryCompressChat(). The client now explicitly resets this cache in startChat() to ensure the new chat instance correctly synchronizes with the model's tool registry.

Related Issues

Part of #17958

How to Validate

Run core package tests:

npm test -w @google/gemini-cli-core -- src/tools/definitions/resolver.test.ts src/tools/tool-registry.test.ts src/tools/read-file.test.ts src/tools/shell.test.ts src/core/client.test.ts

Pre-Merge Checklist

• Updated relevant documentation and README (if needed)
• Added/updated tests (if needed)
• Noted breaking changes (if any)
• Validated on required platforms/methods:
  • Linux
    • npm run

[gemini-code-assist]

Summary of Changes

Hello @aishaneeshah, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces a significant refactoring of the core architecture for tool definitions. The primary objective is to enable tools to dynamically adapt their schemas and descriptions based on the specific AI model they are interacting with, while also centralizing their configuration. This lays down the foundational infrastructure for future enhancements that will allow for more flexible and model-aware tool integrations, ensuring backward compatibility is maintained.

Highlights

• Centralized Tool Definitions: The definitions for read_file and shell tools have been moved to a new coreTools.ts file, promoting better organization and reusability.
• Model-Dependent Schemas: A new getSchema(modelId) method has been introduced across the tool architecture, allowing tool schemas and descriptions to be dynamically adjusted based on the specific AI model in use.
• Late-Binding Tool Updates: Logic has been implemented to update tool definitions with the final resolved model ID just before API requests are made, ensuring that model-specific schemas are correctly applied at the opportune moment.
• Refactored Tool Registry: The ToolRegistry has been updated to pass the modelId when retrieving function declarations, facilitating dynamic resolution of tool schemas based on the active model.

Changelog

• packages/core/src/core/client.ts
  • Modified the setTools method to accept an optional modelId and include logic to prevent redundant updates for the same model.
  • Introduced lastUsedModelId to track the model for which tools were last configured.
  • Updated the run method to call setTools with the resolved modelToUse before making API requests, ensuring late-binding of model-dependent tool descriptions.
• packages/core/src/core/geminiChat.ts
  • Added a lastUsedModel private field to track the model used in the chat.
  • Ensures geminiClient.setTools is called to update tool definitions if the model changes during a chat turn.
• packages/core/src/tools/snapshots/read-file.test.ts.snap
  • Added new snapshot tests for the ReadFileTool's getSchema method, verifying its behavior with and without a modelId.
• packages/core/src/tools/snapshots/shell.test.ts.snap
  • Added new snapshot tests for the ShellTool's getSchema method, verifying its behavior with and without a modelId.
• packages/core/src/tools/definitions/coreTools.ts
  • New file: Centralizes the ToolDefinition for read_file and shell tools.
  • Includes platform-specific logic for generating shell tool descriptions and command descriptions.
• packages/core/src/tools/definitions/resolver.test.ts
  • New file: Contains unit tests for the resolveToolDeclaration function.
• packages/core/src/tools/definitions/resolver.ts
  • New file: Provides the resolveToolDeclaration function, which currently returns the base tool definition, establishing the groundwork for model-specific overrides.
• packages/core/src/tools/definitions/types.ts
  • New file: Defines the ToolDefinition interface, which includes a base FunctionDeclaration.
• packages/core/src/tools/read-file.test.ts
  • Added new test cases for the getSchema method of ReadFileTool to ensure it correctly returns schemas based on modelId.
• packages/core/src/tools/read-file.ts
  • Refactored ReadFileTool to use the centralized READ_FILE_DEFINITION for its base schema.
  • Overrode the getSchema method to utilize resolveToolDeclaration for dynamic schema resolution.
• packages/core/src/tools/shell.test.ts
  • Added new test cases for the getSchema method of ShellTool to ensure it correctly returns schemas based on modelId.
• packages/core/src/tools/shell.ts
  • Removed previously defined getShellToolDescription and getCommandDescription functions, as they are now centralized in coreTools.ts.
  • Refactored ShellTool to use the getShellDefinition from coreTools.ts for its base schema.
  • Overrode the getSchema method to utilize resolveToolDeclaration for dynamic schema resolution.
• packages/core/src/tools/tool-registry.test.ts
  • Added a test to verify that ToolRegistry correctly passes the modelId to a tool's getSchema method when retrieving function declarations.
• packages/core/src/tools/tool-registry.ts
  • Modified getFunctionDeclarations and getFunctionDeclarationsFiltered methods to accept an optional modelId parameter.
  • Updated these methods to call tool.getSchema(modelId) instead of directly accessing tool.schema, enabling model-specific schema retrieval.
• packages/core/src/tools/tools.ts
  • Introduced getSchema(modelId?: string): FunctionDeclaration to the ToolBuilder interface, allowing for model-specific schema retrieval.
  • Deprecated the direct schema: FunctionDeclaration property, replacing it with a readonly schema getter that calls getSchema().
  • Updated the DeclarativeTool abstract class to implement the new getSchema method and the deprecated schema getter.

Activity

• The author has provided a detailed summary of the changes and specific validation steps.
• Tests have been added/updated to cover the new functionality.
• Validation on Linux has been completed successfully.
• This pull request addresses and fixes issue Tool Improvements: Ensure Tool changes are avaialble per model #17958.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces a solid refactoring to support model-dependent tool schemas by centralizing tool definitions and implementing a late-binding mechanism for tool updates. The changes are well-structured and lay a good foundation for future enhancements. I've identified a couple of critical issues related to inconsistent schema generation that could lead to runtime errors, and one high-severity issue regarding redundant state management. Addressing these points will improve the robustness and maintainability of the new architecture.

[gemini-code-assist]

The getSchema method returns objects with different shapes depending on whether modelId is provided.

• If modelId is present, resolveToolDeclaration is called, which returns a FunctionDeclaration with a parameters property, which is correct for the Gemini API.
• If modelId is not present, super.getSchema() is called. This eventually calls DeclarativeTool.getSchema, which returns an object with a parametersJsonSchema property. This is not a valid FunctionDeclaration for the API and will cause issues.

This inconsistency can lead to runtime errors. The method should always return a valid FunctionDeclaration.

You can fix this by always using resolveToolDeclaration, which correctly handles the schema creation. The if statement is not necessary as resolveToolDeclaration can handle a null or undefined modelId.

  override getSchema(modelId?: string) {
    return resolveToolDeclaration(READ_FILE_DEFINITION, modelId);
  }

[gemini-code-assist]

This getSchema method has two issues:

1. Inconsistent Schema Shape: It returns objects with different shapes. When modelId is provided, resolveToolDeclaration returns a valid FunctionDeclaration with a parameters property. When modelId is not provided, super.getSchema() returns an object with a parametersJsonSchema property, which is incorrect for the Gemini API and will cause runtime errors.
2. Stale Data: The super.getSchema() call relies on properties (description, parameterSchema) that were set in the constructor. If the enableInteractiveShell config changes during the tool's lifetime, super.getSchema() will return stale data, while the modelId branch will use fresh data from getShellDefinition.

To fix both issues, you should consistently use resolveToolDeclaration to generate the schema.

  override getSchema(modelId?: string) {
    const definition = getShellDefinition(
      this.config.getEnableInteractiveShell(),
    );
    return resolveToolDeclaration(definition, modelId);
  }

[gemini-code-assist]

The logic to check if the model has changed before calling setTools is redundant. The GeminiClient.setTools method already contains a guard to prevent unnecessary updates if the modelId hasn't changed.

This implementation introduces a duplicate lastUsedModel state in GeminiChat (line 243), which is already managed as lastUsedModelId in GeminiClient. Redundant state can lead to maintenance issues and subtle bugs if they get out of sync.

You can simplify the code and remove the redundant state by removing the if condition here and the lastUsedModel property from GeminiChat.

        await this.config.getGeminiClient().setTools(modelToUse);

[github-actions]

Size Change: +7.61 kB (+0.03%)

Total Size: 23.9 MB

Filename	Size	Change
./bundle/gemini.js	23.9 MB	+7.61 kB (+0.03%)

ℹ️ View Unchanged

Filename	Size
./bundle/sandbox-macos-permissive-closed.sb	1.03 kB
./bundle/sandbox-macos-permissive-open.sb	890 B
./bundle/sandbox-macos-permissive-proxied.sb	1.31 kB
./bundle/sandbox-macos-restrictive-closed.sb	3.29 kB
./bundle/sandbox-macos-restrictive-open.sb	3.36 kB
./bundle/sandbox-macos-restrictive-proxied.sb	3.56 kB

_{compressed-size-action}

[joshualitt]

Do we want to disentangle the schemas from the tools even further? That is, rather than having the tool determine the schema, we store the schemas separately, in a map of tool name to schema or similar. The reason to do this is so we can also support custom schema overrides even in cases where the tool is MCP or similar. Anyway, this point is minor, and I'm also okay with the way you've done it.

[aishaneeshah]

Per discussion, keeping the implementation as is for now, and will disentangle further if we decide it's needed in future.