Coverage for src/lilbee/retrieval/entities/extractor.py: 100%
196 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-09-08 09:20 +0000
1"""Two-phase typed entity extraction.
3Phase 1 (cheap): an LLM reads a stratified sample of chunks and proposes the
4corpus-specific type schema, persisted as machine state inside the index
5before any expensive pass runs (see :mod:`schema`; nothing to review or edit).
7Phase 2 (scales with corpus): each type is found by the cheapest extractor
8that can serve it: compiled regex for identifier-shaped types, spaCy labels
9for the general ones, and an LLM only for types neither can catch. Cost is
10therefore dominated by how many LLM-kind types the schema keeps.
11"""
13from __future__ import annotations
15import logging
16import re
17from collections.abc import Mapping
18from dataclasses import dataclass, field
19from typing import TYPE_CHECKING, Any
21import regex
23from lilbee.core.llm_json import first_json_object, json_reply_format
24from lilbee.providers.base import aux_options
25from lilbee.retrieval.entities.schema import (
26 EntitySchema,
27 EntityType,
28 ExtractorKind,
29 extractor_key,
30)
31from lilbee.retrieval.reasoning import strip_reasoning
33if TYPE_CHECKING:
34 from lilbee.providers.base import LLMProvider
36log = logging.getLogger(__name__)
38INDUCTION_SAMPLE_SIZE = 40
39# Thinking models spend most of their budget reasoning before the JSON
40# appears (the reasoning is stripped afterward); 1200 starved them into
41# emitting nothing parseable.
42INDUCTION_MAX_TOKENS = 4096
43# Chunks per LLM call in phase 2; larger batches save round-trips, smaller
44# ones keep each response comfortably parseable.
45LLM_EXTRACTION_BATCH = 8
46LLM_EXTRACTION_MAX_TOKENS = 800
48_CONFIDENCE = {ExtractorKind.REGEX: 1.0, ExtractorKind.SPACY: 0.8, ExtractorKind.LLM: 0.6}
51@dataclass
52class ExtractionStats:
53 """Mutable state :func:`extract_entities` carries across its calls.
55 Lets the full pass tell a clean zero-entity result from batches the
56 provider failed outright -- the latter must not count as a completed
57 pass, or the schema gets marked applied with rows silently missing.
58 ``timed_out_types`` carries a runaway pattern's name across batches so
59 it is abandoned once per pass instead of per chunk.
60 """
62 llm_batches: int = 0
63 llm_batches_failed: int = 0
64 timed_out_types: set[str] = field(default_factory=set)
67INDUCTION_PROMPT = (
68 "You are designing an entity-extraction schema for a document collection. "
69 "Below are sample passages. Propose the 3-8 entity TYPES most useful for "
70 "counting and cross-referencing in this collection.\n"
71 "Return ONLY a JSON object of the form:\n"
72 '{{"types": [{{"name": "snake_case_name", "kind": "regex|spacy|llm", '
73 '"pattern": "<regex for regex kinds, spaCy label like PERSON/ORG/DATE for '
74 'spacy kinds, empty for llm kinds>", "description": "one line", '
75 '"synonyms": ["words a question would use"]}}]}}\n'
76 "Prefer regex for identifier-shaped types (codes, numbered records), spacy "
77 "for people/organizations/dates, llm only when unavoidable. Regex patterns "
78 "must match identifiers INLINE in running text: describe the identifier "
79 "token itself and never use ^ or $ anchors.\n\n"
80 "Passages:\n{sample}"
81)
83LLM_EXTRACTION_PROMPT = (
84 "Extract entities of these types from each numbered passage.\n"
85 "Types:\n{types}\n"
86 "Return ONLY a JSON object mapping passage number to a list of "
87 '{{"type": ..., "text": ...}} objects; use an empty list when a passage '
88 "has none.\n\nPassages:\n{passages}"
89)
92def normalize_value(text: str) -> str:
93 """Canonical form for grouping: casefold, collapse spaces, and strip
94 leading zeros from purely numeric values so 00482 and 482 group together."""
95 value = " ".join(text.split()).casefold()
96 if value.isdigit():
97 value = value.lstrip("0") or "0"
98 return value
101def _pattern_is_usable(entity_type: EntityType) -> bool:
102 """Whether an induced type's pattern can actually drive its extractor.
104 Only llm kinds legitimately leave the pattern blank (they carry a
105 description instead). An empty regex compiles but then matches at every
106 position of every chunk in the corpus, and an empty spaCy label names
107 nothing.
108 """
109 if entity_type.kind is ExtractorKind.LLM:
110 return True
111 if not entity_type.pattern:
112 log.warning("Dropping induced type %s: empty pattern", entity_type.name)
113 return False
114 if entity_type.kind is ExtractorKind.REGEX:
115 try:
116 _compile_pattern(entity_type.pattern)
117 except regex.error:
118 log.warning("Dropping induced type %s: bad regex", entity_type.name)
119 return False
120 return True
123def induce_schema(sample_texts: list[str], provider: LLMProvider) -> EntitySchema | None:
124 """Phase 1: propose a schema from sampled chunk texts. None on failure."""
125 if not sample_texts:
126 return None
127 sample = "\n---\n".join(t[:600] for t in sample_texts[:INDUCTION_SAMPLE_SIZE])
128 prompt = INDUCTION_PROMPT.format(sample=sample)
129 try:
130 response = provider.chat(
131 [{"role": "user", "content": prompt}],
132 stream=False,
133 # Induction wants one deterministic, well-formed schema, not a
134 # creative sample that parses only some of the time.
135 options=aux_options(
136 INDUCTION_MAX_TOKENS, temperature=0, response_format=json_reply_format()
137 ),
138 )
139 except Exception:
140 log.warning("Entity schema induction failed at the provider", exc_info=True)
141 return None
142 payload = first_json_object(strip_reasoning(response.text))
143 if payload is None:
144 log.warning("Entity schema induction returned no parseable JSON")
145 return None
146 types: list[EntityType] = []
147 seen_extractors: set[tuple[ExtractorKind, str]] = set()
148 for raw in payload.get("types", []):
149 try:
150 entity_type = EntityType.model_validate(raw)
151 except Exception:
152 log.warning("Dropping invalid induced type: %r", raw)
153 continue
154 if not _pattern_is_usable(entity_type):
155 continue
156 # Small models sometimes propose several names for one extractor
157 # (three types sharing a regex triple the table with identical rows);
158 # the first name wins.
159 key = extractor_key(entity_type)
160 if key in seen_extractors:
161 log.warning("Dropping induced type %s: duplicate extractor", entity_type.name)
162 continue
163 seen_extractors.add(key)
164 types.append(entity_type)
165 return EntitySchema(types=types) if types else None
168# Tokens for anchored-pattern matching: identifier-shaped runs of word chars
169# (hyphens allowed inside), so punctuation adjacent to an inline identifier
170# never defeats the match.
171_IDENTIFIER_TOKEN_RE = re.compile(r"[0-9A-Za-z][0-9A-Za-z-]*")
173# Wall-clock budget for one pattern against one chunk. Schema patterns are
174# authored by a small local model, so a pathological one (nested quantifiers)
175# is a realistic input; without a bound it backtracks until sync gives up.
176# A second is orders of magnitude above what a sane identifier pattern costs.
177_PATTERN_MATCH_TIMEOUT_SECONDS = 1.0
180def _compile_pattern(pattern: str) -> regex.Pattern[str]:
181 """Compile a schema pattern on the timeout-capable engine.
183 ``regex`` accepts everything ``re`` does and, unlike the stdlib module,
184 takes a per-match ``timeout``, which is the only way to bound an
185 LLM-authored pattern running over the whole corpus.
186 """
187 return regex.compile(pattern)
190def _extract_regex(entity_type: EntityType, text: str) -> list[str] | None:
191 """Regex mentions, or ``None`` when the match blew its time budget.
193 Schema authors (and the induction model) tend to write ^...$ patterns
194 that describe an identifier's whole shape. Applied to running text those
195 match almost nothing, since identifiers appear inline with adjacent
196 punctuation; a ^...$ pattern therefore full-matches each identifier
197 token instead of the chunk.
198 """
199 pattern = entity_type.pattern
200 inner_text = pattern[1:-1]
201 # Only a single fully-anchored pattern converts; an alternation of
202 # anchored branches (^A$|^B$) would be mangled by stripping the outer
203 # pair, so it falls through to finditer unchanged.
204 anchored = (
205 pattern.startswith("^")
206 and pattern.endswith("$")
207 and "^" not in inner_text
208 and "$" not in inner_text
209 )
210 try:
211 if anchored:
212 inner = _compile_pattern(inner_text)
213 return [
214 token.group(0)
215 for token in _IDENTIFIER_TOKEN_RE.finditer(text)
216 if inner.fullmatch(token.group(0), timeout=_PATTERN_MATCH_TIMEOUT_SECONDS)
217 ]
218 compiled = _compile_pattern(pattern)
219 return [m.group(0) for m in compiled.finditer(text, timeout=_PATTERN_MATCH_TIMEOUT_SECONDS)]
220 except TimeoutError:
221 return None
224def _extract_spacy(types: list[EntityType], text: str, nlp: Any) -> list[tuple[EntityType, str]]:
225 wanted = {t.pattern.upper(): t for t in types}
226 doc = nlp(text)
227 found: list[tuple[EntityType, str]] = []
228 for ent in doc.ents:
229 entity_type = wanted.get(ent.label_)
230 if entity_type is not None:
231 found.append((entity_type, ent.text))
232 return found
235def _extract_llm_batch(
236 types: list[EntityType],
237 texts: list[str],
238 provider: LLMProvider,
239) -> list[list[tuple[EntityType, str]]] | None:
240 """One LLM extraction call over a batch of texts.
242 Returns ``None`` when the provider call itself fails (model down or
243 unloaded), so the caller can tell a failed batch from a batch with no
244 entities. A response that parses to nothing usable is an empty result,
245 not a failure.
246 """
247 by_name = {t.name: t for t in types}
248 type_lines = "\n".join(f"- {t.name}: {t.description or t.name}" for t in types)
249 passages = "\n".join(f"[{i}] {t[:800]}" for i, t in enumerate(texts))
250 prompt = LLM_EXTRACTION_PROMPT.format(types=type_lines, passages=passages)
251 empty: list[list[tuple[EntityType, str]]] = [[] for _ in texts]
252 try:
253 response = provider.chat(
254 [{"role": "user", "content": prompt}],
255 stream=False,
256 options=aux_options(LLM_EXTRACTION_MAX_TOKENS),
257 )
258 except Exception:
259 log.warning("LLM entity extraction failed for a batch", exc_info=True)
260 return None
261 payload = first_json_object(strip_reasoning(response.text))
262 if payload is None:
263 return empty
264 results = empty
265 for key, items in payload.items():
266 try:
267 index = int(key)
268 except (TypeError, ValueError):
269 continue
270 if not (0 <= index < len(texts)) or not isinstance(items, list):
271 continue
272 for item in items:
273 if not isinstance(item, dict):
274 continue
275 entity_type = by_name.get(str(item.get("type", "")))
276 mention = str(item.get("text", "")).strip()
277 if entity_type is not None and mention:
278 results[index].append((entity_type, mention))
279 return results
282def _regex_findings(
283 regex_types: list[EntityType], text: str, timed_out: set[str]
284) -> list[tuple[EntityType, str]]:
285 """Regex-kind findings for one chunk, skipping types that blew their budget.
287 A type whose pattern times out is added to *timed_out* so the caller
288 stops attempting it; the ban lasts as long as that set does.
289 """
290 found: list[tuple[EntityType, str]] = []
291 for entity_type in regex_types:
292 if entity_type.name in timed_out:
293 continue
294 matches = _extract_regex(entity_type, text)
295 if matches is None:
296 timed_out.add(entity_type.name)
297 log.warning(
298 "Entity type %s: its pattern exceeded the %.0fs match budget; "
299 "skipping that type. Its regex is likely pathological.",
300 entity_type.name,
301 _PATTERN_MATCH_TIMEOUT_SECONDS,
302 )
303 continue
304 found.extend((entity_type, m) for m in matches)
305 return found
308# Singles failing in a row before a batch retry concludes the provider is down.
309_SINGLE_RETRY_ABORT = 3
312def _retry_batch_singly(
313 types: list[EntityType],
314 batch: list[Mapping[str, Any]],
315 provider: LLMProvider,
316) -> tuple[list[list[tuple[EntityType, str]]], bool]:
317 """Per-chunk retry of a failed batch: (per-chunk findings, batch failed).
319 A chunk whose single call still fails loses its entities (logged); the
320 batch only counts as failed when nothing succeeds, i.e. the provider
321 itself is down, so the pass retries next sync.
322 """
323 found_all: list[list[tuple[EntityType, str]]] = [[] for _ in batch]
324 any_ok = False
325 consecutive = 0
326 for offset, record in enumerate(batch):
327 single = _extract_llm_batch(types, [record["chunk"]], provider)
328 if single is None:
329 consecutive += 1
330 log.warning(
331 "LLM entity extraction failed for %s#%s; its entities are skipped",
332 record["source"],
333 record["chunk_index"],
334 )
335 if not any_ok and consecutive >= _SINGLE_RETRY_ABORT:
336 break
337 continue
338 consecutive = 0
339 any_ok = True
340 found_all[offset] = single[0]
341 return found_all, not any_ok
344def extract_entities(
345 chunks: list[Mapping[str, Any]],
346 schema: EntitySchema,
347 *,
348 provider: LLMProvider | None = None,
349 nlp: Any = None,
350 stats: ExtractionStats | None = None,
351) -> list[dict]:
352 """Phase 2 over ingest-shaped chunk records; returns entities-table rows.
354 Each chunk dict needs ``chunk`` (text), ``source``, ``chunk_index``, and
355 ``page_start``. Extractor kinds degrade independently: regex always runs,
356 spaCy kinds are skipped without a loaded model, LLM kinds without a
357 provider, so a partial toolchain yields partial extraction, never failure.
358 Pass ``stats`` to observe how many LLM batches ran and how many the
359 provider failed; a failed batch contributes no rows either way.
360 """
361 regex_types = [t for t in schema.types if t.kind is ExtractorKind.REGEX]
362 spacy_types = [t for t in schema.types if t.kind is ExtractorKind.SPACY]
363 llm_types = [t for t in schema.types if t.kind is ExtractorKind.LLM]
365 # A pattern that blows its budget is abandoned rather than retried on
366 # every remaining chunk; with stats the ban holds for the whole pass.
367 timed_out = stats.timed_out_types if stats is not None else set()
369 per_chunk: list[list[tuple[EntityType, str]]] = []
370 for record in chunks:
371 text = record["chunk"]
372 found = _regex_findings(regex_types, text, timed_out)
373 if spacy_types and nlp is not None:
374 found.extend(_extract_spacy(spacy_types, text, nlp))
375 per_chunk.append(found)
377 if llm_types and provider is not None:
378 for start in range(0, len(chunks), LLM_EXTRACTION_BATCH):
379 batch = chunks[start : start + LLM_EXTRACTION_BATCH]
380 batch_found = _extract_llm_batch(llm_types, [r["chunk"] for r in batch], provider)
381 failed = False
382 if batch_found is None:
383 # Retry chunk-by-chunk: one poisoned chunk must not fail the
384 # batch (and with it the whole pass) on every sync.
385 batch_found, failed = _retry_batch_singly(llm_types, batch, provider)
386 if stats is not None:
387 stats.llm_batches += 1
388 if failed:
389 stats.llm_batches_failed += 1
390 for offset, found in enumerate(batch_found):
391 per_chunk[start + offset].extend(found)
393 return [
394 row
395 for record, found in zip(chunks, per_chunk, strict=True)
396 for row in _rows_for_chunk(record, found)
397 ]
400def _rows_for_chunk(record: Mapping[str, Any], found: list[tuple[EntityType, str]]) -> list[dict]:
401 """Deduplicated entities-table rows for one chunk's findings."""
402 rows: list[dict] = []
403 seen: set[tuple[str, str]] = set()
404 for entity_type, mention in found:
405 normalized = normalize_value(mention)
406 if not normalized:
407 continue
408 key = (entity_type.name, normalized)
409 if key in seen:
410 continue
411 seen.add(key)
412 rows.append(
413 {
414 "entity": mention,
415 "type": entity_type.name,
416 "normalized_value": normalized,
417 "source": record["source"],
418 "page": int(record.get("page_start") or 0),
419 "chunk_index": int(record["chunk_index"]),
420 "confidence": _CONFIDENCE[entity_type.kind],
421 }
422 )
423 return rows