fix(agent): include tools schema in context-window budget

This commit is contained in:
matevip 2026-05-02 15:40:03 +08:00
parent c2ebe51f63
commit eaacb3a78f
5 changed files with 101 additions and 9 deletions

View File

@ -303,7 +303,7 @@ public class AgentGraphBuilder {
String reasoningEffort = resolveReasoningEffortForModel(runtimeModel); String reasoningEffort = resolveReasoningEffortForModel(runtimeModel);
CompiledGraph compiledGraph = buildReActGraph(toolSet, chatModel, maxIter, reasoningEffort, runtimeModel, agentId); CompiledGraph compiledGraph = buildReActGraph(toolSet, chatModel, maxIter, reasoningEffort, runtimeModel, agentId);
return new StateGraphReActAgent(chatClient, conversationService, compiledGraph, return new StateGraphReActAgent(chatClient, conversationService, compiledGraph,
chatModel, conversationWindowManager); chatModel, conversationWindowManager, toolSet);
} }
StateGraphPlanExecuteAgent buildPlanExecuteAgent(AgentToolSet toolSet, ModelConfigEntity runtimeModel, int maxIter) { StateGraphPlanExecuteAgent buildPlanExecuteAgent(AgentToolSet toolSet, ModelConfigEntity runtimeModel, int maxIter) {
@ -317,7 +317,7 @@ public class AgentGraphBuilder {
String reasoningEffort = resolveReasoningEffortForModel(runtimeModel); String reasoningEffort = resolveReasoningEffortForModel(runtimeModel);
CompiledGraph graph = buildPlanExecuteGraph(toolSet, chatModel, maxIter, reasoningEffort, runtimeModel, agentId); CompiledGraph graph = buildPlanExecuteGraph(toolSet, chatModel, maxIter, reasoningEffort, runtimeModel, agentId);
return new StateGraphPlanExecuteAgent(chatClient, conversationService, graph, planningService, return new StateGraphPlanExecuteAgent(chatClient, conversationService, graph, planningService,
chatModel, conversationWindowManager); chatModel, conversationWindowManager, toolSet);
} }
CompiledGraph buildPlanExecuteGraph(AgentToolSet toolSet, ChatModel chatModel, int maxIterations, String reasoningEffort) { CompiledGraph buildPlanExecuteGraph(AgentToolSet toolSet, ChatModel chatModel, int maxIterations, String reasoningEffort) {

View File

@ -11,6 +11,7 @@ import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse; import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.ChatOptions; import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt; import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.tool.ToolCallback;
import com.alibaba.cloud.ai.dashscope.chat.DashScopeChatOptions; import com.alibaba.cloud.ai.dashscope.chat.DashScopeChatOptions;
import org.springframework.stereotype.Component; import org.springframework.stereotype.Component;
import vip.mate.agent.prompt.PromptLoader; import vip.mate.agent.prompt.PromptLoader;
@ -116,6 +117,23 @@ public class ConversationWindowManager {
String currentUserMessage, String currentUserMessage,
Integer maxInputTokens, ChatModel chatModel, Integer maxInputTokens, ChatModel chatModel,
String conversationId, Long agentId) { String conversationId, Long agentId) {
return fitToWindow(messages, systemPrompt, currentUserMessage,
maxInputTokens, chatModel, conversationId, agentId, null);
}
/**
* Same as the 7-arg overload but additionally accounts for the tool
* definitions sent on every LLM call. Without {@code toolCallbacks},
* the budget calculation underestimates the actual request size by the
* full size of the tools schema (often several thousand tokens for
* agents bound to multiple MCP servers), making compression fire too
* late and producing HTTP 400 once the request hits the model.
*/
public List<Message> fitToWindow(List<Message> messages, String systemPrompt,
String currentUserMessage,
Integer maxInputTokens, ChatModel chatModel,
String conversationId, Long agentId,
java.util.Collection<ToolCallback> toolCallbacks) {
if (messages == null || messages.isEmpty()) { if (messages == null || messages.isEmpty()) {
return messages; return messages;
} }
@ -127,20 +145,21 @@ public class ConversationWindowManager {
int systemTokens = TokenEstimator.estimateTokens(systemPrompt); int systemTokens = TokenEstimator.estimateTokens(systemPrompt);
int currentMsgTokens = TokenEstimator.estimateTokens(currentUserMessage) + TokenEstimator.PER_MESSAGE_OVERHEAD; int currentMsgTokens = TokenEstimator.estimateTokens(currentUserMessage) + TokenEstimator.PER_MESSAGE_OVERHEAD;
int historyTokens = TokenEstimator.estimateTokens(messages); int historyTokens = TokenEstimator.estimateTokens(messages);
int totalTokens = systemTokens + currentMsgTokens + historyTokens; int toolsTokens = TokenEstimator.estimateToolsTokens(toolCallbacks);
int totalTokens = systemTokens + currentMsgTokens + historyTokens + toolsTokens;
if (totalTokens <= triggerThreshold) { if (totalTokens <= triggerThreshold) {
return messages; return messages;
} }
log.info("[ConversationWindow] 超阈值: {} tokens (system={}, current={}, history={}) > {} 触发阈值 (max={}), conv={}", log.info("[ConversationWindow] 超阈值: {} tokens (system={}, current={}, history={}, tools={}) > {} 触发阈值 (max={}), conv={}",
totalTokens, systemTokens, currentMsgTokens, historyTokens, totalTokens, systemTokens, currentMsgTokens, historyTokens, toolsTokens,
triggerThreshold, effectiveMax, conversationId); triggerThreshold, effectiveMax, conversationId);
evictExpiredEntries(); evictExpiredEntries();
// 可用于历史的 token 预算 = max - system - currentMsg - 安全余量 // 可用于历史的 token 预算 = max - system - currentMsg - tools - 安全余量
int reservedTokens = systemTokens + currentMsgTokens + (int) (effectiveMax * 0.05); int reservedTokens = systemTokens + currentMsgTokens + toolsTokens + (int) (effectiveMax * 0.05);
// RFC-025 Change 1: reserve 硬封顶到 effectiveMax 50% // RFC-025 Change 1: reserve 硬封顶到 effectiveMax 50%
// 小上下文模型Ollama 16K本地 8KsystemTokens + currentMsgTokens 很容易 // 小上下文模型Ollama 16K本地 8KsystemTokens + currentMsgTokens 很容易
// 接近或超过 effectiveMax不封顶会让 historyBudget 变负数导致死循环压缩 // 接近或超过 effectiveMax不封顶会让 historyBudget 变负数导致死循环压缩

View File

@ -1,7 +1,10 @@
package vip.mate.agent.context; package vip.mate.agent.context;
import org.springframework.ai.chat.messages.Message; import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.ai.tool.definition.ToolDefinition;
import java.util.Collection;
import java.util.List; import java.util.List;
/** /**
@ -21,6 +24,13 @@ public final class TokenEstimator {
/** 每条消息的固定开销 tokenrole 标记、分隔符等) */ /** 每条消息的固定开销 tokenrole 标记、分隔符等) */
static final int PER_MESSAGE_OVERHEAD = 4; static final int PER_MESSAGE_OVERHEAD = 4;
/**
* Per-tool wrapper overhead: function/type:object boilerplate, name and
* description framing, parameters key, and JSON braces around the schema.
* Conservative slightly overestimates so budget guards don't underrun.
*/
static final int PER_TOOL_OVERHEAD = 12;
private TokenEstimator() { private TokenEstimator() {
} }
@ -71,6 +81,38 @@ public final class TokenEstimator {
.sum(); .sum();
} }
/**
* Estimate the token cost of the tool definitions sent on every LLM call
* (name + description + JSON inputSchema, plus per-tool wrapper overhead).
* <p>
* A heavily-bound agent (multiple MCP servers, many built-ins) can carry
* several thousand tokens of tool schema on every request leaving them
* out of the context-window budget makes compression decisions fire too
* late and on small models triggers HTTP 400 once the request actually
* goes out.
*/
public static int estimateToolsTokens(Collection<ToolCallback> callbacks) {
if (callbacks == null || callbacks.isEmpty()) {
return 0;
}
int total = 0;
for (ToolCallback cb : callbacks) {
if (cb == null) continue;
ToolDefinition def;
try {
def = cb.getToolDefinition();
} catch (Exception e) {
continue;
}
if (def == null) continue;
total += estimateTokens(def.name())
+ estimateTokens(def.description())
+ estimateTokens(def.inputSchema())
+ PER_TOOL_OVERHEAD;
}
return total;
}
/** /**
* 判断是否为 CJK 字符中日韩统一表意文字 + 常用标点 * 判断是否为 CJK 字符中日韩统一表意文字 + 常用标点
*/ */

View File

@ -53,15 +53,32 @@ public class StateGraphReActAgent extends BaseAgent implements StructuredStreamC
private final CompiledGraph compiledGraph; private final CompiledGraph compiledGraph;
private final org.springframework.ai.chat.model.ChatModel chatModel; private final org.springframework.ai.chat.model.ChatModel chatModel;
private final ConversationWindowManager conversationWindowManager; private final ConversationWindowManager conversationWindowManager;
/**
* Held only so {@link #buildInitialState} can include the tools schema in
* the context-window budget those bytes ride along on every LLM call
* and were previously ignored, making compression decisions fire late.
* Nullable for the legacy 5-arg constructor used by older tests.
*/
private final vip.mate.agent.AgentToolSet toolSet;
public StateGraphReActAgent(ChatClient chatClient, ConversationService conversationService, public StateGraphReActAgent(ChatClient chatClient, ConversationService conversationService,
CompiledGraph compiledGraph, CompiledGraph compiledGraph,
org.springframework.ai.chat.model.ChatModel chatModel, org.springframework.ai.chat.model.ChatModel chatModel,
ConversationWindowManager conversationWindowManager) { ConversationWindowManager conversationWindowManager) {
this(chatClient, conversationService, compiledGraph, chatModel,
conversationWindowManager, null);
}
public StateGraphReActAgent(ChatClient chatClient, ConversationService conversationService,
CompiledGraph compiledGraph,
org.springframework.ai.chat.model.ChatModel chatModel,
ConversationWindowManager conversationWindowManager,
vip.mate.agent.AgentToolSet toolSet) {
super(chatClient, conversationService); super(chatClient, conversationService);
this.compiledGraph = compiledGraph; this.compiledGraph = compiledGraph;
this.chatModel = chatModel; this.chatModel = chatModel;
this.conversationWindowManager = conversationWindowManager; this.conversationWindowManager = conversationWindowManager;
this.toolSet = toolSet;
} }
@Override @Override
@ -372,7 +389,8 @@ public class StateGraphReActAgent extends BaseAgent implements StructuredStreamC
maxInputTokens, maxInputTokens,
chatModel, chatModel,
conversationId, conversationId,
parsedAgentId); parsedAgentId,
toolSet != null ? toolSet.callbacks() : null);
} }
List<Message> messages = new ArrayList<>(historyMessages); List<Message> messages = new ArrayList<>(historyMessages);

