Update backend.py

backend.py

@@ -3,141 +3,203 @@ import tempfile
 import gc
 import logging
 from fastapi import FastAPI, UploadFile, File, HTTPException
-from fastapi.middleware.cors import CORSMiddleware  # <-- ADDED
 from pydantic import BaseModel
 import torch
-from dotenv import load_dotenv
+from dotenv import load_dotenv  # Used to load API key from .env file
+
+# ---------------- Groq API ----------------
 from groq import Groq, APIError
+
+# ---------------- LangChain ----------------
 from langchain_community.document_loaders import PyPDFLoader
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_community.vectorstores import Chroma
 
-# ---------------- Setup ----------------
+# --- Configuration & Setup ---
 logging.basicConfig(level=logging.INFO)
+
+# 1. Load environment variables from .env file
 load_dotenv()
-GROQ_API_KEY = os.getenv("GROQ_API_KEY")
+
+# 2. Load API Key from Environment Variable
+GROQ_API_KEY = os.environ.get("GROQ_API_KEY")
 GROQ_MODEL = "llama-3.1-8b-instant"
+
+# 3. Initialize Groq Client
 client = None
-if GROQ_API_KEY:
+if not GROQ_API_KEY:
+    logging.error(
+        "❌ GROQ_API_KEY is not set in the environment or the .env file. The service will run but cannot answer questions.")
+else:
     try:
         client = Groq(api_key=GROQ_API_KEY)
-        logging.info("✅ Groq client initialized")
+        logging.info("✅ Groq client initialized successfully.")
     except Exception as e:
-        logging.error(f"Groq init failed: {e}")
+        logging.error(f"❌ Failed to initialize Groq client: {e}")
+        client = None
 
 app = FastAPI()
 
-# ==================================================#
-#  CORS Middleware (NEW SECTION)
-# ==================================================
-origins = [
-    "*",  # Allow all origins for deployment on HF Spaces
-]
-
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=origins,
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-# ==================================================#
-
+# Global state for RAG components
 retriever = None
 vectorstore = None
 
+
+# ---------------- Input Schema ----------------
 class Query(BaseModel):
     question: str
 
-# ==================================================#
-#  PDF Upload
+
 # ==================================================
-@app.post("/api/upload")
+#  PDF Upload → Chunk → Embed → Vectorstore
+# ==================================================
+@app.post("/upload")
 async def upload_pdf(file: UploadFile = File(...)):
+    """Handles PDF upload, processing, chunking, embedding, and vectorstore creation."""
     global retriever, vectorstore
+
     if not file.filename.endswith(".pdf"):
         raise HTTPException(400, "Only PDF files allowed")
+
     if not client:
-        raise HTTPException(500, "Groq API key missing")
+        raise HTTPException(500, "Service not fully initialized. Groq API key is missing or invalid.")
+
     path = None
     try:
+        # 1. Save file temporarily
         with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp:
             tmp.write(await file.read())
             path = tmp.name
+
+        logging.info(f"Processing PDF: {path}")
+
+        # 2. Load
         loader = PyPDFLoader(path)
         docs = loader.load()
+
+        # 3. Split
         splitter = RecursiveCharacterTextSplitter(
             chunk_size=800,
             chunk_overlap=50
         )
         chunks = splitter.split_documents(docs)
+
+        # 4. Embeddings (Using CPU-friendly model)
         embeddings = HuggingFaceEmbeddings(
             model_name="sentence-transformers/all-MiniLM-L6-v2",
             model_kwargs={"device": "cpu"},
             encode_kwargs={"normalize_embeddings": True}
         )
+
+        # 5. Clear previous vectorstore to free memory
         if vectorstore:
             del vectorstore
             gc.collect()
+
+        # 6. Create Vectorstore and Retriever
         vectorstore = Chroma.from_documents(chunks, embeddings)
+        # Search for 3 most relevant chunks
         retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
+
+        logging.info(f"PDF processed. Chunks created: {len(chunks)}")
+
         return {"message": "PDF processed", "chunks": len(chunks)}
+
     except Exception as e:
-        raise HTTPException(500, str(e))
+        logging.error(f"Error during PDF processing: {e}")
+        raise HTTPException(500, f"Error: {str(e)}")
     finally:
+        # 7. Cleanup temp file and memory
         if path and os.path.exists(path):
             os.unlink(path)
         gc.collect()
 
-# ==================================================#
-#  Ask Question
+
+# ==================================================
+#  ASK → RETRIEVE → GROQ → ANSWER
 # ==================================================
-@app.post("/api/ask")
+@app.post("/ask")
 async def ask(req: Query):
-    if not retriever:
-        raise HTTPException(400, "Upload PDF first")
+    global retriever
+
+    if client is None:
+        raise HTTPException(500, "Groq client is not initialized. Check API key setup.")
+
+    if retriever is None:
+        raise HTTPException(400, "Upload PDF first to initialize the knowledge base.")
+
     try:
+        # 1. Retrieve relevant chunks (NEW LangChain API)
         docs = retriever.invoke(req.question)
+
         context = "\n\n".join(d.page_content for d in docs)
-        prompt = f"""Use ONLY the context below.If answer not found, say: "I cannot find this in the PDF."CONTEXT:{context}QUESTION: {req.question}ANSWER:"""
+
+        # 2. Build prompt
+        prompt = f"""
+You are a strict RAG Q&A assistant.
+Use ONLY the context provided. If the answer is not found, reply:
+"I cannot find this in the PDF."
+
+---------------- CONTEXT ----------------
+{context}
+-----------------------------------------
+
+QUESTION: {req.question}
+
+FINAL ANSWER:
+"""
+
+        # 3. Call Groq
         response = client.chat.completions.create(
             model=GROQ_MODEL,
             messages=[
-                {"role": "system", "content": "Answer strictly from PDF context"},
+                {"role": "system",
+                 "content": "Use only the PDF content. If answer not found, say: 'I cannot find this in the PDF.'"},
                 {"role": "user", "content": prompt}
             ],
             temperature=0.0
         )
-        return {
-            "answer": response.choices[0].message.content.strip(),
-            "sources": len(docs)
-        }
+
+        answer = response.choices[0].message.content.strip()
+        return {"answer": answer, "sources": len(docs)}
+
     except APIError as e:
-        raise HTTPException(500, str(e))
+        logging.error(f"Groq API Error: {e}")
+        raise HTTPException(500, f"Groq API Error: {str(e)}")
+
+    except Exception as e:
+        logging.error(f"General error in /ask: {e}")
+        raise HTTPException(500, f"General error: {str(e)}")
+
 
-# ==================================================#
-#  Clear Memory
 # ==================================================
-@app.post("/api/clear")
+#  HEALTH & CLEAR
+# ==================================================
+@app.get("/health")
+async def health():
+    """Endpoint for checking service status."""
+    return {
+        "status": "running",
+        "pdf_loaded": retriever is not None,
+        "groq_client_ok": client is not None
+    }
+
+
+@app.post("/clear")
 async def clear():
+    """Clears the current RAG components from memory."""
     global retriever, vectorstore
+
+    # Explicitly clear objects
     if vectorstore:
         del vectorstore
     retriever = None
     vectorstore = None
+
     gc.collect()
+    # Clear CUDA cache if running on a machine with a GPU (good practice)
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
-    return {"message": "Memory cleared"}
 
-# ==================================================#
-#  Health
-# ==================================================
-@app.get("/api/health")
-async def health():
-    return {
-        "status": "running",
-        "pdf_loaded": retriever is not None,
-        "groq_client_ok": client is not None
-    }
+    return {"message": "Memory cleared. Upload a new PDF."}