\ from typing import List, Tuple from sentence_transformers import SentenceTransformer import numpy as np import faiss class SessionRAG: def __init__(self, model_name="sentence-transformers/all-MiniLM-L6-v2"): self.model = SentenceTransformer(model_name) self.docs: List[Tuple[str, str]] = [] self.index = None self.vecs = None def add_docs(self, items: List[Tuple[str, str]]): self.docs.extend(items) texts = [t for _, t in self.docs] if not texts: self.index = None; self.vecs=None; return embs = self.model.encode(texts, convert_to_numpy=True, normalize_embeddings=True).astype(np.float32) self.vecs = embs self.index = faiss.IndexFlatIP(embs.shape[1]) self.index.add(embs) def retrieve(self, query: str, k: int = 6) -> List[str]: if not self.index or self.vecs is None: return [] q = self.model.encode([query], convert_to_numpy=True, normalize_embeddings=True).astype(np.float32) D, I = self.index.search(q, k) out = [] for idx in I[0]: if 0 <= idx < len(self.docs): out.append(self.docs[idx][1]) return out