View File

@ -45,16 +45,28 @@ public class StateGraphPlanExecuteAgent extends BaseAgent implements StructuredS
private final PlanningService planningService; private final PlanningService planningService;
private final org.springframework.ai.chat.model.ChatModel chatModel; private final org.springframework.ai.chat.model.ChatModel chatModel;
private final ConversationWindowManager conversationWindowManager; private final ConversationWindowManager conversationWindowManager;
/** Held only so context-window budget includes the tools schema. Nullable for legacy constructor. */
private final vip.mate.agent.AgentToolSet toolSet;
public StateGraphPlanExecuteAgent(ChatClient chatClient, ConversationService conversationService, public StateGraphPlanExecuteAgent(ChatClient chatClient, ConversationService conversationService,
CompiledGraph compiledGraph, PlanningService planningService, CompiledGraph compiledGraph, PlanningService planningService,
org.springframework.ai.chat.model.ChatModel chatModel, org.springframework.ai.chat.model.ChatModel chatModel,
ConversationWindowManager conversationWindowManager) { ConversationWindowManager conversationWindowManager) {
this(chatClient, conversationService, compiledGraph, planningService,
chatModel, conversationWindowManager, null);
}
public StateGraphPlanExecuteAgent(ChatClient chatClient, ConversationService conversationService,
CompiledGraph compiledGraph, PlanningService planningService,
org.springframework.ai.chat.model.ChatModel chatModel,
ConversationWindowManager conversationWindowManager,
vip.mate.agent.AgentToolSet toolSet) {
super(chatClient, conversationService); super(chatClient, conversationService);
this.compiledGraph = compiledGraph; this.compiledGraph = compiledGraph;
this.planningService = planningService; this.planningService = planningService;
this.chatModel = chatModel; this.chatModel = chatModel;
this.conversationWindowManager = conversationWindowManager; this.conversationWindowManager = conversationWindowManager;
this.toolSet = toolSet;
} }
@Override @Override
@ -254,7 +266,8 @@ public class StateGraphPlanExecuteAgent extends BaseAgent implements StructuredS
maxInputTokens, maxInputTokens,
chatModel, chatModel,
conversationId, conversationId,
parsedAgentId); parsedAgentId,
toolSet != null ? toolSet.callbacks() : null);
} }
List<Message> messages = new ArrayList<>(historyMessages); List<Message> messages = new ArrayList<>(historyMessages);