opensource/mateclaw
2
0
Fork 0
mirror of https://gitee.com/mateos/mateclaw.git synced 2026-09-13 03:13:41 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

fix(wiki): budget the system-prompt KB page listing to the model window

Browse Source 
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.
This commit is contained in:
倪程伟 2026-07-15 14:52:43 +08:00 committed by GitHub
parent 04a9bb9e13
commit fc3d84d6c2
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 94 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
mateclaw-server/src/main/java/vip/mate/agent/AgentGraphBuilder.java
View File

@ -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;
}

37
mateclaw-server/src/main/java/vip/mate/wiki/service/WikiContextService.java
View File

@ -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<WikiKnowledgeBaseEntity> 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 (### <name>) followed by a
// metadata line and its page list. The heading deliberately contains
@ -180,16 +198,20 @@ public class WikiContextService {
List<WikiPageEntity> 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");
}

59
mateclaw-server/src/test/java/vip/mate/wiki/service/WikiContextServiceBudgetTest.java
View File

@ -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<WikiPageEntity> 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);
}
}