Mcp agent v1 2 add _truncation_tool_name #662
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Summary of Changes
Hello @mushenL, I'm Gemini Code Assist1! 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 enhances the LLMAgent by introducing a robust mechanism for handling long tool names. It intelligently truncates tool names before they are presented to the Large Language Model, which can help optimize token usage and prevent potential issues with context window limits. Crucially, it also ensures that the original tool names are correctly restored when the LLM invokes a tool, preserving the intended functionality. Additionally, the project's version has been updated to reflect a release candidate status.
Highlights
- Tool Name Truncation: Implemented a new
_truncation_tool_namemethod within theLLMAgentto automatically shorten tool names that exceed a configurable maximum length (MAX_TOOL_NAME_LENGTH, default 64). This truncation primarily affects the 'server' part of 'server:fun' formatted tool names, appending '...' to indicate truncation. This is crucial for managing token usage and avoiding issues with LLM context windows when tool names are excessively long. - Tool Name Mapping and Restoration: Introduced a
tool_name_mapattribute inLLMAgentto store the mapping between the truncated tool names (sent to the LLM) and their original, full names. This ensures that when the LLM makes a tool call using a truncated name, the system can correctly retrieve and use the original tool name for execution, maintaining functional integrity. - Version Update: The project version has been updated from
1.0.0to1.0.0rc0, indicating a transition to a release candidate or pre-release phase for themcp agent v1 2.
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 in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.
| 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 issue 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 is currently in preview and 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 to provide feedback.
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
-
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. ↩
There was a problem hiding this comment.
Code Review
This pull request introduces a feature to truncate long tool names. The review identifies potential issues in the implementation of _truncation_tool_name, including a critical bug related to tool name collisions, a high-severity issue with state management, and an off-by-one error in length calculation. Specific code suggestions are provided to address these concerns, along with feedback to improve code clarity.
| truncated_server = server_name[:max_server_len] + '...' | ||
| new_name = f'{truncated_server}:{fun_name}' | ||
| # 创建映射关系 | ||
| self.tool_name_map[new_name] = original_name |
There was a problem hiding this comment.
The current truncation logic can cause different tool names to be truncated to the same string (e.g., server_name_long_1:my_func and server_name_long_2:my_func could both become server_...:my_func). This will cause one of the mappings in tool_name_map to be overwritten, leading to subsequent tool calls being routed to the wrong tool.
To prevent this critical issue, check for such collisions and raise an error. This will alert the user that their tool names are not distinct enough for safe truncation.
if new_name in self.tool_name_map and self.tool_name_map[new_name] != original_name:
raise ValueError(
f"Tool name truncation collision: '{original_name}' and "
f"'{self.tool_name_map[new_name]}' both truncate to '{new_name}'. "
'Please provide more distinct tool names to avoid this issue.'
)
self.tool_name_map[new_name] = original_name| @@ -295,6 +296,36 @@ def _handle_stream_message(self, messages: List[Message], | |||
| for message in self.llm.generate(messages, tools=tools): | |||
| yield message | |||
|
|
|||
| def _truncation_tool_name(self, tools: List[Tool]): | |||
| max_length = int(os.environ.get('MAX_TOOL_NAME_LENGTH', 64)) | |||
There was a problem hiding this comment.
The tool_name_map is an instance variable that is populated in _truncation_tool_name but never cleared. If the agent performs multiple operations where the set of available tools changes, this can lead to stale entries persisting in tool_name_map. This could cause incorrect behavior if a new truncated tool name happens to match a stale one.
To prevent this, clear the map at the beginning of this method.
self.tool_name_map = {}
max_length = int(os.environ.get('MAX_TOOL_NAME_LENGTH', 64))| ) | ||
|
|
||
| server_name, fun_name = parts | ||
| max_server_len = max_length - len(fun_name) - 3 |
There was a problem hiding this comment.
There is an off-by-one error in the calculation of max_server_len. The truncated name format is f'{truncated_server}:{fun_name}', and truncated_server includes .... The length of ... is 3 characters, and the colon : is 1 character. Therefore, you need to subtract a total of 4 characters from max_length to account for them, not 3. The current calculation can result in a final tool name that is one character longer than max_length.
max_server_len = max_length - len(fun_name) - 4| if max_length < 4: | ||
| max_length = 64 |
There was a problem hiding this comment.
The check if max_length < 4: uses a magic number 4, and its purpose is unclear. The subsequent check if max_server_len <= 0 on line 315 already correctly handles cases where the name is too short to be truncated. This makes the check on lines 301-302 redundant.
Consider removing these lines for clarity. If a minimum length is truly required, it should be a named constant with a clear explanation of why that specific value is chosen.
ms_agent/agent/llm_agent.py
Outdated
| # 截断服务名并添加 ... | ||
| truncated_server = server_name[:max_server_len] + '...' | ||
| new_name = f'{truncated_server}:{fun_name}' | ||
| # 创建映射关系 |
There was a problem hiding this comment.
The comments in this block are in Chinese, while the rest of the file and project appear to use English. For consistency and maintainability, use English for all comments.
# Truncate the server name and add an ellipsis
truncated_server = server_name[:max_server_len] + '...'
new_name = f'{truncated_server}:{fun_name}'
# Create a mapping from the new name to the original name
Change Summary
Related issue number
Checklist
pre-commit installandpre-commit run --all-filesbefore git commit, and passed lint check.