""" step3_retrieval.py ================== Enterprise-grade retrieval engine for the KCC RAG chatbot. Architecture ------------ 1. Load FAISS index + metadata once at startup (cached in-process) 2. encode_query() → embed user query → float32 unit vector 3. search() → FAISS top-K → fetch matching metadata rows 4. format_context()→ assemble retrieved Q&A pairs as LLM context Design decisions ---------------- - Metadata loaded as PyArrow Table (memory-mapped, zero-copy slicing) - FAISS IVF-PQ with nprobe=64 gives ~95% recall@5 - Query embedding normalised → inner-product == cosine similarity - Duplicate-answer dedup (same KccAns text → keep highest-scoring copy) - Language-agnostic: model handles Hindi, Telugu, Kannada, Marathi, English Usage (standalone test) ----------------------- python step3_retrieval.py "What fertilizer for wheat in Uttar Pradesh?" python step3_retrieval.py --top-k 10 "kharif crops pest control" """ import argparse import re import sqlite3 import sys import time from pathlib import Path from typing import List, Optional from dataclasses import dataclass, field import faiss import numpy as np import pandas as pd import pyarrow as pa import pyarrow.parquet as pq from sentence_transformers import SentenceTransformer # ── ICAR priority retriever ─────────────────────────────────────────────────── try: from icar_retriever import get_icar_retriever as _get_icar_retriever _ICAR_RETRIEVER = _get_icar_retriever() _ICAR_AVAILABLE = True except Exception as _e: print(f"[ICAR] Retriever not available: {_e}") _ICAR_RETRIEVER = None _ICAR_AVAILABLE = False sys.path.insert(0, str(Path(__file__).parent)) import config # ── public data class ───────────────────────────────────────────────────────── @dataclass class RetrievedDoc: """One retrieved document returned by the retrieval engine.""" rank: int score: float # cosine similarity (0–1, higher is better) rerank_score: float = 0.0 # cross-encoder score (higher = more relevant) query: str = "" # original farmer query from the dataset answer: str = "" # KCC expert answer state: str = "" crop: str = "" year: int = 0 @property def score_pct(self) -> str: return f"{self.score * 100:.1f}%" # ── retrieval engine ────────────────────────────────────────────────────────── class KCCRetriever: """ Stateful retrieval engine. Load once, query many times. Thread-safety: encode_query and search are read-only after __init__, so the same instance can safely be called from multiple threads (e.g., Streamlit's multi-user sessions). """ def __init__( self, index_path: Path = config.FAISS_INDEX_FILE, metadata_path: Path = config.METADATA_FILE, model_name: str = config.EMBEDDING_MODEL, top_k: int = config.TOP_K, nprobe: int = 64, device: Optional[str] = None, ): self._top_k = top_k self._loaded = False # ── paths ────────────────────────────────────────────────────────────── if not index_path.exists(): raise FileNotFoundError( f"FAISS index not found: {index_path}\n" "Run step2_embeddings.py first." ) if not metadata_path.exists(): raise FileNotFoundError( f"Metadata not found: {metadata_path}\n" "Run step2_embeddings.py first." ) # ── FAISS ────────────────────────────────────────────────────────────── print(f"[KCCRetriever] Loading FAISS index …", flush=True) t0 = time.perf_counter() self._index = faiss.read_index(str(index_path)) self._index.nprobe = nprobe faiss_load_s = time.perf_counter() - t0 print( f"[KCCRetriever] FAISS loaded: {self._index.ntotal:,} vectors " f"({faiss_load_s:.1f}s)", flush=True, ) # ── Metadata (PyArrow memory-mapped) ─────────────────────────────────── print(f"[KCCRetriever] Loading metadata …", flush=True) t0 = time.perf_counter() # True zero-copy memory map: OS pages data from disk on demand. # This avoids loading 1.1 GB into RAM all at once. self._meta: pa.Table = pq.read_table( metadata_path, columns=["QueryText", "KccAns", "StateName", "Crop", "year"], memory_map=True, ) # Cache column arrays for fast indexed access during search self._col_query = self._meta.column("QueryText") self._col_ans = self._meta.column("KccAns") self._col_state = self._meta.column("StateName") self._col_crop = self._meta.column("Crop") self._col_year = self._meta.column("year") meta_load_s = time.perf_counter() - t0 print( f"[KCCRetriever] Metadata loaded: {self._meta.num_rows:,} rows " f"({meta_load_s:.1f}s)", flush=True, ) # Sanity check assert self._index.ntotal == self._meta.num_rows, ( f"MISMATCH: FAISS {self._index.ntotal:,} vs metadata {self._meta.num_rows:,}. " "Re-run step2_embeddings.py." ) # ── Embedding model ──────────────────────────────────────────────────── import torch if device is None: device = "cuda" if torch.cuda.is_available() else "cpu" print(f"[KCCRetriever] Loading embedding model ({device}) …", flush=True) t0 = time.perf_counter() self._model = SentenceTransformer(model_name, device=device) model_load_s = time.perf_counter() - t0 # Warmup self._model.encode(["warmup"], convert_to_numpy=True, normalize_embeddings=True) print( f"[KCCRetriever] Model ready ({model_load_s:.1f}s)", flush=True, ) # ── Cross-encoder reranker (optional, ~80 MB) ───────────────────────── self._reranker = None if config.RERANKER_MODEL: try: from sentence_transformers import CrossEncoder as _CE print(f"[KCCRetriever] Loading reranker: {config.RERANKER_MODEL} …", flush=True) self._reranker = _CE(config.RERANKER_MODEL) print("[KCCRetriever] Reranker ready.", flush=True) except Exception as e: print(f"[KCCRetriever] Reranker unavailable: {e}", flush=True) # ── BM25 (SQLite FTS5) — disk-based, ~50 MB RAM at query time ───────── self._bm25_conn: Optional[sqlite3.Connection] = None bm25_path = getattr(config, "BM25_INDEX_FILE", None) if bm25_path and Path(bm25_path).exists(): try: self._bm25_conn = sqlite3.connect( f"file:{bm25_path}?mode=ro", uri=True, check_same_thread=False, ) self._bm25_conn.execute("PRAGMA cache_size = -32768") print(f"[KCCRetriever] BM25 index loaded (hybrid search enabled).", flush=True) except Exception as e: print(f"[KCCRetriever] BM25 unavailable: {e}", flush=True) self._bm25_conn = None else: print( "[KCCRetriever] BM25 index not found — " "run step2b_bm25_index.py to enable hybrid search.", flush=True, ) self._loaded = True print("[KCCRetriever] Ready.", flush=True) # ── public API ───────────────────────────────────────────────────────────── @property def index_size(self) -> int: return self._index.ntotal def encode_query(self, query: str) -> np.ndarray: """ Embed a single user query → float32 unit vector (1 × DIM). Normalised so inner-product == cosine similarity in FAISS. """ return self._model.encode( [query.strip()], convert_to_numpy=True, normalize_embeddings=True, show_progress_bar=False, ).astype(np.float32) def search( self, query: str, top_k: Optional[int] = None, deduplicate: bool = True, min_score: float = 0.0, crop_filter: Optional[str] = None, state: str = "", district: str = "", ) -> List[RetrievedDoc]: """ Retrieve the top-K most similar Q&A pairs for a user query. Parameters ---------- query : User's natural-language question (any supported language) top_k : Override default TOP_K from config deduplicate : Drop results with identical KccAns text (keep highest score) min_score : Discard results below this cosine similarity threshold crop_filter : If set, prefer results whose Crop column contains this string (case-insensitive). Falls back to unfiltered results when not enough crop-matched hits are found. state : Optional farmer state (e.g. "Madhya Pradesh"). Boosts retrieved docs from the same state by +0.12. district : Optional farmer district (e.g. "Barwani"). Boosts retrieved docs from the same district by +0.25. Returns ------- List of RetrievedDoc, ranked by similarity (best first). """ k = top_k or self._top_k # Embed query t0 = time.perf_counter() q_vec = self.encode_query(query) embed_ms = (time.perf_counter() - t0) * 1000 # Retrieve more candidates: headroom for quality filter + reranking rerank_n = config.RERANKER_TOP_N if self._reranker else 0 search_k = max(k * 3, rerank_n + k) if deduplicate else k if crop_filter: search_k = max(search_k, k * 20) t1 = time.perf_counter() scores, indices = self._index.search(q_vec, search_k) faiss_ms = (time.perf_counter() - t1) * 1000 # Fetch metadata rows — pre-compute normalized state/district for boosting _state_lower = state.strip().lower() if state else "" _district_lower = district.strip().lower() if district else "" valid_hits = [ (float(s), int(i)) for s, i in zip(scores[0], indices[0]) if i >= 0 and float(s) >= min_score ] crop_lower = crop_filter.lower() if crop_filter else None docs: List[RetrievedDoc] = [] fallback_docs: List[RetrievedDoc] = [] # used if crop filter gives too few seen_answers: set = set() for score, idx in valid_hits: if len(docs) >= k and len(fallback_docs) >= k: break answer_raw = self._col_ans[idx].as_py() answer = str(answer_raw).strip() if answer_raw is not None else "" # Quality filter: skip very short / vague KCC answers if len(answer) < config.MIN_ANSWER_CHARS: continue # Quality filter: skip generic/unhelpful KCC operator responses _GENERIC_SKIP = [ "contact your kvk", "contact nearest kvk", "visit kvk", "kvk se sampark karen", "nazdiki kvk", "apne kshetra ke", "krishi vigyan kendra se", "district agriculture office", "krishi vibhag se sampark", "contact agriculture department", "not available in our system", "uplabdh nahi hai", "jaankari nahi hai", "information not available", "please contact", "please visit", "kindly contact", "apne nazdiki bank", "bank se sampark karen", ] ans_lower = answer.lower() if any(ph in ans_lower for ph in _GENERIC_SKIP): continue if deduplicate and answer in seen_answers: continue seen_answers.add(answer) year_raw = self._col_year[idx].as_py() crop_val = str(self._col_crop[idx].as_py() or "").strip() # ── UPGRADE 1+4: Apply location + recency score boosts ──────────── doc_state = str(self._col_state[idx].as_py() or "").strip() doc_year = int(year_raw) if year_raw is not None else 0 boosted_score = score # Recency boost (UPGRADE 4) if doc_year >= 2023: boosted_score += 0.15 elif doc_year >= 2021: boosted_score += 0.08 elif doc_year >= 2018: boosted_score += 0.03 # Location boost (UPGRADE 1) if _state_lower and doc_state.lower() == _state_lower: boosted_score += 0.12 # District boost: KCC records don't have district col directly, # check if district name appears in the query text if _district_lower and _district_lower in str(self._col_query[idx].as_py() or "").lower(): boosted_score += 0.25 doc = RetrievedDoc( rank = 0, # assigned below score = boosted_score, query = str(self._col_query[idx].as_py() or "").strip(), answer = answer, state = doc_state, crop = crop_val, year = doc_year, ) if crop_lower: if crop_lower in crop_val.lower() and len(docs) < k: docs.append(doc) elif crop_lower not in crop_val.lower() and len(fallback_docs) < k: fallback_docs.append(doc) else: if len(docs) < k: docs.append(doc) # Fill with fallback results if crop filter gave too few if crop_filter and len(docs) < k: docs.extend(fallback_docs[:k - len(docs)]) # ── BM25 hybrid merge ───────────────────────────────────────────────── # Run keyword search and merge with FAISS results. # BM25 catches exact chemical names / variety names / specific doses # that the embedding model may distribute poorly in vector space. if self._bm25_conn is not None: bm25_docs = self.bm25_search( query, top_n = max(k * 2, 20), crop_filter = crop_filter, ) # Merge: add BM25 results not already in FAISS results existing_answers = {d.answer for d in docs} for bd in bm25_docs: if bd.answer not in existing_answers and len(docs) < k * 2: existing_answers.add(bd.answer) docs.append(bd) # ── Cross-encoder reranking ──────────────────────────────────────────── # Re-score top candidates using (query, answer) pair relevance. # Takes ~20ms for 20 pairs on CPU — worth it for quality improvement. if self._reranker and len(docs) > 1: try: pairs = [(query, doc.answer) for doc in docs] scores = self._reranker.predict(pairs) for doc, rs in zip(docs, scores): doc.rerank_score = float(rs) # Re-sort by reranker score (higher = more relevant) docs.sort(key=lambda d: d.rerank_score, reverse=True) docs = docs[:k] # keep top-K after reranking except Exception: pass # fall back to cosine similarity order # Assign final ranks for i, doc in enumerate(docs): doc.rank = i + 1 return docs @staticmethod def _sanitize_fts5(query: str) -> str: """ Sanitize a query string for safe use in SQLite FTS5 MATCH expressions. Removes FTS5 special operators; keeps alphanumeric + Unicode (Hindi/Indian scripts). """ # Strip FTS5 operators that would cause syntax errors cleaned = re.sub(r'["\(\)\*\:\^~]+', ' ', query) # Collapse whitespace cleaned = ' '.join(cleaned.split()) return cleaned def bm25_search( self, query: str, top_n: int = 20, crop_filter: Optional[str] = None, ) -> List[RetrievedDoc]: """ BM25 keyword search via SQLite FTS5. Returns up to top_n results ranked by BM25 relevance. Falls back gracefully to empty list if BM25 index is unavailable or the query contains no searchable tokens. BM25 complements FAISS by catching: - Exact chemical/variety names (e.g. "Chlorpyrifos", "PB-1121") - Specific dose queries ("2ml/liter") - Uncommon tokens the embedding model distributes poorly """ if self._bm25_conn is None: return [] safe_q = self._sanitize_fts5(query) if not safe_q: return [] # Build crop filter clause if requested crop_clause = "" params: list = [safe_q] if crop_filter: crop_clause = " AND crop = ?" params.append(crop_filter) sql = f""" SELECT query_text, answer, state, crop, year, rank FROM kcc_bm25 WHERE kcc_bm25 MATCH ?{crop_clause} ORDER BY rank LIMIT ? """ params.append(top_n) try: cur = self._bm25_conn.cursor() rows = cur.execute(sql, params).fetchall() except sqlite3.OperationalError: # Invalid query syntax — return empty return [] if not rows: # No crop-filtered results → retry without crop filter if crop_filter: return self.bm25_search(query, top_n=top_n, crop_filter=None) return [] # Convert BM25 rank (negative float, closer to 0 = better) to a # normalised score in [0, 1]. We cap raw rank at -100 to avoid # extremely bad matches distorting the scale. raw_ranks = [r[5] for r in rows] best = min(raw_ranks) # most negative = best worst = max(raw_ranks, default=best) span = (worst - best) or 1.0 # avoid div-by-zero docs: List[RetrievedDoc] = [] seen_answers: set = set() for i, (qt, ans, state, crop, year, rank) in enumerate(rows): ans = str(ans or "").strip() if len(ans) < config.MIN_ANSWER_CHARS: continue if ans in seen_answers: continue seen_answers.add(ans) # Normalise: best rank → 1.0, worst → 0.05 (never 0 so it can merge) score = 0.05 + 0.95 * (1.0 - (rank - best) / span) docs.append(RetrievedDoc( rank = i + 1, score = round(score, 4), query = str(qt or "").strip(), answer = ans, state = str(state or "").strip(), crop = str(crop or "").strip(), year = int(year) if year else 0, )) return docs def format_context(self, docs: List[RetrievedDoc]) -> str: """ Format retrieved docs as a concise context block for the LLM prompt. Each retrieved Q&A pair is numbered and trimmed to keep the prompt short. """ if not docs: return "No relevant past Q&A found." MAX_ANSWER_CHARS = 600 # truncate very long answers lines = ["RELEVANT PAST FARMER Q&A FROM KCC DATABASE:"] for doc in docs: ans = doc.answer[:MAX_ANSWER_CHARS] if len(doc.answer) > MAX_ANSWER_CHARS: ans += "…" lines.append( f"\n[{doc.rank}] Similarity: {doc.score_pct} | " f"State: {doc.state} | Crop: {doc.crop} | Year: {doc.year}" f"\n Q: {doc.query}" f"\n A: {ans}" ) return "\n".join(lines) # ── KCC Golden Set Retriever (UPGRADE 5) ───────────────────────────────────── class KCCGoldenRetriever: """ Fast lookup for pre-verified golden Q&A pairs from the KCC corpus. Checked BEFORE FAISS — if a golden match is found with high confidence, it is prepended to the context with a [VERIFIED ANSWER] tag. Uses TF-IDF cosine similarity for matching (no GPU required, <5ms per query). """ SIMILARITY_THRESHOLD = 0.85 # minimum cosine sim to use golden answer def __init__(self, golden_path: str): import json as _json from pathlib import Path as _Path self._entries: list = [] self._vectorizer = None self._tfidf_matrix = None gp = _Path(golden_path) if not gp.exists(): print(f"[KCCGoldenRetriever] Golden set not found at {golden_path} — skipping", flush=True) return with open(gp, "r", encoding="utf-8") as f: data = _json.load(f) self._entries = data.get("entries", []) if not self._entries: print("[KCCGoldenRetriever] Empty golden set", flush=True) return # Build TF-IDF index try: from sklearn.feature_extraction.text import TfidfVectorizer from sklearn.metrics.pairwise import cosine_similarity as _cos_sim self._cos_sim = _cos_sim queries = [e["query"] for e in self._entries] self._vectorizer = TfidfVectorizer( ngram_range=(1, 2), max_features=20000, sublinear_tf=True, min_df=1, ) self._tfidf_matrix = self._vectorizer.fit_transform(queries) print(f"[KCCGoldenRetriever] Loaded {len(self._entries)} golden entries", flush=True) except ImportError: print("[KCCGoldenRetriever] sklearn not available — golden retriever disabled", flush=True) self._vectorizer = None def lookup(self, query: str, top_k: int = 3) -> list: """ TF-IDF similarity search against golden set. Returns list of dicts with keys: query, answer, crop, state, year, score Only returns matches above SIMILARITY_THRESHOLD. """ if self._vectorizer is None or self._tfidf_matrix is None or not self._entries: return [] try: q_vec = self._vectorizer.transform([query]) sims = self._cos_sim(q_vec, self._tfidf_matrix).flatten() top_indices = sims.argsort()[::-1][:top_k] results = [] for idx in top_indices: score = float(sims[idx]) if score >= self.SIMILARITY_THRESHOLD: entry = dict(self._entries[idx]) entry["score"] = score results.append(entry) return results except Exception as e: print(f"[KCCGoldenRetriever] lookup error: {e}", flush=True) return [] @property def size(self) -> int: return len(self._entries) # ── module-level golden retriever singleton ─────────────────────────────────── _golden_retriever: Optional["KCCGoldenRetriever"] = None def get_golden_retriever(golden_path: Optional[str] = None) -> "KCCGoldenRetriever": """Return the module-level golden retriever singleton (create on first call).""" global _golden_retriever if _golden_retriever is None: if golden_path is None: from pathlib import Path as _P golden_path = str(_P(__file__).parent / "kcc_golden_set.json") _golden_retriever = KCCGoldenRetriever(golden_path) return _golden_retriever # ── module-level singleton (lazy) ───────────────────────────────────────────── _retriever: Optional[KCCRetriever] = None def get_retriever(**kwargs) -> KCCRetriever: """Return the module-level retriever singleton (create on first call).""" global _retriever if _retriever is None: _retriever = KCCRetriever(**kwargs) return _retriever # ── standalone test ─────────────────────────────────────────────────────────── def _standalone_demo(query: str, top_k: int) -> None: print("=" * 70) print("KCC RAG — Step 3: Retrieval Demo") print("=" * 70) retriever = get_retriever(top_k=top_k) print(f"\nQuery: {query!r}") print(f"Index size: {retriever.index_size:,} vectors\n") t0 = time.perf_counter() docs = retriever.search(query, top_k=top_k) dt = (time.perf_counter() - t0) * 1000 print(f"Retrieved {len(docs)} results in {dt:.1f} ms\n") print(retriever.format_context(docs)) print("=" * 70) if __name__ == "__main__": parser = argparse.ArgumentParser( description="KCC RAG – Step 3: Retrieval engine test" ) parser.add_argument("query", help="Query to search for") parser.add_argument("--top-k", type=int, default=config.TOP_K) args = parser.parse_args() _standalone_demo(args.query, args.top_k)