mirror of
https://gitee.com/mateos/mateclaw.git
synced 2026-09-13 11:13:43 +08:00
feat(tool): 工具 schema 预算阈值门——超出窗口预算时按使用频度自动降级到扩展目录,enable_tool 可找回
This commit is contained in:
matevip
parent
56737e197f
commit
bf0d64e46a
@ -52,6 +52,7 @@ import vip.mate.skill.runtime.SkillCatalogRenderer;
|
||||
import vip.mate.skill.service.SkillService;
|
||||
import vip.mate.system.service.SystemSettingService;
|
||||
import vip.mate.tool.ToolRegistry;
|
||||
import vip.mate.tool.disclosure.ToolUsageRecencyTracker;
|
||||
import vip.mate.memory.spi.MemoryManager;
|
||||
import vip.mate.workspace.document.WorkspaceFileService;
|
||||
import vip.mate.tool.guard.service.ToolGuardService;
|
||||
@ -103,6 +104,7 @@ public class AgentGraphBuilder {
|
||||
private final ModelProviderService modelProviderService;
|
||||
private final ModelContextWindowResolver contextWindowResolver;
|
||||
private final PrefixBudgetPlanner prefixBudgetPlanner;
|
||||
private final ToolUsageRecencyTracker toolUsageRecencyTracker;
|
||||
private final vip.mate.llm.service.ModelCapabilityService modelCapabilityService;
|
||||
private final ProviderRouter providerRouter;
|
||||
private final PlanningService planningService;
|
||||
@ -428,10 +430,17 @@ public class AgentGraphBuilder {
|
||||
// in ReasoningNode; Plan-Execute keeps advertising every tool (it has no
|
||||
// action node to record enable_tool), so baking the catalog there would
|
||||
// describe an enable_tool flow that can never take effect.
|
||||
// Auto-demotion is likewise ReAct-only: hiding a tool from Plan-Execute
|
||||
// would remove it with no enable_tool path to recover it.
|
||||
boolean isPlanExecute = "plan_execute".equals(entity.getAgentType());
|
||||
Set<String> autoDemotedTools = Set.of();
|
||||
if (!isPlanExecute) {
|
||||
if (prefixBudgetPlan.enabled()) {
|
||||
autoDemotedTools = toolDisclosureService.computeAutoDemotions(
|
||||
toolSet, prefixBudgetPlan.toolSchemaBudgetTokens());
|
||||
}
|
||||
String extensionCatalog = toolDisclosureService.renderExtensionCatalog(
|
||||
toolSet, effectiveMaxInputTokens);
|
||||
toolSet, effectiveMaxInputTokens, autoDemotedTools);
|
||||
if (extensionCatalog != null && !extensionCatalog.isBlank()) {
|
||||
enhancedPrompt = enhancedPrompt + extensionCatalog;
|
||||
}
|
||||
@ -452,7 +461,7 @@ public class AgentGraphBuilder {
|
||||
entity.getName(), maxIter, toolSet.size(), protocol.getId());
|
||||
} else {
|
||||
agent = buildReActAgent(toolSet, runtimeModel, maxIter, entity.getId(), skillCatalogRenderer,
|
||||
prefixBudgetPlan);
|
||||
prefixBudgetPlan, autoDemotedTools);
|
||||
// StateGraph 路径下工具调用由 ActionNode 控制，始终启用
|
||||
toolCallingEnabled = true;
|
||||
log.info("Built StateGraph ReAct agent: {} (maxIterations={}, tools={}, protocol={})",
|
||||
@ -539,17 +548,17 @@ public class AgentGraphBuilder {
|
||||
|
||||
StateGraphReActAgent buildReActAgent(AgentToolSet toolSet, ModelConfigEntity runtimeModel,
|
||||
int maxIter, Long agentId, SkillCatalogRenderer skillCatalogRenderer) {
|
||||
return buildReActAgent(toolSet, runtimeModel, maxIter, agentId, skillCatalogRenderer, null);
|
||||
return buildReActAgent(toolSet, runtimeModel, maxIter, agentId, skillCatalogRenderer, null, Set.of());
|
||||
}
|
||||
|
||||
StateGraphReActAgent buildReActAgent(AgentToolSet toolSet, ModelConfigEntity runtimeModel,
|
||||
int maxIter, Long agentId, SkillCatalogRenderer skillCatalogRenderer,
|
||||
PrefixBudgetPlan prefixBudgetPlan) {
|
||||
PrefixBudgetPlan prefixBudgetPlan, Set<String> autoDemotedTools) {
|
||||
ChatModel chatModel = buildRuntimeChatModel(runtimeModel);
|
||||
ChatClient chatClient = ChatClient.create(chatModel);
|
||||
String reasoningEffort = resolveReasoningEffortForModel(runtimeModel);
|
||||
CompiledGraph compiledGraph = buildReActGraph(toolSet, chatModel, maxIter, reasoningEffort,
|
||||
runtimeModel, agentId, skillCatalogRenderer, prefixBudgetPlan);
|
||||
runtimeModel, agentId, skillCatalogRenderer, prefixBudgetPlan, autoDemotedTools);
|
||||
return new StateGraphReActAgent(chatClient, conversationService, compiledGraph,
|
||||
chatModel, conversationWindowManager, toolSet);
|
||||
}
|
||||
@ -616,6 +625,7 @@ public class AgentGraphBuilder {
|
||||
// LLM mis-calls a skill name as a tool, the response tells it
|
||||
// the right invocation pattern instead of a dead-end error.
|
||||
executor.setSkillRuntimeService(skillRuntimeService);
|
||||
executor.setUsageRecencyTracker(toolUsageRecencyTracker);
|
||||
// Optional: route child-agent denied-tool audit events through
|
||||
// the audit pipeline. Null when audit is not wired (legacy / test).
|
||||
if (auditEventService != null) {
|
||||
@ -876,13 +886,13 @@ public class AgentGraphBuilder {
|
||||
String reasoningEffort, ModelConfigEntity primaryModelConfig,
|
||||
Long agentId, SkillCatalogRenderer skillCatalogRenderer) {
|
||||
return buildReActGraph(toolSet, chatModel, maxIterations, reasoningEffort,
|
||||
primaryModelConfig, agentId, skillCatalogRenderer, null);
|
||||
primaryModelConfig, agentId, skillCatalogRenderer, null, Set.of());
|
||||
}
|
||||
|
||||
CompiledGraph buildReActGraph(AgentToolSet toolSet, ChatModel chatModel, int maxIterations,
|
||||
String reasoningEffort, ModelConfigEntity primaryModelConfig,
|
||||
Long agentId, SkillCatalogRenderer skillCatalogRenderer,
|
||||
PrefixBudgetPlan prefixBudgetPlan) {
|
||||
PrefixBudgetPlan prefixBudgetPlan, Set<String> autoDemotedTools) {
|
||||
try {
|
||||
List<vip.mate.llm.failover.FallbackEntry> fallbackChain = buildFallbackChain(primaryModelConfig, agentId);
|
||||
NodeStreamingChatHelper streamingHelper = new NodeStreamingChatHelper(
|
||||
@ -905,6 +915,7 @@ public class AgentGraphBuilder {
|
||||
// LLM mis-calls a skill name as a tool, the response tells it
|
||||
// the right invocation pattern instead of a dead-end error.
|
||||
executor.setSkillRuntimeService(skillRuntimeService);
|
||||
executor.setUsageRecencyTracker(toolUsageRecencyTracker);
|
||||
// Optional: route child-agent denied-tool audit events through
|
||||
// the audit pipeline. Null when audit is not wired (legacy / test).
|
||||
if (auditEventService != null) {
|
||||
@ -920,6 +931,7 @@ public class AgentGraphBuilder {
|
||||
streamingHelper, conversationWindowManager, streamTracker, 0, wikiContextService,
|
||||
skillCatalogRenderer, toolDisclosureService, progressLedgerService);
|
||||
reasoningNode.setPrefixBudgetPlan(prefixBudgetPlan);
|
||||
reasoningNode.setAutoDemotedTools(autoDemotedTools);
|
||||
ActionNode actionNode = new ActionNode(executor, streamTracker);
|
||||
ObservationProcessor observationProcessor = new ObservationProcessor(graphObservationProperties);
|
||||
ObservationNode observationNode = new ObservationNode(observationProcessor, streamTracker);
|
||||
|
||||
@ -19,7 +19,8 @@ public record PrefixBudgetPlan(
|
||||
int wikiTokens,
|
||||
int skillCatalogTokens,
|
||||
int extensionCatalogTokens,
|
||||
int ledgerTokens) {
|
||||
int ledgerTokens,
|
||||
int toolSchemaBudgetTokens) {
|
||||
|
||||
/** Window-size tier. Small windows tighten the injection ratio. */
|
||||
public enum Profile {
|
||||
@ -35,6 +36,6 @@ public record PrefixBudgetPlan(
|
||||
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);
|
||||
Integer.MAX_VALUE, Integer.MAX_VALUE, Integer.MAX_VALUE, Integer.MAX_VALUE);
|
||||
}
|
||||
}
|
||||
|
||||
@ -72,7 +72,8 @@ public class PrefixBudgetPlanner {
|
||||
(int) (injectionBudget * shares.getWiki() / sum),
|
||||
(int) (injectionBudget * shares.getSkill() / sum),
|
||||
(int) (injectionBudget * shares.getExtensionCatalog() / sum),
|
||||
(int) (injectionBudget * shares.getLedger() / sum));
|
||||
(int) (injectionBudget * shares.getLedger() / sum),
|
||||
(int) (effectiveMax * properties.getToolSchemaRatio()));
|
||||
|
||||
if (profile != PrefixBudgetPlan.Profile.NORMAL) {
|
||||
log.info("[PrefixBudget] 窗口 {} tokens 进入 {} 档:注入预算 {} tokens"
|
||||
|
||||
@ -7,6 +7,7 @@ import org.springframework.ai.chat.messages.ToolResponseMessage;
|
||||
import org.springframework.ai.chat.model.ToolContext;
|
||||
import org.springframework.ai.tool.ToolCallback;
|
||||
import vip.mate.tool.builtin.ToolExecutionContext;
|
||||
import vip.mate.tool.disclosure.ToolUsageRecencyTracker;
|
||||
import vip.mate.agent.AgentToolSet;
|
||||
import vip.mate.agent.GraphEventPublisher;
|
||||
import vip.mate.agent.context.ChatOrigin;
|
||||
@ -251,6 +252,13 @@ public class ToolExecutionExecutor {
|
||||
*/
|
||||
private vip.mate.skill.runtime.SkillRuntimeService skillRuntimeService;
|
||||
|
||||
/** Optional recency feed for budget-driven tool-disclosure demotion. */
|
||||
private ToolUsageRecencyTracker usageRecencyTracker;
|
||||
|
||||
public void setUsageRecencyTracker(ToolUsageRecencyTracker tracker) {
|
||||
this.usageRecencyTracker = tracker;
|
||||
}
|
||||
|
||||
public void setSkillRuntimeService(vip.mate.skill.runtime.SkillRuntimeService s) {
|
||||
this.skillRuntimeService = s;
|
||||
}
|
||||
@ -885,6 +893,12 @@ public class ToolExecutionExecutor {
|
||||
ToolExecutionContext.clear();
|
||||
}
|
||||
|
||||
// Recency feed for budget-driven disclosure demotion: recently used
|
||||
// tools keep their advertised schema, never-used ones demote first.
|
||||
if (usageRecencyTracker != null) {
|
||||
usageRecencyTracker.recordUse(toolName);
|
||||
}
|
||||
|
||||
int rawLen = result != null ? result.length() : 0;
|
||||
// RFC-052: returnDirect tools bypass spill / truncation / LLM context.
|
||||
// Their full text goes to the user verbatim and is never persisted to
|
||||
|
||||
@ -363,6 +363,19 @@ public class ReasoningNode implements NodeAction {
|
||||
this.prefixBudgetPlan = prefixBudgetPlan;
|
||||
}
|
||||
|
||||
/**
|
||||
* Core-tier tools auto-demoted to the extension catalog because the
|
||||
* advertised schemas exceeded the window's tool-schema budget. Decided
|
||||
* once at agent-build time (kept stable for prompt caching); the baked
|
||||
* extension catalog lists them so {@code enable_tool} can surface any of
|
||||
* them back.
|
||||
*/
|
||||
private Set<String> autoDemotedTools = Set.of();
|
||||
|
||||
public void setAutoDemotedTools(Set<String> autoDemotedTools) {
|
||||
this.autoDemotedTools = autoDemotedTools == null ? Set.of() : autoDemotedTools;
|
||||
}
|
||||
|
||||
public ReasoningNode(ChatModel chatModel, AgentToolSet toolSet, String reasoningEffort,
|
||||
NodeStreamingChatHelper streamingHelper,
|
||||
ConversationWindowManager conversationWindowManager,
|
||||
@ -728,7 +741,8 @@ public class ReasoningNode implements NodeAction {
|
||||
// so an enable_tool call earlier in this loop takes effect immediately.
|
||||
// Falls back to the full tool set when no disclosure service is wired.
|
||||
List<ToolCallback> activeCallbacks = (toolDisclosureService != null && toolSet != null)
|
||||
? toolDisclosureService.split(toolSet, accessor.enabledExtensionTools()).activeCallbacks()
|
||||
? toolDisclosureService.split(toolSet, accessor.enabledExtensionTools(), autoDemotedTools)
|
||||
.activeCallbacks()
|
||||
: toolCallbacks;
|
||||
|
||||
ChatOptions options = buildChatOptions(effectiveReasoning, activeCallbacks);
|
||||
|
||||
@ -47,6 +47,14 @@ public class PrefixBudgetProperties {
|
||||
*/
|
||||
private double compactTriggerRatioOverride = 0.85;
|
||||
|
||||
/**
|
||||
* Fraction of the effective window the advertised tool schemas may
|
||||
* occupy. When the core tool set estimates above this, the least
|
||||
* recently used demotable tools are auto-moved to the extension catalog
|
||||
* (recoverable via {@code enable_tool}) until the set fits.
|
||||
*/
|
||||
private double toolSchemaRatio = 0.25;
|
||||
|
||||
/** Relative shares of the injection budget. Normalized at plan time. */
|
||||
private Shares shares = new Shares();
|
||||
|
||||
|
||||
@ -6,6 +6,7 @@ import org.springframework.ai.tool.ToolCallback;
|
||||
import org.springframework.beans.factory.annotation.Value;
|
||||
import org.springframework.stereotype.Service;
|
||||
import vip.mate.agent.AgentToolSet;
|
||||
import vip.mate.agent.context.TokenEstimator;
|
||||
import vip.mate.tool.ToolRegistry;
|
||||
import vip.mate.tool.mcp.model.McpServerEntity;
|
||||
import vip.mate.tool.mcp.service.McpServerService;
|
||||
@ -15,6 +16,7 @@ import vip.mate.tool.service.AvailableToolService;
|
||||
import vip.mate.tool.service.ToolService;
|
||||
|
||||
import java.util.ArrayList;
|
||||
import java.util.Comparator;
|
||||
import java.util.LinkedHashMap;
|
||||
import java.util.LinkedHashSet;
|
||||
import java.util.List;
|
||||
@ -53,6 +55,7 @@ public class DefaultToolDisclosureService implements ToolDisclosureService {
|
||||
private final McpServerService mcpServerService;
|
||||
private final AvailableToolService availableToolService;
|
||||
private final ToolRegistry toolRegistry;
|
||||
private final ToolUsageRecencyTracker usageRecencyTracker;
|
||||
|
||||
@Value("${mateclaw.tools.disclosure.mode:progressive}")
|
||||
private String disclosureMode;
|
||||
@ -97,17 +100,25 @@ public class DefaultToolDisclosureService implements ToolDisclosureService {
|
||||
|
||||
@Override
|
||||
public ToolDisclosureSplit split(AgentToolSet baseSet, Set<String> enabledExtensions) {
|
||||
return split(baseSet, enabledExtensions, Set.of());
|
||||
}
|
||||
|
||||
@Override
|
||||
public ToolDisclosureSplit split(AgentToolSet baseSet, Set<String> enabledExtensions,
|
||||
Set<String> autoDemoted) {
|
||||
List<ToolCallback> all = baseSet == null ? List.of() : baseSet.callbacks();
|
||||
if (legacyMode()) {
|
||||
return new ToolDisclosureSplit(all, List.of());
|
||||
}
|
||||
Set<String> enabled = enabledExtensions == null ? Set.of() : enabledExtensions;
|
||||
Set<String> demoted = autoDemoted == null ? Set.of() : autoDemoted;
|
||||
List<ToolCallback> active = new ArrayList<>(all.size());
|
||||
List<ToolCallback> extensionCatalog = new ArrayList<>();
|
||||
for (ToolCallback cb : all) {
|
||||
if (resolveTier(cb) == DisclosureTier.EXTENSION) {
|
||||
String name = cb.getToolDefinition().name();
|
||||
if (resolveTier(cb) == DisclosureTier.EXTENSION || demoted.contains(name)) {
|
||||
extensionCatalog.add(cb);
|
||||
if (enabled.contains(cb.getToolDefinition().name())) {
|
||||
if (enabled.contains(name)) {
|
||||
active.add(cb);
|
||||
}
|
||||
} else {
|
||||
@ -117,12 +128,79 @@ public class DefaultToolDisclosureService implements ToolDisclosureService {
|
||||
return new ToolDisclosureSplit(active, extensionCatalog);
|
||||
}
|
||||
|
||||
/**
|
||||
* {@inheritDoc}
|
||||
*
|
||||
* <p>Protection set: {@link #ALWAYS_CORE} meta-tools and builtin tools
|
||||
* with an explicit {@code disclosure_tier = core} row. MCP tools remain
|
||||
* demotable — the server-level tier cannot distinguish an explicit core
|
||||
* choice from the default, and MCP schemas are typically the heaviest
|
||||
* part of the advertisement.
|
||||
*/
|
||||
@Override
|
||||
public Set<String> computeAutoDemotions(AgentToolSet baseSet, Integer budgetTokens) {
|
||||
if (legacyMode() || baseSet == null || budgetTokens == null
|
||||
|| budgetTokens <= 0 || budgetTokens == Integer.MAX_VALUE) {
|
||||
return Set.of();
|
||||
}
|
||||
List<ToolCallback> core = split(baseSet, Set.of()).activeCallbacks();
|
||||
int coreTokens = TokenEstimator.estimateToolsTokens(core);
|
||||
if (coreTokens <= budgetTokens) {
|
||||
return Set.of();
|
||||
}
|
||||
Snapshot snap = snapshot();
|
||||
List<ToolCallback> candidates = core.stream()
|
||||
.filter(cb -> isDemotable(cb.getToolDefinition().name(), snap))
|
||||
.sorted(demotionOrder())
|
||||
.toList();
|
||||
Set<String> demoted = new LinkedHashSet<>();
|
||||
int remainingTokens = coreTokens;
|
||||
for (ToolCallback cb : candidates) {
|
||||
if (remainingTokens <= budgetTokens) {
|
||||
break;
|
||||
}
|
||||
remainingTokens -= TokenEstimator.estimateToolsTokens(List.of(cb));
|
||||
demoted.add(cb.getToolDefinition().name());
|
||||
}
|
||||
if (!demoted.isEmpty()) {
|
||||
log.info("[ToolDisclosure] 工具 schema 估算 {} tokens 超出预算 {}——已将 {} 个最少使用的工具"
|
||||
+ "降级到扩展目录(enable_tool 可找回): {}",
|
||||
coreTokens, budgetTokens, demoted.size(), demoted);
|
||||
}
|
||||
return demoted;
|
||||
}
|
||||
|
||||
private boolean isDemotable(String toolName, Snapshot snap) {
|
||||
if (toolName == null || ALWAYS_CORE.contains(toolName)) {
|
||||
return false;
|
||||
}
|
||||
// An explicit core row is an operator decision — never override it.
|
||||
return snap.builtinTierByName.get(toolName) != DisclosureTier.CORE;
|
||||
}
|
||||
|
||||
/** Never-used tools demote first, then least recently used; name-tiebreak keeps builds deterministic. */
|
||||
private Comparator<ToolCallback> demotionOrder() {
|
||||
return Comparator
|
||||
.<ToolCallback, Long>comparing(cb -> {
|
||||
Long lastUsed = usageRecencyTracker == null
|
||||
? null : usageRecencyTracker.lastUsedAt(cb.getToolDefinition().name());
|
||||
return lastUsed == null ? Long.MIN_VALUE : lastUsed;
|
||||
})
|
||||
.thenComparing(cb -> cb.getToolDefinition().name());
|
||||
}
|
||||
|
||||
@Override
|
||||
public String renderExtensionCatalog(AgentToolSet baseSet, Integer maxInputTokens) {
|
||||
return renderExtensionCatalog(baseSet, maxInputTokens, Set.of());
|
||||
}
|
||||
|
||||
@Override
|
||||
public String renderExtensionCatalog(AgentToolSet baseSet, Integer maxInputTokens,
|
||||
Set<String> autoDemoted) {
|
||||
if (legacyMode() || baseSet == null) {
|
||||
return "";
|
||||
}
|
||||
List<ToolCallback> extension = split(baseSet, Set.of()).extensionCatalog();
|
||||
List<ToolCallback> extension = split(baseSet, Set.of(), autoDemoted).extensionCatalog();
|
||||
if (extension.isEmpty()) {
|
||||
return "";
|
||||
}
|
||||
|
||||
@ -31,6 +31,29 @@ public interface ToolDisclosureService {
|
||||
*/
|
||||
ToolDisclosureSplit split(AgentToolSet baseSet, Set<String> enabledExtensions);
|
||||
|
||||
/**
|
||||
* Budget-aware variant: tools in {@code autoDemoted} are treated as
|
||||
* extension tier for this split even when their resolved tier is core.
|
||||
* The demotion set is decided once per agent build (see
|
||||
* {@link #computeAutoDemotions}) so the runtime split, the baked catalog
|
||||
* and the prompt-cache prefix stay consistent with each other.
|
||||
*/
|
||||
default ToolDisclosureSplit split(AgentToolSet baseSet, Set<String> enabledExtensions,
|
||||
Set<String> autoDemoted) {
|
||||
return split(baseSet, enabledExtensions);
|
||||
}
|
||||
|
||||
/**
|
||||
* Decide which core-tier tools to auto-demote so the advertised tool
|
||||
* schemas fit {@code budgetTokens} (estimated). Ranking: never-used tools
|
||||
* first, then least recently used; meta-tools and explicitly configured
|
||||
* core tools are never demoted. Empty when the set already fits, when
|
||||
* {@code budgetTokens} is null, or in legacy disclosure mode.
|
||||
*/
|
||||
default Set<String> computeAutoDemotions(AgentToolSet baseSet, Integer budgetTokens) {
|
||||
return Set.of();
|
||||
}
|
||||
|
||||
/**
|
||||
* Render the {@code ## Extension Tools} system-prompt segment for the
|
||||
* agent's extension tools, or an empty string when there are none / when
|
||||
@ -38,6 +61,15 @@ public interface ToolDisclosureService {
|
||||
*/
|
||||
String renderExtensionCatalog(AgentToolSet baseSet, Integer maxInputTokens);
|
||||
|
||||
/**
|
||||
* Budget-aware variant: auto-demoted tools are listed in the catalog too,
|
||||
* so the model can discover and {@code enable_tool} them back.
|
||||
*/
|
||||
default String renderExtensionCatalog(AgentToolSet baseSet, Integer maxInputTokens,
|
||||
Set<String> autoDemoted) {
|
||||
return renderExtensionCatalog(baseSet, maxInputTokens);
|
||||
}
|
||||
|
||||
/** Drop the cached tier snapshot so the next resolve re-reads the DB. */
|
||||
void invalidate();
|
||||
|
||||
|
||||
@ -0,0 +1,33 @@
|
||||
package vip.mate.tool.disclosure;
|
||||
|
||||
import org.springframework.stereotype.Component;
|
||||
|
||||
import java.util.Map;
|
||||
import java.util.concurrent.ConcurrentHashMap;
|
||||
|
||||
/**
|
||||
* In-memory recency signal for tool usage. Feeds the budget-driven
|
||||
* auto-demotion ranking: never-used tools demote first, then the least
|
||||
* recently used ones.
|
||||
*
|
||||
* <p>Deliberately process-local and unpersisted — this is an advisory
|
||||
* ranking, not an audit trail. A restart resets everything to "never used",
|
||||
* which merely makes the first demotion pass alphabetical.
|
||||
*/
|
||||
@Component
|
||||
public class ToolUsageRecencyTracker {
|
||||
|
||||
private final Map<String, Long> lastUsedAtMs = new ConcurrentHashMap<>();
|
||||
|
||||
/** Record a successful execution of {@code toolName}. */
|
||||
public void recordUse(String toolName) {
|
||||
if (toolName != null && !toolName.isBlank()) {
|
||||
lastUsedAtMs.put(toolName, System.currentTimeMillis());
|
||||
}
|
||||
}
|
||||
|
||||
/** @return epoch millis of the last recorded use, or null when never used. */
|
||||
public Long lastUsedAt(String toolName) {
|
||||
return toolName == null ? null : lastUsedAtMs.get(toolName);
|
||||
}
|
||||
}
|
||||
@ -12,6 +12,7 @@ import java.util.List;
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
import static org.mockito.ArgumentMatchers.anyString;
|
||||
import static org.mockito.ArgumentMatchers.eq;
|
||||
import static org.mockito.ArgumentMatchers.isNull;
|
||||
import static org.mockito.Mockito.mock;
|
||||
import static org.mockito.Mockito.verify;
|
||||
import static org.mockito.Mockito.when;
|
||||
@ -38,7 +39,7 @@ class ReasoningNodePtlPromptTest {
|
||||
@Test
|
||||
void prefixIncludesSystemRuntimeAndWikiSegments() {
|
||||
WikiContextService wikiContextService = mock(WikiContextService.class);
|
||||
when(wikiContextService.buildRelevantContext(eq(42L), anyString()))
|
||||
when(wikiContextService.buildRelevantContext(eq(42L), anyString(), isNull()))
|
||||
.thenReturn(WIKI_RELEVANT_TEXT);
|
||||
|
||||
ReasoningNode node = newNode(wikiContextService);
|
||||
@ -67,7 +68,7 @@ class ReasoningNodePtlPromptTest {
|
||||
// that two independent calls with the same inputs produce
|
||||
// structurally identical output.
|
||||
WikiContextService wikiContextService = mock(WikiContextService.class);
|
||||
when(wikiContextService.buildRelevantContext(eq(42L), anyString()))
|
||||
when(wikiContextService.buildRelevantContext(eq(42L), anyString(), isNull()))
|
||||
.thenReturn(WIKI_RELEVANT_TEXT);
|
||||
|
||||
ReasoningNode node = newNode(wikiContextService);
|
||||
@ -125,7 +126,7 @@ class ReasoningNodePtlPromptTest {
|
||||
@Test
|
||||
void blankWikiResultSkipsWikiSegment() {
|
||||
WikiContextService wikiContextService = mock(WikiContextService.class);
|
||||
when(wikiContextService.buildRelevantContext(eq(42L), anyString()))
|
||||
when(wikiContextService.buildRelevantContext(eq(42L), anyString(), isNull()))
|
||||
.thenReturn(" "); // blank → drop the layer
|
||||
|
||||
ReasoningNode node = newNode(wikiContextService);
|
||||
|
||||
@ -7,6 +7,7 @@ import org.springframework.ai.tool.ToolCallback;
|
||||
import org.springframework.ai.tool.annotation.Tool;
|
||||
import org.springframework.test.util.ReflectionTestUtils;
|
||||
import vip.mate.agent.AgentToolSet;
|
||||
import vip.mate.agent.context.TokenEstimator;
|
||||
import vip.mate.tool.ToolRegistry;
|
||||
import vip.mate.tool.mcp.model.McpServerEntity;
|
||||
import vip.mate.tool.mcp.service.McpServerService;
|
||||
@ -88,7 +89,7 @@ class ToolDisclosureServiceTest {
|
||||
lenient().when(ms.listAll()).thenReturn(servers);
|
||||
lenient().when(as.listAvailable()).thenReturn(available);
|
||||
lenient().when(tr.getEnabledToolSet()).thenReturn(globalSet());
|
||||
return new DefaultToolDisclosureService(ts, ms, as, tr);
|
||||
return new DefaultToolDisclosureService(ts, ms, as, tr, new ToolUsageRecencyTracker());
|
||||
}
|
||||
|
||||
@Test
|
||||
@ -202,4 +203,105 @@ class ToolDisclosureServiceTest {
|
||||
private static List<String> names(List<ToolCallback> cbs) {
|
||||
return cbs.stream().map(c -> c.getToolDefinition().name()).toList();
|
||||
}
|
||||
|
||||
// ==================== budget-driven auto-demotion ====================
|
||||
|
||||
/** Three plain core tools for demotion-ranking tests. */
|
||||
static class ManyCoreTools {
|
||||
@Tool(description = "core tool a")
|
||||
public String tool_a() { return ""; }
|
||||
|
||||
@Tool(description = "core tool b")
|
||||
public String tool_b() { return ""; }
|
||||
|
||||
@Tool(description = "core tool c")
|
||||
public String tool_c() { return ""; }
|
||||
}
|
||||
|
||||
private static AgentToolSet manyCoreSet() {
|
||||
return AgentToolSet.fromCallbacks(List.of(new ManyCoreTools()),
|
||||
List.of(ToolCallbacks.from(new ManyCoreTools())));
|
||||
}
|
||||
|
||||
@Test
|
||||
@DisplayName("no demotion when the core schemas fit the budget, or when budget is absent")
|
||||
void noDemotionWhenBudgetFits() {
|
||||
var svc = service(List.of(), List.of(), List.of());
|
||||
AgentToolSet set = manyCoreSet();
|
||||
assertTrue(svc.computeAutoDemotions(set, Integer.MAX_VALUE).isEmpty());
|
||||
assertTrue(svc.computeAutoDemotions(set, null).isEmpty());
|
||||
assertTrue(svc.computeAutoDemotions(set, 1_000_000).isEmpty());
|
||||
}
|
||||
|
||||
@Test
|
||||
@DisplayName("tiny budget demotes every demotable tool, alphabetical when nothing was ever used")
|
||||
void tinyBudgetDemotesAll() {
|
||||
var svc = service(List.of(), List.of(), List.of());
|
||||
var demoted = svc.computeAutoDemotions(manyCoreSet(), 1);
|
||||
assertEquals(Set.of("tool_a", "tool_b", "tool_c"), demoted);
|
||||
}
|
||||
|
||||
@Test
|
||||
@DisplayName("budget one tool short demotes exactly the first never-used candidate")
|
||||
void partialDemotionTakesFirstCandidate() {
|
||||
var svc = service(List.of(), List.of(), List.of());
|
||||
AgentToolSet set = manyCoreSet();
|
||||
int coreTokens = TokenEstimator.estimateToolsTokens(svc.split(set, Set.of()).activeCallbacks());
|
||||
var demoted = svc.computeAutoDemotions(set, coreTokens - 1);
|
||||
assertEquals(Set.of("tool_a"), demoted);
|
||||
}
|
||||
|
||||
@Test
|
||||
@DisplayName("recently used tools demote last")
|
||||
void recencyProtectsRecentlyUsed() {
|
||||
ToolUsageRecencyTracker tracker = new ToolUsageRecencyTracker();
|
||||
tracker.recordUse("tool_a");
|
||||
ToolService ts = mock(ToolService.class);
|
||||
McpServerService ms = mock(McpServerService.class);
|
||||
AvailableToolService as = mock(AvailableToolService.class);
|
||||
ToolRegistry tr = mock(ToolRegistry.class);
|
||||
lenient().when(ts.listTools()).thenReturn(List.of());
|
||||
lenient().when(ms.listAll()).thenReturn(List.of());
|
||||
lenient().when(as.listAvailable()).thenReturn(List.of());
|
||||
lenient().when(tr.getEnabledToolSet()).thenReturn(globalSet());
|
||||
var svc = new DefaultToolDisclosureService(ts, ms, as, tr, tracker);
|
||||
|
||||
AgentToolSet set = manyCoreSet();
|
||||
int coreTokens = TokenEstimator.estimateToolsTokens(svc.split(set, Set.of()).activeCallbacks());
|
||||
// One tool over budget: the never-used tool_b (alphabetically first
|
||||
// among never-used) demotes, the recently used tool_a survives.
|
||||
var demoted = svc.computeAutoDemotions(set, coreTokens - 1);
|
||||
assertEquals(Set.of("tool_b"), demoted);
|
||||
}
|
||||
|
||||
@Test
|
||||
@DisplayName("explicit core DB row and meta-tools are never demoted")
|
||||
void explicitCoreProtected() {
|
||||
var svc = service(List.of(toolRow("tool_a", "builtin", "core")), List.of(), List.of());
|
||||
var demoted = svc.computeAutoDemotions(manyCoreSet(), 1);
|
||||
assertEquals(Set.of("tool_b", "tool_c"), demoted);
|
||||
}
|
||||
|
||||
@Test
|
||||
@DisplayName("auto-demoted tools behave as extension in split and can be enabled back")
|
||||
void splitHonorsAutoDemotions() {
|
||||
var svc = service(List.of(), List.of(), List.of());
|
||||
AgentToolSet set = manyCoreSet();
|
||||
|
||||
var split = svc.split(set, Set.of(), Set.of("tool_b"));
|
||||
assertEquals(List.of("tool_a", "tool_c"), names(split.activeCallbacks()));
|
||||
assertEquals(List.of("tool_b"), names(split.extensionCatalog()));
|
||||
|
||||
var enabledBack = svc.split(set, Set.of("tool_b"), Set.of("tool_b"));
|
||||
assertTrue(names(enabledBack.activeCallbacks()).contains("tool_b"));
|
||||
}
|
||||
|
||||
@Test
|
||||
@DisplayName("catalog rendering lists auto-demoted tools for discoverability")
|
||||
void catalogListsAutoDemoted() {
|
||||
var svc = service(List.of(), List.of(), List.of());
|
||||
String catalog = svc.renderExtensionCatalog(manyCoreSet(), 8192, Set.of("tool_b"));
|
||||
assertTrue(catalog.contains("tool_b"));
|
||||
assertFalse(catalog.contains("| `tool_a`"), "non-demoted core tools stay out of the catalog");
|
||||
}
|
||||
}
|
||||
|
||||
Loading…
Reference in New Issue
Block a user