From 7fcdb57615d21d0e1fe849532aacb13767f8bae9 Mon Sep 17 00:00:00 2001 From: YueGuobin Date: Thu, 5 Mar 2026 01:50:48 +0800 Subject: [PATCH] feat(context-manager): improve context window management with tool token accounting Enhanced the context window management system to properly account for tool definition tokens when trimming messages. The key changes include: - Updated `trim_messages_for_context` function to accept `tool_tokens` parameter - Modified token budget allocation logic to subtract tool tokens before message trimming - Added detailed documentation explaining the token budget distribution between messages and tool definitions - Implemented prioritized trimming strategy that preserves system messages and recent conversation history - Added boundary case handling for scenarios where system messages or tools exceed available budget The improvements ensure more accurate context window management by accounting for the ~1000-2000 tokens typically consumed by tool definitions that LangChain automatically includes in LLM requests. --- .../gns3-copilot/context-window-management.md | 87 ++++++++- .../gns3_copilot/agent/context_manager.py | 183 +++++++++++++----- .../tools_v2/config_tools_nornir.py | 4 +- .../tools_v2/display_tools_nornir.py | 4 +- .../gns3_copilot/tools_v2/gns3_start_node.py | 4 +- .../tools_v2/vpcs_tools_telnetlib3.py | 4 +- 6 files changed, 220 insertions(+), 66 deletions(-) diff --git a/docs/gns3-copilot/context-window-management.md b/docs/gns3-copilot/context-window-management.md index 914237259..69f5a8c7a 100644 --- a/docs/gns3-copilot/context-window-management.md +++ b/docs/gns3-copilot/context-window-management.md @@ -241,15 +241,90 @@ pip install tiktoken>=0.8.0 - `balanced`: 使用 75% 限制（默认） - `aggressive`: 使用 85% 限制（最大化输入） -**`trim_messages_for_context(messages, model_name, strategy) -> list`** -- 使用 LangChain 的 `trim_messages` 裁剪消息 +**`trim_messages_for_context(messages, model_name, strategy, tool_tokens) -> list`** +- 使用 tiktoken 裁剪消息 - 保留最近的消息 - 始终保留系统消息 +- **关键**：考虑工具定义占用的 tokens -**`prepare_context_messages(state_messages, system_prompt, topology_context, model_name, trim_strategy) -> list`** -- GNS3 Copilot 的主要入口点 -- 构建完整上下文（系统提示 + 拓扑信息 + 消息历史） -- 自动裁剪以适应模型限制 +### 2. 裁剪逻辑详解 + +#### 2.1 Token 预算分配 + +当调用 LLM 时，发送的内容包含两部分： + +``` +发送给 LLM 的完整请求: +┌─────────────────────────────────────────────────────────────┐ +│ 1. Messages (我们管理的) │ +│ ├─ SystemMessage: system prompt │ +│ ├─ SystemMessage: topology info │ +│ └─ HumanMessage: 用户消息 / 历史消息 │ +├─────────────────────────────────────────────────────────────┤ +│ 2. Tool Definitions (LangChain 自动添加，不在消息中) │ +│ ├─ Tool 1 schema (name, description, parameters) │ +│ ├─ Tool 2 schema │ +│ └─ ... (约 1000-2000 tokens per tool) │ +└─────────────────────────────────────────────────────────────┘ +``` + +#### 2.2 裁剪流程 + +``` +第一步：计算输入预算 +┌─────────────────────────────────────────────────────────────┐ +│ context_limit: 16,000 tokens │ +│ strategy: conservative (60%) │ +│ │ +│ 输入预算 = 16,000 × 0.6 = 9,600 tokens │ +└─────────────────────────────────────────────────────────────┘ + ↓ +第二步：减去工具定义 +┌─────────────────────────────────────────────────────────────┐ +│ 输入预算: 9,600 tokens │ +│ 工具定义: 1,862 tokens │ +│ │ +│ 可用于消息 = 9,600 - 1,862 = 7,738 tokens │ +└─────────────────────────────────────────────────────────────┘ + ↓ +第三步：分离系统消息和对话历史 +┌─────────────────────────────────────────────────────────────┐ +│ 可用于消息: 7,738 tokens │ +│ │ +│ System Messages (保留): │ +│ - System prompt: ~1,000 tokens │ +│ - Topology info: ~3,000 tokens │ +│ = 4,000 tokens │ +│ │ +│ 可用于历史 = 7,738 - 4,000 = 3,738 tokens │ +└─────────────────────────────────────────────────────────────┘ + ↓ +第四步：裁剪对话历史 +┌─────────────────────────────────────────────────────────────┐ +│ 历史消息: 5,000 tokens → 需要裁剪到 3,738 tokens │ +│ │ +│ 丢弃旧消息，直到满足限制 │ +└─────────────────────────────────────────────────────────────┘ +``` + +#### 2.3 裁剪优先级 + +系统按以下优先级保留内容： + +| 优先级 | 内容 | 说明 | +|--------|------|------| +| 1️⃣ | System Prompt | 永远保留（第1个） | +| 2️⃣ | Topology Info | 永远保留（第2个） | +| 3️⃣ | 最新用户消息 | 至少保留最后1条 | +| 4️⃣ | 旧对话历史 | 按时间顺序丢弃 | + +#### 2.4 边界情况处理 + +| 情况 | 处理方式 | +|------|----------| +| System > 预算 | 只保留 system prompt，丢弃所有其他消息 | +| Tools > 预算 | 警告日志，建议增加 context_limit | +| 历史全被裁剪 | 保留最后1条用户消息 | ### 2. 集成到 GNS3 Copilot diff --git a/gns3server/agent/gns3_copilot/agent/context_manager.py b/gns3server/agent/gns3_copilot/agent/context_manager.py index 82dec7493..f33776d1f 100644 --- a/gns3server/agent/gns3_copilot/agent/context_manager.py +++ b/gns3server/agent/gns3_copilot/agent/context_manager.py @@ -281,13 +281,17 @@ def trim_messages_for_context( llm_config: dict[str, Any] | None = None, strategy: Literal["conservative", "balanced", "aggressive"] = "balanced", preserve_system: bool = True, + tool_tokens: int = 0, ) -> list[Any]: """ Trim messages to fit within model's context window. - This function uses tiktoken for accurate token counting and LangChain's - trim_messages utility to intelligently reduce message history while - preserving conversation flow. + This function uses tiktoken for accurate token counting and intelligently + reduces message history while preserving conversation flow. + + IMPORTANT: Tool definitions are sent separately by LangChain and count towards + the context limit. This function accounts for tool tokens when making + trimming decisions. Args: messages: List of LangChain messages (HumanMessage, AIMessage, etc.) @@ -295,13 +299,14 @@ def trim_messages_for_context( llm_config: Optional LLM config dict from database (may contain context_limit) strategy: How aggressively to use the context window preserve_system: Whether to always preserve system messages + tool_tokens: Token count for tool definitions (these are sent separately by LangChain) Returns: list: Trimmed list of messages that fit within context limit Examples: >>> messages = [HumanMessage("Hello"), AIMessage("Hi there!")] - >>> trimmed = trim_messages_for_context(messages, "gpt-4o") + >>> trimmed = trim_messages_for_context(messages, "gpt-4o", tool_tokens=1000) >>> len(trimmed) <= len(messages) True """ @@ -314,58 +319,99 @@ def trim_messages_for_context( # Calculate usable tokens (reserve space for output) max_tokens = calculate_max_tokens(model_limit, strategy) + # Account for tool tokens - these are sent separately by LangChain + # and count towards the context limit + available_for_messages = max_tokens - tool_tokens + + if available_for_messages < 0: + logger.warning( + "Tool definitions (%d tokens) exceed input budget (%d tokens). " + "Consider reducing context_limit or using fewer tools.", + tool_tokens, max_tokens + ) + available_for_messages = 0 + # Check if trimming is needed current_tokens = count_messages_tokens(messages) - if current_tokens <= max_tokens: + if current_tokens <= available_for_messages: logger.debug( - "Messages fit in context: %d / %d tokens", - current_tokens, max_tokens + "Messages fit in context: %d / %d tokens (available: %d, tools: %d)", + current_tokens, max_tokens, available_for_messages, tool_tokens ) return messages logger.info( - "Trimming messages: %d → %d tokens (model: %s)", - current_tokens, max_tokens, model_name + "Trimming messages: %d → %d tokens (budget: %d, tools: %d)", + current_tokens, available_for_messages, max_tokens, tool_tokens ) - # Trim messages using LangChain's utility - try: - from langchain_core.messages.utils import count_tokens_approximately - trimmed = trim_messages( - messages, - strategy="last", # Keep most recent messages - max_tokens=max_tokens, - token_counter=count_tokens_approximately, # Required: token counting function - include_system=preserve_system, # Keep system messages if requested - start_on="human", # Ensure we start with a human message - end_on=("human", "tool", "ai"), # End on human/tool/ai messages - ) + # Manually separate and trim to ensure system messages are preserved + # This is more reliable than using trim_messages with include_system + system_msgs = [m for m in messages if isinstance(m, SystemMessage)] + other_msgs = [m for m in messages if not isinstance(m, SystemMessage)] - logger.info( - "Trimmed %d → %d messages", - len(messages), len(trimmed) - ) + # Calculate tokens for system messages (these will always be preserved) + system_tokens = count_messages_tokens(system_msgs) - return trimmed - - except Exception as e: - logger.error("Failed to trim messages: %s", e, exc_info=True) - - # Fallback: simple slicing (keep last N messages) - # Estimate average tokens per message (~100 tokens) - fallback_msg_count = max(1, max_tokens // 100) + # Calculate available tokens for non-system messages (after tools and system) + available_for_other = available_for_messages - system_tokens + if available_for_other <= 0: + # Not enough space for system messages - keep only system messages logger.warning( - "Using fallback trimming: keeping last %d messages", - fallback_msg_count + "System messages (%d tokens) exceed available space (%d tokens), truncating to system only", + system_tokens, available_for_messages ) + return system_msgs[:1] if system_msgs else messages[-1:] - # Always preserve system messages - system_msgs = [m for m in messages if isinstance(m, SystemMessage)] - other_msgs = [m for m in messages if not isinstance(m, SystemMessage)] + # Trim non-system messages to fit available space + trimmed_other = _trim_to_token_limit(other_msgs, available_for_other) - return system_msgs + other_msgs[-fallback_msg_count:] + # Combine system messages with trimmed conversation + trimmed = system_msgs + trimmed_other + + logger.info( + "Trimmed %d → %d messages (system: %d, history: %d → %d)", + len(messages), len(trimmed), + len(system_msgs), len(other_msgs), len(trimmed_other) + ) + + return trimmed + + +def _trim_to_token_limit(messages: list[Any], max_tokens: int) -> list[Any]: + """ + Trim messages to fit within token limit using tiktoken. + + Iteratively removes oldest messages until under token limit. + Always keeps at least the most recent message. + + Args: + messages: List of messages to trim + max_tokens: Maximum tokens allowed + + Returns: + Trimmed list of messages + """ + if not messages: + return messages + + current_tokens = count_messages_tokens(messages) + + if current_tokens <= max_tokens: + return messages + + # Iteratively remove oldest messages + trimmed = list(messages) + while trimmed and count_messages_tokens(trimmed) > max_tokens: + trimmed.pop(0) + + # Ensure at least one message remains + if not trimmed and messages: + trimmed = [messages[-1]] + + return trimmed def get_token_usage_summary( @@ -472,11 +518,21 @@ def prepare_context_messages( # Build base context (system + topology) context_messages = [SystemMessage(content=system_prompt)] + # Calculate token breakdown + system_prompt_tokens = count_tokens_accurately(system_prompt) + topology_tokens = 0 + if topology_context: + topology_context_full = f"Current Topology:\n{topology_context}" + topology_tokens = count_tokens_accurately(topology_context_full) context_messages.append( - SystemMessage(content=f"Current Topology:\n{topology_context}") + SystemMessage(content=topology_context_full) ) + # Calculate conversation history tokens + history_tokens = count_messages_tokens(state_messages) + total_context_tokens = system_prompt_tokens + topology_tokens + history_tokens + # Combine with conversation history full_messages = context_messages + state_messages @@ -487,20 +543,43 @@ def prepare_context_messages( llm_config=llm_config, strategy=trim_strategy, preserve_system=True, # Always keep system prompts + tool_tokens=tool_tokens, # Account for tool definitions ) - # Log summary (including tool tokens) - summary = get_token_usage_summary(trimmed_messages, model_name, llm_config, tool_tokens) - logger.info( - "Context prepared: %d msgs, ~%d tokens (messages) + %d tokens (tools) = %d total / %dK limit (%.1f%%), strategy=%s", - summary["message_count"], - summary["estimated_tokens"], - summary["tool_tokens"], - summary["total_tokens"], - summary["model_limit_k"], - summary["total_usage_percentage"], - trim_strategy - ) + # Recalculate after trimming + trimmed_history_tokens = count_messages_tokens([m for m in trimmed_messages if not isinstance(m, SystemMessage)]) + trimmed_total_tokens = system_prompt_tokens + topology_tokens + trimmed_history_tokens + + # Log summary with detailed breakdown + model_limit_k = get_model_context_limit(model_name, llm_config) // 1000 + + if trimmed_history_tokens < history_tokens: + # Trimming happened + logger.info( + "Context prepared (trimmed): system=%d + topology=%d + history=%d→%d + tools=%d = %d total / %dK limit (%.1f%%), strategy=%s", + system_prompt_tokens, + topology_tokens, + history_tokens, + trimmed_history_tokens, + tool_tokens, + trimmed_total_tokens, + model_limit_k, + (trimmed_total_tokens / (model_limit_k * 1000)) * 100, + trim_strategy + ) + else: + # No trimming + logger.info( + "Context prepared: system=%d + topology=%d + history=%d + tools=%d = %d total / %dK limit (%.1f%%), strategy=%s", + system_prompt_tokens, + topology_tokens, + history_tokens, + tool_tokens, + total_context_tokens, + model_limit_k, + (total_context_tokens / (model_limit_k * 1000)) * 100, + trim_strategy + ) return trimmed_messages diff --git a/gns3server/agent/gns3_copilot/tools_v2/config_tools_nornir.py b/gns3server/agent/gns3_copilot/tools_v2/config_tools_nornir.py index 496536daf..f4e5df1d0 100644 --- a/gns3server/agent/gns3_copilot/tools_v2/config_tools_nornir.py +++ b/gns3server/agent/gns3_copilot/tools_v2/config_tools_nornir.py @@ -71,7 +71,7 @@ class ExecuteMultipleDeviceConfigCommands(BaseTool): Input should be a JSON object containing project_id and device configurations. Example input: { - "project_id": "f32ebf3d-ef8c-4910-b0d6-566ed828cd24", + "project_id": "", "device_configs": [ { "device_name": "R-1", @@ -463,7 +463,7 @@ if __name__ == "__main__": # example tool_input with project_id input_paras = json.dumps( { - "project_id": "f32ebf3d-ef8c-4910-b0d6-566ed828cd24", + "project_id": "", "device_configs": [ { "device_name": "R-1", diff --git a/gns3server/agent/gns3_copilot/tools_v2/display_tools_nornir.py b/gns3server/agent/gns3_copilot/tools_v2/display_tools_nornir.py index c7ce08404..13b805814 100644 --- a/gns3server/agent/gns3_copilot/tools_v2/display_tools_nornir.py +++ b/gns3server/agent/gns3_copilot/tools_v2/display_tools_nornir.py @@ -95,7 +95,7 @@ class ExecuteMultipleDeviceCommands(BaseTool): **Input Format:** { - "project_id": "f32ebf3d-ef8c-4910-b0d6-566ed828cd24", + "project_id": "", "device_configs": [ { "device_name": "R-1", @@ -476,7 +476,7 @@ if __name__ == "__main__": # Example usage with new format device_commands = json.dumps( { - "project_id": "f32ebf3d-ef8c-4910-b0d6-566ed828cd24", + "project_id": "", "device_configs": [ { "device_name": "R-2", diff --git a/gns3server/agent/gns3_copilot/tools_v2/gns3_start_node.py b/gns3server/agent/gns3_copilot/tools_v2/gns3_start_node.py index 4a539d0fc..f07846cf6 100644 --- a/gns3server/agent/gns3_copilot/tools_v2/gns3_start_node.py +++ b/gns3server/agent/gns3_copilot/tools_v2/gns3_start_node.py @@ -355,7 +355,7 @@ if __name__ == "__main__": print("=== Testing single node startup ===") test_input_single = json.dumps( { - "project_id": "f32ebf3d-ef8c-4910-b0d6-566ed828cd24", # Replace with actual project UUID + "project_id": "", # Replace with actual project UUID "node_ids": [ "fbeda109-9a74-4d8c-a749-cc3847911a90" ], # Replace with actual node UUID @@ -369,7 +369,7 @@ if __name__ == "__main__": print("\n=== Testing multiple nodes startup ===") test_input_multiple = json.dumps( { - "project_id": "f32ebf3d-ef8c-4910-b0d6-566ed828cd24", # Replace with actual project UUID + "project_id": "", # Replace with actual project UUID "node_ids": [ "fbeda109-9a74-4d8c-a749-cc3847911a90", # Replace with actual node UUIDs "another-node-uuid-here", diff --git a/gns3server/agent/gns3_copilot/tools_v2/vpcs_tools_telnetlib3.py b/gns3server/agent/gns3_copilot/tools_v2/vpcs_tools_telnetlib3.py index 8d4ea969c..27b86ab52 100644 --- a/gns3server/agent/gns3_copilot/tools_v2/vpcs_tools_telnetlib3.py +++ b/gns3server/agent/gns3_copilot/tools_v2/vpcs_tools_telnetlib3.py @@ -65,7 +65,7 @@ class VPCSMultiCommands(BaseTool): **Input Format:** { - "project_id": "f32ebf3d-ef8c-4910-b0d6-566ed828cd24", + "project_id": "", "device_configs": [ { "device_name": "PC1", @@ -407,7 +407,7 @@ if __name__ == "__main__": # Example usage command_groups = json.dumps( { - "project_id": "f32ebf3d-ef8c-4910-b0d6-566ed828cd24", + "project_id": "", "device_configs": [ { "device_name": "PC1",