From 4aa36afb97da75232ec615fd0164739e8c22532c Mon Sep 17 00:00:00 2001 From: rodrigedilson-ia Date: Tue, 22 Sep 2026 11:24:52 -0300 Subject: [PATCH] fix(compiler): cap concept/entity brief lists in the plan prompt MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `_read_concept_briefs` and `_read_entity_briefs` read every page in the KB and the result is rebuilt into the concepts-plan prompt for each compiled document. Their size is therefore O(KB), with no ceiling — so a large KB eventually pushes that call past the model's context window. The failure mode is unkind: it lands as a hard error partway through a recompile, not as degraded output. We hit it on a KB of ~260 documents, where seven consecutive `recompile` runs exited non-zero, all of them at the plan step, with prompts measured between 200k and 356k tokens against a 200k limit. Because it scales with the KB, every retry was guaranteed to fail the same way. This adds a per-list character budget (`briefs_budget_chars`, default 120_000 — roughly 30k tokens each, so a KB has to grow well past a few hundred pages before anything is trimmed). Set it to 0 to restore the previous uncapped behavior. Two details worth calling out: - **Ranking is by source count, not alphabetical.** That count is already the cross-document recurrence signal the plan call uses for create-vs-update, so it is also the right thing to keep when the list has to be cut. Ties break alphabetically, because prompt caching depends on byte-identical prefixes and an unstable order would silently cost cache hits. - **The trim is announced in the list.** A shortened list that reads as complete is worse than a long one: the planner uses these briefs to decide create-vs-update, so a silently dropped concept comes back as a duplicate page for something the KB already has. The trailing marker states how many were omitted and tells the model to prefer `update` when unsure. In the degenerate case where the budget cannot fit even one line, the output says the list was omitted rather than returning an empty one — an empty list reads as an empty KB, which would push the planner to recreate everything. Small KBs are byte-identical to before (covered by a test), so prompt caches are not invalidated for existing users. --- openkb/agent/compiler.py | 144 +++++++++++++++++++++++++++++++++++---- openkb/config.py | 5 ++ tests/test_compiler.py | 124 +++++++++++++++++++++++++++++++++ 3 files changed, 261 insertions(+), 12 deletions(-) diff --git a/openkb/agent/compiler.py b/openkb/agent/compiler.py index d0c9f878..f0e513dd 100644 --- a/openkb/agent/compiler.py +++ b/openkb/agent/compiler.py @@ -724,7 +724,103 @@ def _resolve_description(fm: dict) -> str: return "" -def _read_concept_briefs(wiki_dir: Path) -> str: +DEFAULT_BRIEFS_BUDGET_CHARS = 120_000 +"""Per-list character cap for the concept/entity briefs in the plan prompt. + +Roughly 30k tokens each, so a KB has to grow well past a few hundred pages +before anything is trimmed — the cap is a ceiling against unbounded growth, not +a working limit. Overridable with the ``briefs_budget_chars`` config key; set it +to 0 to restore the previous uncapped behavior. + +Why a cap at all: ``_read_concept_briefs`` and ``_read_entity_briefs`` read +*every* page in the KB and the result is rebuilt into the plan prompt for each +compiled document. The size is therefore O(KB), and a large KB eventually +exceeds the model's context window — which surfaces as a hard failure partway +through a recompile, not as degraded output. +""" + + +def _resolve_briefs_budget(config: dict) -> int: + """Read ``briefs_budget_chars`` from config, falling back to the default. + + A non-integer or negative value falls back rather than raising: a typo in a + config file should not abort a compile, and a negative budget has no sane + reading. ``0`` is meaningful and passes through — it disables trimming. + """ + raw = config.get("briefs_budget_chars", DEFAULT_BRIEFS_BUDGET_CHARS) + if isinstance(raw, bool) or not isinstance(raw, int) or raw < 0: + logger.warning( + "briefs_budget_chars=%r is not a non-negative integer — using %d", + raw, + DEFAULT_BRIEFS_BUDGET_CHARS, + ) + return DEFAULT_BRIEFS_BUDGET_CHARS + return raw + + +def _n_sources(fm_dict: dict) -> int: + """How many documents a page cites, from its ``sources:`` frontmatter list. + + This is the cross-document recurrence signal already used in the entity + brief line. It doubles as the salience ranking when the brief list has to + be trimmed to fit a budget: a concept seen in many documents is the one the + planner most needs to know about, because it is the one most likely to be + updated rather than created. + """ + sources = fm_dict.get("sources") + return len(sources) if isinstance(sources, list) else 0 + + +def _fit_briefs(ranked: list[tuple[int, str]], budget_chars: int, noun: str) -> str: + """Join brief lines under a character budget, ranked by salience. + + ``ranked`` is ``[(n_sources, line), ...]``. Lines are emitted most-cited + first, then alphabetically (the caller sorts), until the budget is spent. + + **The truncation is announced in the returned text.** A shortened list that + looks complete is worse than a long one: the planner reads these briefs to + decide create-vs-update, so a silently dropped concept comes back as a + duplicate page for something the KB already has. The trailing marker tells + the model the list is partial, so "not listed" stops meaning "not present". + + A budget of 0 or less disables trimming, which keeps the previous behavior + available for callers that want the full list. + """ + if budget_chars <= 0: + return "\n".join(line for _, line in ranked) or "(none yet)" + + kept: list[str] = [] + used = 0 + for _, line in ranked: + if used + len(line) + 1 > budget_chars: + break + kept.append(line) + used += len(line) + 1 + + dropped = len(ranked) - len(kept) + if not dropped: + return "\n".join(kept) or "(none yet)" + if not kept: + # Budget smaller than a single line. Say so rather than return an empty + # list that reads like an empty KB. + return f"(list omitted: {len(ranked)} {noun} exceed the brief budget)" + + logger.info( + "brief list trimmed to fit budget: kept %d of %d %s (%d chars)", + len(kept), + len(ranked), + noun, + budget_chars, + ) + kept.append( + f"- (… {dropped} more {noun} not listed — this list is truncated to the " + f"{len(kept) - 1} most-cited; absence here does NOT mean the page is " + f"missing, so prefer 'update' over 'create' when unsure)" + ) + return "\n".join(kept) + + +def _read_concept_briefs(wiki_dir: Path, budget_chars: int = 0) -> str: """Read existing concept pages and return compact one-line summaries. For each concept, reads the ``description:`` field (falling back to legacy @@ -732,6 +828,13 @@ def _read_concept_briefs(wiki_dir: Path) -> str: truncating the first 150 chars of the body (newlines collapsed to spaces). Formats each as ``- {slug}: {description}``. + With ``budget_chars`` > 0 the list is capped at that many characters, most- + cited concepts first (see :func:`_fit_briefs`). The cap exists because this + list grows with the KB and lands in the plan prompt on every compiled + document: an unbounded list eventually exceeds the model's context window, + and the failure arrives as a hard error mid-compile rather than as degraded + output. + Returns "(none yet)" if the concepts directory is missing or empty. """ concepts_dir = wiki_dir / "concepts" @@ -742,7 +845,7 @@ def _read_concept_briefs(wiki_dir: Path) -> str: if not md_files: return "(none yet)" - lines: list[str] = [] + ranked: list[tuple[int, str]] = [] for path in md_files: text = path.read_text(encoding="utf-8") fm_dict = frontmatter.parse(text) @@ -752,17 +855,22 @@ def _read_concept_briefs(wiki_dir: Path) -> str: body = parts[1] if parts is not None else text brief = body.strip().replace("\n", " ")[:150] if brief: - lines.append(f"- {path.stem}: {brief}") + ranked.append((_n_sources(fm_dict), f"- {path.stem}: {brief}")) - return "\n".join(lines) or "(none yet)" + # Most-cited first; alphabetical within a tier so the prompt stays stable + # across runs (prompt caching depends on byte-identical prefixes). + ranked.sort(key=lambda item: (-item[0], item[1])) + return _fit_briefs(ranked, budget_chars, "concepts") -def _read_entity_briefs(wiki_dir: Path) -> str: +def _read_entity_briefs(wiki_dir: Path, budget_chars: int = 0) -> str: """Read existing entity pages as compact lines for the plan call. Formats each as ``- {slug} ({type}, {n} sources) — {brief}``. The source count is the cross-document recurrence signal the LLM uses to decide - create-vs-update and salience. Returns "(none yet)" when empty. + create-vs-update and salience — and, with ``budget_chars`` > 0, the ranking + used to decide what stays when the list is capped. Returns "(none yet)" + when empty. """ entities_dir = wiki_dir / "entities" if not entities_dir.exists(): @@ -772,21 +880,22 @@ def _read_entity_briefs(wiki_dir: Path) -> str: if not md_files: return "(none yet)" - lines: list[str] = [] + ranked: list[tuple[int, str]] = [] for path in md_files: text = path.read_text(encoding="utf-8") fm_dict = frontmatter.parse(text) brief = _resolve_description(fm_dict) etype = str(fm_dict.get("type") or "").strip().lower() or "other" - n_sources = len(fm_dict["sources"]) if isinstance(fm_dict.get("sources"), list) else 0 + n_sources = _n_sources(fm_dict) if not brief: parts = frontmatter.split(text) body = parts[1] if parts is not None else text brief = body.strip().replace("\n", " ")[:150] suffix = f" — {brief}" if brief else "" - lines.append(f"- {path.stem} ({etype}, {n_sources} sources){suffix}") + ranked.append((n_sources, f"- {path.stem} ({etype}, {n_sources} sources){suffix}")) - return "\n".join(lines) or "(none yet)" + ranked.sort(key=lambda item: (-item[0], item[1])) + return _fit_briefs(ranked, budget_chars, "entities") def _iter_h2_headings(lines: list[str]) -> list[tuple[int, str]]: @@ -1603,6 +1712,7 @@ async def _compile_concepts( doc_type: str = "short", rewrite_summary: bool = False, entity_types: list[str] | None = None, + briefs_budget_chars: int | None = None, bundle=None, ) -> None: """Shared Steps 2-4: concepts plan → generate/update → index. @@ -1624,8 +1734,14 @@ async def _compile_concepts( valid_types = frozenset(entity_types) # --- Step 2: Get concepts plan (A cached) --- - concept_briefs = _read_concept_briefs(wiki_dir) - entity_briefs = _read_entity_briefs(wiki_dir) + # Both lists grow with the KB and are rebuilt into the plan prompt for every + # compiled document, so their combined size is what eventually pushes this + # call past the model's context window. The budget caps each one; see + # ``_fit_briefs`` for why the trim is announced rather than silent. + if briefs_budget_chars is None: + briefs_budget_chars = DEFAULT_BRIEFS_BUDGET_CHARS + concept_briefs = _read_concept_briefs(wiki_dir, briefs_budget_chars) + entity_briefs = _read_entity_briefs(wiki_dir, briefs_budget_chars) # Second cache breakpoint: end of the assistant summary message. Covers # (system + doc + summary) for the plan call and every concept call. @@ -2220,6 +2336,7 @@ async def compile_short_doc( config = resolve_effective_config(kb_dir)[0] language: str = config.get("language", "en") entity_types = resolve_entity_types(config) + briefs_budget = _resolve_briefs_budget(config) wiki_dir = kb_dir / "wiki" schema_md = get_agents_md(wiki_dir) @@ -2280,6 +2397,7 @@ async def compile_short_doc( doc_type="short", rewrite_summary=True, entity_types=entity_types, + briefs_budget_chars=briefs_budget, bundle=bundle, ) finally: @@ -2308,6 +2426,7 @@ async def compile_long_doc( config = resolve_effective_config(kb_dir)[0] language: str = config.get("language", "en") entity_types = resolve_entity_types(config) + briefs_budget = _resolve_briefs_budget(config) wiki_dir = kb_dir / "wiki" schema_md = get_agents_md(wiki_dir) @@ -2364,6 +2483,7 @@ async def compile_long_doc( doc_brief=doc_description, doc_type="pageindex", entity_types=entity_types, + briefs_budget_chars=briefs_budget, bundle=bundle, ) finally: diff --git a/openkb/config.py b/openkb/config.py index 95ca8691..c2e274d2 100644 --- a/openkb/config.py +++ b/openkb/config.py @@ -36,6 +36,11 @@ # global/KB list overrides it wholesale; resolve_entity_types cleans the # effective value on read. "entity_types": list(DEFAULT_ENTITY_TYPES), + # Per-list character cap for the concept/entity briefs in the compile plan + # prompt. Those lists are O(KB) and are rebuilt for every compiled document, + # so without a ceiling a large KB eventually exceeds the model's context + # window mid-recompile. 0 disables trimming. + "briefs_budget_chars": 120_000, } GLOBAL_CONFIG_DIR = Path.home() / ".config" / "openkb" diff --git a/tests/test_compiler.py b/tests/test_compiler.py index 95a57cc4..dbe4ba57 100644 --- a/tests/test_compiler.py +++ b/tests/test_compiler.py @@ -9,6 +9,7 @@ from openkb.agent.compiler import ( _ENTITY_TYPE_LIST, + DEFAULT_BRIEFS_BUDGET_CHARS, _add_related_link, _backlink_concepts, _backlink_entities, @@ -23,6 +24,7 @@ _read_entity_briefs, _read_wiki_context, _remove_source_from_frontmatter, + _resolve_briefs_budget, _sanitize_concept_name, _update_index, _write_concept, @@ -2830,3 +2832,125 @@ def test_concept_update_malformed_frontmatter_rebuilds(self, tmp_path): assert 'type: "Concept"' in text # Must have a properly closed frontmatter block (two '---' occurrences). assert text.count("---") >= 2 + + +class TestBriefsBudget: + """The concept/entity brief lists are O(KB) and are rebuilt into the plan + prompt for every compiled document. Without a ceiling a large KB eventually + exceeds the model's context window, and the failure arrives as a hard error + partway through a recompile rather than as degraded output. + """ + + @staticmethod + def _concept(wiki, slug, description, n_sources=1): + d = wiki / "concepts" + d.mkdir(parents=True, exist_ok=True) + sources = ", ".join(f'"summaries/doc{i}.md"' for i in range(n_sources)) + (d / f"{slug}.md").write_text( + f'---\ntype: "Concept"\nsources: [{sources}]\n' + f'description: "{description}"\n---\n\n# {slug}\n', + encoding="utf-8", + ) + + def test_under_budget_is_unchanged(self, tmp_path): + """Small KBs must see byte-identical output — the cap is a ceiling, not + a working limit, and the plan prompt is prompt-cached on exact prefixes. + """ + wiki = tmp_path / "wiki" + self._concept(wiki, "alpha", "first") + self._concept(wiki, "beta", "second") + + assert _read_concept_briefs(wiki, 120_000) == _read_concept_briefs(wiki, 0) + + def test_budget_keeps_most_cited_first(self, tmp_path): + """Source count is the recurrence signal the planner uses for + create-vs-update, so it is also the right thing to keep when trimming. + """ + wiki = tmp_path / "wiki" + self._concept(wiki, "rare", "x" * 60, n_sources=1) + self._concept(wiki, "common", "y" * 60, n_sources=9) + + out = _read_concept_briefs(wiki, 90) + + assert "common" in out + assert "- rare:" not in out + + def test_truncation_is_announced(self, tmp_path): + """A shortened list that looks complete is worse than a long one: the + planner reads these briefs to decide create-vs-update, so a silently + dropped concept comes back as a duplicate page for something the KB + already has. + """ + wiki = tmp_path / "wiki" + for i in range(12): + self._concept(wiki, f"concept-{i:02d}", "z" * 80, n_sources=i) + + out = _read_concept_briefs(wiki, 300) + + assert "not listed" in out + assert "prefer 'update' over 'create'" in out + + def test_budget_zero_disables_trimming(self, tmp_path): + wiki = tmp_path / "wiki" + for i in range(30): + self._concept(wiki, f"c{i:02d}", "w" * 100) + + out = _read_concept_briefs(wiki, 0) + + assert len(out.splitlines()) == 30 + assert "not listed" not in out + + def test_budget_smaller_than_one_line_says_so(self, tmp_path): + """An empty string would read as an empty KB, which is a different + thing entirely and would push the planner to recreate everything.""" + wiki = tmp_path / "wiki" + self._concept(wiki, "alpha", "x" * 200) + + out = _read_concept_briefs(wiki, 5) + + assert "list omitted" in out + assert out != "(none yet)" + + def test_ordering_is_stable_within_a_tier(self, tmp_path): + """Prompt caching depends on byte-identical prefixes, so equal-salience + entries must not reorder between runs.""" + wiki = tmp_path / "wiki" + for slug in ("gamma", "alpha", "beta"): + self._concept(wiki, slug, "same", n_sources=2) + + primeira = _read_concept_briefs(wiki, 120_000) + + assert primeira == _read_concept_briefs(wiki, 120_000) + assert primeira.index("alpha") < primeira.index("beta") < primeira.index("gamma") + + def test_entity_briefs_respect_the_same_budget(self, tmp_path): + wiki = tmp_path / "wiki" + d = wiki / "entities" + d.mkdir(parents=True) + for i in range(10): + (d / f"e{i:02d}.md").write_text( + f'---\ntype: "organization"\nsources: ["summaries/d{i}.md"]\n' + f'description: "{"q" * 90}"\n---\n\n# e{i:02d}\n', + encoding="utf-8", + ) + + out = _read_entity_briefs(wiki, 250) + + assert "not listed" in out + assert len(out.splitlines()) < 10 + + +class TestResolveBriefsBudget: + """A typo in a config file must not abort a compile.""" + + def test_default_when_absent(self): + assert _resolve_briefs_budget({}) == DEFAULT_BRIEFS_BUDGET_CHARS + + def test_explicit_zero_passes_through(self): + """0 is meaningful — it disables trimming — so it must not be treated + as 'unset'.""" + assert _resolve_briefs_budget({"briefs_budget_chars": 0}) == 0 + + @pytest.mark.parametrize("bad", ["lots", -1, None, 1.5, True]) + def test_invalid_falls_back(self, bad): + assert _resolve_briefs_budget({"briefs_budget_chars": bad}) == DEFAULT_BRIEFS_BUDGET_CHARS