rag_pipeline.py

"""
RAG pipeline: loads the FAISS index and builds a ConversationalRetrievalChain
backed by Groq (llama-3.3-70b-versatile). Returns the answer and source chunks.
"""

import os
from pathlib import Path

from dotenv import load_dotenv
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from langchain_classic.chains import ConversationalRetrievalChain
from langchain_classic.memory import ConversationBufferMemory
from langchain_groq import ChatGroq

load_dotenv()

EMBEDDING_MODEL = "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
INDEX_DIR = Path("faiss_index")
GROQ_MODEL = "llama-3.3-70b-versatile"
TOP_K = 4

SYSTEM_PROMPT = (
    "You are a helpful multilingual assistant that supports English, Arabic, and Malay. "
    "Always reply in the same language the user used. "
    "Base your answer strictly on the provided context. "
    "If the context does not contain enough information, say so honestly."
)


def load_vectorstore() -> FAISS:
    if not INDEX_DIR.exists():
        raise FileNotFoundError(
            "FAISS index not found. Run `python ingest.py` first."
        )
    embeddings = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL)
    return FAISS.load_local(
        str(INDEX_DIR), embeddings, allow_dangerous_deserialization=True
    )


class RAGPipeline:
    def __init__(self):
        vectorstore = load_vectorstore()
        retriever = vectorstore.as_retriever(search_kwargs={"k": TOP_K})

        llm = ChatGroq(
            model=GROQ_MODEL,
            api_key=os.environ["GROQ_API_KEY"],
            max_tokens=1024,
        )

        memory = ConversationBufferMemory(
            memory_key="chat_history",
            return_messages=True,
            output_key="answer",
        )

        self.chain = ConversationalRetrievalChain.from_llm(
            llm=llm,
            retriever=retriever,
            memory=memory,
            return_source_documents=True,
            output_key="answer",
            combine_docs_chain_kwargs={"prompt": _build_prompt()},
        )

    def ask(self, question: str) -> dict:
        """
        Returns:
            {
                "answer": str,
                "sources": [{"content": str, "source": str}, ...]
            }
        """
        result = self.chain.invoke({"question": question})
        sources = [
            {
                "content": doc.page_content,
                "source": doc.metadata.get("source", "unknown"),
            }
            for doc in result.get("source_documents", [])
        ]
        return {"answer": result["answer"], "sources": sources}


def _build_prompt():
    from langchain_core.prompts import PromptTemplate

    template = (
        SYSTEM_PROMPT
        + "\n\nContext:\n{context}\n\nQuestion: {question}\n\nAnswer:"
    )
    return PromptTemplate(input_variables=["context", "question"], template=template)