Update backend.py

backend.py

@@ -5,201 +5,160 @@ import logging
 from fastapi import FastAPI, UploadFile, File, HTTPException
 from pydantic import BaseModel
 import torch
-from dotenv import load_dotenv  # Used to load API key from .env file
+from dotenv import load_dotenv
 
-# ---------------- 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
 
-# --- Configuration & Setup ---
+# ---------------- Setup ----------------
 logging.basicConfig(level=logging.INFO)
-
-# 1. Load environment variables from .env file
 load_dotenv()
 
-# 2. Load API Key from Environment Variable
-GROQ_API_KEY = os.environ.get("GROQ_API_KEY")
+GROQ_API_KEY = os.getenv("GROQ_API_KEY")
 GROQ_MODEL = "llama-3.1-8b-instant"
 
-# 3. Initialize Groq Client
 client = None
-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:
+if GROQ_API_KEY:
     try:
         client = Groq(api_key=GROQ_API_KEY)
-        logging.info("✅ Groq client initialized successfully.")
+        logging.info("✅ Groq client initialized")
     except Exception as e:
-        logging.error(f"❌ Failed to initialize Groq client: {e}")
-        client = None
+        logging.error(f"Groq init failed: {e}")
 
 app = FastAPI()
 
-# Global state for RAG components
 retriever = None
 vectorstore = None
 
 
-# ---------------- Input Schema ----------------
 class Query(BaseModel):
     question: str
 
 
 # ==================================================
-#  PDF Upload → Chunk → Embed → Vectorstore
+#  PDF Upload
 # ==================================================
-@app.post("/upload")
+@app.post("/api/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, "Service not fully initialized. Groq API key is missing or invalid.")
+        raise HTTPException(500, "Groq API key missing")
 
     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:
-        logging.error(f"Error during PDF processing: {e}")
-        raise HTTPException(500, f"Error: {str(e)}")
+        raise HTTPException(500, str(e))
+
     finally:
-        # 7. Cleanup temp file and memory
         if path and os.path.exists(path):
             os.unlink(path)
         gc.collect()
 
 
 # ==================================================
-#  ASK → RETRIEVE → GROQ → ANSWER
+#  Ask Question
 # ==================================================
-@app.post("/ask")
+@app.post("/api/ask")
 async def ask(req: Query):
-    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.")
+    if not retriever:
+        raise HTTPException(400, "Upload PDF first")
 
     try:
-        # 1. Retrieve relevant chunks (NEW LangChain API)
         docs = retriever.invoke(req.question)
-
         context = "\n\n".join(d.page_content for d in docs)
 
-        # 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."
+Use ONLY the context below.
+If answer not found, say: "I cannot find this in the PDF."
 
----------------- CONTEXT ----------------
+CONTEXT:
 {context}
------------------------------------------
 
 QUESTION: {req.question}
 
-FINAL ANSWER:
+ANSWER:
 """
 
-        # 3. Call Groq
         response = client.chat.completions.create(
             model=GROQ_MODEL,
             messages=[
-                {"role": "system",
-                 "content": "Use only the PDF content. If answer not found, say: 'I cannot find this in the PDF.'"},
+                {"role": "system", "content": "Answer strictly from PDF context"},
                 {"role": "user", "content": prompt}
             ],
             temperature=0.0
         )
 
-        answer = response.choices[0].message.content.strip()
-        return {"answer": answer, "sources": len(docs)}
+        return {
+            "answer": response.choices[0].message.content.strip(),
+            "sources": len(docs)
+        }
 
     except APIError as 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)}")
+        raise HTTPException(500, str(e))
 
 
 # ==================================================
-#  HEALTH & CLEAR
+#  Clear Memory
 # ==================================================
-@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")
+@app.post("/api/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. Upload a new PDF."}
+    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
+    }