Coverage for src/lilbee/wiki/lazy.py: 100%

1"""Generate one wiki page on demand, from the index rather than from a build. 

2 

3A build walks every source document and spends a call on each. This path walks 

4one subject: it takes that subject's chunks across every source naming it and 

5writes a single page. The cost is one call for the page someone actually asked 

6for, and the evidence is wider than a build's, which assigns each entity to the 

7one source mentioning it most and never sees the rest. 

8""" 

9 

10from __future__ import annotations 

11 

12import functools 

13import logging 

14from typing import TYPE_CHECKING 

15 

16from lilbee.app.services import get_services 

17from lilbee.core.config import cfg 

18from lilbee.core.text import clean_label_for_display 

19from lilbee.runtime.progress import ( 

20 DetailedProgressCallback, 

21 EventType, 

22 WikiPageEvent, 

23 WikiPhase, 

24 WikiPhaseEvent, 

25 noop_callback, 

26) 

27 

28from .batch import hash_existing_sources 

29from .citations import resolve_multi_source_citations 

30from .generation import rewrite_links_across_wiki 

31from .page import ( 

32 chunks_to_text, 

33 generate_page, 

34 truncate_chunks_to_budget, 

35) 

36from .shared import WIKI_BUILD_LOCK 

37from .stubs import WikiStub, load_stub_index 

38 

39if TYPE_CHECKING: 

40 import threading 

41 from pathlib import Path 

42 

43 from lilbee.core.config import Config 

44 from lilbee.data.store import CitationRecord, SearchChunk, Store 

45 

46 from .citations import ParsedCitation 

47 from .stats import BuildStats 

48 

49log = logging.getLogger(__name__) 

50 

51 

52class UnknownStubError(LookupError): 

53 """Raised when a slug names no entry in the wiki index.""" 

54 

55 

56def _chunks_for_stub(stub: WikiStub, store: Store) -> tuple[dict[str, list[SearchChunk]], int]: 

57 """The stub's chunks grouped by source, plus how many refs went unresolved. 

58 

59 Refs are looked up per source and filtered by index, so a source that was 

60 re-ingested with fewer chunks contributes what it still has instead of 

61 failing the page. 

62 """ 

63 wanted: dict[str, set[int]] = {} 

64 for source, index in stub.chunk_refs: 

65 wanted.setdefault(source, set()).add(index) 

66 if not wanted: 

67 # Subtracting a re-indexed source can empty the refs while leaving real 

68 # evidence: the cap may have given every ref to the source that went. 

69 # The recorded sources are the truth, so fall back to reading them. 

70 wanted = {source: set() for source in stub.sources} 

71 

72 # Most-mentioning source first, so when the context budget truncates it 

73 # drops the documents with least to say rather than the alphabetically 

74 # unlucky ones. 

75 counts = dict(stub.source_mentions) 

76 order = sorted(wanted, key=lambda name: (-counts.get(name, 0), name)) 

77 

78 by_source: dict[str, list[SearchChunk]] = {} 

79 resolved = 0 

80 for source in order: 

81 available = {c.chunk_index: c for c in store.get_chunks_by_source(source)} 

82 indexes = wanted[source] or set(available) 

83 kept = [available[i] for i in sorted(indexes) if i in available] 

84 resolved += len(kept) 

85 if kept: 

86 by_source[source] = kept 

87 return by_source, max(0, len(stub.chunk_refs) - resolved) 

88 

89 

90def _resolve(slug: str, stubs: dict[str, WikiStub]) -> WikiStub | None: 

91 """Find a stub by bare slug or by the subdir-qualified form. 

92 

93 Every surface shows pages as ``entities/ford``, so that is what a user 

94 types, while the index is keyed by the bare slug. Both resolve. 

95 """ 

96 direct = stubs.get(slug) 

97 if direct is not None: 

98 return direct 

99 return next((stub for stub in stubs.values() if stub.wiki_slug == slug), None) 

100 

101 

102def resolve_stub(slug: str, config: Config | None = None) -> WikiStub | None: 

103 """The index entry for *slug*, by bare or subdir-qualified form.""" 

104 return _resolve(slug, load_stub_index(config or cfg)) 

105 

106 

107def generate_stub_page( 

108 slug: str, 

109 store: Store, 

110 config: Config | None = None, 

111 *, 

112 stats: BuildStats | None = None, 

113 on_progress: DetailedProgressCallback = noop_callback, 

114 cancel: threading.Event | None = None, 

115) -> Path | None: 

116 """Write the page for one indexed subject. Returns its path, or None. 

117 

118 Runs the same citation verification, faithfulness gate, and drafts 

119 quarantine a build does; nothing here bypasses them. Holds the wiki mutex, 

120 so a page generated from the browse tree cannot interleave with a build. 

121 Emits the same wiki_phase/wiki_page events a build does; a *cancel* set 

122 before the model call skips it. 

123 """ 

124 if config is None: 

125 config = cfg 

126 with WIKI_BUILD_LOCK: 

127 stubs = load_stub_index(config) 

128 stub = _resolve(slug, stubs) 

129 if stub is None: 

130 raise UnknownStubError(f"no indexed page named {slug!r}") 

131 

132 chunks_by_source, unresolved = _chunks_for_stub(stub, store) 

133 if not chunks_by_source: 

134 log.warning("No chunks remain for %s; the index is stale", slug) 

135 return None 

136 if unresolved: 

137 log.info("%d of %s's indexed chunks are gone", unresolved, slug) 

138 

139 on_progress(EventType.WIKI_PHASE, WikiPhaseEvent(phase=WikiPhase.GENERATE, total=1)) 

140 if cancel is not None and cancel.is_set(): 

141 log.info("Generation of %s cancelled before the model call", slug) 

142 return None 

143 

144 source_names = sorted(chunks_by_source) 

145 all_chunks = [c for chunks in chunks_by_source.values() for c in chunks] 

146 source_list = "\n".join(f"- {name}" for name in source_names) 

147 display = clean_label_for_display(stub.label) 

148 render = functools.partial( 

149 config.wiki_entity_page_prompt.format, topic=display, source_list=source_list 

150 ) 

151 all_chunks = truncate_chunks_to_budget(all_chunks, config, len(render(chunks_text=""))) 

152 prompt = render(chunks_text=chunks_to_text(all_chunks)) 

153 source_hashes = hash_existing_sources(source_names) 

154 

155 def resolver(parsed: list[ParsedCitation]) -> list[CitationRecord]: 

156 return resolve_multi_source_citations( 

157 parsed, source_names, source_hashes, chunks_by_source 

158 ) 

159 

160 path = generate_page( 

161 label=stub.label, 

162 prompt=prompt, 

163 chunks=all_chunks, 

164 citation_resolver=resolver, 

165 page_type=stub.subdir, 

166 slug=stub.slug, 

167 source_names=source_names, 

168 provider=get_services().provider, 

169 store=store, 

170 config=config, 

171 stats=stats, 

172 # These documents mention the subject; the page does not replace 

173 # them. Pruning here would delete every document that named it. 

174 supersedes_sources=False, 

175 ) 

176 if path is not None: 

177 # The link pass a build runs; without it the page has no [[links]] 

178 # and sits alone in the graph. No entities: the surface map is then 

179 # built from the pages on disk, this one included, so links go both 

180 # ways. 

181 rewrite_links_across_wiki([], config) 

182 on_progress( 

183 EventType.WIKI_PAGE, 

184 WikiPageEvent(label=stub.label, pages=1, current=1, total=1), 

185 ) 

186 return path