From fc3d84d6c2e35e9aa7912e07bdc93b0e2f0333bf Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=80=AA=E7=A8=8B=E4=BC=9F?= Date: Wed, 15 Jul 2026 14:52:43 +0800 Subject: [PATCH] fix(wiki): budget the system-prompt KB page listing to the model window MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit buildWikiContext enumerated an agent's bound knowledge-base pages into the system prompt capped only by maxContextChars (default 10000, sized for large cloud models). On a small-context model a large KB therefore consumed a big fixed slice of the window on every turn — the "tool token estimate fills the context" report in #521 (the growth lands in the system-prompt bucket, not the tool-schema bucket; wiki tool schemas are fixed-size and do not scale with file count). Add a budgeted buildWikiContext(agentId, budgetTokens) overload mirroring buildRelevantContext: the page enumeration also stops once the estimated token total exceeds the budget, appending the existing '... and more (use wiki_list_pages)' hint. AgentGraphBuilder passes the same prefix budget it already applies to the memory block; the legacy Integer.MAX_VALUE path keeps chars-only behavior for large models. Tests cover null-budget (all pages), token-budget truncation, and zero-budget skip. --- .../vip/mate/agent/AgentGraphBuilder.java | 9 ++- .../mate/wiki/service/WikiContextService.java | 37 ++++++++++-- .../service/WikiContextServiceBudgetTest.java | 59 ++++++++++++++++++- 3 files changed, 94 insertions(+), 11 deletions(-) diff --git a/mateclaw-server/src/main/java/vip/mate/agent/AgentGraphBuilder.java b/mateclaw-server/src/main/java/vip/mate/agent/AgentGraphBuilder.java index 2c669e0d..45a47bd8 100644 --- a/mateclaw-server/src/main/java/vip/mate/agent/AgentGraphBuilder.java +++ b/mateclaw-server/src/main/java/vip/mate/agent/AgentGraphBuilder.java @@ -1743,8 +1743,13 @@ public class AgentGraphBuilder { """; } - // Wiki 知识库上下文注入 - String wikiContext = wikiContextService.buildWikiContext(entity.getId()); + // Wiki 知识库上下文注入。Share the same prefix budget as the memory + // block so a large KB's page listing can't consume a fixed + // maxContextChars-sized slice of a small model's window every turn + // (issue #521). Integer.MAX_VALUE (the unbudgeted default path) keeps + // the legacy chars-only cap for large cloud models. + Integer wikiBudgetTokens = memoryBudgetTokens == Integer.MAX_VALUE ? null : memoryBudgetTokens; + String wikiContext = wikiContextService.buildWikiContext(entity.getId(), wikiBudgetTokens); return basePrompt + ABOUT_YOU_BLOCK + toolGuidance + searchGuidance + wikiContext; } diff --git a/mateclaw-server/src/main/java/vip/mate/wiki/service/WikiContextService.java b/mateclaw-server/src/main/java/vip/mate/wiki/service/WikiContextService.java index 29965c7b..f2c2adca 100644 --- a/mateclaw-server/src/main/java/vip/mate/wiki/service/WikiContextService.java +++ b/mateclaw-server/src/main/java/vip/mate/wiki/service/WikiContextService.java @@ -151,9 +151,24 @@ public class WikiContextService { * Build full wiki context for agent system prompt. */ public String buildWikiContext(Long agentId) { + return buildWikiContext(agentId, null); + } + + /** + * Budgeted variant of the system-prompt wiki listing. In addition to the + * absolute {@code maxContextChars} cap (sized for large cloud models), the + * enumerated page list may not exceed {@code budgetTokens} (estimated), so + * a large KB cannot consume a fixed ~{@code maxContextChars}-sized slice of + * a small local model's context window on every turn. A null budget keeps + * the previous chars-only behavior; a non-positive budget skips injection. + */ + public String buildWikiContext(Long agentId, Integer budgetTokens) { if (!properties.isEnabled()) { return ""; } + if (budgetTokens != null && budgetTokens <= 0) { + return ""; + } List kbs = kbService.listByAgentId(agentId); if (kbs.isEmpty()) { @@ -166,6 +181,9 @@ public class WikiContextService { int totalChars = 0; int maxChars = properties.getMaxContextChars(); + // Running estimate of what has been appended, so the page enumeration + // (the part that scales with KB file count) can respect budgetTokens. + int totalTokens = TokenEstimator.estimateTokens(sb.toString()); // Each KB renders as a HEADING-ONLY block (### ) followed by a // metadata line and its page list. The heading deliberately contains @@ -180,16 +198,20 @@ public class WikiContextService { List pages = pageService.listSummaries(kb.getId()); if (pages.isEmpty()) continue; - // Heading: pure KB name. This is what `kbName` expects verbatim. - sb.append("### ").append(kb.getName()).append("\n"); + // Heading: pure KB name (what `kbName` expects verbatim) plus a + // metadata line — built as one string so its tokens are budgeted too. + StringBuilder heading = new StringBuilder(); + heading.append("### ").append(kb.getName()).append("\n"); // Metadata line: page count first (easy to scan), then optional // description. Lives on its own line so it can't be confused for // part of the name. - sb.append(pages.size()).append(" pages"); + heading.append(pages.size()).append(" pages"); if (kb.getDescription() != null && !kb.getDescription().isBlank()) { - sb.append(" — ").append(kb.getDescription()); + heading.append(" — ").append(kb.getDescription()); } - sb.append("\n\n"); + heading.append("\n\n"); + sb.append(heading); + totalTokens += TokenEstimator.estimateTokens(heading.toString()); boolean compact = pages.size() > 20; @@ -204,12 +226,15 @@ public class WikiContextService { } line += "\n"; } - if (totalChars + line.length() > maxChars) { + int lineTokens = TokenEstimator.estimateTokens(line); + if (totalChars + line.length() > maxChars + || (budgetTokens != null && totalTokens + lineTokens > budgetTokens)) { sb.append("- ... and more (use wiki_list_pages to see all)\n"); break; } sb.append(line); totalChars += line.length(); + totalTokens += lineTokens; } sb.append("\n"); } diff --git a/mateclaw-server/src/test/java/vip/mate/wiki/service/WikiContextServiceBudgetTest.java b/mateclaw-server/src/test/java/vip/mate/wiki/service/WikiContextServiceBudgetTest.java index e5047a6e..655074f0 100644 --- a/mateclaw-server/src/test/java/vip/mate/wiki/service/WikiContextServiceBudgetTest.java +++ b/mateclaw-server/src/test/java/vip/mate/wiki/service/WikiContextServiceBudgetTest.java @@ -7,6 +7,7 @@ import org.mockito.Mockito; import vip.mate.wiki.WikiProperties; import vip.mate.wiki.dto.PageSearchResult; import vip.mate.wiki.model.WikiKnowledgeBaseEntity; +import vip.mate.wiki.model.WikiPageEntity; import java.util.List; @@ -17,12 +18,15 @@ import static org.mockito.ArgumentMatchers.anyLong; import static org.mockito.ArgumentMatchers.anyString; /** - * Tests for the token-budgeted knowledge-base relevance injection in - * {@link WikiContextService}. + * Tests for the token-budgeted knowledge-base injection in + * {@link WikiContextService} — both the per-query relevance block + * ({@code buildRelevantContext}) and the system-prompt page listing + * ({@code buildWikiContext}). */ class WikiContextServiceBudgetTest { private WikiKnowledgeBaseService kbService; + private WikiPageService pageService; private HybridRetriever hybridRetriever; private WikiProperties properties; private WikiContextService service; @@ -31,7 +35,7 @@ class WikiContextServiceBudgetTest { void setUp() { kbService = Mockito.mock(WikiKnowledgeBaseService.class); hybridRetriever = Mockito.mock(HybridRetriever.class); - WikiPageService pageService = Mockito.mock(WikiPageService.class); + pageService = Mockito.mock(WikiPageService.class); properties = new WikiProperties(); service = new WikiContextService(kbService, pageService, hybridRetriever, properties); @@ -84,4 +88,53 @@ class WikiContextServiceBudgetTest { assertEquals("", result); Mockito.verifyNoInteractions(hybridRetriever); } + + // ---- buildWikiContext (system-prompt page listing) budget (issue #521) ---- + + private void stubKbWithPages(int pageCount) { + WikiKnowledgeBaseEntity kb = new WikiKnowledgeBaseEntity(); + kb.setId(7L); + kb.setName("产品知识库"); + Mockito.when(kbService.listByAgentId(anyLong())).thenReturn(List.of(kb)); + + List pages = new java.util.ArrayList<>(); + for (int i = 0; i < pageCount; i++) { + WikiPageEntity p = new WikiPageEntity(); + p.setSlug("page-" + i); + p.setTitle("知识库页面标题" + i); + pages.add(p); + } + Mockito.when(pageService.listSummaries(anyLong())).thenReturn(pages); + } + + @Test + @DisplayName("null budget lists all pages (chars-only legacy behavior)") + void wikiContextNullBudgetListsAll() { + stubKbWithPages(50); + String result = service.buildWikiContext(1L, null); + assertTrue(result.contains("page-0")); + assertTrue(result.contains("page-49")); + assertFalse(result.contains("... and more")); + } + + @Test + @DisplayName("token budget truncates the page listing and appends the search hint") + void wikiContextBudgetTruncates() { + stubKbWithPages(50); + // Each compact line ("- page-N: 知识库页面标题N\n") is ~10+ CJK tokens; a + // small budget must cut the listing well before all 50 pages. + String result = service.buildWikiContext(1L, 80); + assertTrue(result.contains("page-0")); + assertFalse(result.contains("page-49")); + assertTrue(result.contains("... and more (use wiki_list_pages to see all)")); + } + + @Test + @DisplayName("zero or negative budget skips the wiki block entirely") + void wikiContextZeroBudgetSkips() { + String result = service.buildWikiContext(1L, 0); + assertEquals("", result); + Mockito.verifyNoInteractions(kbService); + Mockito.verifyNoInteractions(pageService); + } }