rag_pipeline.py

import tempfile
from langchain_chroma import Chroma
from langchain_community.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
import os
from langchain_huggingface import HuggingFaceEmbeddings
import re
from groq import Groq
from dotenv import load_dotenv
load_dotenv()

client = Groq(api_key=os.environ["GROQ_API_KEY"])
CHROMA_DIR = "./chroma_db"
# Embedding model (free HuggingFace)
embedding_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2", model_kwargs={"trust_remote_code": True})

def process_file(file_bytes, filename, file_id):
    ext = filename.split('.')[-1].lower()
    with tempfile.NamedTemporaryFile(delete=False, suffix=f".{ext}") as tmp:
        tmp.write(file_bytes)
        tmp_path = tmp.name

    loader = PyPDFLoader(tmp_path) if ext == 'pdf' else None
    docs = loader.load()

    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
    chunks = text_splitter.split_documents(docs)

    vectorstore = Chroma.from_documents(
        documents=chunks,
        embedding=embedding_model,
        persist_directory=f"{CHROMA_DIR}/{file_id}"
    )
    retriever = vectorstore.as_retriever(search_type="similarity", search_kwargs={"k": 4})
    os.unlink(tmp_path)
    return retriever


def answer_query(question, context, explain_like_5=False):
    system_prompt = (
        "You are a helpful assistant answering user queries based on provided document chunks.\n"
        "Only use the given context. If the answer is not found, respond with 'I don't know.'"
    )
    if explain_like_5:
        system_prompt += "\nExplain the answer in a simple way, like you're talking to a 5-year-old."

    # Step 2: Send to LLM
    print("context\======================================================================", context)
    print("question =====================================================================", question)
    response = client.chat.completions.create(
        model="llama-3.3-70b-versatile",
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": f"Context:\n{context}\n\nQuestion:\n{question}"}
        ]
    )
    return response.choices[0].message.content