Coverage for src/lilbee/retrieval/entities/schema.py: 100%

1"""Typed entity extraction: the schema model and the entities table. 

2 

3The extraction taxonomy is induced from the corpus, not fixed: a general NER 

4tag set has no notion of the identifier types a specific corpus carries. 

5Sync induces a schema from a corpus sample and applies it automatically; the 

6schema is machine state, persisted inside the LanceDB index so it travels 

7with the data and needs no management. There is nothing to review or edit: 

8if induction quality needs improving, the fix belongs in induction itself. 

9""" 

10 

11from __future__ import annotations 

12 

13import json 

14import logging 

15from enum import Enum 

16from typing import TYPE_CHECKING 

17 

18import pyarrow as pa 

19from pydantic import BaseModel, Field, field_validator 

20 

21from lilbee.retrieval.language import noun_variants 

22 

23if TYPE_CHECKING: 

24 from lilbee.data.store import Store 

25 from lilbee.data.store.types import EntitySchemaState 

26 

27log = logging.getLogger(__name__) 

28 

29 

30class ExtractorKind(Enum): 

31 """How a type's mentions are found, cheapest first.""" 

32 

33 REGEX = "regex" 

34 SPACY = "spacy" 

35 LLM = "llm" 

36 

37 

38class EntityType(BaseModel): 

39 """One induced type: how to find it and what to call it in questions.""" 

40 

41 name: str = Field(min_length=1, max_length=64) 

42 kind: ExtractorKind 

43 # REGEX kinds compile this; SPACY kinds name a spaCy label (PERSON, ORG, 

44 # DATE, ...); LLM kinds carry a one-line description for the prompt. 

45 pattern: str = "" 

46 description: str = "" 

47 # Question nouns that mean this type ("part number", "part numbers"). 

48 synonyms: list[str] = Field(default_factory=list) 

49 

50 @field_validator("name") 

51 @classmethod 

52 def _slugify(cls, v: str) -> str: 

53 slug = "_".join(v.strip().lower().split()) 

54 if not slug.replace("_", "").isalnum(): 

55 raise ValueError(f"type name must be alphanumeric words: {v!r}") 

56 return slug 

57 

58 

59class EntitySchema(BaseModel): 

60 """The induced extraction contract for one corpus.""" 

61 

62 types: list[EntityType] 

63 

64 def type_for_noun(self, noun: str) -> EntityType | None: 

65 """Resolve a question noun (singular or plural) to a type, if any.""" 

66 candidates = noun_variants(noun) 

67 for entity_type in self.types: 

68 names = {entity_type.name, entity_type.name.replace("_", " ")} 

69 names.update(entity_type.synonyms) 

70 expanded: set[str] = set() 

71 for name in names: 

72 expanded |= noun_variants(name) 

73 if candidates & expanded: 

74 return entity_type 

75 return None 

76 

77 

78def extractor_key(entity_type: EntityType) -> tuple[ExtractorKind, str]: 

79 """What a type actually extracts, ignoring the name it was given. 

80 

81 Two types with the same kind and the same pattern (or, for LLM kinds, 

82 the same description) find the same mentions, so they are the same type 

83 wearing different names. Induction dedupes by this, and schema evolution 

84 uses it to tell a genuinely new type from a rename of a known one. 

85 """ 

86 signature = entity_type.pattern.strip() or entity_type.description.strip() 

87 return entity_type.kind, signature.casefold() 

88 

89 

90def merge_schemas(existing: EntitySchema, induced: EntitySchema) -> EntitySchema: 

91 """Existing types plus any genuinely new ones from *induced*. 

92 

93 Union, not replace: existing types keep their names and positions so 

94 answers stay stable across a re-induction, and a type the fresh sample 

95 happens not to cover is never silently dropped (the documents it was 

96 induced from are still in the index). Only extractors the schema has 

97 never seen are appended. 

98 """ 

99 known = {extractor_key(t) for t in existing.types} 

100 known_names = {t.name for t in existing.types} 

101 additions = [ 

102 t for t in induced.types if extractor_key(t) not in known and t.name not in known_names 

103 ] 

104 if not additions: 

105 return existing 

106 return EntitySchema(types=[*existing.types, *additions]) 

107 

108 

109def parse_schema(state: EntitySchemaState) -> EntitySchema | None: 

110 """The schema carried by a persisted row, or ``None`` when unreadable. 

111 

112 Unreadable is logged, not raised: a corrupt row is machine state gone 

113 wrong, so the lifecycle re-induces automatically instead of failing sync. 

114 """ 

115 try: 

116 return EntitySchema.model_validate_json(state["schema_json"]) 

117 except Exception: 

118 log.warning("Persisted entity schema is unreadable; re-inducing", exc_info=True) 

119 return None 

120 

121 

122def load_schema(store: Store) -> EntitySchema | None: 

123 """Read the persisted schema from the index, or ``None`` when never induced.""" 

124 state = store.entity_schema_state() 

125 return parse_schema(state) if state is not None else None 

126 

127 

128def save_schema(schema: EntitySchema, store: Store, *, applied: bool, source_count: int) -> None: 

129 """Persist the schema into the index (stable key order). 

130 

131 ``source_count`` is the document count the schema was induced from, and 

132 is the baseline the next sync compares against to decide the corpus has 

133 drifted; it is required so a save can never silently record a baseline 

134 of zero, which would read as a tiny corpus and re-induce every sync. 

135 """ 

136 payload = json.dumps(schema.model_dump(mode="json"), sort_keys=True) 

137 store.save_entity_schema(payload, applied=applied, source_count=source_count) 

138 

139 

140def _entities_schema() -> pa.Schema: 

141 return pa.schema( 

142 [ 

143 pa.field("entity", pa.utf8()), 

144 pa.field("type", pa.utf8()), 

145 pa.field("normalized_value", pa.utf8()), 

146 pa.field("source", pa.utf8()), 

147 pa.field("page", pa.int32()), 

148 pa.field("chunk_index", pa.int32()), 

149 pa.field("confidence", pa.float32()), 

150 ] 

151 )