Coverage for astrocyte/pipeline/reranking.py: 91%
170 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-10-29 22:10 +0000
1"""Basic reranking — keyword and entity-aware scoring for Phase 1.
3Sync, pure computation — Rust migration candidate.
4Phase 2: cross-encoder or LLM-based reranking.
5"""
7from __future__ import annotations
9import json
10from dataclasses import replace
11from string import punctuation
12from typing import TYPE_CHECKING
14from astrocyte.mip.schema import RerankSpec
15from astrocyte.pipeline.fusion import ScoredItem
16from astrocyte.types import Message
18if TYPE_CHECKING:
19 from astrocyte.provider import LLMProvider
21COMMON_QUESTION_WORDS: set[str] = {
22 "what",
23 "when",
24 "where",
25 "who",
26 "why",
27 "how",
28 "which",
29 "did",
30 "does",
31 "do",
32 "is",
33 "are",
34 "was",
35 "were",
36 "has",
37 "have",
38 "can",
39 "could",
40 "would",
41 "should",
42 "will",
43 "tell",
44 "describe",
45}
47KEYWORD_OVERLAP_WEIGHT = 0.05
48PROPER_NOUN_WEIGHT = 0.10
49QUERY_INTERACTION_WEIGHT = 0.35
50QUERY_PHRASE_WEIGHT = 0.08
51QUERY_NAME_MATCH_WEIGHT = 0.25
52QUERY_NAME_MISS_PENALTY = 0.20
53COMPILED_LAYER_WEIGHT = 0.30
54OBSERVATION_LAYER_WEIGHT = 0.20
55SESSION_DIVERSITY_PENALTY = 0.08
56# Observation proof-count boost: each additional confirming memory adds this
57# to the score, capped at OBSERVATION_PROOF_CAP × weight. A 5-evidence
58# observation gets a +0.10 bonus over a single-evidence raw memory.
59OBSERVATION_PROOF_WEIGHT = 0.025
60OBSERVATION_PROOF_CAP = 4 # clamp at 4 additional proofs
62# Characters allowed inside proper names (apostrophes and hyphens).
63# Straight apostrophe, left/right single quotation marks, and hyphen.
64NAME_CONNECTOR_CHARS = ("'", "\u2018", "\u2019", "-")
67def _tokenize_terms(text: str) -> list[str]:
68 """Tokenize text consistently for keyword and item matching.
70 - Split on whitespace
71 - Strip leading/trailing punctuation
72 - Lowercase
73 - Drop empty tokens
74 """
75 return [t for t in (w.strip(punctuation).lower() for w in text.split()) if t]
78def _content_terms(text: str) -> set[str]:
79 return {t for t in _tokenize_terms(text) if t not in COMMON_QUESTION_WORDS and len(t) > 2}
82def _is_name_token(token: str) -> bool:
83 """Check if a token looks like a proper name, allowing apostrophes and hyphens.
85 Accepts: "Alice", "O'Brien", "Mary-Ann", "Jean-Paul"
86 Rejects: "--", "'hello", "123", ""
87 """
88 if not token:
89 return False
90 # Must start and end with a letter
91 if not token[0].isalpha() or not token[-1].isalpha():
92 return False
93 # Interior characters must be letters, apostrophes, or hyphens
94 for ch in token:
95 if not (ch.isalpha() or ch in NAME_CONNECTOR_CHARS):
96 return False
97 return True
100def basic_rerank(
101 items: list[ScoredItem],
102 query: str,
103 *,
104 mip_rerank: RerankSpec | None = None,
105) -> list[ScoredItem]:
106 """Rerank items using keyword overlap and proper-noun boosting.
108 Adds bonuses to items whose text contains:
109 - General query terms (``keyword_weight`` per matching term)
110 - Proper nouns / names from the query (``proper_noun_weight`` per match)
112 Defaults come from module constants (``KEYWORD_OVERLAP_WEIGHT`` /
113 ``PROPER_NOUN_WEIGHT``); a MIP ``RerankSpec`` from the active routing
114 decision can override either weight on a per-call basis without mutating
115 module state. ``None`` fields fall through to the default.
117 This is a heuristic — production systems should use cross-encoders.
118 """
119 if not items or not query:
120 return items
122 keyword_weight = (
123 mip_rerank.keyword_weight
124 if mip_rerank is not None and mip_rerank.keyword_weight is not None
125 else KEYWORD_OVERLAP_WEIGHT
126 )
127 proper_noun_weight = (
128 mip_rerank.proper_noun_weight
129 if mip_rerank is not None and mip_rerank.proper_noun_weight is not None
130 else PROPER_NOUN_WEIGHT
131 )
133 # Tokenize query once; filter common question words for overlap scoring
134 query_terms = {t for t in _tokenize_terms(query) if t not in COMMON_QUESTION_WORDS}
136 # Detect proper nouns from all words.
137 # Matches: Title Case ("Alice"), ALL CAPS ("USA"), and lowercase names that
138 # appear as query terms but aren't common words (caught by _is_name_token).
139 proper_nouns: set[str] = set()
140 for w in query.split():
141 cleaned = w.strip(punctuation)
142 if not cleaned or cleaned.lower() in COMMON_QUESTION_WORDS:
143 continue
144 is_proper = (
145 cleaned.istitle() # "Alice"
146 or (cleaned.isupper() and len(cleaned) >= 2) # "USA", "AI"
147 )
148 if is_proper and _is_name_token(cleaned):
149 proper_nouns.add(cleaned.lower())
151 # Pre-compute tokenized terms for all items to avoid repeated work.
152 item_terms_by_item = [(item, set(_tokenize_terms(item.text))) for item in items]
154 return sorted(
155 (
156 # ``replace`` rather than a field-by-field rebuild: these
157 # transforms change only the score, and re-enumerating fields is
158 # how ``occurred_at`` was silently dropped here while
159 # ``retained_at`` survived — every recalled memory came back
160 # undated, which is unanswerable for temporal questions and wrong
161 # on the AML ``created_at`` contract. Enumerating fields means the
162 # next field added to ScoredItem gets dropped too.
163 replace(
164 item,
165 score=item.score
166 + len(query_terms & item_terms) * keyword_weight
167 + len(proper_nouns & item_terms) * proper_noun_weight
168 + _observation_proof_boost(item),
169 )
170 for item, item_terms in item_terms_by_item
171 ),
172 key=lambda x: x.score,
173 reverse=True,
174 )
177def cross_encoder_like_rerank(
178 items: list[ScoredItem],
179 query: str,
180) -> list[ScoredItem]:
181 """Final precision rerank using query-item interaction features.
183 This is a deterministic local stand-in for a cross-encoder: it scores each
184 query/memory pair jointly, then applies entity/person and memory-layer
185 signals. It is intentionally cheap enough to run before ``reflect()``
186 synthesis, where precision matters more than broad candidate coverage.
187 """
188 if not items or not query:
189 return items
191 query_terms = _content_terms(query)
192 query_names = _proper_names(query)
193 query_bigrams = _bigrams(_tokenize_terms(query))
195 scored: list[ScoredItem] = []
196 for item in items:
197 item_terms = _content_terms(item.text)
198 item_names = _candidate_names(item)
199 overlap = len(query_terms & item_terms) / max(len(query_terms), 1)
200 phrase_hits = len(query_bigrams & _bigrams(_tokenize_terms(item.text)))
202 score = item.score
203 score += overlap * QUERY_INTERACTION_WEIGHT
204 score += min(phrase_hits, 3) * QUERY_PHRASE_WEIGHT
205 score += _layer_boost(item)
207 if query_names:
208 if query_names & item_names:
209 score += QUERY_NAME_MATCH_WEIGHT
210 elif item_names:
211 score -= QUERY_NAME_MISS_PENALTY
213 scored.append(
214 replace(item, score=max(score, 0.0))
215 )
217 return sorted(scored, key=lambda x: x.score, reverse=True)
220async def llm_pairwise_rerank(
221 items: list[ScoredItem],
222 query: str,
223 llm_provider: LLMProvider,
224 *,
225 top_n: int = 30,
226 keep_n: int | None = None,
227) -> list[ScoredItem]:
228 """Use the configured LLM as a lightweight listwise reranker.
230 The prompt asks for candidate IDs in descending relevance order. If the LLM
231 response is missing or malformed, the deterministic local reranker is used.
232 """
233 if not items or not query:
234 return items
235 candidates = cross_encoder_like_rerank(items[:top_n], query)
236 remainder = items[top_n:]
237 keep = keep_n or len(candidates)
239 prompt_lines = [
240 "Rank the memory candidates by how directly they answer the query.",
241 "Penalize wrong-person or wrong-premise candidates.",
242 'Return JSON only: {"ranked_ids": ["id1", "id2"]}.',
243 "",
244 f"Query: {query}",
245 "",
246 "Candidates:",
247 ]
248 for item in candidates:
249 prompt_lines.append(f"- id={item.id} score={item.score:.4f}: {item.text[:500]}")
251 try:
252 completion = await llm_provider.complete(
253 [
254 Message(role="system", content="You are a strict memory reranker."),
255 Message(role="user", content="\n".join(prompt_lines)),
256 ],
257 max_tokens=512,
258 temperature=0.0,
259 )
260 ranked_ids = _parse_ranked_ids(completion.text)
261 except Exception:
262 ranked_ids = []
264 if not ranked_ids:
265 return apply_context_diversity(candidates, query)[:keep] + remainder
267 by_id = {item.id: item for item in candidates}
268 ordered: list[ScoredItem] = []
269 seen: set[str] = set()
270 for item_id in ranked_ids:
271 item = by_id.get(item_id)
272 if item is not None and item_id not in seen:
273 ordered.append(item)
274 seen.add(item_id)
275 ordered.extend(item for item in candidates if item.id not in seen)
276 return apply_context_diversity(ordered, query)[:keep] + remainder
279def apply_context_diversity(items: list[ScoredItem], query: str) -> list[ScoredItem]:
280 """Softly penalize repeated sessions unless the query asks for aggregation."""
281 if not items:
282 return items
283 if _is_aggregate_query(query):
284 return items
286 session_counts: dict[str, int] = {}
287 diversified: list[ScoredItem] = []
288 for item in items:
289 session = str((item.metadata or {}).get("session_id") or "")
290 penalty = 0.0
291 if session:
292 seen = session_counts.get(session, 0)
293 penalty = seen * SESSION_DIVERSITY_PENALTY
294 session_counts[session] = seen + 1
295 diversified.append(
296 replace(item, score=max(item.score - penalty, 0.0))
297 )
298 return sorted(diversified, key=lambda item: item.score, reverse=True)
301def _proper_names(text: str) -> set[str]:
302 names: set[str] = set()
303 for word in text.split():
304 cleaned = word.strip(punctuation)
305 if cleaned and cleaned.istitle() and _is_name_token(cleaned):
306 names.add(cleaned.lower())
307 return names
310def _candidate_names(item: ScoredItem) -> set[str]:
311 names = _proper_names(item.text)
312 metadata = item.metadata or {}
313 for key in ("locomo_speakers", "locomo_persons", "speakers", "person"):
314 raw = metadata.get(key)
315 if not raw:
316 continue
317 for part in str(raw).replace("|", ",").split(","):
318 cleaned = part.strip().lower()
319 if cleaned:
320 names.add(cleaned)
321 return names
324def _parse_ranked_ids(text: str) -> list[str]:
325 raw = text.strip()
326 if raw.startswith("```"):
327 raw = "\n".join(line for line in raw.splitlines() if not line.strip().startswith("```"))
328 start = raw.find("{")
329 end = raw.rfind("}")
330 if start >= 0 and end > start:
331 raw = raw[start : end + 1]
332 try:
333 parsed = json.loads(raw)
334 except json.JSONDecodeError:
335 return []
336 ids = parsed.get("ranked_ids") if isinstance(parsed, dict) else None
337 if not isinstance(ids, list):
338 return []
339 return [str(item) for item in ids if item]
342def _is_aggregate_query(query: str) -> bool:
343 query_l = (query or "").lower()
344 return any(term in query_l for term in ("how many", "what are", "list", "all ", "which "))
347def _bigrams(tokens: list[str]) -> set[tuple[str, str]]:
348 return set(zip(tokens, tokens[1:], strict=False))
351def _layer_boost(item: ScoredItem) -> float:
352 if item.fact_type == "wiki" or item.memory_layer == "compiled":
353 return COMPILED_LAYER_WEIGHT
354 if item.fact_type == "observation" or item.memory_layer == "observation":
355 return OBSERVATION_LAYER_WEIGHT + _observation_proof_boost(item)
356 return 0.0
359def _observation_proof_boost(item: ScoredItem) -> float:
360 """Additive boost for observation items proportional to their proof count.
362 A single-evidence observation (``_obs_proof_count=1``) gets +0.0.
363 Each additional corroborating memory adds ``OBSERVATION_PROOF_WEIGHT``,
364 capped at ``OBSERVATION_PROOF_CAP`` extra proofs (+0.10 total).
365 Raw memories (no ``_obs_proof_count``) are unaffected.
366 """
367 if item.fact_type != "observation" or not item.metadata:
368 return 0.0
369 proof = item.metadata.get("_obs_proof_count", 1)
370 try:
371 extra = max(0, int(proof) - 1)
372 except (TypeError, ValueError):
373 return 0.0
374 return min(extra, OBSERVATION_PROOF_CAP) * OBSERVATION_PROOF_WEIGHT