Coverage for src/lilbee/retrieval/entities/lifecycle.py: 100%
121 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-14 11:46 +0000
1"""Automatic entity-extraction lifecycle, run by sync.
3Turning ``entity_extraction`` on is the whole user interaction. Sync does
4the rest: the first run samples the indexed chunks, induces a schema, and
5extracts across the whole index; later runs extract new files at ingest
6(see ``pipeline.build_entity_records``). The schema is machine state
7persisted inside the index, so it travels with the data and there is
8nothing for the user to manage.
10The taxonomy keeps up with the corpus on its own. A library that grows a
11new kind of document would otherwise keep answering with the type set
12induced on day one, so sync re-induces from a fresh sample once the corpus
13has grown materially since the schema was written, and unions in any type
14the schema has never seen. Re-induction that proposes nothing new costs one
15sampled chat call and no extraction pass.
17Extraction is idempotent: every full pass clears previously extracted
18rows first, so an interrupted pass, or a schema that gained a type, can
19never double-count.
20"""
22from __future__ import annotations
24import logging
25from typing import TYPE_CHECKING
27from lilbee.core.config import CHUNKS_TABLE, ENTITIES_TABLE, active_config
28from lilbee.retrieval.entities.extractor import (
29 INDUCTION_SAMPLE_SIZE,
30 ExtractionStats,
31 extract_entities,
32 induce_schema,
33)
34from lilbee.retrieval.entities.schema import (
35 EntitySchema,
36 ExtractorKind,
37 merge_schemas,
38 parse_schema,
39 save_schema,
40)
42if TYPE_CHECKING:
43 from lilbee.data.store import Store
44 from lilbee.data.store.types import EntitySchemaState
45 from lilbee.runtime.cancellation import CancelSignal
47log = logging.getLogger(__name__)
49# Chunks per extraction batch during a full pass; bounds the boxed Python
50# rows in flight (the scan itself is projected and stays columnar).
51_BACKFILL_BATCH = 2000
53# Corpus growth that makes the induced taxonomy worth revisiting, as a
54# multiple of the document count at induction. Half again as many documents
55# is enough to have introduced a family the first sample never saw; smaller
56# ratios would re-induce on ordinary drip-feed additions, whose new files are
57# already extracted at ingest under the existing schema.
58_REINDUCE_GROWTH_FACTOR = 1.5
59# Below this many documents, growth ratios are meaningless (2 -> 3 documents
60# is 1.5x): a corpus this small re-induces on any growth, since the first
61# sample necessarily saw very little.
62_REINDUCE_MIN_SOURCES = 10
65def ensure_entities(cancel: CancelSignal | None = None) -> None:
66 """Bring extracted entities in line with the schema; no-op when off.
68 Failure never fails the sync: a missing chat model defers induction to
69 the next sync with a log line, and a cancelled pass leaves the schema
70 unapplied so the next sync redoes the (idempotent) full pass.
71 """
72 # Read through the scope: under the library API the active config is the
73 # caller's, not the process-global cfg (which may say the feature is off).
74 if not active_config().entity_extraction:
75 return
76 from lilbee.app.services import get_services
78 store = get_services().store
79 # One read of the schema row: it carries the taxonomy, whether a full
80 # pass ever completed under it, and the corpus size it was induced from.
81 state = store.entity_schema_state()
82 schema = parse_schema(state) if state is not None else None
83 if state is None or schema is None:
84 # Never induced (or the persisted row is unreadable): induce afresh.
85 schema = _induce(store)
86 if schema is None:
87 return
88 else:
89 evolved = _evolve(store, schema, state)
90 if evolved is not None:
91 schema = evolved
92 elif state["applied"]:
93 return # up to date; new files were covered at ingest
94 else:
95 log.info("Entity extraction pass incomplete; redoing the full pass")
96 if _full_pass(store, schema, cancel):
97 store.mark_entity_schema_applied()
100def _evolve(store: Store, schema: EntitySchema, state: EntitySchemaState) -> EntitySchema | None:
101 """The schema grown to cover a drifted corpus, or ``None`` to leave it be.
103 Returns a schema only when re-induction actually found a type the
104 current one lacks, so a corpus that grew without changing character
105 costs one chat call and no extraction pass. The new document count is
106 recorded either way, so a re-induction that adds nothing does not run
107 again on the next sync.
108 """
109 from lilbee.app.services import get_services
111 sources = store.count_sources()
112 at_induction = state["source_count"]
113 if not _corpus_drifted(at_induction, sources):
114 return None
115 # A drifted corpus has chunks to sample by construction; an empty sample
116 # would fall through to induce_schema, which declines it like any other
117 # unusable input, so it needs no separate branch here.
118 texts = _sample_chunks(store, INDUCTION_SAMPLE_SIZE)
119 induced = induce_schema(texts, get_services().provider)
120 if induced is None:
121 log.warning("Entity schema re-induction produced nothing usable; retrying next sync")
122 return None
123 merged = merge_schemas(schema, induced)
124 if len(merged.types) == len(schema.types):
125 # Nothing new: record the corpus size so this does not re-run, and
126 # leave the applied flag alone (extraction is already current).
127 save_schema(schema, store, applied=state["applied"], source_count=sources)
128 log.info("Corpus grew but the entity taxonomy still fits; no re-extraction")
129 return None
130 added = [t.name for t in merged.types[len(schema.types) :]]
131 log.info(
132 "Corpus grew from %d to %d documents; entity schema gained %s. Re-extracting",
133 at_induction,
134 sources,
135 ", ".join(added),
136 )
137 save_schema(merged, store, applied=False, source_count=sources)
138 return merged
141def _corpus_drifted(at_induction: int, now: int) -> bool:
142 """Whether the corpus has grown enough to revisit the taxonomy."""
143 if now <= at_induction:
144 return False
145 if at_induction < _REINDUCE_MIN_SOURCES:
146 return True
147 return now >= at_induction * _REINDUCE_GROWTH_FACTOR
150def _induce(store: Store) -> EntitySchema | None:
151 """First-run schema induction from the indexed chunks. None defers."""
152 from lilbee.app.services import get_services
154 texts = _sample_chunks(store, INDUCTION_SAMPLE_SIZE)
155 if not texts:
156 log.info("Entity extraction is on but nothing is indexed yet; deferring")
157 return None
158 schema = induce_schema(texts, get_services().provider)
159 if schema is None:
160 log.warning("Entity schema induction produced nothing usable; retrying next sync")
161 return None
162 save_schema(schema, store, applied=False, source_count=store.count_sources())
163 log.info("Induced entity schema (%d types); extracting across the index", len(schema.types))
164 return schema
167def _sample_chunks(store: Store, limit: int) -> list[str]:
168 """Stratified sample: chunks read evenly across the table so a corpus
169 with several document families contributes all of them."""
170 table = store.open_table(CHUNKS_TABLE)
171 if table is None:
172 return []
173 total = table.count_rows()
174 if total == 0:
175 return []
176 # Pick the target indices outright rather than striding. A stride wide
177 # enough to span the table also thins the sample (total=41 yielded 21 rows,
178 # not 40), while a narrow one never reaches the tail -- the most recently
179 # ingested documents, exactly what re-induction exists to catch. Spreading
180 # across the closed range fills the budget and hits both ends at any size.
181 wanted = {round(i * (total - 1) / max(1, limit - 1)) for i in range(limit)}
182 texts: list[str] = []
183 # Projection is pushed into LanceDB: a bare to_arrow() would drag every
184 # column, embedding vectors included, into memory before selecting.
185 arrow = table.search().select(["chunk"]).limit(None).to_arrow()
186 index = 0
187 for batch in arrow.to_batches(max_chunksize=_BACKFILL_BATCH):
188 for text in batch.column("chunk").to_pylist():
189 if index in wanted:
190 texts.append(text)
191 if len(texts) >= limit:
192 return texts
193 index += 1
194 return texts
197def _full_pass(store: Store, schema: EntitySchema, cancel: CancelSignal | None) -> bool:
198 """Extract entities across every stored chunk. True when it completed.
200 Clears previously extracted rows first so the pass is idempotent. A pass
201 with failed LLM batches (chat model down or erroring) does NOT count as
202 completed: rows are missing for those chunks, and marking the schema
203 applied would make every later sync return early and never retry.
204 """
205 table = store.open_table(CHUNKS_TABLE)
206 if table is None:
207 return False
208 store.clear_table(ENTITIES_TABLE, "entity IS NOT NULL")
209 nlp = None
210 if any(t.kind is ExtractorKind.SPACY for t in schema.types):
211 from lilbee.retrieval.concepts import concepts_available
213 if concepts_available():
214 from lilbee.retrieval.concepts.nlp import load_spacy_pipeline
216 try:
217 nlp = load_spacy_pipeline()
218 except ImportError:
219 log.warning("spaCy model unavailable; skipping spaCy entity types")
220 provider = None
221 if any(t.kind is ExtractorKind.LLM for t in schema.types):
222 from lilbee.app.services import get_services
224 provider = get_services().provider
225 written = 0
226 stats = ExtractionStats()
227 columns = ["chunk", "source", "chunk_index", "page_start"]
228 # Projection pushed into LanceDB; see _sample_chunks.
229 arrow = table.search().select(columns).limit(None).to_arrow()
230 for batch in arrow.to_batches(max_chunksize=_BACKFILL_BATCH):
231 if cancel is not None and cancel.is_set():
232 log.info("Entity extraction cancelled; the next sync restarts the pass")
233 return False
234 records = batch.to_pylist()
235 rows = extract_entities(records, schema, provider=provider, nlp=nlp, stats=stats)
236 written += store.add_entities(rows)
237 if stats.llm_batches_failed:
238 log.warning(
239 "Entity extraction pass had %d of %d LLM batches fail; the next sync redoes the pass",
240 stats.llm_batches_failed,
241 stats.llm_batches,
242 )
243 return False
244 log.info("Entity extraction complete: %d rows across the index", written)
245 return True