Coverage for src/lilbee/wiki/stubs.py: 100%
240 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-11-06 08:17 +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
23from lilbee.data.store import SourceType
25from .entity_extractor import EntityKind, get_entity_extractor
26from .entity_extractor.factory import effective_entity_mode
27from .generation import _corpus_chunks
28from .shared import (
29 WIKI_BUILD_LOCK,
30 WikiSubdir,
31 atomic_write_text,
32 is_pending_marker_text,
33)
35if TYPE_CHECKING:
36 from pathlib import Path
38 from lilbee.core.config import Config
39 from lilbee.data.store import SearchChunk, Store
41 from .entity_extractor import EntityExtractor, ExtractedEntity
43log = logging.getLogger(__name__)
45# The index file is a derived cache: the authoritative mention evidence lives in
46# the store's _wiki_mentions table, and this file is the corpus-wide >=floor view
47# rebuilt from it on every refresh. It lives beside the pages because it is read
48# on every browse and a wipe should remove it with everything else.
49STUB_INDEX_FILENAME = "stubs.json"
51# Schema marker, so a shape change is detected rather than parsed as garbage.
52# Version 4 names each source by its id. A file of another version is
53# aggregated again from the store's mention table, which needs no extraction;
54# with no mention rows it reads as no index at all and the next sync rebuilds it.
55_INDEX_VERSION = 4
57_NO_EXTRACTOR = "Wiki entity extractor unavailable; leaving the stub index unchanged"
60@dataclass(frozen=True)
61class WikiStub:
62 """One entity the corpus names, with the evidence a page would be built from."""
64 slug: str
65 label: str
66 kind: EntityKind
67 type_hint: str
68 # Per-source mention counts by source id, sorted, summed by ``mentions`` for
69 # the floor. Kept per source rather than as a single total because the
70 # store aggregate produces them per source, and because chunk refs are
71 # capped, so counting those would undercount a subject named more often
72 # than the cap.
73 source_mentions: tuple[tuple[str, int], ...]
74 chunk_refs: tuple[tuple[str, int], ...]
76 @property
77 def sources(self) -> tuple[str, ...]:
78 """The id of every document naming this subject."""
79 return tuple(source for source, _ in self.source_mentions)
81 @property
82 def mentions(self) -> int:
83 """How often the corpus names this subject, across all its sources."""
84 return sum(count for _, count in self.source_mentions)
86 @property
87 def subdir(self) -> WikiSubdir:
88 """Where this stub's page lands once generated."""
89 return WikiSubdir.CONCEPTS if self.kind is EntityKind.CONCEPT else WikiSubdir.ENTITIES
91 @property
92 def wiki_slug(self) -> str:
93 """The browse slug, ``<subdir>/<slug>``."""
94 return f"{self.subdir}/{self.slug}"
96 def to_dict(self) -> dict[str, Any]:
97 """Serialize for the on-disk index."""
98 return {
99 "slug": self.slug,
100 "label": self.label,
101 "kind": self.kind.value,
102 "type_hint": self.type_hint,
103 "source_mentions": [[source, count] for source, count in self.source_mentions],
104 "chunk_refs": [[source, index] for source, index in self.chunk_refs],
105 }
108def _stub_from_dict(raw: dict[str, Any]) -> WikiStub | None:
109 """Rebuild a stub from the index, or None when the row is unusable.
111 The index is a plain file a user can edit or a partial write can truncate,
112 so a bad row is dropped rather than failing the whole browse.
113 """
114 try:
115 return WikiStub(
116 slug=str(raw["slug"]),
117 label=str(raw["label"]),
118 kind=EntityKind(raw["kind"]),
119 type_hint=str(raw.get("type_hint", "")),
120 source_mentions=tuple((str(s), int(c)) for s, c in raw["source_mentions"]),
121 chunk_refs=tuple((str(s), int(i)) for s, i in raw.get("chunk_refs", [])),
122 )
123 except (KeyError, TypeError, ValueError):
124 log.warning("Dropping unreadable wiki stub row: %r", raw)
125 return None
128def stub_index_path(config: Config | None = None) -> Path:
129 """Where the index file lives."""
130 if config is None:
131 config = cfg
132 return config.data_root / config.wiki_dir / STUB_INDEX_FILENAME
135def _read_stub_index(config: Config | None = None) -> dict[str, WikiStub] | None:
136 """The index as stored, or None when it cannot be read.
138 None and an empty dict mean different things to an incremental refresh:
139 one has nothing to build on and must rebuild, the other is a corpus that
140 genuinely indexes to nothing and must not re-scan on every sync.
141 """
142 path = stub_index_path(config)
143 if not path.is_file():
144 return None
145 try:
146 payload = json.loads(path.read_text(encoding="utf-8"))
147 except (OSError, json.JSONDecodeError, UnicodeDecodeError):
148 log.warning("Wiki stub index unreadable", exc_info=True)
149 return None
150 if not isinstance(payload, dict) or payload.get("version") != _INDEX_VERSION:
151 log.info("Wiki stub index version mismatch")
152 return _index_from_mentions(config or cfg)
153 rows = payload.get("stubs")
154 if not isinstance(rows, list):
155 # Every index this code writes carries a stubs list, so one without it
156 # is damaged rather than describing a corpus that names nothing.
157 log.warning("Wiki stub index has no stubs list")
158 return None
159 parsed = [_stub_from_dict(row) for row in rows if isinstance(row, dict)]
160 if rows and not any(stub is not None for stub in parsed):
161 # It claims entries and none survived parsing. Reporting that as empty
162 # would let an incremental refresh build on it and drop everything the
163 # current sync did not touch.
164 log.warning("Wiki stub index holds no usable rows")
165 return None
166 return {stub.slug: stub for stub in parsed if stub is not None}
169def _index_from_mentions(config: Config) -> dict[str, WikiStub] | None:
170 """The index aggregated from the store's mention rows and saved; None with no rows."""
171 store = get_services().store
172 with WIKI_BUILD_LOCK:
173 if not store.has_wiki_mentions():
174 return None
175 rows = store.wiki_mention_rows()
176 stubs = _stubs_from_mention_rows(rows, config, config.wiki_stub_max_chunk_refs)
177 save_stub_index(stubs, config)
178 return stubs
181def load_stub_index(config: Config | None = None) -> dict[str, WikiStub]:
182 """Read the index, keyed by slug. Empty when absent or unreadable."""
183 return _read_stub_index(config) or {}
186def save_stub_index(stubs: dict[str, WikiStub], config: Config | None = None) -> None:
187 """Write the index atomically, in slug order so the file is diffable."""
188 if config is None:
189 config = cfg
190 payload = {
191 "version": _INDEX_VERSION,
192 "stubs": [stubs[slug].to_dict() for slug in sorted(stubs)],
193 }
194 atomic_write_text(stub_index_path(config), json.dumps(payload, indent=2, sort_keys=False))
197def _recut_refs(stub: WikiStub, cap: int) -> WikiStub:
198 """Re-apply the cap across the stub's refs, most-mentioning source first."""
199 if len(stub.chunk_refs) <= cap:
200 return stub
201 by_source: dict[str, list[int]] = {}
202 for source, index in stub.chunk_refs:
203 by_source.setdefault(source, []).append(index)
204 counts = dict(stub.source_mentions)
205 ordered = sorted(by_source.items(), key=lambda kv: (-counts.get(kv[0], 0), kv[0]))
206 refs: list[tuple[str, int]] = []
207 for source, indexes in ordered:
208 for index in sorted(set(indexes)):
209 if len(refs) >= cap:
210 break
211 refs.append((source, index))
212 return WikiStub(
213 slug=stub.slug,
214 label=stub.label,
215 kind=stub.kind,
216 type_hint=stub.type_hint,
217 source_mentions=stub.source_mentions,
218 chunk_refs=tuple(refs),
219 )
222def _mention_rows_by_source(
223 entities: list[ExtractedEntity], cap: int
224) -> dict[str, list[dict[str, Any]]]:
225 """Split extracted entities into per-(subject, source) store rows.
227 One extraction pass yields entities whose refs span every source; the store
228 keeps them per source so a later sync can replace one source's evidence
229 without re-reading the rest. ``mention_count`` is the true per-source count;
230 the indices are capped, since the aggregate re-caps across sources anyway.
231 """
232 by_source: dict[str, list[dict[str, Any]]] = {}
233 for entity in entities:
234 per_source: dict[str, list[int]] = {}
235 for ref in entity.chunk_refs:
236 per_source.setdefault(ref.source, []).append(ref.chunk_index)
237 for source, indices in per_source.items():
238 unique = sorted(set(indices))
239 by_source.setdefault(source, []).append(
240 {
241 "slug": entity.slug,
242 "label": entity.label,
243 "kind": entity.kind.value,
244 "type_hint": entity.type_hint,
245 "source": source,
246 "mention_count": len(unique),
247 "chunk_indices": unique[:cap],
248 }
249 )
250 return by_source
253def _stubs_from_mention_rows(
254 rows: list[dict[str, Any]], config: Config, cap: int
255) -> dict[str, WikiStub]:
256 """Aggregate mention rows into the stubs whose corpus-wide count clears the
257 floor.
259 The floor is judged over the count summed across every source, so a subject
260 below it in each separately-synced document still qualifies once all of its
261 rows are present. This is the one place the floor is applied.
262 """
263 floor = config.wiki_entity_min_mentions
264 by_slug: dict[str, list[dict[str, Any]]] = {}
265 for row in rows:
266 by_slug.setdefault(row["slug"], []).append(row)
267 stubs: dict[str, WikiStub] = {}
268 for slug, slug_rows in by_slug.items():
269 if sum(r["mention_count"] for r in slug_rows) < floor:
270 continue
271 # The source that names the subject most often labels its page.
272 canonical = max(slug_rows, key=lambda r: (r["mention_count"], r["source"]))
273 source_mentions = tuple(sorted((r["source_id"], r["mention_count"]) for r in slug_rows))
274 refs = tuple(
275 (r["source_id"], index)
276 for r in sorted(slug_rows, key=lambda r: r["source"])
277 for index in r["chunk_indices"]
278 )
279 stub = WikiStub(
280 slug=slug,
281 label=canonical["label"],
282 kind=EntityKind(canonical["kind"]),
283 type_hint=canonical["type_hint"],
284 source_mentions=source_mentions,
285 chunk_refs=refs,
286 )
287 stubs[slug] = _recut_refs(stub, cap)
288 return stubs
291def _write_source_mentions(
292 store: Store,
293 extractor: EntityExtractor,
294 chunks: list[SearchChunk],
295 sources: set[str],
296 cap: int,
297) -> set[str]:
298 """Extract *chunks* and replace the mention rows for every source in
299 *sources*, returning the slugs those sources now name.
301 A source that named something and now names nothing is still replaced, with
302 an empty set, so its stale rows go.
303 """
304 rows_by_source = _mention_rows_by_source(extractor.extract(chunks), cap)
305 affected: set[str] = set()
306 for source in sources:
307 rows = rows_by_source.get(source, [])
308 store.replace_wiki_mentions_for_source(source, rows)
309 affected.update(row["slug"] for row in rows)
310 return affected
313def _sources_of(stubs: dict[str, WikiStub]) -> set[str]:
314 """The id of every source *stubs* list under a subject."""
315 return {source for stub in stubs.values() for source in stub.sources}
318def _gone_sources(store: Store, stubs: dict[str, WikiStub]) -> set[str]:
319 """The ids *stubs* list that no listed source row holds.
321 A removed source that a citation names keeps its row, held out of every listing.
322 """
323 listed = _sources_of(stubs)
324 paths = store.source_paths(listed)
325 unlisted = store.missing_sources(paths.values())
326 return {i for i in listed if i not in paths or paths[i] in unlisted}
329def shown_sources(stub: WikiStub, store: Store) -> list[str]:
330 """The paths of the documents naming *stub*, sorted."""
331 return sorted(store.source_paths(stub.sources).values())
334def _extractor_can_run(config: Config) -> bool:
335 """Whether the backend of the extractor *config* names can run. Logs nothing.
337 It asks for the mode that runs, so the factory has no other mode to report.
338 """
339 mode = effective_entity_mode(config.wiki_entity_mode)
340 return get_entity_extractor(mode, get_services().provider, config).available(quiet=True)
343def _available_extractor(config: Config) -> EntityExtractor | None:
344 """The extractor *config* names, or None when its backend cannot run.
346 It extracts at a floor of one: the store keeps every mention, and
347 _stubs_from_mention_rows judges the real floor over the corpus-wide aggregate.
348 """
349 pass_config = config.model_copy(update={"wiki_entity_min_mentions": 1})
350 extractor = get_entity_extractor(config.wiki_entity_mode, get_services().provider, pass_config)
351 return extractor if extractor.available() else None
354def _rebuild(
355 store: Store,
356 config: Config,
357 extractor: EntityExtractor,
358 sources: set[str] | None,
359 gone: set[str] | None = None,
360) -> dict[str, WikiStub]:
361 """Extract *sources*, or the whole corpus for None, and persist the index.
363 The subjects that list one of *sources* or one of the ids in *gone* are
364 aggregated again. Caller holds ``WIKI_BUILD_LOCK``.
365 """
366 cap = config.wiki_stub_max_chunk_refs
367 # A cold or file-only-migrated store has no rows to aggregate, so an
368 # incremental pass would derive an empty index. Rebuild in full to seed
369 # the table; incremental passes maintain it from there.
370 if sources is None or not store.has_wiki_mentions():
371 store.clear_wiki_mentions()
372 all_sources = {record["filename"] for record in store.get_sources()}
373 _write_source_mentions(store, extractor, _corpus_chunks(store), all_sources, cap)
374 stubs = _stubs_from_mention_rows(store.wiki_mention_rows(), config, cap)
375 else:
376 previous = load_stub_index(config)
377 chunks = [c for name in sorted(sources) for c in store.get_chunks_by_source(name)]
378 affected = _write_source_mentions(store, extractor, chunks, sources, cap)
379 listed = set(store.source_ids(sources, SourceType.DOCUMENT).values()) | (gone or set())
380 affected |= {slug for slug, stub in previous.items() if set(stub.sources) & listed}
381 recomputed = _stubs_from_mention_rows(store.wiki_mention_rows(slugs=affected), config, cap)
382 stubs = {slug: stub for slug, stub in previous.items() if slug not in affected}
383 stubs.update(recomputed)
384 save_stub_index(stubs, config)
385 log.info("Wiki stub index: %d entities across the corpus", len(stubs))
386 return stubs
389def refresh_stub_index(
390 store: Store,
391 config: Config | None = None,
392 *,
393 sources: set[str] | None = None,
394) -> dict[str, WikiStub]:
395 """Rebuild the index from the corpus and persist it.
397 The store's ``_wiki_mentions`` table is the source of truth. ``sources=None``
398 re-extracts the whole corpus; otherwise only those sources are re-extracted
399 and their mention rows replaced. The stub index is then the corpus-wide
400 aggregate of that table filtered to the mention floor, so a subject below
401 the floor in each separately-synced document still appears once all its
402 evidence is present.
404 Spends no LLM call. Holds the wiki mutex because it writes into the wiki
405 directory alongside builds and prunes.
406 """
407 if config is None:
408 config = cfg
409 with WIKI_BUILD_LOCK:
410 extractor = _available_extractor(config)
411 if extractor is None:
412 # An unavailable backend extracts nothing, which is indistinguishable
413 # from a corpus that names nothing. Rebuilding on that would drop the
414 # whole index; leave both the store and the file untouched.
415 log.warning(_NO_EXTRACTOR)
416 return load_stub_index(config)
417 return _rebuild(store, config, extractor, sources)
420def _owed_sources(store: Store, config: Config, since: str) -> tuple[set[str], set[str]]:
421 """The sources written at or after *since*, and the ids the index lists that left the store."""
422 return set(store.sources_ingested_since(since)), _gone_sources(store, load_stub_index(config))
425def _extractor_for_the_pass(config: Config, named: set[str]) -> EntityExtractor | None:
426 """The extractor of a pass with something to give, or None with a line only for *named*."""
427 # With no file of its own to give, a sync first asks whether an extractor can run.
428 if not named and not _extractor_can_run(config):
429 log.debug(_NO_EXTRACTOR)
430 return None
431 extractor = _available_extractor(config)
432 if extractor is None:
433 log.warning(_NO_EXTRACTOR)
434 return extractor
437def _refresh_corpus(store: Store, config: Config, named: set[str]) -> None:
438 """Give the index the whole corpus: no record covers any source row."""
439 extractor = _extractor_for_the_pass(config, named)
440 if extractor is None:
441 return
442 checkpoint = store.sources_checkpoint()
443 _rebuild(store, config, extractor, None)
444 store.record_wiki_refreshed(checkpoint)
447def refresh_owed_sources(store: Store, config: Config, named: set[str]) -> set[str]:
448 """Give the index *named* and every source the store holds that no refresh has read.
450 The store records the checkpoint of the last refresh that completed. With
451 the source table at that version nothing is owed and nothing is read. With
452 the version moved and no row stamped since, nothing is given and no
453 extractor is built or asked: the refresh records the new version. With no
454 record the whole corpus is owed. A refresh with no extractor records
455 nothing, so the sources stay owed. Returns the sources it gave the index,
456 none for the whole corpus.
457 """
458 covered = store.wiki_refreshed()
459 if covered is not None and covered.version == store.sources_version():
460 return set()
461 with WIKI_BUILD_LOCK:
462 # The read takes an index of an older version to this one, extractor or not.
463 load_stub_index(config)
464 if covered is None:
465 _refresh_corpus(store, config, named)
466 return set()
467 # Before the read of the owed rows: a row committed after it is owed to the next refresh.
468 checkpoint = store.sources_checkpoint()
469 written, gone = _owed_sources(store, config, covered.time)
470 owed = written | named
471 if owed or gone:
472 extractor = _extractor_for_the_pass(config, named)
473 if extractor is None:
474 return set()
475 _rebuild(store, config, extractor, owed, gone)
476 store.record_wiki_refreshed(checkpoint)
477 return owed
480def _page_exists(stub: WikiStub, wiki_root: Path) -> bool:
481 """Whether this subject already has a page, published or awaiting review.
483 A page the faithfulness gate routed to drafts/ has been written and is
484 waiting on a human. Offering it as unwritten invites generating it again
485 and again, each time for another LLM call, while the draft sits there.
487 An archived page does not count. Prune retires a page when the sources it
488 cited are deleted, while the subject can still be named by documents that
489 survive, and nothing lists or restores archive/. Suppressing the stub would
490 make the subject unreachable from either pane; listing it costs nothing,
491 since generation only ever runs when a user asks for it.
492 """
493 if (wiki_root / f"{stub.wiki_slug}.md").is_file():
494 return True
495 draft = wiki_root / WikiSubdir.DRAFTS / f"{stub.slug}.md"
496 return draft.is_file() and not _is_placeholder(draft)
499def drop_removed_sources(store: Store, config: Config | None = None) -> None:
500 """Forget each document the index lists that no source row holds.
502 Re-aggregates the affected slugs from the store, whose rows for the removed
503 sources went with their chunks, so it costs no extraction pass. Without it a
504 removed document keeps its subjects in the browse tree forever: its skip
505 marker keeps it out of later syncs, so no refresh ever revisits them.
506 """
507 if config is None:
508 config = cfg
509 cap = config.wiki_stub_max_chunk_refs
510 with WIKI_BUILD_LOCK:
511 previous = load_stub_index(config)
512 gone = _gone_sources(store, previous)
513 affected = {slug for slug, stub in previous.items() if set(stub.sources) & gone}
514 if not affected:
515 return
516 # A subject only the removed sources named drops out; one still named
517 # elsewhere keeps its surviving evidence.
518 recomputed = _stubs_from_mention_rows(store.wiki_mention_rows(slugs=affected), config, cap)
519 stubs = {slug: stub for slug, stub in previous.items() if slug not in affected}
520 stubs.update(recomputed)
521 if stubs != previous:
522 save_stub_index(stubs, config)
525def _is_placeholder(draft: Path) -> bool:
526 """Whether a draft is a PENDING marker rather than written content.
528 A marker records that generation failed to produce the section, so the
529 subject still has no page and must stay listed as unwritten.
530 """
531 try:
532 text = draft.read_text(encoding="utf-8")
533 except OSError:
534 return False
535 return is_pending_marker_text(text)
538def ungenerated_stubs(stubs: dict[str, WikiStub], wiki_root: Path) -> list[WikiStub]:
539 """The stubs nothing has written a page for yet, in slug order."""
540 return [stub for _, stub in sorted(stubs.items()) if not _page_exists(stub, wiki_root)]