session_rag.py

\
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