rag/search.py

import os
import faiss
import pickle
from sentence_transformers import SentenceTransformer

BASE_DIR = os.path.dirname(os.path.dirname(__file__))
VECTOR_DIR = os.path.join(BASE_DIR, "rag", "vectorstore")

index = faiss.read_index(os.path.join(VECTOR_DIR, "index.faiss"))

with open(os.path.join(VECTOR_DIR, "meta.pkl"), "rb") as f:
    documents = pickle.load(f)

model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")


def search_context(query, top_k=3):
    q_emb = model.encode([query])
    D, I = index.search(q_emb, top_k)
    return "\n".join([documents[i] for i in I[0]])