Coverage for collectra/tasks/irn_resolvers.py: 0.00%
190 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-16 09:56 +0000
1"""IRN (Internal Reference Number) resolver tasks for Collectra workflows.
3This module provides an IRN resolver that uses semantic similarity search
4(via Chroma + OpenAI embeddings) plus an LLM verification step to map
5free-text query inputs to canonical IRN identifiers loaded from an xlsx
6source file.
8Public API:
9 IRNResolver: Task class that resolves a query to an IRN string.
10 normalise_column_name: Helper that normalises raw spreadsheet column names.
11 render_card: Helper that renders a per-row card from a template.
12"""
14__all__ = ["IRNResolver", "normalise_column_name", "render_card"]
16import hashlib
17import math
18import re
19from pathlib import Path
21import pandas as pd
22from dotenv import load_dotenv
23from langchain_chroma import Chroma
24from langchain_core.documents import Document
25from langchain_core.messages import HumanMessage
27from collectra.tasks.canonicalisers import LLMCanonicaliser
28from collectra.types.texts import Text
30from ..logger import get_logger
32load_dotenv()
34logger = get_logger(__name__)
37NO_MATCH_INSTRUCTION = (
38 "If none of the candidates is clearly the same entity referenced in the input, "
39 'respond with empty string (""). Do NOT guess. Respond with ONLY the IRN integer, '
40 "or empty string."
41)
44class _SafeDict(dict):
45 """A dict subclass that returns "" for missing keys when used with format_map."""
47 def __missing__(self, key):
48 return ""
51def normalise_column_name(raw: str) -> str:
52 """Normalise a raw spreadsheet column name into a snake_case identifier.
54 Rules applied in order:
55 1. Strip a single trailing parenthetical ``(...)`` together with surrounding
56 whitespace.
57 2. Strip a trailing ``:``.
58 3. Lowercase.
59 4. Replace runs of whitespace and ``/`` with ``_``.
60 5. Drop any character not in ``[a-z0-9_]``.
61 6. Collapse runs of ``_`` and strip leading/trailing ``_``.
62 """
63 s = str(raw)
65 # 1. Strip a single trailing parenthetical (...) with surrounding whitespace.
66 s = re.sub(r"\s*\([^()]*\)\s*$", "", s)
68 # 2. Strip a trailing ':'.
69 s = re.sub(r":\s*$", "", s)
70 s = s.rstrip(":")
72 # 3. Lowercase.
73 s = s.lower()
75 # 4. Replace runs of whitespace AND '/' with '_'.
76 s = re.sub(r"[\s/]+", "_", s)
78 # 5. Drop any character not in [a-z0-9_].
79 s = re.sub(r"[^a-z0-9_]", "", s)
81 # 6. Collapse runs of '_' and strip leading/trailing '_'.
82 s = re.sub(r"_+", "_", s)
83 s = s.strip("_")
85 return s
88def _is_missing(value) -> bool:
89 """Return True if value is None, NaN, or pandas-NA."""
90 if value is None:
91 return True
92 try:
93 if pd.isna(value):
94 return True
95 except (TypeError, ValueError):
96 pass
97 if isinstance(value, float):
98 try:
99 if math.isnan(value):
100 return True
101 except (TypeError, ValueError):
102 pass
103 return False
106def render_card(template: str, row: dict) -> str:
107 """Render a card from a template and a row mapping.
109 Replaces None/NaN values with "", applies ``template.format(**row)``,
110 then collapses runs of separator-only segments (``|`` or ``,`` with
111 surrounding whitespace) and strips leading/trailing separators.
112 """
113 cleaned = {k: ("" if _is_missing(v) else str(v)) for k, v in row.items()}
114 rendered = template.format_map(_SafeDict(cleaned))
116 # Collapse runs of (\s*[|,]\s*){2,} into the LAST single \s*[|,]\s* in the
117 # run. We reconstruct the "last single" as: optional single leading space
118 # (if the run begins with whitespace), the last separator character, and
119 # optional single trailing space (if the run ends with whitespace). This
120 # avoids the greedy-matching artefact where earlier iterations of the
121 # capture group consume whitespace that should belong to the last match.
122 sep_chars = "|,"
124 def _collapse(match: re.Match) -> str:
125 run = match.group(0)
126 # Last separator character in the run.
127 last_sep = None
128 for ch in reversed(run):
129 if ch in sep_chars:
130 last_sep = ch
131 break
132 if last_sep is None:
133 return run
134 leading = " " if run and run[0].isspace() else ""
135 trailing = " " if run and run[-1].isspace() else ""
136 return f"{leading}{last_sep}{trailing}"
138 rendered = re.sub(r"(\s*[|,]\s*){2,}", _collapse, rendered)
140 # Strip leading/trailing whitespace and any leading/trailing | or , plus
141 # surrounding whitespace.
142 rendered = re.sub(r"^(\s*[|,]\s*)+", "", rendered)
143 rendered = re.sub(r"(\s*[|,]\s*)+$", "", rendered)
144 rendered = rendered.strip()
146 return rendered
149class IRNResolver(LLMCanonicaliser):
150 """Resolve a free-text query to an IRN identifier.
152 Builds one Chroma sub-index per entry in ``index_filters`` from rows of
153 a source xlsx file. At query time, tries each filter in order, retrieves
154 the top-``count`` candidates by embedding distance, applies acceptance
155 rules (``score_floor`` and optional ``score_gap_eps``), then asks the LLM
156 to either pick a candidate IRN or return empty string.
158 Attributes:
159 source (Path): Path to the xlsx source file.
160 card_template (str): str.format template using normalised column names.
161 query_template (str): str.format template using input Text names.
162 irn_column (str): Name of the IRN column (default "IRN").
163 count (int): Number of candidates to retrieve per filter (default 20).
164 embedding_model (str): Embedding model name.
165 score_floor (float): Maximum acceptable Chroma distance for top-1.
166 score_gap_eps (float): Minimum gap between top-1 and top-2 distances.
167 index_filters (list): Pandas-query strings (or None) - one per sub-index.
168 """
170 def __init__(self, name: str, model: str, **kwargs):
171 # Required kwargs (validate before super-init since parent doesn't know
172 # about them).
173 if "source" not in kwargs:
174 raise ValueError("IRNResolver requires a 'source' kwarg (path to xlsx).")
175 if "card_template" not in kwargs:
176 raise ValueError("IRNResolver requires a 'card_template' kwarg.")
177 if "query_template" not in kwargs:
178 raise ValueError("IRNResolver requires a 'query_template' kwarg.")
180 # The parent LLMCanonicaliser.__init__ wraps `entities` in
181 # Path(...) and then checks .exists()/.suffix. We don't use the
182 # parent's entity machinery (retrieve_entities and run are overridden),
183 # so we inject a Path-safe sentinel. Using "." makes Path(".") valid
184 # and exists() True; the parent then assigns self.entities = Path(".")
185 # without triggering file reads.
186 kwargs.setdefault("entities", ".")
187 super().__init__(name, model, **kwargs)
189 self.source: Path = Path(kwargs["source"])
190 self.card_template: str = kwargs["card_template"]
191 self.query_template: str = kwargs["query_template"]
193 self.irn_column: str = kwargs.get("irn_column", "IRN")
194 self.count: int = kwargs.get("count", 20)
195 self.embedding_model: str = kwargs.get(
196 "embedding_model", "text-embedding-3-large"
197 )
198 self.score_floor: float = kwargs.get("score_floor", 0.5)
199 self.score_gap_eps: float = kwargs.get("score_gap_eps", 0.0)
200 self.index_filters: list = kwargs.get("index_filters", [None])
202 # Augment preamble once with the no-match instruction.
203 if self.preamble:
204 self.preamble = f"{self.preamble}\n\n{NO_MATCH_INSTRUCTION}"
205 else:
206 self.preamble = NO_MATCH_INSTRUCTION
208 # Build/load one Chroma vector store per filter.
209 from langchain_openai import OpenAIEmbeddings
211 self._embeddings = OpenAIEmbeddings(model=self.embedding_model)
212 self._vector_stores: list[Chroma] = []
213 self._build_indices()
215 # ------------------------------------------------------------------
216 # Index building / caching
217 # ------------------------------------------------------------------
218 def _source_hash(self) -> str:
219 with open(self.source, "rb") as f:
220 return hashlib.sha256(f.read()).hexdigest()
222 def _index_path_for(self, i: int) -> Path:
223 return self.source.parent / f"{self.source.stem}_irn_index_{i}"
225 def _collection_name_for(self, i: int) -> str:
226 # Chroma collection names: 3–512 chars, [a-zA-Z0-9._-], must start/end alnum.
227 safe_stem = re.sub(r"[^a-zA-Z0-9._-]", "_", self.source.stem)
228 return f"{safe_stem}_irn_collection_{i}"
230 def _build_indices(self) -> None:
231 current_hash = self._source_hash()
232 df_raw_cache: pd.DataFrame | None = None
234 for i, filter_str in enumerate(self.index_filters):
235 index_path = self._index_path_for(i)
236 sidecar = index_path / ".source_hash"
237 collection_name = self._collection_name_for(i)
239 reuse = (
240 index_path.exists()
241 and index_path.is_dir()
242 and sidecar.exists()
243 and sidecar.read_text().strip() == current_hash
244 )
246 if reuse:
247 logger.info(
248 "Reusing existing IRN index at %s for filter %r",
249 index_path,
250 filter_str,
251 )
252 vector_store = Chroma(
253 collection_name=collection_name,
254 embedding_function=self._embeddings,
255 persist_directory=str(index_path),
256 )
257 self._vector_stores.append(vector_store)
258 continue
260 # Need to (re)build this index.
261 if df_raw_cache is None:
262 df_raw_cache = pd.read_excel(self.source)
263 self.df = df_raw_cache.rename(
264 columns={c: normalise_column_name(c) for c in df_raw_cache.columns}
265 )
267 # Filter against the raw spreadsheet column names so users can write
268 # filters like "Species.isna()" using the actual xlsx headers.
269 df = df_raw_cache
270 if filter_str is not None:
271 df = df.query(filter_str)
272 df = df.rename(columns={c: normalise_column_name(c) for c in df.columns})
274 normalised_irn_column = normalise_column_name(self.irn_column)
275 id_col = (
276 self.irn_column
277 if self.irn_column in df.columns
278 else normalised_irn_column
279 )
280 if id_col not in df.columns:
281 raise ValueError(
282 f"IRN column {self.irn_column!r} (normalised {normalised_irn_column!r}) "
283 f"not found in source columns: {list(df.columns)}"
284 )
286 ids: list[str] = []
287 documents: list[Document] = []
288 seen: set[str] = set()
289 for _, row in df.iterrows():
290 row_dict = row.to_dict()
291 irn_value = row_dict.get(id_col)
292 if _is_missing(irn_value):
293 continue
294 irn_str = str(irn_value).strip()
295 if not irn_str or irn_str in seen:
296 continue
297 card = render_card(self.card_template, row_dict)
298 if not card:
299 continue
300 ids.append(irn_str)
301 documents.append(Document(page_content=card))
302 seen.add(irn_str)
304 index_path.mkdir(parents=True, exist_ok=True)
305 vector_store = Chroma(
306 collection_name=collection_name,
307 embedding_function=self._embeddings,
308 persist_directory=str(index_path),
309 )
310 if documents:
311 vector_store.add_documents(documents=documents, ids=ids)
313 sidecar.write_text(current_hash)
314 logger.info(
315 "Built IRN index at %s with %d rows for filter %r",
316 index_path,
317 len(documents),
318 filter_str,
319 )
320 self._vector_stores.append(vector_store)
322 # ------------------------------------------------------------------
323 # Retrieval / run
324 # ------------------------------------------------------------------
325 def retrieve_entities(
326 self, query_data: str, filter_index: int = 0
327 ) -> list[tuple[str, str]]:
328 """Retrieve top-``count`` (id, page_content) pairs for the given filter."""
329 vector_store = self._vector_stores[filter_index]
330 results = vector_store.similarity_search_with_score(query_data, k=self.count)
331 out: list[tuple[str, str]] = []
332 for doc, _score in results:
333 out.append((str(doc.id), doc.page_content))
334 return out
336 def _render_query(self, inputs_dict: dict) -> str:
337 return self.query_template.format_map(_SafeDict(inputs_dict))
339 def _try_filter(
340 self,
341 filter_index: int,
342 query_text: str,
343 args: tuple,
344 ) -> str | None:
345 """Try one filter; return matched IRN string or None."""
346 vector_store = self._vector_stores[filter_index]
347 results = vector_store.similarity_search_with_score(query_text, k=self.count)
348 if not results:
349 return None
350 if results[0][1] > self.score_floor:
351 return None
352 if (
353 self.score_gap_eps > 0
354 and len(results) >= 2
355 and (results[1][1] - results[0][1]) < self.score_gap_eps
356 ):
357 return None
359 candidate_ids: list[str] = [str(doc.id) for doc, _ in results]
360 candidates_string = "\n".join(
361 f"IRN {doc.id}: {doc.page_content}" for doc, _ in results
362 )
364 # Mirror LLMCanonicaliser.run for the LLM call portion.
365 prompt, pattern = self.pre_run()
366 messages = self.replace_inputs(
367 pattern, prompt, *args, entities=candidates_string
368 )
369 self.messages.append(HumanMessage(content=messages))
370 self.check_inputs(*args)
371 response = self.invoke()
373 # Parse response: strip whitespace and surrounding quotes.
374 cleaned = response.strip()
375 if cleaned and cleaned[0] in ("'", '"') and cleaned[-1] == cleaned[0]:
376 cleaned = cleaned[1:-1].strip()
378 if not cleaned or cleaned.lower() in ("none", "n/a"):
379 return None
381 if cleaned in candidate_ids:
382 return cleaned
383 return None
385 def run(self, *args: Text) -> Text:
386 """Resolve the inputs to an IRN string by trying each filter in order."""
387 inputs_dict = {arg.name: str(arg.data) for arg in args}
388 query_text = self._render_query(inputs_dict)
390 for filter_index in range(len(self._vector_stores)):
391 matched = self._try_filter(filter_index, query_text, args)
392 if matched is not None:
393 return Text(name=self.get_output_name(), data=matched)
395 logger.info("IRN resolver %s returned empty (no acceptable match)", self.name)
396 return Text(name=self.get_output_name(), data="")