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