rag.py

import os
from dotenv import load_dotenv
import shutil
import uuid

from langchain_community.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter

from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.vectorstores import Chroma

from langchain_google_genai import ChatGoogleGenerativeAI

from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser


# Load environment variables
load_dotenv()

# Chroma database directory
DB_DIRECTORY = "chroma_db"


####################################
# Create Vector Store
####################################
def create_vector_store(pdf_path):

    global vector_db

    loader = PyPDFLoader(pdf_path)
    documents = loader.load()

    splitter = RecursiveCharacterTextSplitter(
        chunk_size=1000,
        chunk_overlap=200
    )

    chunks = splitter.split_documents(documents)

    embeddings = HuggingFaceEmbeddings(
        model_name="sentence-transformers/all-MiniLM-L6-v2"
    )

    vector_db = Chroma.from_documents(
        documents=chunks,
        embedding=embeddings,
        collection_name=str(uuid.uuid4())
    )

    return vector_db

    
####################################
# Create RAG Chain
####################################
def get_chain():

    global vector_db

    # Check whether a PDF has been uploaded
    if vector_db is None:
        raise Exception(
            "No vector database found. Please upload a PDF first."
        )

    # Retriever
    retriever = vector_db.as_retriever(
        search_kwargs={"k": 3}
    )

    # Gemini LLM
    llm = ChatGoogleGenerativeAI(
        model="gemini-2.5-flash",
        google_api_key=os.getenv("GOOGLE_API_KEY"),
        temperature=0.2
    )

    # Prompt
    prompt = ChatPromptTemplate.from_template(
        """
You are a helpful AI assistant.

Answer ONLY using the information provided in the context.

If the answer cannot be found in the context, reply exactly:

"I could not find the answer in the uploaded document."

Context:
{context}

Question:
{question}

Answer:
"""
    )

    # Build chain
    chain = (
        {
            "context": retriever,
            "question": lambda x: x
        }
        | prompt
        | llm
        | StrOutputParser()
    )

    return chain
####################################
# Ask Question
####################################
def ask_question(question):

    chain = get_chain()

    response = chain.invoke(question)

    return response