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);
CompiledGraph compiledGraph = buildReActGraph(toolSet, chatModel, maxIter, reasoningEffort, runtimeModel, agentId);
return new StateGraphReActAgent(chatClient, conversationService, compiledGraph,
chatModel, conversationWindowManager);
chatModel, conversationWindowManager, toolSet);
}
StateGraphPlanExecuteAgent buildPlanExecuteAgent(AgentToolSet toolSet, ModelConfigEntity runtimeModel, int maxIter) {
@ -317,7 +317,7 @@ public class AgentGraphBuilder {
String reasoningEffort = resolveReasoningEffortForModel(runtimeModel);
CompiledGraph graph = buildPlanExecuteGraph(toolSet, chatModel, maxIter, reasoningEffort, runtimeModel, agentId);
return new StateGraphPlanExecuteAgent(chatClient, conversationService, graph, planningService,
chatModel, conversationWindowManager);
chatModel, conversationWindowManager, toolSet);
}
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.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.tool.ToolCallback;
import com.alibaba.cloud.ai.dashscope.chat.DashScopeChatOptions;
import org.springframework.stereotype.Component;
import vip.mate.agent.prompt.PromptLoader;
@ -116,6 +117,23 @@ public class ConversationWindowManager {
String currentUserMessage,
Integer maxInputTokens, ChatModel chatModel,
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()) {
return messages;
}
@ -127,20 +145,21 @@ public class ConversationWindowManager {
int systemTokens = TokenEstimator.estimateTokens(systemPrompt);
int currentMsgTokens = TokenEstimator.estimateTokens(currentUserMessage) + TokenEstimator.PER_MESSAGE_OVERHEAD;
int historyTokens = TokenEstimator.estimateTokens(messages);
int totalTokens = systemTokens + currentMsgTokens + historyTokens;
int toolsTokens = TokenEstimator.estimateToolsTokens(toolCallbacks);
int totalTokens = systemTokens + currentMsgTokens + historyTokens + toolsTokens;
if (totalTokens <= triggerThreshold) {
return messages;
}
log.info("[ConversationWindow] 超阈值: {} tokens (system={}, current={}, history={}) > {} 触发阈值 (max={}), conv={}",
totalTokens, systemTokens, currentMsgTokens, historyTokens,
log.info("[ConversationWindow] 超阈值: {} tokens (system={}, current={}, history={}, tools={}) > {} 触发阈值 (max={}), conv={}",
totalTokens, systemTokens, currentMsgTokens, historyTokens, toolsTokens,
triggerThreshold, effectiveMax, conversationId);
evictExpiredEntries();
// 可用于历史的 token 预算 = max - system - currentMsg - 安全余量
int reservedTokens = systemTokens + currentMsgTokens + (int) (effectiveMax * 0.05);
// 可用于历史的 token 预算 = max - system - currentMsg - tools - 安全余量
int reservedTokens = systemTokens + currentMsgTokens + toolsTokens + (int) (effectiveMax * 0.05);
// RFC-025 Change 1: reserve 硬封顶到 effectiveMax 50%
// 小上下文模型Ollama 16K本地 8KsystemTokens + currentMsgTokens 很容易
// 接近或超过 effectiveMax不封顶会让 historyBudget 变负数导致死循环压缩

View File

@ -1,7 +1,10 @@
package vip.mate.agent.context;
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;
/**
@ -21,6 +24,13 @@ public final class TokenEstimator {
/** 每条消息的固定开销 tokenrole 标记、分隔符等) */
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() {
}
@ -71,6 +81,38 @@ public final class TokenEstimator {
.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 字符中日韩统一表意文字 + 常用标点
*/

View File

@ -53,15 +53,32 @@ public class StateGraphReActAgent extends BaseAgent implements StructuredStreamC
private final CompiledGraph compiledGraph;
private final org.springframework.ai.chat.model.ChatModel chatModel;
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,
CompiledGraph compiledGraph,
org.springframework.ai.chat.model.ChatModel chatModel,
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);
this.compiledGraph = compiledGraph;
this.chatModel = chatModel;
this.conversationWindowManager = conversationWindowManager;
this.toolSet = toolSet;
}
@Override
@ -372,7 +389,8 @@ public class StateGraphReActAgent extends BaseAgent implements StructuredStreamC
maxInputTokens,
chatModel,
conversationId,
parsedAgentId);
parsedAgentId,
toolSet != null ? toolSet.callbacks() : null);
}
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 org.springframework.ai.chat.model.ChatModel chatModel;
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,
CompiledGraph compiledGraph, PlanningService planningService,
org.springframework.ai.chat.model.ChatModel chatModel,
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);
this.compiledGraph = compiledGraph;
this.planningService = planningService;
this.chatModel = chatModel;
this.conversationWindowManager = conversationWindowManager;
this.toolSet = toolSet;
}
@Override
@ -254,7 +266,8 @@ public class StateGraphPlanExecuteAgent extends BaseAgent implements StructuredS
maxInputTokens,
chatModel,
conversationId,
parsedAgentId);
parsedAgentId,
toolSet != null ? toolSet.callbacks() : null);
}
List<Message> messages = new ArrayList<>(historyMessages);