Coverage for src/lilbee/wiki/generation.py: 100%
182 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-14 11:46 +0000
1"""Top-level wiki build orchestrators.
3Two public entry points live here:
5- :func:`build_wiki` produces entity and LLM-curated concept pages
6 per source, runs the one-time legacy-concept-page archival first,
7 then rewrites ``[[link]]`` slugs across all wiki content subdirs.
8- :func:`generate_synthesis_pages` produces cross-source synthesis
9 pages from concept clusters spanning 3+ documents.
11Both reuse the per-source batch path and the single-page pipeline
12from :mod:`lilbee.wiki.synthesis` and :mod:`lilbee.wiki.page`.
13"""
15from __future__ import annotations
17import logging
18import threading
19from pathlib import Path
20from typing import TypedDict
22from lilbee.app.services import get_services
23from lilbee.core.config import Config, cfg
24from lilbee.data.store import SearchChunk, Store
25from lilbee.providers.base import LLMProvider
26from lilbee.retrieval.clustering import SourceClusterer
27from lilbee.runtime.progress import (
28 DetailedProgressCallback,
29 EventType,
30 WikiPageEvent,
31 WikiPhase,
32 WikiPhaseEvent,
33 noop_callback,
34)
35from lilbee.wiki.batch import archive_legacy_concept_pages
36from lilbee.wiki.entity_extractor import EntityKind, ExtractedEntity, get_entity_extractor
37from lilbee.wiki.index import append_wiki_log, update_wiki_index
38from lilbee.wiki.links import apply_rewriter, compile_rewriter
39from lilbee.wiki.shared import (
40 MIN_CLUSTER_SOURCES,
41 WIKI_BUILD_LOCK,
42 WIKI_CONTENT_SUBDIRS,
43 WikiLogAction,
44 WikiSubdir,
45 atomic_write_text,
46)
47from lilbee.wiki.stats import BuildStats, BuildStatsDict
48from lilbee.wiki.synthesis import (
49 generate_source_batch,
50 generate_synthesis_page,
51 group_entities_by_primary_source,
52)
54log = logging.getLogger(__name__)
56_ENTITY_LIKE_SUBDIRS: tuple[str, ...] = (WikiSubdir.CONCEPTS, WikiSubdir.ENTITIES)
59def _generate_for_cluster(
60 label: str,
61 sources: frozenset[str],
62 provider: LLMProvider,
63 store: Store,
64 config: Config,
65 stats: BuildStats,
66) -> Path | None:
67 """Gather chunks for a cluster and generate a synthesis page."""
68 source_names = sorted(sources)
69 chunks_by_source: dict[str, list] = {}
70 for name in source_names:
71 chunks = store.get_chunks_by_source(name)
72 if chunks:
73 chunks_by_source[name] = chunks
75 if len(chunks_by_source) < MIN_CLUSTER_SOURCES:
76 return None
78 return generate_synthesis_page(
79 label, source_names, chunks_by_source, provider, store, config, stats
80 )
83def generate_synthesis_pages(
84 provider: LLMProvider,
85 store: Store,
86 clusterer: SourceClusterer,
87 config: Config | None = None,
88 on_progress: DetailedProgressCallback = noop_callback,
89 stats: BuildStats | None = None,
90 cancel: threading.Event | None = None,
91) -> list[Path]:
92 """Generate synthesis pages for source clusters spanning 3+ documents.
94 Setting *cancel* stops at the next cluster boundary so a disconnected
95 client does not leave the run holding the wiki build mutex.
96 """
97 if config is None:
98 config = cfg
99 stats = BuildStats.ensure(stats)
101 clusters = clusterer.get_clusters(min_sources=MIN_CLUSTER_SOURCES)
102 if not clusters:
103 log.info("No source clusters span %d+ sources, skipping synthesis", MIN_CLUSTER_SOURCES)
104 return []
106 on_progress(EventType.WIKI_PHASE, WikiPhaseEvent(phase=WikiPhase.GENERATE, total=len(clusters)))
107 pages: list[Path] = []
108 for index, cluster in enumerate(clusters, start=1):
109 if cancel is not None and cancel.is_set():
110 log.info("Wiki synthesis cancelled after %d of %d clusters", index - 1, len(clusters))
111 break
112 page = _generate_for_cluster(cluster.label, cluster.sources, provider, store, config, stats)
113 if page is not None:
114 pages.append(page)
115 on_progress(
116 EventType.WIKI_PAGE,
117 WikiPageEvent(
118 label=cluster.label,
119 pages=0 if page is None else 1,
120 current=index,
121 total=len(clusters),
122 ),
123 )
125 log.info("Generated %d synthesis pages", len(pages))
126 return pages
129def _all_sources_in_scope(
130 grouped: dict[str, list[ExtractedEntity]],
131 store: Store,
132 config: Config,
133 extract_concepts: bool,
134) -> set[str]:
135 """Union of sources with entities and (when enabled) eligible for concept curation.
137 Seed the union with every entity's primary source (the keys of
138 ``grouped``). When ``extract_concepts`` is True AND
139 ``wiki_batch_min_chunks`` is satisfied, add any source in the store
140 that passes the floor. This gives concept-only sources (no extracted
141 entities) their chance at curation while keeping zero-entity short
142 sources skipped entirely.
143 """
144 sources: set[str] = set(grouped)
145 if not extract_concepts:
146 return sources
147 try:
148 records = store.get_sources()
149 except Exception as exc:
150 log.warning("get_sources failed; sticking to entity-grouped sources: %s", exc)
151 return sources
152 for record in records:
153 name = record.get("filename", "") if isinstance(record, dict) else ""
154 if not name:
155 continue
156 if name in sources:
157 continue
158 chunk_count = record.get("chunk_count", 0) if isinstance(record, dict) else 0
159 if chunk_count >= config.wiki_batch_min_chunks:
160 sources.add(name)
161 return sources
164def _entity_surface_map(entities: list[ExtractedEntity]) -> dict[str, str]:
165 """Build the surface-form -> slug map for the ``[[link]]`` rewriter.
167 Includes both the entity's human label (e.g. *"Henry Ford"*) and
168 the slug-with-hyphens-as-spaces variant (*"henry ford"*) so the
169 rewriter catches either form in body text.
170 """
171 mapping: dict[str, str] = {}
172 for entity in entities:
173 mapping[entity.label] = entity.slug
174 spaced = entity.slug.replace("-", " ")
175 if spaced and spaced != entity.label:
176 mapping[spaced] = entity.slug
177 return mapping
180def _augment_surface_map_with_existing_pages(
181 surface_to_slug: dict[str, str], wiki_root: Path
182) -> None:
183 """Add slugs for pages already on disk so an incremental rebuild of
184 one concept still links to its unchanged neighbors. **Mutates
185 surface_to_slug in place.** Only enriches the map with the
186 hyphen-to-space surface form because frontmatter labels aren't
187 read here; body prose typically uses the spaced form so this
188 covers the common case.
189 """
190 for subdir in _ENTITY_LIKE_SUBDIRS:
191 subdir_path = wiki_root / subdir
192 if not subdir_path.is_dir():
193 continue
194 for md_path in subdir_path.rglob("*.md"):
195 slug = md_path.stem
196 spaced = slug.replace("-", " ")
197 surface_to_slug.setdefault(spaced, slug)
200def rewrite_links_across_wiki(
201 entities: list[ExtractedEntity], config: Config, wiki_root: Path | None = None
202) -> None:
203 """Rewrite ``[[slug]]`` links on every page under ``wiki/`` content subdirs.
205 A page never receives a link to itself: the rewriter takes the
206 owning slug and drops it inside its match callback, so the
207 surface map is shared unmodified across every page in the walk
208 (no O(M) dict rebuild per file). The map is augmented with
209 slugs from the existing on-disk corpus so a touched page still
210 links to untouched neighbors. The alternation regex + lookup are
211 compiled once per build and reused across pages.
212 """
213 surface_to_slug = _entity_surface_map(entities)
214 # Callers that were handed a root use it: accept publishes into the root it
215 # was given, which is not always the one the config points at.
216 if wiki_root is None:
217 wiki_root = config.data_root / config.wiki_dir
218 _augment_surface_map_with_existing_pages(surface_to_slug, wiki_root)
219 rewriter = compile_rewriter(surface_to_slug)
220 if rewriter is None:
221 return
223 for subdir in WIKI_CONTENT_SUBDIRS:
224 subdir_path = wiki_root / subdir
225 if not subdir_path.is_dir():
226 continue
227 is_entity_subdir = subdir in _ENTITY_LIKE_SUBDIRS
228 for md_path in subdir_path.rglob("*.md"):
229 owning_slug = md_path.stem if is_entity_subdir else None
230 original = md_path.read_text(encoding="utf-8")
231 rewritten = apply_rewriter(original, rewriter, skip_slug=owning_slug)
232 if rewritten != original:
233 atomic_write_text(md_path, rewritten)
236def build_wiki(
237 entities: list[ExtractedEntity],
238 provider: LLMProvider,
239 store: Store,
240 config: Config | None = None,
241 *,
242 extract_concepts: bool = True,
243 on_progress: DetailedProgressCallback = noop_callback,
244 stats: BuildStats | None = None,
245 cancel: threading.Event | None = None,
246) -> list[Path]:
247 """Produce entity and LLM-curated concept pages per source.
249 Per-entity / per-concept fan-out is collapsed into a per-source
250 batched call: for each source in ``entities``' chunk refs, one LLM
251 call identifies 3-5 concepts AND writes a wiki section for every
252 pre-extracted entity belonging to that source. Output sections are
253 split, citation-verified, embedding-scored, and landed under
254 ``wiki/entities/`` or ``wiki/concepts/`` depending on kind.
256 ``extract_concepts=False`` (used by the incremental-ingest hook)
257 drops the concept-curation paragraph from the prompt so a
258 touched source does not churn concept slugs.
260 A one-time archive migration runs first (idempotently, gated by
261 ``{data_dir}/.phase-d-migrated``), moving legacy concept pages
262 under ``wiki/archive/concepts/`` and unwrapping stale
263 ``[[archived-slug]]`` links across the remaining pages.
264 """
265 if config is None:
266 config = cfg
267 stats = BuildStats.ensure(stats)
268 wiki_root = config.data_root / config.wiki_dir
269 archive_legacy_concept_pages(wiki_root, config.data_dir, store, config)
271 grouped = group_entities_by_primary_source(entities)
272 all_sources = _all_sources_in_scope(grouped, store, config, extract_concepts)
273 written_concept_slugs: dict[str, str] = {}
274 pages: list[Path] = []
276 on_progress(
277 EventType.WIKI_PHASE, WikiPhaseEvent(phase=WikiPhase.GENERATE, total=len(all_sources))
278 )
279 for index, source in enumerate(sorted(all_sources), start=1):
280 if cancel is not None and cancel.is_set():
281 log.info("Wiki build cancelled after %d of %d sources", index - 1, len(all_sources))
282 break
283 source_entities = grouped.get(source, [])
284 chunks = store.get_chunks_by_source(source)
285 chunk_count = len(chunks)
286 source_extract = extract_concepts and chunk_count >= config.wiki_batch_min_chunks
287 source_pages: list[Path] = []
288 if source_entities or source_extract:
289 source_pages = generate_source_batch(
290 source=source,
291 entities=source_entities,
292 chunks=chunks,
293 provider=provider,
294 store=store,
295 config=config,
296 extract_concepts=source_extract,
297 written_concept_slugs=written_concept_slugs,
298 stats=stats,
299 )
300 pages.extend(source_pages)
301 else:
302 log.info(
303 "Skipping source %s: %d entities, %d chunks, min=%d, extract=%s",
304 source,
305 len(source_entities),
306 chunk_count,
307 config.wiki_batch_min_chunks,
308 source_extract,
309 )
310 on_progress(
311 EventType.WIKI_PAGE,
312 WikiPageEvent(
313 label=source,
314 pages=len(source_pages),
315 current=index,
316 total=len(all_sources),
317 ),
318 )
320 rewrite_links_across_wiki(entities, config)
321 log.info("Generated %d batched wiki pages", len(pages))
322 return pages
325class WikiBuildSummary(TypedDict):
326 """Result of a full wiki build/update."""
328 paths: list[str]
329 entities: int
330 count: int
331 stats: BuildStatsDict
334class WikiEntityCandidate(TypedDict):
335 """One NER entity a build would write a page for."""
337 slug: str
338 label: str
339 kind: EntityKind
340 type_hint: str
341 mentions: int
342 sources: list[str]
345DRY_RUN_CONCEPT_NOTE = (
346 "LLM-curated concepts are not part of a dry run. "
347 "Run the build to see which concepts the LLM proposes."
348)
351def _corpus_chunks(store: Store) -> list[SearchChunk]:
352 """Every chunk of every ingested source, in source order."""
353 chunks: list[SearchChunk] = []
354 for record in store.get_sources():
355 chunks.extend(store.get_chunks_by_source(record["filename"]))
356 return chunks
359def preview_build_entities(config: Config) -> list[WikiEntityCandidate]:
360 """Entity candidates a build would cover, extracted with no LLM call.
362 The dry run every surface shares. Concepts come from the per-source
363 batched call, so they are absent here: see :data:`DRY_RUN_CONCEPT_NOTE`.
364 """
365 svc = get_services()
366 extractor = get_entity_extractor(config.wiki_entity_mode, svc.provider, config)
367 return [
368 WikiEntityCandidate(
369 slug=entity.slug,
370 label=entity.label,
371 kind=entity.kind,
372 type_hint=entity.type_hint,
373 mentions=len(entity.chunk_refs),
374 sources=sorted({ref.source for ref in entity.chunk_refs}),
375 )
376 for entity in extractor.extract(_corpus_chunks(svc.store))
377 ]
380def run_full_build(
381 config: Config | None = None,
382 on_progress: DetailedProgressCallback = noop_callback,
383 cancel: threading.Event | None = None,
384) -> WikiBuildSummary:
385 """Extract entities and build wiki pages for every ingested source.
387 Holds the wiki build mutex for the whole run, so a build started from any
388 surface waits for one already in flight instead of interleaving writes.
389 Setting *cancel* stops the run at the next source boundary and releases the
390 mutex; without it a client that disconnects mid-stream leaves a build
391 holding the mutex until it finishes the whole corpus.
392 """
393 if config is None:
394 config = cfg
395 stats = BuildStats()
396 with WIKI_BUILD_LOCK:
397 svc = get_services()
398 on_progress(EventType.WIKI_PHASE, WikiPhaseEvent(phase=WikiPhase.EXTRACT))
399 extractor = get_entity_extractor(config.wiki_entity_mode, svc.provider, config)
400 entities = extractor.extract(_corpus_chunks(svc.store))
401 pages = build_wiki(
402 entities,
403 svc.provider,
404 svc.store,
405 config,
406 extract_concepts=config.wiki_extract_concepts,
407 on_progress=on_progress,
408 stats=stats,
409 cancel=cancel,
410 )
411 on_progress(EventType.WIKI_PHASE, WikiPhaseEvent(phase=WikiPhase.INDEX))
412 update_wiki_index(config)
413 append_wiki_log(
414 WikiLogAction.BUILD,
415 f"{len(pages)} pages from {len(entities)} records; {stats.summary_line()}",
416 config,
417 )
418 return {
419 "paths": [str(p) for p in pages],
420 "entities": len(entities),
421 "count": len(pages),
422 "stats": stats.as_dict(),
423 }
426class WikiSynthesizeSummary(TypedDict):
427 """Result of running synthesis-page generation."""
429 paths: list[str]
430 count: int
431 stats: BuildStatsDict
434def run_full_synthesize(
435 config: Config | None = None,
436 on_progress: DetailedProgressCallback = noop_callback,
437 cancel: threading.Event | None = None,
438) -> WikiSynthesizeSummary:
439 """Generate synthesis pages for cross-source clusters.
441 Shares the wiki build mutex with :func:`run_full_build` so synthesis and a
442 build never write the same tree at once.
443 """
444 if config is None:
445 config = cfg
446 stats = BuildStats()
447 svc = get_services()
448 with WIKI_BUILD_LOCK:
449 paths = generate_synthesis_pages(
450 svc.provider, svc.store, svc.clusterer, config, on_progress, stats, cancel
451 )
452 append_wiki_log(
453 WikiLogAction.SYNTHESIZE,
454 f"{len(paths)} synthesis pages; {stats.summary_line()}",
455 config,
456 )
457 return {
458 "paths": [str(p) for p in paths],
459 "count": len(paths),
460 "stats": stats.as_dict(),
461 }