Coverage for src/asketmc_bot/rag_filter.py: 82%
140 statements
« prev ^ index » next coverage.py v7.5.1, created at 2026-04-08 17:14 +0000
1from __future__ import annotations
3import asyncio
4import hashlib
5import logging
6from collections import OrderedDict
7from typing import FrozenSet, List, Optional, Sequence, Set, Tuple
9from llama_index.core.schema import NodeWithScore
11from asketmc_bot import config as cfg
13__all__ = ["get_filtered_nodes", "build_context", "purge_filter_cache"]
15rag_log = logging.getLogger("asketmc.rag")
18def _cfg(name: str, default):
19 return getattr(cfg, name, default)
22TOP_K: int = int(_cfg("TOP_K", 8))
23FILTER_ALPHA: float = float(_cfg("FILTER_ALPHA", 0.5))
24SCORE_RELATIVE_THRESHOLD: float = float(_cfg("SCORE_RELATIVE_THRESHOLD", 0.7))
25LEMMA_MATCH_RATIO: float = float(_cfg("LEMMA_MATCH_RATIO", 0.10))
26CONTEXT_BONUS_PER_INTERSECTION: float = float(_cfg("CONTEXT_BONUS_PER_INTERSECTION", 0.1))
27FILTER_CACHE_MAX_SIZE: int = int(_cfg("FILTER_CACHE_MAX_SIZE", 256))
29# Use a stable per-process salt to reduce cross-run cache-key collisions if ids repeat.
30_CACHE_SALT = str(id(object()))
33def _safe_score(v: Optional[float]) -> float:
34 try:
35 return float(v) if v is not None else 0.0
36 except Exception:
37 return 0.0
40def _node_id(n: NodeWithScore) -> str:
41 node = n.node
42 return (
43 getattr(node, "node_id", None)
44 or getattr(node, "id_", None)
45 or node.metadata.get("node_id")
46 or node.metadata.get("doc_id")
47 or node.metadata.get("file_name")
48 or f"node@{id(node)}"
49 )
52def _metadata_lemmas(n: NodeWithScore) -> FrozenSet[str]:
53 lem = n.node.metadata.get("lemmas", [])
54 try:
55 return frozenset(str(x) for x in lem)
56 except Exception:
57 return frozenset()
60def _cache_key(qlem: FrozenSet[str], raw_nodes: Sequence[NodeWithScore]) -> str:
61 # Bound the key size and ensure deterministic ordering.
62 q_part = " ".join(sorted(qlem))
63 fp_items: List[str] = []
64 for n in raw_nodes[:64]:
65 fp_items.append(f"{_node_id(n)}:{_safe_score(n.score):.6f}")
66 fp = "|".join(fp_items)
67 material = f"{_CACHE_SALT}||{q_part}||{fp}"
68 return hashlib.sha256(material.encode("utf-8")).hexdigest()
71def _content_from_node(n: NodeWithScore) -> str:
72 node = n.node
73 try:
74 if hasattr(node, "get_content"): 74 ↛ 76line 74 didn't jump to line 76, because the condition on line 74 was always true
75 return node.get_content(metadata_mode="none") or ""
76 return getattr(node, "text", "") or ""
77 except Exception:
78 return ""
81_FILTER_CACHE: "OrderedDict[str, Tuple[NodeWithScore, ...]]" = OrderedDict()
82_FILTER_CACHE_LOCK = asyncio.Lock()
85def _cache_get(key: str) -> Optional[List[NodeWithScore]]:
86 tpl = _FILTER_CACHE.get(key)
87 if tpl is None:
88 return None
89 _FILTER_CACHE.move_to_end(key, last=True)
90 return list(tpl)
93def _cache_put(key: str, nodes: List[NodeWithScore]) -> None:
94 _FILTER_CACHE[key] = tuple(nodes)
95 _FILTER_CACHE.move_to_end(key, last=True)
96 while len(_FILTER_CACHE) > FILTER_CACHE_MAX_SIZE: 96 ↛ 97line 96 didn't jump to line 97, because the condition on line 96 was never true
97 _FILTER_CACHE.popitem(last=False)
100def purge_filter_cache() -> None:
101 _FILTER_CACHE.clear()
104async def _filter_nodes(raw_nodes: List[NodeWithScore], qlem: FrozenSet[str]) -> List[NodeWithScore]:
105 if not raw_nodes:
106 return []
108 scores = [_safe_score(n.score) for n in raw_nodes]
109 max_score = max(scores) if scores else 1.0
110 if max_score <= 0.0: 110 ↛ 111line 110 didn't jump to line 111, because the condition on line 110 was never true
111 max_score = 1.0
113 alpha = FILTER_ALPHA
114 qlen = len(qlem) or 1
116 strict: List[Tuple[NodeWithScore, float]] = []
117 for n in raw_nodes:
118 s = _safe_score(n.score)
119 lemmas = _metadata_lemmas(n)
120 inter = len(qlem & lemmas)
121 rel_score = s / max_score
122 ratio = inter / qlen
123 weight = s + alpha * ratio
124 if (rel_score >= SCORE_RELATIVE_THRESHOLD) or (ratio >= LEMMA_MATCH_RATIO): 124 ↛ 117line 124 didn't jump to line 117, because the condition on line 124 was always true
125 strict.append((n, weight))
127 if strict: 127 ↛ 132line 127 didn't jump to line 132, because the condition on line 127 was always true
128 strict.sort(key=lambda x: x[1], reverse=True)
129 return [n for n, _ in strict[:TOP_K]]
131 # Fallback: require at least 1 overlapping lemma
132 fallback: List[Tuple[NodeWithScore, float]] = []
133 for n in raw_nodes:
134 s = _safe_score(n.score)
135 lemmas = _metadata_lemmas(n)
136 inter = len(qlem & lemmas)
137 if inter <= 0:
138 continue
139 ratio = inter / qlen
140 weight = s + alpha * ratio
141 fallback.append((n, weight))
143 fallback.sort(key=lambda x: x[1], reverse=True)
144 return [n for n, _ in fallback[:TOP_K]]
147async def get_filtered_nodes(raw_nodes: List[NodeWithScore], qlem: FrozenSet[str]) -> List[NodeWithScore]:
148 key = _cache_key(qlem, raw_nodes)
149 async with _FILTER_CACHE_LOCK:
150 cached = _cache_get(key)
151 if cached is not None:
152 return cached
154 nodes = await _filter_nodes(raw_nodes, qlem)
156 async with _FILTER_CACHE_LOCK:
157 _cache_put(key, nodes)
159 return list(nodes)
162def build_context(nodes: List[NodeWithScore], qlem: FrozenSet[str], char_limit: int) -> str:
163 beta = CONTEXT_BONUS_PER_INTERSECTION
164 scored: List[Tuple[NodeWithScore, float]] = []
165 for n in nodes:
166 s = _safe_score(n.score)
167 inter = len(qlem & _metadata_lemmas(n))
168 weight = s + beta * float(inter)
169 scored.append((n, weight))
171 scored.sort(key=lambda x: x[1], reverse=True)
173 parts: List[str] = []
174 seen_hashes: Set[str] = set()
175 total = 0
177 sep = "\n---\n"
178 sep_len = len(sep)
179 limit = max(0, int(char_limit))
181 for n, _ in scored:
182 txt = (_content_from_node(n) or "").strip()
183 if not txt:
184 continue
186 h = hashlib.sha256(txt.encode("utf-8")).hexdigest()
187 if h in seen_hashes:
188 continue
190 prospective = total + (sep_len if parts else 0) + len(txt)
191 if prospective > limit:
192 break
194 seen_hashes.add(h)
195 if parts:
196 parts.append(sep)
197 total += sep_len
198 parts.append(txt)
199 total += len(txt)
201 return "".join(parts)