mirror of
https://gitee.com/mateos/mateclaw.git
synced 2026-09-13 03:13:41 +08:00
feat(agent): prefix 注入块统一 token 预算——记忆/Wiki 注入随模型有效窗口缩放,身份 prompt 超大告警
This commit is contained in:
matevip
parent
67fd74f7fb
commit
56737e197f
@ -38,6 +38,9 @@ import vip.mate.llm.chatmodel.ReasoningEffortResolver;
|
||||
import vip.mate.llm.model.ModelConfigEntity;
|
||||
import vip.mate.llm.model.ModelFamily;
|
||||
import vip.mate.llm.model.ModelProtocol;
|
||||
import vip.mate.agent.context.PrefixBudgetPlan;
|
||||
import vip.mate.agent.context.PrefixBudgetPlanner;
|
||||
import vip.mate.agent.context.TokenEstimator;
|
||||
import vip.mate.llm.model.ModelProviderEntity;
|
||||
import vip.mate.llm.probe.ModelContextWindowResolver;
|
||||
import vip.mate.llm.routing.ProviderModelRef;
|
||||
@ -99,6 +102,7 @@ public class AgentGraphBuilder {
|
||||
private final ModelConfigService modelConfigService;
|
||||
private final ModelProviderService modelProviderService;
|
||||
private final ModelContextWindowResolver contextWindowResolver;
|
||||
private final PrefixBudgetPlanner prefixBudgetPlanner;
|
||||
private final vip.mate.llm.service.ModelCapabilityService modelCapabilityService;
|
||||
private final ProviderRouter providerRouter;
|
||||
private final PlanningService planningService;
|
||||
@ -396,7 +400,22 @@ public class AgentGraphBuilder {
|
||||
}
|
||||
}
|
||||
|
||||
String enhancedPrompt = buildEnhancedPrompt(entity, builtinSearchEnabled);
|
||||
// Prefix injection budget: optional blocks (memory / wiki / skill
|
||||
// catalog / extension catalog / ledger) share a token budget scaled
|
||||
// to the model's effective window. The agent's own prompt and the
|
||||
// tool schemas are never truncated — they are subtracted from the
|
||||
// budget so the optional blocks absorb the squeeze.
|
||||
int basePromptTokens = TokenEstimator.estimateTokens(entity.getSystemPrompt());
|
||||
int toolSchemaTokens = TokenEstimator.estimateToolsTokens(toolSet.callbacks());
|
||||
PrefixBudgetPlan prefixBudgetPlan = prefixBudgetPlanner.plan(
|
||||
effectiveMaxInputTokens, basePromptTokens, toolSchemaTokens);
|
||||
if (basePromptTokens > prefixBudgetPlan.effectiveMaxTokens() / 2) {
|
||||
log.warn("Agent {} 的身份 prompt 约 {} tokens,已超过模型有效窗口 {} 的一半——"
|
||||
+ "系统不会截断用户自写的身份 prompt,请自行精简,否则小上下文模型可能无法响应",
|
||||
entity.getId(), basePromptTokens, prefixBudgetPlan.effectiveMaxTokens());
|
||||
}
|
||||
|
||||
String enhancedPrompt = buildEnhancedPrompt(entity, builtinSearchEnabled, prefixBudgetPlan.memoryTokens());
|
||||
|
||||
// Runtime skill-catalog renderer — captures this agent's bound skills,
|
||||
// effective tool allowlist, model window and workspace; invoked each
|
||||
@ -432,7 +451,8 @@ public class AgentGraphBuilder {
|
||||
log.info("Built StateGraph Plan-Execute agent: {} (maxIterations={}, tools={}, protocol={})",
|
||||
entity.getName(), maxIter, toolSet.size(), protocol.getId());
|
||||
} else {
|
||||
agent = buildReActAgent(toolSet, runtimeModel, maxIter, entity.getId(), skillCatalogRenderer);
|
||||
agent = buildReActAgent(toolSet, runtimeModel, maxIter, entity.getId(), skillCatalogRenderer,
|
||||
prefixBudgetPlan);
|
||||
// StateGraph 路径下工具调用由 ActionNode 控制，始终启用
|
||||
toolCallingEnabled = true;
|
||||
log.info("Built StateGraph ReAct agent: {} (maxIterations={}, tools={}, protocol={})",
|
||||
@ -519,11 +539,17 @@ public class AgentGraphBuilder {
|
||||
|
||||
StateGraphReActAgent buildReActAgent(AgentToolSet toolSet, ModelConfigEntity runtimeModel,
|
||||
int maxIter, Long agentId, SkillCatalogRenderer skillCatalogRenderer) {
|
||||
return buildReActAgent(toolSet, runtimeModel, maxIter, agentId, skillCatalogRenderer, null);
|
||||
}
|
||||
|
||||
StateGraphReActAgent buildReActAgent(AgentToolSet toolSet, ModelConfigEntity runtimeModel,
|
||||
int maxIter, Long agentId, SkillCatalogRenderer skillCatalogRenderer,
|
||||
PrefixBudgetPlan prefixBudgetPlan) {
|
||||
ChatModel chatModel = buildRuntimeChatModel(runtimeModel);
|
||||
ChatClient chatClient = ChatClient.create(chatModel);
|
||||
String reasoningEffort = resolveReasoningEffortForModel(runtimeModel);
|
||||
CompiledGraph compiledGraph = buildReActGraph(toolSet, chatModel, maxIter, reasoningEffort,
|
||||
runtimeModel, agentId, skillCatalogRenderer);
|
||||
runtimeModel, agentId, skillCatalogRenderer, prefixBudgetPlan);
|
||||
return new StateGraphReActAgent(chatClient, conversationService, compiledGraph,
|
||||
chatModel, conversationWindowManager, toolSet);
|
||||
}
|
||||
@ -849,6 +875,14 @@ public class AgentGraphBuilder {
|
||||
CompiledGraph buildReActGraph(AgentToolSet toolSet, ChatModel chatModel, int maxIterations,
|
||||
String reasoningEffort, ModelConfigEntity primaryModelConfig,
|
||||
Long agentId, SkillCatalogRenderer skillCatalogRenderer) {
|
||||
return buildReActGraph(toolSet, chatModel, maxIterations, reasoningEffort,
|
||||
primaryModelConfig, agentId, skillCatalogRenderer, null);
|
||||
}
|
||||
|
||||
CompiledGraph buildReActGraph(AgentToolSet toolSet, ChatModel chatModel, int maxIterations,
|
||||
String reasoningEffort, ModelConfigEntity primaryModelConfig,
|
||||
Long agentId, SkillCatalogRenderer skillCatalogRenderer,
|
||||
PrefixBudgetPlan prefixBudgetPlan) {
|
||||
try {
|
||||
List<vip.mate.llm.failover.FallbackEntry> fallbackChain = buildFallbackChain(primaryModelConfig, agentId);
|
||||
NodeStreamingChatHelper streamingHelper = new NodeStreamingChatHelper(
|
||||
@ -885,6 +919,7 @@ public class AgentGraphBuilder {
|
||||
supportsReasoningEffort,
|
||||
streamingHelper, conversationWindowManager, streamTracker, 0, wikiContextService,
|
||||
skillCatalogRenderer, toolDisclosureService, progressLedgerService);
|
||||
reasoningNode.setPrefixBudgetPlan(prefixBudgetPlan);
|
||||
ActionNode actionNode = new ActionNode(executor, streamTracker);
|
||||
ObservationProcessor observationProcessor = new ObservationProcessor(graphObservationProperties);
|
||||
ObservationNode observationNode = new ObservationNode(observationProcessor, streamTracker);
|
||||
@ -1473,6 +1508,10 @@ public class AgentGraphBuilder {
|
||||
""";
|
||||
|
||||
private String buildEnhancedPrompt(AgentEntity entity, boolean builtinSearchEnabled) {
|
||||
return buildEnhancedPrompt(entity, builtinSearchEnabled, Integer.MAX_VALUE);
|
||||
}
|
||||
|
||||
private String buildEnhancedPrompt(AgentEntity entity, boolean builtinSearchEnabled, int memoryBudgetTokens) {
|
||||
// The agent's own systemPrompt encodes its identity (role / goal /
|
||||
// backstory). The memory block from workspace files (AGENTS.md, SOUL.md,
|
||||
// PROFILE.md, MEMORY.md, ...) augments that identity with durable
|
||||
@ -1481,7 +1520,7 @@ public class AgentGraphBuilder {
|
||||
// dropped the identity prompt, so editor-side identity changes never
|
||||
// reached runtime if the agent had any workspace files.
|
||||
String identityPrompt = entity.getSystemPrompt() != null ? entity.getSystemPrompt().trim() : "";
|
||||
String memoryPrompt = memoryManager.buildSystemPromptBlock(entity.getId());
|
||||
String memoryPrompt = memoryManager.buildSystemPromptBlock(entity.getId(), memoryBudgetTokens);
|
||||
StringBuilder basePromptBuilder = new StringBuilder();
|
||||
if (!identityPrompt.isEmpty()) {
|
||||
basePromptBuilder.append(identityPrompt);
|
||||
|
||||
@ -0,0 +1,40 @@
|
||||
package vip.mate.agent.context;
|
||||
|
||||
/**
|
||||
* Per-agent-build token budget for the prompt prefix's optional injection
|
||||
* blocks. Produced once by {@link PrefixBudgetPlanner} when the agent graph
|
||||
* is assembled (the inputs — effective window, base prompt, tool schemas —
|
||||
* are all stable per build) and handed to each injection site.
|
||||
*
|
||||
* <p>{@code enabled == false} means budgeting is switched off: every budget
|
||||
* field holds {@link Integer#MAX_VALUE} and consumers keep their existing
|
||||
* absolute caps untouched.
|
||||
*/
|
||||
public record PrefixBudgetPlan(
|
||||
boolean enabled,
|
||||
int effectiveMaxTokens,
|
||||
Profile profile,
|
||||
int injectionBudgetTokens,
|
||||
int memoryTokens,
|
||||
int wikiTokens,
|
||||
int skillCatalogTokens,
|
||||
int extensionCatalogTokens,
|
||||
int ledgerTokens) {
|
||||
|
||||
/** Window-size tier. Small windows tighten the injection ratio. */
|
||||
public enum Profile {
|
||||
/** Regular window — budget shares rarely bind (absolute caps are smaller). */
|
||||
NORMAL,
|
||||
/** Window below the compact threshold — tightened injection ratio. */
|
||||
COMPACT,
|
||||
/** Window below the minimal threshold — injection cut to the bone. */
|
||||
MINIMAL
|
||||
}
|
||||
|
||||
/** Budgeting disabled — unlimited budgets, previous behavior. */
|
||||
public static PrefixBudgetPlan unlimited(int effectiveMaxTokens) {
|
||||
return new PrefixBudgetPlan(false, effectiveMaxTokens, Profile.NORMAL,
|
||||
Integer.MAX_VALUE, Integer.MAX_VALUE, Integer.MAX_VALUE,
|
||||
Integer.MAX_VALUE, Integer.MAX_VALUE, Integer.MAX_VALUE);
|
||||
}
|
||||
}
|
||||
@ -0,0 +1,105 @@
|
||||
package vip.mate.agent.context;
|
||||
|
||||
import lombok.RequiredArgsConstructor;
|
||||
import lombok.extern.slf4j.Slf4j;
|
||||
import org.springframework.boot.context.properties.EnableConfigurationProperties;
|
||||
import org.springframework.stereotype.Component;
|
||||
import vip.mate.config.ConversationWindowProperties;
|
||||
import vip.mate.config.PrefixBudgetProperties;
|
||||
|
||||
/**
|
||||
* Computes the {@link PrefixBudgetPlan} for one agent build: how many tokens
|
||||
* each optional prefix injection block (memory / wiki / skill catalog /
|
||||
* extension catalog / progress ledger) may spend, scaled to the model's
|
||||
* effective context window.
|
||||
*
|
||||
* <pre>
|
||||
* injectionBudget = max(0, effectiveMax × ratio(profile)
|
||||
* − basePromptTokens − toolSchemaTokens)
|
||||
* block budget = injectionBudget × normalizedShare(block)
|
||||
* </pre>
|
||||
*
|
||||
* The agent's own prompt and the tool schemas are never truncated here —
|
||||
* they are subtracted from the injection budget so the optional blocks
|
||||
* absorb the squeeze. On large windows the shares far exceed each block's
|
||||
* absolute cap, so behavior is byte-identical to the pre-budget code.
|
||||
*/
|
||||
@Slf4j
|
||||
@Component
|
||||
@RequiredArgsConstructor
|
||||
@EnableConfigurationProperties(PrefixBudgetProperties.class)
|
||||
public class PrefixBudgetPlanner {
|
||||
|
||||
private final PrefixBudgetProperties properties;
|
||||
private final ConversationWindowProperties windowProperties;
|
||||
|
||||
/**
|
||||
* @param effectiveMaxInputTokens the model's effective window (explicit
|
||||
* config or probed); null/0 falls back to
|
||||
* the global default
|
||||
* @param basePromptTokens estimated tokens of the agent's own
|
||||
* identity prompt (before memory/guidance)
|
||||
* @param toolSchemaTokens estimated tokens of the advertised tool
|
||||
* schemas
|
||||
*/
|
||||
public PrefixBudgetPlan plan(Integer effectiveMaxInputTokens, int basePromptTokens, int toolSchemaTokens) {
|
||||
int effectiveMax = (effectiveMaxInputTokens != null && effectiveMaxInputTokens > 0)
|
||||
? effectiveMaxInputTokens : windowProperties.getDefaultMaxInputTokens();
|
||||
if (!properties.isEnabled()) {
|
||||
return PrefixBudgetPlan.unlimited(effectiveMax);
|
||||
}
|
||||
|
||||
PrefixBudgetPlan.Profile profile = profileFor(effectiveMax);
|
||||
double ratio = switch (profile) {
|
||||
case NORMAL -> properties.getInjectionRatio();
|
||||
case COMPACT -> properties.getCompactInjectionRatio();
|
||||
case MINIMAL -> properties.getMinimalInjectionRatio();
|
||||
};
|
||||
|
||||
int injectionBudget = Math.max(0,
|
||||
(int) (effectiveMax * ratio) - Math.max(0, basePromptTokens) - Math.max(0, toolSchemaTokens));
|
||||
|
||||
PrefixBudgetProperties.Shares shares = properties.getShares();
|
||||
double sum = shares.getMemory() + shares.getWiki() + shares.getSkill()
|
||||
+ shares.getExtensionCatalog() + shares.getLedger();
|
||||
if (sum <= 0) {
|
||||
sum = 1.0;
|
||||
}
|
||||
|
||||
PrefixBudgetPlan plan = new PrefixBudgetPlan(
|
||||
true, effectiveMax, profile, injectionBudget,
|
||||
(int) (injectionBudget * shares.getMemory() / sum),
|
||||
(int) (injectionBudget * shares.getWiki() / sum),
|
||||
(int) (injectionBudget * shares.getSkill() / sum),
|
||||
(int) (injectionBudget * shares.getExtensionCatalog() / sum),
|
||||
(int) (injectionBudget * shares.getLedger() / sum));
|
||||
|
||||
if (profile != PrefixBudgetPlan.Profile.NORMAL) {
|
||||
log.info("[PrefixBudget] 窗口 {} tokens 进入 {} 档:注入预算 {} tokens"
|
||||
+ "(memory={}, wiki={}, skill={}, extCatalog={}, ledger={})",
|
||||
effectiveMax, profile, injectionBudget,
|
||||
plan.memoryTokens(), plan.wikiTokens(), plan.skillCatalogTokens(),
|
||||
plan.extensionCatalogTokens(), plan.ledgerTokens());
|
||||
}
|
||||
return plan;
|
||||
}
|
||||
|
||||
/** Compaction trigger ratio for this window size (small windows fill up before compacting). */
|
||||
public double compactTriggerRatioFor(int effectiveMaxTokens, double defaultRatio) {
|
||||
if (!properties.isEnabled()) {
|
||||
return defaultRatio;
|
||||
}
|
||||
return profileFor(effectiveMaxTokens) == PrefixBudgetPlan.Profile.NORMAL
|
||||
? defaultRatio : properties.getCompactTriggerRatioOverride();
|
||||
}
|
||||
|
||||
private PrefixBudgetPlan.Profile profileFor(int effectiveMax) {
|
||||
if (effectiveMax < properties.getMinimalThresholdTokens()) {
|
||||
return PrefixBudgetPlan.Profile.MINIMAL;
|
||||
}
|
||||
if (effectiveMax < properties.getCompactThresholdTokens()) {
|
||||
return PrefixBudgetPlan.Profile.COMPACT;
|
||||
}
|
||||
return PrefixBudgetPlan.Profile.NORMAL;
|
||||
}
|
||||
}
|
||||
@ -22,6 +22,7 @@ import vip.mate.llm.chatmodel.ThinkingLevelHolder;
|
||||
import vip.mate.agent.graph.NodeStreamingChatHelper;
|
||||
import vip.mate.agent.context.ConversationWindowManager;
|
||||
import vip.mate.agent.context.LoopBudgetConfig;
|
||||
import vip.mate.agent.context.PrefixBudgetPlan;
|
||||
import vip.mate.agent.context.LoopMessageBudgeter;
|
||||
import vip.mate.agent.context.RuntimeContextInjector;
|
||||
import vip.mate.agent.context.TokenEstimator;
|
||||
@ -350,6 +351,18 @@ public class ReasoningNode implements NodeAction {
|
||||
*/
|
||||
private final vip.mate.agent.progress.ProgressLedgerService progressLedgerService;
|
||||
|
||||
/**
|
||||
* Token budget for the optional prefix injection blocks, computed at
|
||||
* agent-build time against the model's effective context window. Null
|
||||
* when the graph was assembled without budgeting (tests, legacy paths) —
|
||||
* all injection sites then keep their previous absolute-cap behavior.
|
||||
*/
|
||||
private PrefixBudgetPlan prefixBudgetPlan;
|
||||
|
||||
public void setPrefixBudgetPlan(PrefixBudgetPlan prefixBudgetPlan) {
|
||||
this.prefixBudgetPlan = prefixBudgetPlan;
|
||||
}
|
||||
|
||||
public ReasoningNode(ChatModel chatModel, AgentToolSet toolSet, String reasoningEffort,
|
||||
NodeStreamingChatHelper streamingHelper,
|
||||
ConversationWindowManager conversationWindowManager,
|
||||
@ -470,6 +483,12 @@ public class ReasoningNode implements NodeAction {
|
||||
* otherwise the documented fallback.
|
||||
*/
|
||||
private int loopContextWindowTokens() {
|
||||
// Per-model effective window (explicit config or probed) beats the
|
||||
// global default — the loop budgeter is otherwise blind to small
|
||||
// local models and never trims for them.
|
||||
if (prefixBudgetPlan != null && prefixBudgetPlan.effectiveMaxTokens() > 0) {
|
||||
return prefixBudgetPlan.effectiveMaxTokens();
|
||||
}
|
||||
if (conversationWindowManager != null) {
|
||||
int v = conversationWindowManager.getDefaultMaxInputTokens();
|
||||
if (v > 0) return v;
|
||||
@ -1119,7 +1138,9 @@ public class ReasoningNode implements NodeAction {
|
||||
if (!projectRecalled && wikiContextService != null && agentIdStr != null && !agentIdStr.isEmpty()) {
|
||||
try {
|
||||
Long parsedAgentId = Long.parseLong(agentIdStr);
|
||||
String wikiRelevant = wikiContextService.buildRelevantContext(parsedAgentId, userMsg);
|
||||
Integer wikiBudgetTokens = (prefixBudgetPlan != null && prefixBudgetPlan.enabled())
|
||||
? prefixBudgetPlan.wikiTokens() : null;
|
||||
String wikiRelevant = wikiContextService.buildRelevantContext(parsedAgentId, userMsg, wikiBudgetTokens);
|
||||
if (wikiRelevant != null && !wikiRelevant.isBlank()) {
|
||||
prefix.add(new UserMessage(wikiRelevant));
|
||||
}
|
||||
|
||||
@ -0,0 +1,61 @@
|
||||
package vip.mate.config;
|
||||
|
||||
import lombok.Data;
|
||||
import org.springframework.boot.context.properties.ConfigurationProperties;
|
||||
|
||||
/**
|
||||
* Budget configuration for the prompt prefix's optional injection blocks
|
||||
* (memory, wiki relevance, skill catalog, extension-tool catalog, progress
|
||||
* ledger).
|
||||
*
|
||||
* <p>Each block already has its own absolute cap (chars), but those caps are
|
||||
* mutually blind and sized for large cloud models — stacked together they
|
||||
* easily exceed a local 8k/16k window on the very first request, when there
|
||||
* is no history to compact. This budget scales every block against the
|
||||
* model's effective context window instead: the enforced limit is always
|
||||
* {@code min(block's own cap, its share of the injection budget)}, so large
|
||||
* windows behave exactly as before while small windows shrink each block
|
||||
* proportionally.
|
||||
*/
|
||||
@Data
|
||||
@ConfigurationProperties(prefix = "mateclaw.context.prefix-budget")
|
||||
public class PrefixBudgetProperties {
|
||||
|
||||
/** Kill switch — false restores the previous per-block-absolute-cap behavior. */
|
||||
private boolean enabled = true;
|
||||
|
||||
/** Fraction of the effective window granted to prefix injection blocks (normal profile). */
|
||||
private double injectionRatio = 0.35;
|
||||
|
||||
/** Injection ratio when the window is below {@link #compactThresholdTokens}. */
|
||||
private double compactInjectionRatio = 0.25;
|
||||
|
||||
/** Injection ratio when the window is below {@link #minimalThresholdTokens}. */
|
||||
private double minimalInjectionRatio = 0.15;
|
||||
|
||||
/** Windows below this enter the compact profile. */
|
||||
private int compactThresholdTokens = 32768;
|
||||
|
||||
/** Windows below this enter the minimal profile. */
|
||||
private int minimalThresholdTokens = 8192;
|
||||
|
||||
/**
|
||||
* Compaction trigger ratio used for compact / minimal profiles instead of
|
||||
* the global {@code compactTriggerRatio}. Small windows should be used up
|
||||
* before summarizing — compacting at 75% of an 8k window wastes what
|
||||
* little room there is.
|
||||
*/
|
||||
private double compactTriggerRatioOverride = 0.85;
|
||||
|
||||
/** Relative shares of the injection budget. Normalized at plan time. */
|
||||
private Shares shares = new Shares();
|
||||
|
||||
@Data
|
||||
public static class Shares {
|
||||
private double memory = 0.35;
|
||||
private double wiki = 0.30;
|
||||
private double skill = 0.20;
|
||||
private double extensionCatalog = 0.10;
|
||||
private double ledger = 0.05;
|
||||
}
|
||||
}
|
||||
@ -3,6 +3,7 @@ package vip.mate.memory.spi;
|
||||
import io.micrometer.core.instrument.MeterRegistry;
|
||||
import lombok.extern.slf4j.Slf4j;
|
||||
import org.springframework.stereotype.Component;
|
||||
import vip.mate.agent.context.TokenEstimator;
|
||||
import vip.mate.memory.MemoryProperties;
|
||||
import vip.mate.memory.spi.decorator.MetricsMemoryProvider;
|
||||
import vip.mate.memory.spi.decorator.RetryableMemoryProvider;
|
||||
@ -81,13 +82,41 @@ public class MemoryManager {
|
||||
* Called once at agent build time (snapshot frozen for session).
|
||||
*/
|
||||
public String buildSystemPromptBlock(Long agentId) {
|
||||
return buildSystemPromptBlock(agentId, Integer.MAX_VALUE);
|
||||
}
|
||||
|
||||
/**
|
||||
* Budgeted variant: providers keep their own per-block caps, but the
|
||||
* combined output additionally may not exceed {@code budgetTokens}
|
||||
* (estimated). Provider order is priority order — once the budget is
|
||||
* spent, later providers are dropped whole and a partially fitting block
|
||||
* is truncated at a line boundary. Small local context windows need this:
|
||||
* the individual caps are sized for large cloud models and stack up past
|
||||
* an 8k/16k window on their own.
|
||||
*/
|
||||
public String buildSystemPromptBlock(Long agentId, int budgetTokens) {
|
||||
List<String> blocks = new ArrayList<>();
|
||||
int usedTokens = 0;
|
||||
for (MemoryProvider provider : providers) {
|
||||
try {
|
||||
String block = provider.systemPromptBlock(agentId);
|
||||
if (block != null && !block.isBlank()) {
|
||||
blocks.add(block);
|
||||
if (block == null || block.isBlank()) {
|
||||
continue;
|
||||
}
|
||||
int blockTokens = TokenEstimator.estimateTokens(block);
|
||||
if (usedTokens + blockTokens <= budgetTokens) {
|
||||
blocks.add(block);
|
||||
usedTokens += blockTokens;
|
||||
continue;
|
||||
}
|
||||
int remaining = budgetTokens - usedTokens;
|
||||
String truncated = truncateToTokenBudget(block, remaining);
|
||||
if (!truncated.isBlank()) {
|
||||
blocks.add(truncated + "\n\n[memory truncated to fit the model context window]");
|
||||
}
|
||||
log.info("[MemoryManager] Memory block budget {} tokens reached at provider '{}' — "
|
||||
+ "remaining providers dropped from the system prompt", budgetTokens, provider.id());
|
||||
break;
|
||||
} catch (Exception e) {
|
||||
log.warn("[MemoryManager] Provider '{}' systemPromptBlock() failed: {}",
|
||||
provider.id(), e.getMessage());
|
||||
@ -96,6 +125,27 @@ public class MemoryManager {
|
||||
return String.join("\n\n", blocks);
|
||||
}
|
||||
|
||||
/** Trim to the last full line that fits the token budget; empty when nothing fits. */
|
||||
private static String truncateToTokenBudget(String block, int budgetTokens) {
|
||||
if (budgetTokens <= 0) {
|
||||
return "";
|
||||
}
|
||||
StringBuilder kept = new StringBuilder();
|
||||
int usedTokens = 0;
|
||||
for (String line : block.split("\n", -1)) {
|
||||
int lineTokens = TokenEstimator.estimateTokens(line) + 1;
|
||||
if (usedTokens + lineTokens > budgetTokens) {
|
||||
break;
|
||||
}
|
||||
if (kept.length() > 0) {
|
||||
kept.append('\n');
|
||||
}
|
||||
kept.append(line);
|
||||
usedTokens += lineTokens;
|
||||
}
|
||||
return kept.toString();
|
||||
}
|
||||
|
||||
// ==================== Prefetch / Recall ====================
|
||||
|
||||
/**
|
||||
|
||||
@ -3,6 +3,7 @@ package vip.mate.wiki.service;
|
||||
import lombok.RequiredArgsConstructor;
|
||||
import lombok.extern.slf4j.Slf4j;
|
||||
import org.springframework.stereotype.Service;
|
||||
import vip.mate.agent.context.TokenEstimator;
|
||||
import vip.mate.wiki.WikiProperties;
|
||||
import vip.mate.wiki.dto.PageSearchResult;
|
||||
import vip.mate.wiki.model.WikiKnowledgeBaseEntity;
|
||||
@ -47,9 +48,22 @@ public class WikiContextService {
|
||||
* returns snippet + reason instead of just summary.
|
||||
*/
|
||||
public String buildRelevantContext(Long agentId, String userMessage) {
|
||||
return buildRelevantContext(agentId, userMessage, null);
|
||||
}
|
||||
|
||||
/**
|
||||
* Budgeted variant: in addition to the absolute {@code maxContextChars}
|
||||
* cap, the injected block may not exceed {@code budgetTokens} (estimated).
|
||||
* Null budget keeps the previous chars-only behavior. Small local context
|
||||
* windows need this — the chars cap is sized for large cloud models.
|
||||
*/
|
||||
public String buildRelevantContext(Long agentId, String userMessage, Integer budgetTokens) {
|
||||
if (!properties.isEnabled() || userMessage == null || userMessage.isBlank()) {
|
||||
return "";
|
||||
}
|
||||
if (budgetTokens != null && budgetTokens <= 0) {
|
||||
return "";
|
||||
}
|
||||
|
||||
// Skip retrieval for continuation / acknowledgement turns — observed
|
||||
// in production: a user reply of "继续" produced top-5 hits dominated
|
||||
@ -94,15 +108,26 @@ public class WikiContextService {
|
||||
"e.g. 「来源：[[页面标题]]」or「(来源：页面标题)」.]\n\n");
|
||||
int totalChars = 0;
|
||||
int maxChars = properties.getMaxContextChars();
|
||||
int totalTokens = TokenEstimator.estimateTokens(sb.toString());
|
||||
boolean anyEntry = false;
|
||||
|
||||
for (PageSearchResult hit : hits) {
|
||||
String entry = buildContextEntry(hit);
|
||||
if (totalChars + entry.length() > maxChars) {
|
||||
int entryTokens = TokenEstimator.estimateTokens(entry);
|
||||
if (totalChars + entry.length() > maxChars
|
||||
|| (budgetTokens != null && totalTokens + entryTokens > budgetTokens)) {
|
||||
sb.append("- ... (use wiki_search_pages for more)\n");
|
||||
break;
|
||||
}
|
||||
sb.append(entry);
|
||||
totalChars += entry.length();
|
||||
totalTokens += entryTokens;
|
||||
anyEntry = true;
|
||||
}
|
||||
// A budget too small for even one entry yields a header-only block —
|
||||
// pure overhead. Skip the injection entirely; wiki tools stay usable.
|
||||
if (!anyEntry && budgetTokens != null) {
|
||||
return "";
|
||||
}
|
||||
sb.append("</wiki-relevant>");
|
||||
return sb.toString();
|
||||
|
||||
@ -0,0 +1,121 @@
|
||||
package vip.mate.agent.context;
|
||||
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
import org.junit.jupiter.api.DisplayName;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import vip.mate.config.ConversationWindowProperties;
|
||||
import vip.mate.config.PrefixBudgetProperties;
|
||||
|
||||
import static org.junit.jupiter.api.Assertions.*;
|
||||
|
||||
/**
|
||||
* Unit tests for {@link PrefixBudgetPlanner} — profile selection, share
|
||||
* allocation, and the disabled/rollback path.
|
||||
*/
|
||||
class PrefixBudgetPlannerTest {
|
||||
|
||||
private PrefixBudgetProperties properties;
|
||||
private ConversationWindowProperties windowProperties;
|
||||
private PrefixBudgetPlanner planner;
|
||||
|
||||
@BeforeEach
|
||||
void setUp() {
|
||||
properties = new PrefixBudgetProperties();
|
||||
windowProperties = new ConversationWindowProperties();
|
||||
planner = new PrefixBudgetPlanner(properties, windowProperties);
|
||||
}
|
||||
|
||||
@Test
|
||||
@DisplayName("large window → NORMAL profile, budget = max*ratio − base − tools")
|
||||
void normalProfile() {
|
||||
PrefixBudgetPlan plan = planner.plan(128000, 2000, 3000);
|
||||
assertEquals(PrefixBudgetPlan.Profile.NORMAL, plan.profile());
|
||||
assertEquals(128000, plan.effectiveMaxTokens());
|
||||
assertEquals((int) (128000 * 0.35) - 2000 - 3000, plan.injectionBudgetTokens());
|
||||
assertTrue(plan.enabled());
|
||||
}
|
||||
|
||||
@Test
|
||||
@DisplayName("16k window → COMPACT profile with tightened ratio")
|
||||
void compactProfile() {
|
||||
PrefixBudgetPlan plan = planner.plan(16384, 1000, 1000);
|
||||
assertEquals(PrefixBudgetPlan.Profile.COMPACT, plan.profile());
|
||||
assertEquals((int) (16384 * 0.25) - 2000, plan.injectionBudgetTokens());
|
||||
}
|
||||
|
||||
@Test
|
||||
@DisplayName("4k window → MINIMAL profile")
|
||||
void minimalProfile() {
|
||||
PrefixBudgetPlan plan = planner.plan(4096, 500, 500);
|
||||
assertEquals(PrefixBudgetPlan.Profile.MINIMAL, plan.profile());
|
||||
assertEquals(Math.max(0, (int) (4096 * 0.15) - 1000), plan.injectionBudgetTokens());
|
||||
}
|
||||
|
||||
@Test
|
||||
@DisplayName("oversized base prompt + tools clamp the budget to zero, never negative")
|
||||
void budgetNeverNegative() {
|
||||
PrefixBudgetPlan plan = planner.plan(8192, 9000, 5000);
|
||||
assertEquals(0, plan.injectionBudgetTokens());
|
||||
assertEquals(0, plan.memoryTokens());
|
||||
assertEquals(0, plan.wikiTokens());
|
||||
}
|
||||
|
||||
@Test
|
||||
@DisplayName("shares split the injection budget and sum to at most the budget")
|
||||
void sharesSplitBudget() {
|
||||
PrefixBudgetPlan plan = planner.plan(128000, 0, 0);
|
||||
int budget = plan.injectionBudgetTokens();
|
||||
// ±1 token tolerance: share division is floating point.
|
||||
assertEquals(budget * 0.35, plan.memoryTokens(), 1.0);
|
||||
assertEquals(budget * 0.30, plan.wikiTokens(), 1.0);
|
||||
assertEquals(budget * 0.20, plan.skillCatalogTokens(), 1.0);
|
||||
assertEquals(budget * 0.10, plan.extensionCatalogTokens(), 1.0);
|
||||
assertEquals(budget * 0.05, plan.ledgerTokens(), 1.0);
|
||||
int sum = plan.memoryTokens() + plan.wikiTokens() + plan.skillCatalogTokens()
|
||||
+ plan.extensionCatalogTokens() + plan.ledgerTokens();
|
||||
assertTrue(sum <= budget);
|
||||
}
|
||||
|
||||
@Test
|
||||
@DisplayName("shares not summing to 1 are normalized")
|
||||
void sharesNormalized() {
|
||||
properties.getShares().setMemory(2.0);
|
||||
properties.getShares().setWiki(2.0);
|
||||
properties.getShares().setSkill(0.0);
|
||||
properties.getShares().setExtensionCatalog(0.0);
|
||||
properties.getShares().setLedger(0.0);
|
||||
PrefixBudgetPlan plan = planner.plan(128000, 0, 0);
|
||||
assertEquals(plan.injectionBudgetTokens() / 2, plan.memoryTokens());
|
||||
assertEquals(plan.injectionBudgetTokens() / 2, plan.wikiTokens());
|
||||
assertEquals(0, plan.skillCatalogTokens());
|
||||
}
|
||||
|
||||
@Test
|
||||
@DisplayName("disabled → unlimited plan, previous behavior")
|
||||
void disabledYieldsUnlimited() {
|
||||
properties.setEnabled(false);
|
||||
PrefixBudgetPlan plan = planner.plan(8192, 100, 100);
|
||||
assertFalse(plan.enabled());
|
||||
assertEquals(Integer.MAX_VALUE, plan.memoryTokens());
|
||||
assertEquals(Integer.MAX_VALUE, plan.wikiTokens());
|
||||
assertEquals(8192, plan.effectiveMaxTokens());
|
||||
}
|
||||
|
||||
@Test
|
||||
@DisplayName("null/zero effective window falls back to the global default")
|
||||
void nullWindowFallsBackToGlobalDefault() {
|
||||
PrefixBudgetPlan plan = planner.plan(null, 0, 0);
|
||||
assertEquals(windowProperties.getDefaultMaxInputTokens(), plan.effectiveMaxTokens());
|
||||
assertEquals(PrefixBudgetPlan.Profile.NORMAL, plan.profile());
|
||||
}
|
||||
|
||||
@Test
|
||||
@DisplayName("small windows compact later (0.85), large windows keep the default ratio")
|
||||
void compactTriggerRatioAdapts() {
|
||||
assertEquals(0.75, planner.compactTriggerRatioFor(128000, 0.75));
|
||||
assertEquals(0.85, planner.compactTriggerRatioFor(16384, 0.75));
|
||||
assertEquals(0.85, planner.compactTriggerRatioFor(4096, 0.75));
|
||||
properties.setEnabled(false);
|
||||
assertEquals(0.75, planner.compactTriggerRatioFor(4096, 0.75));
|
||||
}
|
||||
}
|
||||
@ -0,0 +1,96 @@
|
||||
package vip.mate.memory;
|
||||
|
||||
import io.micrometer.core.instrument.MeterRegistry;
|
||||
import org.junit.jupiter.api.DisplayName;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.springframework.beans.factory.ObjectProvider;
|
||||
import vip.mate.memory.spi.MemoryManager;
|
||||
import vip.mate.memory.spi.MemoryProvider;
|
||||
|
||||
import java.util.List;
|
||||
|
||||
import static org.junit.jupiter.api.Assertions.*;
|
||||
|
||||
/**
|
||||
* Tests for the token-budgeted system-prompt block assembly in
|
||||
* {@link MemoryManager} — provider order is priority order, later providers
|
||||
* drop whole once the budget is spent, a partially fitting block truncates
|
||||
* at a line boundary.
|
||||
*/
|
||||
class MemoryManagerBudgetTest {
|
||||
|
||||
/** CJK chars estimate ≈ 1 token each, which makes budgets easy to reason about. */
|
||||
private static final String BLOCK_A = "甲".repeat(300);
|
||||
private static final String BLOCK_B = "乙".repeat(300);
|
||||
|
||||
private static MemoryProvider provider(String id, String block) {
|
||||
return new MemoryProvider() {
|
||||
@Override
|
||||
public String id() {
|
||||
return id;
|
||||
}
|
||||
|
||||
@Override
|
||||
public String systemPromptBlock(Long agentId) {
|
||||
return block;
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
private static MemoryManager manager(MemoryProvider... providers) {
|
||||
ObjectProvider<MeterRegistry> noRegistry = new ObjectProvider<>() {
|
||||
@Override
|
||||
public MeterRegistry getObject(Object... args) {
|
||||
throw new UnsupportedOperationException();
|
||||
}
|
||||
|
||||
@Override
|
||||
public MeterRegistry getIfAvailable() {
|
||||
return null;
|
||||
}
|
||||
};
|
||||
return new MemoryManager(List.of(providers), new MemoryProperties(), noRegistry);
|
||||
}
|
||||
|
||||
@Test
|
||||
@DisplayName("unbudgeted call joins every provider block — previous behavior")
|
||||
void unbudgetedJoinsAll() {
|
||||
MemoryManager manager = manager(provider("a", BLOCK_A), provider("b", BLOCK_B));
|
||||
String result = manager.buildSystemPromptBlock(1L);
|
||||
assertTrue(result.contains(BLOCK_A));
|
||||
assertTrue(result.contains(BLOCK_B));
|
||||
}
|
||||
|
||||
@Test
|
||||
@DisplayName("budget exhausted after the first block → later provider dropped whole")
|
||||
void budgetDropsLaterProviders() {
|
||||
MemoryManager manager = manager(provider("a", BLOCK_A), provider("b", BLOCK_B));
|
||||
// 300 CJK chars ≈ 300 tokens; 350 fits block A but not A+B.
|
||||
String result = manager.buildSystemPromptBlock(1L, 350);
|
||||
assertTrue(result.contains(BLOCK_A));
|
||||
assertFalse(result.contains("乙"));
|
||||
}
|
||||
|
||||
@Test
|
||||
@DisplayName("single multi-line block over budget truncates at a line boundary with a marker")
|
||||
void oversizedBlockTruncatesAtLineBoundary() {
|
||||
String multiLine = ("行".repeat(100) + "\n").repeat(10).trim();
|
||||
MemoryManager manager = manager(provider("a", multiLine));
|
||||
String result = manager.buildSystemPromptBlock(1L, 350);
|
||||
assertTrue(result.contains("[memory truncated to fit the model context window]"));
|
||||
// Only whole lines are kept — every kept content line is the full 100-char line.
|
||||
for (String line : result.split("\n")) {
|
||||
if (line.startsWith("行")) {
|
||||
assertEquals(100, line.length());
|
||||
}
|
||||
}
|
||||
assertTrue(result.length() < multiLine.length());
|
||||
}
|
||||
|
||||
@Test
|
||||
@DisplayName("zero budget yields an empty block, not an exception")
|
||||
void zeroBudgetYieldsEmpty() {
|
||||
MemoryManager manager = manager(provider("a", BLOCK_A));
|
||||
assertEquals("", manager.buildSystemPromptBlock(1L, 0));
|
||||
}
|
||||
}
|
||||
@ -0,0 +1,87 @@
|
||||
package vip.mate.wiki.service;
|
||||
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
import org.junit.jupiter.api.DisplayName;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.mockito.Mockito;
|
||||
import vip.mate.wiki.WikiProperties;
|
||||
import vip.mate.wiki.dto.PageSearchResult;
|
||||
import vip.mate.wiki.model.WikiKnowledgeBaseEntity;
|
||||
|
||||
import java.util.List;
|
||||
|
||||
import static org.junit.jupiter.api.Assertions.*;
|
||||
import static org.mockito.ArgumentMatchers.any;
|
||||
import static org.mockito.ArgumentMatchers.anyInt;
|
||||
import static org.mockito.ArgumentMatchers.anyLong;
|
||||
import static org.mockito.ArgumentMatchers.anyString;
|
||||
|
||||
/**
|
||||
* Tests for the token-budgeted knowledge-base relevance injection in
|
||||
* {@link WikiContextService}.
|
||||
*/
|
||||
class WikiContextServiceBudgetTest {
|
||||
|
||||
private WikiKnowledgeBaseService kbService;
|
||||
private HybridRetriever hybridRetriever;
|
||||
private WikiProperties properties;
|
||||
private WikiContextService service;
|
||||
|
||||
@BeforeEach
|
||||
void setUp() {
|
||||
kbService = Mockito.mock(WikiKnowledgeBaseService.class);
|
||||
hybridRetriever = Mockito.mock(HybridRetriever.class);
|
||||
WikiPageService pageService = Mockito.mock(WikiPageService.class);
|
||||
properties = new WikiProperties();
|
||||
service = new WikiContextService(kbService, pageService, hybridRetriever, properties);
|
||||
|
||||
WikiKnowledgeBaseEntity kb = new WikiKnowledgeBaseEntity();
|
||||
kb.setId(7L);
|
||||
Mockito.when(kbService.resolvePrimaryKb(anyLong())).thenReturn(kb);
|
||||
}
|
||||
|
||||
private void stubHits(int count, int excerptChars) {
|
||||
List<PageSearchResult> hits = new java.util.ArrayList<>();
|
||||
for (int i = 0; i < count; i++) {
|
||||
hits.add(PageSearchResult.of("page-" + i, "页面" + i, null,
|
||||
"摘".repeat(excerptChars), List.of("keyword"), null, 1.0));
|
||||
}
|
||||
Mockito.when(hybridRetriever.search(any(), anyString(), anyString(), anyInt()))
|
||||
.thenReturn(hits);
|
||||
}
|
||||
|
||||
@Test
|
||||
@DisplayName("null budget keeps the chars-only behavior — all hits injected")
|
||||
void nullBudgetKeepsAllHits() {
|
||||
stubHits(3, 200);
|
||||
String result = service.buildRelevantContext(1L, "如何配置数据库连接", null);
|
||||
assertTrue(result.contains("page-0"));
|
||||
assertTrue(result.contains("page-2"));
|
||||
}
|
||||
|
||||
@Test
|
||||
@DisplayName("token budget cuts the tail hits and appends the search hint")
|
||||
void budgetCutsTailHits() {
|
||||
stubHits(3, 400); // each entry ≈ 400+ tokens (CJK)
|
||||
String result = service.buildRelevantContext(1L, "如何配置数据库连接", 600);
|
||||
assertTrue(result.contains("page-0"));
|
||||
assertFalse(result.contains("page-2"));
|
||||
assertTrue(result.contains("use wiki_search_pages for more"));
|
||||
}
|
||||
|
||||
@Test
|
||||
@DisplayName("budget too small for even one entry → injection skipped entirely")
|
||||
void tinyBudgetSkipsInjection() {
|
||||
stubHits(3, 400);
|
||||
String result = service.buildRelevantContext(1L, "如何配置数据库连接", 50);
|
||||
assertEquals("", result);
|
||||
}
|
||||
|
||||
@Test
|
||||
@DisplayName("zero or negative budget short-circuits without retrieval")
|
||||
void zeroBudgetShortCircuits() {
|
||||
String result = service.buildRelevantContext(1L, "如何配置数据库连接", 0);
|
||||
assertEquals("", result);
|
||||
Mockito.verifyNoInteractions(hybridRetriever);
|
||||
}
|
||||
}
|
||||
Loading…
Reference in New Issue
Block a user