RAG app upload

.env

@@ -0,0 +1,2 @@
+OPENAI_API_KEY="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"
+GROQ_API_KEY="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"

Dockerfile

@@ -0,0 +1,10 @@
+FROM python:3.10-slim
+
+WORKDIR /app
+COPY ./app /app
+
+RUN pip install --upgrade pip
+RUN pip install -r requirements.txt
+
+EXPOSE 7860
+CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860"]

main.py

@@ -0,0 +1,81 @@
+from rag_pipeline import process_file, answer_query
+from pydantic import BaseModel
+
+class QueryRequest(BaseModel):
+    file_id: str
+    question: str
+    page: int
+    explainLike5: bool = False
+
+from fastapi import FastAPI, UploadFile, File
+from fastapi.middleware.cors import CORSMiddleware
+from dotenv import load_dotenv
+from fastapi import Body
+import uuid
+import os
+from langchain_chroma import Chroma
+from langchain_huggingface import HuggingFaceEmbeddings
+import re
+
+load_dotenv()
+
+CHROMA_DIR = "./chroma_db"
+embedding_model = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
+
+
+app = FastAPI()
+BASE_DIR = "files"
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+file_store = {}
+@app.get("/test")
+async def test():
+    return {"message": "hello world!"}
+
+@app.post("/upload")
+async def upload(file: UploadFile = File(...)):
+    content = await file.read()
+    file_id = str(uuid.uuid4())
+    safe_filename = file.filename.replace(" ", "_")
+    full_filename = f"{file_id}_{safe_filename}"
+    save_path = os.path.join(BASE_DIR, full_filename)
+    os.makedirs(BASE_DIR, exist_ok=True)
+    with open(save_path, "wb") as f:
+        f.write(content)
+    retriever = process_file(content, safe_filename, file_id)
+    file_store[file_id] = retriever
+    return {"message": "File processed", "file_id": file_id}
+
+@app.post("/query")
+async def query_endpoint(request = Body(...)):
+    file_id = request.get("file_id")
+    question = request.get("question")
+    selected_text = request.get("selectedText")
+    explain_like_5 = request.get("explainLike5", False)
+
+    if not file_id or not question:
+        return {"error": "Missing file_id or question"}
+
+    retriever_path = f"{CHROMA_DIR}/{file_id}"
+    # Load retriever from disk
+    if not os.path.exists(retriever_path):
+        return {"error": "Vectorstore for this file_id not found."}
+
+    vectorstore = Chroma(
+        embedding_function=embedding_model,
+        persist_directory=retriever_path
+    )
+    retriever = vectorstore.as_retriever(search_type="similarity", search_kwargs={"k": 4})
+    retrieved_docs = retriever.invoke(selected_text or question)
+    retrieved_context = "\n\n".join(
+        re.sub(r"\s+", " ", doc.page_content.strip()) for doc in retrieved_docs
+    )
+    
+    combined_context = f"User selected this:\n\"{selected_text}\"\n\nRelated parts from the document:\n{retrieved_context}"
+    print("Combined context", combined_context)
+    answer = answer_query(question, combined_context, explain_like_5)
+    return {"answer": answer}

rag_pipeline.py

@@ -0,0 +1,57 @@
+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="all-MiniLM-L6-v2")
+
+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

requirments.txt

@@ -0,0 +1,11 @@
+fastapi
+uvicorn
+python-multipart
+pydantic
+langchain
+langchain-community
+langchain-chroma
+huggingface-hub
+sentence-transformers
+groq
+python-dotenv