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.
2026-08-27 20:40:13 +03:00 · 2026-03-05 01:50:48 +08:00 · 2026-03-05 01:50:48 +08:00 · 7fcdb57615
commit 7fcdb57615
parent f8aa653ef9
6 changed files with 220 additions and 66 deletions
--- 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

--- 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

--- 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": "<PROJECT_UUID>",
            "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": "<PROJECT_UUID>",
            "device_configs": [
                {
                    "device_name": "R-1",
--- 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": "<PROJECT_UUID>",
            "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": "<PROJECT_UUID>",
            "device_configs": [
                {
                    "device_name": "R-2",
--- 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": "<PROJECT_UUID>",  # 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": "<PROJECT_UUID>",  # Replace with actual project UUID
            "node_ids": [
                "fbeda109-9a74-4d8c-a749-cc3847911a90",  # Replace with actual node UUIDs
                "another-node-uuid-here",
--- 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": "<PROJECT_UUID>",
            "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": "<PROJECT_UUID>",
            "device_configs": [
                {
                    "device_name": "PC1",