opensource/mateclaw
2
0
Fork 0
mirror of https://gitee.com/mateos/mateclaw.git synced 2026-11-02 14:11:49 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

fix: unify prompt budgets on AgentScope 2.0.4

Browse Source
This commit is contained in:
matevip 2026-10-20 04:03:19 -04:00
parent 576eb8055a
commit ddbd37b51b
5 changed files with 143 additions and 20 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
mateclaw-server/src/main/java/vip/mate/agent/context/BudgetedCompactionModel.java
View File

@ -3,7 +3,6 @@ package vip.mate.agent.context;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.model.tool.ToolCallingChatOptions;
import vip.mate.agent.runtime.state.RunState;
import vip.mate.llm.chatmodel.ChatModelDefaults;
import java.util.concurrent.CancellationException;
@ -44,13 +43,7 @@ public final class BudgetedCompactionModel implements ChatModel {
if (value != null && value>0) usage.merge(key,value,Math::addExact);
}
public static int outputReserveTokens(ChatOptions options,int fallback) {
Integer maximum = options == null ? null : options.getMaxTokens();
if (options instanceof org.springframework.ai.openai.OpenAiChatOptions openAi
&& openAi.getMaxCompletionTokens() != null) {
maximum = maximum == null ? openAi.getMaxCompletionTokens()
: Math.max(maximum,openAi.getMaxCompletionTokens());
}
return maximum == null ? Math.max(0,fallback) : Math.max(0,maximum);
return PromptBudget.outputReserveTokens(options, null, fallback);
}
public BudgetedCompactionModel(ChatModel delegate,long allowedCalls,int windowTokens,int fallbackOutputTokens) {
this.delegate=delegate; this.allowedCalls=allowedCalls;
@ -63,12 +56,7 @@ public final class BudgetedCompactionModel implements ChatModel {
}
private void checkPrompt(Prompt prompt) {
if (windowTokens <= 0) return;
ChatOptions defaults = ChatModelDefaults.optionsFor(delegate);
ChatOptions options = prompt.getOptions() != null ? prompt.getOptions() : defaults;
int output = outputReserveTokens(options, outputReserveTokens(defaults,fallbackOutputTokens));
int tools = options instanceof ToolCallingChatOptions toolOptions
? TokenEstimator.estimateToolsTokens(toolOptions.getToolCallbacks()) : 0;
int total = TokenEstimator.estimateTokens(prompt.getInstructions()) + tools + output;
long total = PromptBudget.estimateTokens(prompt, ChatModelDefaults.optionsFor(delegate), fallbackOutputTokens);
if (total > windowTokens) {
blockedPrompt = "摘要 Prompt 估算 " + total + " tokens 超过上下文窗口 " + windowTokens
+ " tokens；未调用摘要模型。";

52
mateclaw-server/src/main/java/vip/mate/agent/context/PromptBudget.java Normal file
View File

@ -0,0 +1,52 @@
package vip.mate.agent.context;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.model.tool.ToolCallingChatOptions;
import org.springframework.ai.openai.OpenAiChatOptions;
import org.springframework.ai.tool.ToolCallback;
import java.util.List;
/** Estimates the effective request without changing the provider's prompt or options. */
public final class PromptBudget {
private PromptBudget() {}
public static long estimateTokens(Prompt prompt, ChatOptions defaults, int fallbackOutputTokens) {
ChatOptions runtime = prompt.getOptions();
ChatOptions effective = effectiveOptions(runtime, defaults);
// Native combineWith treats empty lists as absent. Provider adapters explicitly replace
// callbacks when the runtime list is non-null, including a deliberately empty list.
List<ToolCallback> callbacks = runtime instanceof ToolCallingChatOptions tools
&& tools.getToolCallbacks() != null ? tools.getToolCallbacks()
: effective instanceof ToolCallingChatOptions tools ? tools.getToolCallbacks() : List.of();
return (long) TokenEstimator.estimateTokens(prompt.getInstructions())
+ TokenEstimator.estimateToolsTokens(callbacks)
+ effectiveOutputReserve(effective, runtime, fallbackOutputTokens);
}
public static int outputReserveTokens(ChatOptions runtime, ChatOptions defaults, int fallback) {
return effectiveOutputReserve(effectiveOptions(runtime, defaults), runtime, fallback);
}
private static int effectiveOutputReserve(ChatOptions effective, ChatOptions runtime, int fallback) {
// A generic defaults builder cannot represent provider-specific runtime fields. Preserve
// the runtime's declared maximum as a conservative floor, as the existing guard did.
return Math.max(outputReserveTokens(effective, fallback), outputReserveTokens(runtime, 0));
}
private static ChatOptions effectiveOptions(ChatOptions runtime, ChatOptions defaults) {
if (defaults == null) return runtime;
if (runtime == null) return defaults;
return defaults.mutate().combineWith(runtime.mutate()).build();
}
private static int outputReserveTokens(ChatOptions options, int fallback) {
Integer maximum = options == null ? null : options.getMaxTokens();
if (options instanceof OpenAiChatOptions openAi && openAi.getMaxCompletionTokens() != null) {
maximum = maximum == null ? openAi.getMaxCompletionTokens()
: Math.max(maximum, openAi.getMaxCompletionTokens());
}
return maximum == null ? Math.max(0, fallback) : Math.max(0, maximum);
}
}

11
mateclaw-server/src/main/java/vip/mate/agent/graph/node/ReasoningNode.java
View File

@ -1023,8 +1023,7 @@ public class ReasoningNode implements NodeAction {
// budgeting retains legacy behavior.
boolean mandatoryPrefixOverflow = loopReservedPrefixTokens + protectedHistoryTokens
> loopContextWindowTokens();
int assembledPromptEstimateTokens = TokenEstimator.estimateTokens(promptMessages)
+ toolSchemaEstimateTokens + outputReserveTokens;
long assembledPromptEstimateTokens = vip.mate.agent.context.PromptBudget.estimateTokens(prompt, null, outputReserveTokens);
boolean promptOverflow = (holdContent || prefixBudgetPlan != null && prefixBudgetPlan.enabled())
&& assembledPromptEstimateTokens > loopContextWindowTokens();
boolean localPromptRejected = promptOverflow;
@ -1149,8 +1148,8 @@ public class ReasoningNode implements NodeAction {
// Reuse the SAME non-history prefix — wiki/runtime context preserved.
List<Message> retryPromptMessages = new ArrayList<>(nonHistoryPrefix);
retryPromptMessages.addAll(compactedMessages);
int retryEstimateTokens = TokenEstimator.estimateTokens(retryPromptMessages)
+ toolSchemaEstimateTokens + outputReserveTokens;
long retryEstimateTokens = vip.mate.agent.context.PromptBudget.estimateTokens(
new Prompt(retryPromptMessages, options), null, outputReserveTokens);
if ((holdContent || prefixBudgetPlan != null && prefixBudgetPlan.enabled())
&& retryEstimateTokens > loopContextWindowTokens()) {
localPromptRejected = true;
@ -1234,8 +1233,8 @@ public class ReasoningNode implements NodeAction {
List<Message> nudgeHistory = LOOP_BUDGETER.budget(messages, nudgeConfig).messages();
List<Message> nudgedMessages = new ArrayList<>(nudgePrefix);
nudgedMessages.addAll(nudgeHistory);
int nudgeEstimate = TokenEstimator.estimateTokens(nudgedMessages)
+ toolSchemaEstimateTokens + outputReserveTokens;
long nudgeEstimate = vip.mate.agent.context.PromptBudget.estimateTokens(
new Prompt(nudgedMessages, options), null, outputReserveTokens);
if ((holdContent || prefixBudgetPlan != null && prefixBudgetPlan.enabled())
&& nudgeEstimate > loopContextWindowTokens()) {
localPromptRejected = true;

26
mateclaw-server/src/test/java/vip/mate/agent/context/BudgetedCompactionModelTest.java
View File

@ -12,6 +12,32 @@ import static org.mockito.Mockito.*;
import static org.mockito.ArgumentMatchers.*;
import static vip.mate.agent.graph.state.MateClawStateKeys.*;
class BudgetedCompactionModelTest {
@Test void veryLargeOutputReserveCannotOverflowTotalAndBypassGuard() {
var delegate=mock(ChatModel.class);
var prompt=new Prompt("summary",ChatOptions.builder().maxTokens(Integer.MAX_VALUE).build());
var guarded=new BudgetedCompactionModel(delegate,2,24000,1024);
assertThrows(BudgetedCompactionModel.SummaryPromptOverflow.class,()->guarded.call(prompt));
assertEquals(0,guarded.calls());verify(delegate,never()).call(any(Prompt.class));
}
@Test void inheritedCompletionReserveIsCountedWhenRuntimeOverridesOnlyMaxTokens() {
var delegate=mock(ChatModel.class);
when(delegate.getDefaultOptions()).thenReturn(org.springframework.ai.openai.OpenAiChatOptions.builder().maxCompletionTokens(30000).build());
var prompt=new Prompt("summary",org.springframework.ai.openai.OpenAiChatOptions.builder().maxTokens(512).build());
var guarded=new BudgetedCompactionModel(delegate,2,24000,1024);
assertThrows(BudgetedCompactionModel.SummaryPromptOverflow.class,()->guarded.call(prompt));
assertEquals(0,guarded.calls());verify(delegate,never()).call(any(Prompt.class));
}
@Test void inheritedDefaultToolSchemasAreCountedBeforeStreamingProviderInvocation() {
var delegate=mock(ChatModel.class);var tool=mock(org.springframework.ai.tool.ToolCallback.class);
when(delegate.stream(any(Prompt.class))).thenReturn(reactor.core.publisher.Flux.empty());
when(tool.getToolDefinition()).thenReturn(org.springframework.ai.tool.definition.ToolDefinition.builder()
.name("large").description("界".repeat(1000)).inputSchema("{\"type\":\"object\"}").build());
when(delegate.getDefaultOptions()).thenReturn(org.springframework.ai.openai.OpenAiChatOptions.builder().toolCallbacks(List.of(tool)).maxTokens(64).build());
var prompt=new Prompt("summary",org.springframework.ai.openai.OpenAiChatOptions.builder().maxTokens(64).build());
var guarded=new BudgetedCompactionModel(delegate,2,1000,64);
assertThrows(BudgetedCompactionModel.SummaryPromptOverflow.class,()->guarded.stream(prompt).blockLast());
assertEquals(0,guarded.calls());verify(delegate,never()).stream(any(Prompt.class));
}
@Test void exactInvocationAndAvailableUsageAccumulateWithoutMutatingInputState() {
var delegate=mock(ChatModel.class);
when(delegate.call(any(Prompt.class))).thenReturn(new ChatResponse(List.of(new Generation(new AssistantMessage("summary"))),ChatResponseMetadata.builder().usage(new DefaultUsage(19,11,30,null,7L,3L)).build()));

58
mateclaw-server/src/test/java/vip/mate/agent/context/PromptBudgetTest.java Normal file
View File

@ -0,0 +1,58 @@
package vip.mate.agent.context;
import org.junit.jupiter.api.Test;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.openai.OpenAiChatOptions;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.ai.tool.definition.ToolDefinition;
import java.util.List;
import static org.junit.jupiter.api.Assertions.*;
import static org.mockito.Mockito.*;
class PromptBudgetTest {
@Test void explicitEmptyCallbacksAndLowerOutputOverrideDefaultsWithoutMutatingOptions() {
var tool=mock(ToolCallback.class);
var defaults=OpenAiChatOptions.builder().maxTokens(4096).toolCallbacks(List.of(tool)).build();
var runtime=OpenAiChatOptions.builder().maxTokens(64).toolCallbacks(List.of()).build();
var prompt=new Prompt("summary",runtime);
assertEquals(TokenEstimator.estimateTokens(prompt.getInstructions())+64,
PromptBudget.estimateTokens(prompt,defaults,1024));
assertEquals(4096,defaults.getMaxTokens());
assertEquals(List.of(tool),defaults.getToolCallbacks());
assertEquals(64,runtime.getMaxTokens());
assertTrue(runtime.getToolCallbacks().isEmpty());
assertSame(runtime,prompt.getOptions());
verifyNoInteractions(tool);
}
@Test void inheritedCompletionFieldAndRuntimeOverrideFollowNativeOptionMerge() {
var defaults=OpenAiChatOptions.builder().maxCompletionTokens(30000).build();
var partial=OpenAiChatOptions.builder().maxTokens(512).build();
assertEquals(30000,PromptBudget.outputReserveTokens(partial,defaults,1024));
var explicit=partial.mutate().maxCompletionTokens(128).build();
assertNull(explicit.getMaxTokens(),"Native builder clears maxTokens when setting maxCompletionTokens");
assertEquals(128,PromptBudget.outputReserveTokens(explicit,defaults,1024));
assertEquals(30000,defaults.getMaxCompletionTokens());
assertNull(partial.getMaxCompletionTokens());
}
@Test void providerSpecificRuntimeOutputFieldSurvivesGenericDefaultBuilder() {
var defaults=ChatOptions.builder().maxTokens(512).build();
var runtime=OpenAiChatOptions.builder().maxCompletionTokens(30000).build();
assertEquals(30000,PromptBudget.outputReserveTokens(runtime,defaults,1024));
assertTrue(PromptBudget.estimateTokens(new Prompt("small",runtime),defaults,1024)>30000);
}
@Test void nullAndGenericRuntimeOptionsInheritDefaultToolSchemas() {
var tool=mock(ToolCallback.class);
when(tool.getToolDefinition()).thenReturn(ToolDefinition.builder()
.name("tool").description("界".repeat(1000)).inputSchema("{\"type\":\"object\"}").build());
var defaults=OpenAiChatOptions.builder().maxTokens(64).toolCallbacks(List.of(tool)).build();
assertTrue(PromptBudget.estimateTokens(new Prompt("small"),defaults,64)>1000);
assertTrue(PromptBudget.estimateTokens(new Prompt("small",ChatOptions.builder().maxTokens(64).build()),defaults,64)>1000);
}
}