Coverage for src/lilbee/wiki/synthesis.py: 100%
168 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-09-17 10:02 +0000
1"""Cross-source synthesis pages and per-source batched generation.
3Two related orchestrators live here:
5- ``generate_synthesis_page`` and friends produce a single
6 cross-source page from a concept cluster spanning 3+ documents.
7- ``generate_source_batch`` issues one LLM call per source that
8 emits sections for every pre-extracted entity plus 3-5 LLM-curated
9 concepts; the response is split into per-section bodies and each
10 section is finalized via :func:`finalize_section`.
12The shared output-parsing helpers (``_split_batched_output``,
13``_prefix_heading``, ``match_label``) cover both paths.
14"""
16from __future__ import annotations
18import functools
19import logging
20import re
21from collections.abc import Callable
22from pathlib import Path
24import yaml
26from lilbee.core.config import Config
27from lilbee.core.text import clean_label_for_display, make_slug
28from lilbee.data.store import CitationRecord, SearchChunk, Store
29from lilbee.providers.base import LLMProvider
30from lilbee.retrieval.reasoning import strip_reasoning
31from lilbee.wiki.batch import (
32 finalize_section,
33 hash_existing_sources,
34 match_label,
35)
36from lilbee.wiki.citations import (
37 ParsedCitation,
38 parse_wiki_citations,
39 resolve_multi_source_citations,
40)
41from lilbee.wiki.entity_extractor import EntityKind, ExtractedEntity
42from lilbee.wiki.page import (
43 build_wiki_messages,
44 chunks_to_text,
45 generate_page,
46 truncate_chunks_to_budget,
47 wiki_generation_options,
48)
49from lilbee.wiki.persistence import write_pending_marker
50from lilbee.wiki.shared import (
51 PENDING_PARSE_MARKER_PREFIX,
52 PendingKind,
53 WikiSubdir,
54)
55from lilbee.wiki.stats import BuildStats
57log = logging.getLogger(__name__)
59# Section headers the batch parser recognizes: H1 (``# Name``) or H2
60# (``## Name``). The name capture is anchored to the rest of the line so
61# labels like ``## Brake System (hydraulic)`` still parse. Bold lines are
62# not headers: a mid-body ``**emphasis**`` would otherwise truncate its
63# section and open a bogus one.
64_SECTION_HEADER_RE = re.compile(r"^##?[ \t]+(?P<name>[^\n]+?)[ \t]*$", re.MULTILINE)
66# Machine-readable concept declaration the batched prompt requires when
67# concept curation is on. Only a declared name may open a concept section.
68_CONCEPT_DECLARATION_LABEL = "CONCEPTS"
69_CONCEPT_DECLARATION_PREFIX = f"{_CONCEPT_DECLARATION_LABEL}:"
70_CONCEPT_DECLARATION_SEPARATOR = ";"
72# Markdown decoration models wrap the declaration line in: blockquote and
73# heading markers, bold/italic runs, inline code. Stripped from the line and
74# from each parsed label. The labels stay on one line, so only spaces and tabs
75# are tolerated around the colon.
76_CONCEPT_DECLARATION_RE = re.compile(
77 rf"^[ \t>#*_`]*{_CONCEPT_DECLARATION_LABEL}[*_`]*[ \t]*:[ \t]*(?P<labels>[^\n]+)$",
78 re.MULTILINE | re.IGNORECASE,
79)
80_CONCEPT_LABEL_DECORATION = " \t*_`"
82# Cap on the published concept names fed back into the batch prompt for reuse.
83# The list is a nudge toward established names, not a complete index, and an
84# uncapped one eats the whole chunk budget on a large wiki.
85_MAX_REUSE_CONCEPT_LABELS = 50
88def generate_synthesis_page(
89 topic: str,
90 source_names: list[str],
91 chunks_by_source: dict[str, list[SearchChunk]],
92 provider: LLMProvider,
93 store: Store,
94 config: Config,
95 stats: BuildStats | None = None,
96) -> Path | None:
97 """Generate a single synthesis page for a concept cluster.
98 Returns the path to the generated page, or None on failure.
99 """
100 all_chunks = [c for cs in chunks_by_source.values() for c in cs]
101 if not all_chunks:
102 log.warning("No chunks for synthesis topic %r, skipping", topic)
103 return None
105 source_list = "\n".join(f"- {name}" for name in sorted(source_names))
106 template = config.wiki_synthesis_prompt
107 display_topic = clean_label_for_display(topic)
108 render = functools.partial(template.format, topic=display_topic, source_list=source_list)
109 # Budget against the prompt as it renders: the source list is a per-call
110 # substitution the raw template does not carry.
111 all_chunks = truncate_chunks_to_budget(all_chunks, config, len(render(chunks_text="")))
112 prompt = render(chunks_text=chunks_to_text(all_chunks))
113 slug = make_slug(topic)
115 source_hashes = hash_existing_sources(source_names)
117 def resolver(parsed: list[ParsedCitation]) -> list[CitationRecord]:
118 return resolve_multi_source_citations(parsed, source_names, source_hashes, chunks_by_source)
120 return generate_page(
121 label=topic,
122 prompt=prompt,
123 chunks=all_chunks,
124 citation_resolver=resolver,
125 page_type=WikiSubdir.SYNTHESIS,
126 slug=slug,
127 source_names=source_names,
128 provider=provider,
129 store=store,
130 config=config,
131 stats=stats,
132 )
135def _split_batched_output(
136 text: str,
137 expected_entity_labels: set[str],
138 expected_concept_labels: set[str] | None = None,
139) -> dict[str, tuple[EntityKind, str]]:
140 """Parse the batched LLM response into per-label bodies.
142 Splits on H1/H2 headers, then binds each header to an expected entity
143 label or to a concept the response declared, via :func:`match_label`:
144 entities win ties, but an exact match in either set beats any substring
145 match. A header matching neither is dropped as noise. Labels whose section
146 could not be recovered are simply absent from the result; the caller
147 loops over the expected sets to write their PENDING markers.
148 """
149 # A ``## CONCEPTS: a; b`` declaration parses as an H2 header, and its text
150 # substring-matches the labels it declares. Excised from the preamble only:
151 # a body line starting with "concepts:" is section prose, not a declaration.
152 preamble_end = _declaration_scope_end(text)
153 text = _CONCEPT_DECLARATION_RE.sub("", text[:preamble_end]) + text[preamble_end:]
154 concepts = expected_concept_labels or set()
155 recovered: dict[str, tuple[EntityKind, str]] = {}
156 matches = list(_SECTION_HEADER_RE.finditer(text))
157 for i, match in enumerate(matches):
158 name = match.group("name").strip()
159 start = match.end()
160 end = matches[i + 1].start() if i + 1 < len(matches) else len(text)
161 body = text[start:end].strip()
162 if not body:
163 continue
164 lowered = name.lower()
165 kind_label = match_label(
166 lowered,
167 [(expected_entity_labels, EntityKind.ENTITY), (concepts, EntityKind.CONCEPT)],
168 )
169 if kind_label is None:
170 log.info("Dropping section %r: matches no expected entity or declared concept", name)
171 continue
172 kind, label = kind_label
173 if label in recovered:
174 log.info("Dropping section %r: label %r already has a section", name, label)
175 continue
176 recovered[label] = (kind, _prefix_heading(label, body))
177 return recovered
180def _declaration_scope_end(text: str) -> int:
181 """End of the preamble a CONCEPTS declaration may occupy.
183 The prompt requires the declaration before any section, so the scope ends
184 at the first section header. A declaration rendered as a header is itself
185 part of the preamble and does not end the scope.
186 """
187 for match in _SECTION_HEADER_RE.finditer(text):
188 if not _CONCEPT_DECLARATION_RE.search(match.group(0)):
189 return match.start()
190 return len(text)
193def _parse_declared_concepts(text: str) -> set[str]:
194 """Read the ``CONCEPTS: a; b; c`` line the batched prompt requires.
196 An absent declaration means the response curated no concepts, so every
197 non-entity header in it is noise. Only the preamble before the first
198 section header is searched, per the prompt contract that the declaration
199 comes before any section; body prose starting with "concepts:" is content.
200 """
201 match = _CONCEPT_DECLARATION_RE.search(text[: _declaration_scope_end(text)])
202 if match is None:
203 return set()
204 parts = match.group("labels").split(_CONCEPT_DECLARATION_SEPARATOR)
205 stripped = (part.strip(_CONCEPT_LABEL_DECORATION) for part in parts)
206 return {part for part in stripped if part}
209def _prefix_heading(label: str, body: str) -> str:
210 """Rebuild the section's ``# Label`` H1 from the label it was bound to.
212 Splitting consumes the model's own header line. Rebuilding the heading
213 from the matched label rather than the model's wording keeps the
214 title/body coherence gate reading the label the page is filed under,
215 including when the header matched only as a substring of it.
216 """
217 return f"# {clean_label_for_display(label)}\n\n{body}"
220def _existing_concept_labels(wiki_root: Path) -> list[str]:
221 """Published concept slugs as spaced names, so rebuilds reuse established names.
223 Capped at :data:`_MAX_REUSE_CONCEPT_LABELS`, taken in slug order so the same
224 wiki renders the same prompt on every run.
225 """
226 concepts_dir = wiki_root / WikiSubdir.CONCEPTS
227 if not concepts_dir.is_dir():
228 return []
229 labels = sorted({path.stem.replace("-", " ") for path in concepts_dir.rglob("*.md")})
230 return labels[:_MAX_REUSE_CONCEPT_LABELS]
233def _concept_instruction(existing_concepts: list[str]) -> str:
234 """Concept-curation paragraph, including the declaration-line contract.
236 Lives in code rather than in the writable prompt template: the declaration
237 line is what the parser enforces, so overriding the template from settings
238 cannot silently break section recovery.
239 """
240 reuse = ""
241 if existing_concepts:
242 reuse = (
243 "Reuse these existing concept names verbatim when they fit: "
244 f"{', '.join(existing_concepts)}.\n\n"
245 )
246 separator = f"{_CONCEPT_DECLARATION_SEPARATOR} "
247 return (
248 "First, identify 3-5 CONCEPTS: abstract topics or domain terms "
249 "from the source that deserve a standalone wiki page. Do NOT include "
250 "pronouns, articles, or generic nouns.\n\n"
251 f"{reuse}"
252 f"Declare them on a single line as `{_CONCEPT_DECLARATION_PREFIX} "
253 f"first{separator}second{separator}third` before writing anything else. "
254 "A section whose heading is neither a declared concept nor a listed "
255 "entity is discarded.\n\n"
256 "Then write a wiki section for each of the concepts you identified, "
257 "PLUS one section for each NER ENTITY listed below.\n\n"
258 )
261def _build_batch_prompt(
262 source: str,
263 entities: list[ExtractedEntity],
264 chunks_text: str,
265 concept_instruction: str,
266 config: Config,
267) -> str:
268 """Render :attr:`Config.wiki_entity_batch_prompt` for one source call.
270 An empty ``concept_instruction`` (the incremental-ingest hook) leaves the
271 LLM writing entity sections only. Keeps the per-source batched call the
272 single entry point whether or not concepts are requested.
273 """
274 entity_labels = ", ".join(clean_label_for_display(e.label) for e in entities) or "(none)"
275 return config.wiki_entity_batch_prompt.format(
276 source=source,
277 entity_list=entity_labels,
278 chunks_text=chunks_text,
279 concept_instruction=concept_instruction,
280 )
283def group_entities_by_primary_source(
284 entities: list[ExtractedEntity],
285) -> dict[str, list[ExtractedEntity]]:
286 """Group entities under the source that mentions them most.
288 Primary source = source with the highest chunk-ref count;
289 lexicographic tiebreak. An entity with no refs is dropped
290 silently (defensive: extractor always attaches refs, but a
291 future extractor might not).
292 """
293 grouped: dict[str, list[ExtractedEntity]] = {}
294 for entity in entities:
295 if not entity.chunk_refs:
296 continue
297 counts: dict[str, int] = {}
298 for ref in entity.chunk_refs:
299 counts[ref.source] = counts.get(ref.source, 0) + 1
300 primary = min(counts.items(), key=lambda kv: (-kv[1], kv[0]))[0]
301 grouped.setdefault(primary, []).append(entity)
302 return grouped
305def generate_source_batch(
306 source: str,
307 entities: list[ExtractedEntity],
308 chunks: list[SearchChunk],
309 provider: LLMProvider,
310 store: Store,
311 config: Config,
312 *,
313 extract_concepts: bool,
314 written_concept_slugs: dict[str, str],
315 stats: BuildStats | None = None,
316) -> list[Path]:
317 """Issue one LLM call for *source* and finalize every recovered section.
319 Returns the list of page paths written (entities + concepts
320 combined). Every expected entity that produced no page leaves a
321 PENDING marker under ``wiki/drafts/`` so the next build can retry,
322 whether the parser missed its section, a downstream gate rejected
323 it, or the LLM call itself failed. Concept slugs already written by
324 an earlier source produce a PENDING-COLLISION marker on the losing
325 side (see :func:`divert_concept_collision`).
327 ``written_concept_slugs`` is the per-build ledger of
328 slug → first_source. Callers share one dict across the per-source
329 loop. The second source to propose a slug is the one that gets
330 diverted to a collision marker.
331 """
332 stats = BuildStats.ensure(stats)
333 if not chunks:
334 return []
335 wiki_root = config.data_root / config.wiki_dir
336 drafts_dir = wiki_root / WikiSubdir.DRAFTS
337 concept_instruction = (
338 _concept_instruction(_existing_concept_labels(wiki_root)) if extract_concepts else ""
339 )
340 render = functools.partial(
341 _build_batch_prompt,
342 source,
343 entities,
344 concept_instruction=concept_instruction,
345 config=config,
346 )
347 # Budget against the prompt as it renders: the concept instruction and the
348 # entity list are per-call substitutions the raw template does not carry.
349 budgeted = truncate_chunks_to_budget(chunks, config, len(render("")))
350 prompt = render(chunks_to_text(budgeted))
351 text = _request_batch_sections(source, prompt, provider, config)
352 if text is None:
353 _write_pending_markers(entities, set(), source, drafts_dir, stats)
354 return []
356 declared_concepts = _parse_declared_concepts(text) if extract_concepts else set()
357 parsed = _split_batched_output(text, {e.label for e in entities}, declared_concepts)
359 source_names = [source]
360 finalize = functools.partial(
361 finalize_section,
362 chunks=budgeted,
363 citation_resolver=functools.partial(
364 resolve_multi_source_citations,
365 source_names=source_names,
366 source_hashes=hash_existing_sources(source_names),
367 chunks_by_source={source: budgeted},
368 ),
369 source_names=source_names,
370 store=store,
371 config=config,
372 source=source,
373 written_concept_slugs=written_concept_slugs,
374 drafts_dir=drafts_dir,
375 # Citation definitions live in the trailing block of the WHOLE response,
376 # not inside any one section body. Parse once and replay for every
377 # section, so pages other than the last still resolve their footnotes.
378 shared_parsed_citations=parse_wiki_citations(text),
379 scoring_chunks_by_label=_entity_scoring_chunks(entities, budgeted),
380 stats=stats,
381 )
382 pages, written_labels = _finalize_sections(parsed, finalize)
383 _write_pending_markers(entities, written_labels, source, drafts_dir, stats)
384 return pages
387def _entity_scoring_chunks(
388 entities: list[ExtractedEntity],
389 chunks: list[SearchChunk],
390) -> dict[str, list[SearchChunk]]:
391 """Map each entity label to the budgeted chunks its extraction refs name.
393 Labels whose refs all fell outside the budget are absent, so
394 :func:`finalize_section` falls back to the whole-source pool.
395 """
396 by_ref = {(c.source, c.chunk_index): c for c in chunks}
397 scoped: dict[str, list[SearchChunk]] = {}
398 for entity in entities:
399 selected = [
400 by_ref[(ref.source, ref.chunk_index)]
401 for ref in entity.chunk_refs
402 if (ref.source, ref.chunk_index) in by_ref
403 ]
404 if selected:
405 scoped[entity.label] = selected
406 return scoped
409def _request_batch_sections(
410 source: str,
411 prompt: str,
412 provider: LLMProvider,
413 config: Config,
414) -> str | None:
415 """Issue the batched LLM call; None when it raised or came back empty."""
416 messages = build_wiki_messages(prompt, provider, config)
417 try:
418 response = provider.chat(messages, stream=False, options=wiki_generation_options(config))
419 except Exception as exc:
420 log.warning("Batched LLM call failed for source %s: %s", source, exc)
421 return None
422 text = strip_reasoning(response.text).strip()
423 if not text:
424 log.warning("Batched LLM call returned empty response for source %s", source)
425 return None
426 return text
429def _finalize_sections(
430 parsed: dict[str, tuple[EntityKind, str]],
431 finalize: Callable[..., Path | None],
432) -> tuple[list[Path], set[str]]:
433 """Finalize each recovered section; return the pages written and labels covered.
435 A label counts as covered only once ``finalize`` returns a path, so a
436 section dropped by the citation or slug gate still earns a PENDING marker.
437 """
438 pages: list[Path] = []
439 written: set[str] = set()
440 for header_label, (kind, body) in parsed.items():
441 page = finalize(header_label=header_label, kind=kind, body=body)
442 if page is not None:
443 pages.append(page)
444 written.add(header_label)
445 return pages, written
448def _write_pending_markers(
449 entities: list[ExtractedEntity],
450 written_labels: set[str],
451 source: str,
452 drafts_dir: Path,
453 stats: BuildStats,
454) -> None:
455 """Write a PENDING-PARSE marker for every expected entity that produced no page.
457 An entity whose slug is occupied by a draft awaiting review keeps that
458 draft, so no marker is written and none is counted.
459 """
460 for entity in entities:
461 if entity.label in written_labels:
462 continue
463 marker = (
464 f"{PENDING_PARSE_MARKER_PREFIX} for source {source}, "
465 f"entity/concept {entity.label} - "
466 "run wiki build again or manually accept via wiki drafts accept -->"
467 )
468 # Route through ``yaml.safe_dump`` so a label or source containing a
469 # colon, quote, or newline does not produce a frontmatter block that
470 # ``parse_frontmatter`` silently drops.
471 frontmatter_body = yaml.safe_dump(
472 {
473 "pending_source": source,
474 "pending_label": entity.label,
475 "pending_kind": PendingKind.PARSE.value,
476 },
477 sort_keys=False,
478 )
479 path = write_pending_marker(
480 drafts_dir, entity.slug, marker, f"---\n{frontmatter_body}---\n"
481 )
482 if path is None:
483 continue
484 stats.record_pending_marker()
485 log.info("Wrote PENDING-PARSE marker for %s -> %s", entity.slug, path)