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