Coverage for src/lilbee/wiki/stubs.py: 100%
176 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-14 11:46 +0000
1"""The wiki's page index: every page that could exist, generated or not.
3NER extraction names every entity in the corpus and its chunk refs without
4spending an LLM call, so the browse tree can list every page as soon as a sync
5finishes and a body can be generated only when someone asks for that page. A
6build costs one call per source document, which scales with library size rather
7than with what the reader opens; this is the half of that cost that is free.
9Only entities are enumerable this way. LLM-curated concepts come from the
10per-source batched call, so they keep arriving published from an explicit
11``lilbee wiki build``.
12"""
14from __future__ import annotations
16import json
17import logging
18from dataclasses import dataclass
19from typing import TYPE_CHECKING, Any
21from lilbee.app.services import get_services
22from lilbee.core.config import cfg
24from .entity_extractor import EntityKind, get_entity_extractor
25from .generation import _corpus_chunks
26from .shared import (
27 WIKI_BUILD_LOCK,
28 WikiSubdir,
29 atomic_write_text,
30 is_pending_marker_text,
31)
33if TYPE_CHECKING:
34 from pathlib import Path
36 from lilbee.core.config import Config
37 from lilbee.data.store import SearchChunk, Store
39 from .entity_extractor import EntityExtractor, ExtractedEntity
41log = logging.getLogger(__name__)
43# The index file is a derived cache: the authoritative mention evidence lives in
44# the store's _wiki_mentions table, and this file is the corpus-wide >=floor view
45# rebuilt from it on every refresh. It lives beside the pages because it is read
46# on every browse and a wipe should remove it with everything else.
47STUB_INDEX_FILENAME = "stubs.json"
49# Schema marker, so a future shape change can be detected rather than parsed
50# as garbage. A mismatch is treated as no index at all and the next sync
51# rebuilds it.
52# Bumped to 3 when the index became a derived view of the store mention table:
53# an older file reads as absent so the first refresh rebuilds in full and seeds
54# the table.
55_INDEX_VERSION = 3
58@dataclass(frozen=True)
59class WikiStub:
60 """One entity the corpus names, with the evidence a page would be built from."""
62 slug: str
63 label: str
64 kind: EntityKind
65 type_hint: str
66 # Per-source mention counts, sorted by source, summed by ``mentions`` for
67 # the floor. Kept per source rather than as a single total because the
68 # store aggregate produces them per source, and because chunk refs are
69 # capped, so counting those would undercount a subject named more often
70 # than the cap.
71 source_mentions: tuple[tuple[str, int], ...]
72 chunk_refs: tuple[tuple[str, int], ...]
74 @property
75 def sources(self) -> tuple[str, ...]:
76 """Every document naming this subject."""
77 return tuple(source for source, _ in self.source_mentions)
79 @property
80 def mentions(self) -> int:
81 """How often the corpus names this subject, across all its sources."""
82 return sum(count for _, count in self.source_mentions)
84 @property
85 def subdir(self) -> WikiSubdir:
86 """Where this stub's page lands once generated."""
87 return WikiSubdir.CONCEPTS if self.kind is EntityKind.CONCEPT else WikiSubdir.ENTITIES
89 @property
90 def wiki_slug(self) -> str:
91 """The browse slug, ``<subdir>/<slug>``."""
92 return f"{self.subdir}/{self.slug}"
94 def to_dict(self) -> dict[str, Any]:
95 """Serialize for the on-disk index."""
96 return {
97 "slug": self.slug,
98 "label": self.label,
99 "kind": self.kind.value,
100 "type_hint": self.type_hint,
101 "source_mentions": [[source, count] for source, count in self.source_mentions],
102 "chunk_refs": [[source, index] for source, index in self.chunk_refs],
103 }
106def _stub_from_dict(raw: dict[str, Any]) -> WikiStub | None:
107 """Rebuild a stub from the index, or None when the row is unusable.
109 The index is a plain file a user can edit or a partial write can truncate,
110 so a bad row is dropped rather than failing the whole browse.
111 """
112 try:
113 return WikiStub(
114 slug=str(raw["slug"]),
115 label=str(raw["label"]),
116 kind=EntityKind(raw["kind"]),
117 type_hint=str(raw.get("type_hint", "")),
118 source_mentions=tuple((str(s), int(c)) for s, c in raw["source_mentions"]),
119 chunk_refs=tuple((str(s), int(i)) for s, i in raw.get("chunk_refs", [])),
120 )
121 except (KeyError, TypeError, ValueError):
122 log.warning("Dropping unreadable wiki stub row: %r", raw)
123 return None
126def stub_index_path(config: Config | None = None) -> Path:
127 """Where the index file lives."""
128 if config is None:
129 config = cfg
130 return config.data_root / config.wiki_dir / STUB_INDEX_FILENAME
133def _read_stub_index(config: Config | None = None) -> dict[str, WikiStub] | None:
134 """The index as stored, or None when it cannot be read.
136 None and an empty dict mean different things to an incremental refresh:
137 one has nothing to build on and must rebuild, the other is a corpus that
138 genuinely indexes to nothing and must not re-scan on every sync.
139 """
140 path = stub_index_path(config)
141 if not path.is_file():
142 return None
143 try:
144 payload = json.loads(path.read_text(encoding="utf-8"))
145 except (OSError, json.JSONDecodeError, UnicodeDecodeError):
146 log.warning("Wiki stub index unreadable", exc_info=True)
147 return None
148 if not isinstance(payload, dict) or payload.get("version") != _INDEX_VERSION:
149 log.info("Wiki stub index version mismatch")
150 return None
151 rows = payload.get("stubs")
152 if not isinstance(rows, list):
153 # Every index this code writes carries a stubs list, so one without it
154 # is damaged rather than describing a corpus that names nothing.
155 log.warning("Wiki stub index has no stubs list")
156 return None
157 parsed = [_stub_from_dict(row) for row in rows if isinstance(row, dict)]
158 if rows and not any(stub is not None for stub in parsed):
159 # It claims entries and none survived parsing. Reporting that as empty
160 # would let an incremental refresh build on it and drop everything the
161 # current sync did not touch.
162 log.warning("Wiki stub index holds no usable rows")
163 return None
164 return {stub.slug: stub for stub in parsed if stub is not None}
167def load_stub_index(config: Config | None = None) -> dict[str, WikiStub]:
168 """Read the index, keyed by slug. Empty when absent or unreadable."""
169 return _read_stub_index(config) or {}
172def save_stub_index(stubs: dict[str, WikiStub], config: Config | None = None) -> None:
173 """Write the index atomically, in slug order so the file is diffable."""
174 if config is None:
175 config = cfg
176 payload = {
177 "version": _INDEX_VERSION,
178 "stubs": [stubs[slug].to_dict() for slug in sorted(stubs)],
179 }
180 atomic_write_text(stub_index_path(config), json.dumps(payload, indent=2, sort_keys=False))
183def _recut_refs(stub: WikiStub, cap: int) -> WikiStub:
184 """Re-apply the cap across the stub's refs, most-mentioning source first."""
185 if len(stub.chunk_refs) <= cap:
186 return stub
187 by_source: dict[str, list[int]] = {}
188 for source, index in stub.chunk_refs:
189 by_source.setdefault(source, []).append(index)
190 counts = dict(stub.source_mentions)
191 ordered = sorted(by_source.items(), key=lambda kv: (-counts.get(kv[0], 0), kv[0]))
192 refs: list[tuple[str, int]] = []
193 for source, indexes in ordered:
194 for index in sorted(set(indexes)):
195 if len(refs) >= cap:
196 break
197 refs.append((source, index))
198 return WikiStub(
199 slug=stub.slug,
200 label=stub.label,
201 kind=stub.kind,
202 type_hint=stub.type_hint,
203 source_mentions=stub.source_mentions,
204 chunk_refs=tuple(refs),
205 )
208def _mention_rows_by_source(
209 entities: list[ExtractedEntity], cap: int
210) -> dict[str, list[dict[str, Any]]]:
211 """Split extracted entities into per-(subject, source) store rows.
213 One extraction pass yields entities whose refs span every source; the store
214 keeps them per source so a later sync can replace one source's evidence
215 without re-reading the rest. ``mention_count`` is the true per-source count;
216 the indices are capped, since the aggregate re-caps across sources anyway.
217 """
218 by_source: dict[str, list[dict[str, Any]]] = {}
219 for entity in entities:
220 per_source: dict[str, list[int]] = {}
221 for ref in entity.chunk_refs:
222 per_source.setdefault(ref.source, []).append(ref.chunk_index)
223 for source, indices in per_source.items():
224 unique = sorted(set(indices))
225 by_source.setdefault(source, []).append(
226 {
227 "slug": entity.slug,
228 "label": entity.label,
229 "kind": entity.kind.value,
230 "type_hint": entity.type_hint,
231 "source": source,
232 "mention_count": len(unique),
233 "chunk_indices": unique[:cap],
234 }
235 )
236 return by_source
239def _stubs_from_mention_rows(
240 rows: list[dict[str, Any]], config: Config, cap: int
241) -> dict[str, WikiStub]:
242 """Aggregate mention rows into the stubs whose corpus-wide count clears the
243 floor.
245 The floor is judged over the count summed across every source, so a subject
246 below it in each separately-synced document still qualifies once all of its
247 rows are present. This is the one place the floor is applied.
248 """
249 floor = config.wiki_entity_min_mentions
250 by_slug: dict[str, list[dict[str, Any]]] = {}
251 for row in rows:
252 by_slug.setdefault(row["slug"], []).append(row)
253 stubs: dict[str, WikiStub] = {}
254 for slug, slug_rows in by_slug.items():
255 if sum(r["mention_count"] for r in slug_rows) < floor:
256 continue
257 # The source that names the subject most often labels its page.
258 canonical = max(slug_rows, key=lambda r: (r["mention_count"], r["source"]))
259 source_mentions = tuple(sorted((r["source"], r["mention_count"]) for r in slug_rows))
260 refs = tuple(
261 (r["source"], index)
262 for r in sorted(slug_rows, key=lambda r: r["source"])
263 for index in r["chunk_indices"]
264 )
265 stub = WikiStub(
266 slug=slug,
267 label=canonical["label"],
268 kind=EntityKind(canonical["kind"]),
269 type_hint=canonical["type_hint"],
270 source_mentions=source_mentions,
271 chunk_refs=refs,
272 )
273 stubs[slug] = _recut_refs(stub, cap)
274 return stubs
277def _write_source_mentions(
278 store: Store,
279 extractor: EntityExtractor,
280 chunks: list[SearchChunk],
281 sources: set[str],
282 cap: int,
283) -> set[str]:
284 """Extract *chunks* and replace the mention rows for every source in
285 *sources*, returning the slugs those sources now name.
287 A source that named something and now names nothing is still replaced, with
288 an empty set, so its stale rows go.
289 """
290 rows_by_source = _mention_rows_by_source(extractor.extract(chunks), cap)
291 affected: set[str] = set()
292 for source in sources:
293 rows = rows_by_source.get(source, [])
294 store.replace_wiki_mentions_for_source(source, rows)
295 affected.update(row["slug"] for row in rows)
296 return affected
299def refresh_stub_index(
300 store: Store,
301 config: Config | None = None,
302 *,
303 sources: set[str] | None = None,
304) -> dict[str, WikiStub]:
305 """Rebuild the index from the corpus and persist it.
307 The store's ``_wiki_mentions`` table is the source of truth. ``sources=None``
308 re-extracts the whole corpus; otherwise only those sources are re-extracted
309 and their mention rows replaced. The stub index is then the corpus-wide
310 aggregate of that table filtered to the mention floor, so a subject below
311 the floor in each separately-synced document still appears once all its
312 evidence is present.
314 Spends no LLM call. Holds the wiki mutex because it writes into the wiki
315 directory alongside builds and prunes.
316 """
317 if config is None:
318 config = cfg
319 cap = config.wiki_stub_max_chunk_refs
320 with WIKI_BUILD_LOCK:
321 # Floor forced to one for extraction: the store keeps every mention so
322 # the corpus-wide floor can be judged over the aggregate rather than one
323 # sync's slice. _stubs_from_mention_rows applies the real floor.
324 pass_config = config.model_copy(update={"wiki_entity_min_mentions": 1})
325 extractor = get_entity_extractor(
326 config.wiki_entity_mode, get_services().provider, pass_config
327 )
328 if not extractor.available():
329 # An unavailable backend extracts nothing, which is indistinguishable
330 # from a corpus that names nothing. Rebuilding on that would drop the
331 # whole index; leave both the store and the file untouched.
332 log.warning("Wiki entity extractor unavailable; leaving the stub index unchanged")
333 return load_stub_index(config)
335 # A cold or file-only-migrated store has no rows to aggregate, so an
336 # incremental pass would derive an empty index. Rebuild in full to seed
337 # the table; incremental passes maintain it from there.
338 if sources is None or not store.has_wiki_mentions():
339 store.clear_wiki_mentions()
340 all_sources = {record["filename"] for record in store.get_sources()}
341 _write_source_mentions(store, extractor, _corpus_chunks(store), all_sources, cap)
342 stubs = _stubs_from_mention_rows(store.wiki_mention_rows(), config, cap)
343 else:
344 previous = load_stub_index(config)
345 chunks = [c for name in sorted(sources) for c in store.get_chunks_by_source(name)]
346 affected = _write_source_mentions(store, extractor, chunks, sources, cap)
347 affected |= {slug for slug, stub in previous.items() if set(stub.sources) & sources}
348 recomputed = _stubs_from_mention_rows(
349 store.wiki_mention_rows(slugs=affected), config, cap
350 )
351 stubs = {slug: stub for slug, stub in previous.items() if slug not in affected}
352 stubs.update(recomputed)
353 save_stub_index(stubs, config)
354 log.info("Wiki stub index: %d entities across the corpus", len(stubs))
355 return stubs
358def _page_exists(stub: WikiStub, wiki_root: Path) -> bool:
359 """Whether this subject already has a page, published or awaiting review.
361 A page the faithfulness gate routed to drafts/ has been written and is
362 waiting on a human. Offering it as unwritten invites generating it again
363 and again, each time for another LLM call, while the draft sits there.
365 An archived page does not count. Prune retires a page when the sources it
366 cited are deleted, while the subject can still be named by documents that
367 survive, and nothing lists or restores archive/. Suppressing the stub would
368 make the subject unreachable from either pane; listing it costs nothing,
369 since generation only ever runs when a user asks for it.
370 """
371 if (wiki_root / f"{stub.wiki_slug}.md").is_file():
372 return True
373 draft = wiki_root / WikiSubdir.DRAFTS / f"{stub.slug}.md"
374 return draft.is_file() and not _is_placeholder(draft)
377def drop_sources_from_index(names: set[str], config: Config | None = None) -> None:
378 """Forget documents the library no longer has.
380 Re-aggregates the affected slugs from the store, whose rows for the removed
381 sources went with their chunks, so it costs no extraction pass. Without it a
382 removed document keeps its subjects in the browse tree forever: its skip
383 marker keeps it out of later syncs, so no refresh ever revisits them.
384 """
385 if not names:
386 return
387 if config is None:
388 config = cfg
389 cap = config.wiki_stub_max_chunk_refs
390 with WIKI_BUILD_LOCK:
391 previous = load_stub_index(config)
392 affected = {slug for slug, stub in previous.items() if set(stub.sources) & names}
393 if not affected:
394 return
395 # The removed sources' mention rows went with their chunks, so
396 # re-aggregating the affected slugs reflects the smaller corpus: a
397 # subject only those sources named drops out, one still named elsewhere
398 # keeps its surviving evidence.
399 store = get_services().store
400 recomputed = _stubs_from_mention_rows(store.wiki_mention_rows(slugs=affected), config, cap)
401 stubs = {slug: stub for slug, stub in previous.items() if slug not in affected}
402 stubs.update(recomputed)
403 if stubs != previous:
404 save_stub_index(stubs, config)
407def _is_placeholder(draft: Path) -> bool:
408 """Whether a draft is a PENDING marker rather than written content.
410 A marker records that generation failed to produce the section, so the
411 subject still has no page and must stay listed as unwritten.
412 """
413 try:
414 text = draft.read_text(encoding="utf-8")
415 except OSError:
416 return False
417 return is_pending_marker_text(text)
420def ungenerated_stubs(stubs: dict[str, WikiStub], wiki_root: Path) -> list[WikiStub]:
421 """The stubs nothing has written a page for yet, in slug order."""
422 return [stub for _, stub in sorted(stubs.items()) if not _page_exists(stub, wiki_root)]