Update main.py

main.py

@@ -1,104 +1,159 @@
+import os
+import zipfile
+import logging
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
-from langchain.document_loaders import WikipediaLoader
-from langchain_groq import ChatGroq
-import os
-
-app = FastAPI(title="Quiz Generator API")
 
-# in-memory store for the last quiz + context
-STORE = {
-    "quiz": None,      # str
-    "context": None,   # str
-}
-
-# Replace with your actual Groq API key
-GROQ_API_KEY = os.getenv('api_key')
+from langchain_community.vectorstores import FAISS
+from langchain_huggingface import HuggingFaceEmbeddings
+from langchain_groq import ChatGroq
+from langchain.chains import RetrievalQA
+from langchain.prompts import PromptTemplate
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+)
+logger = logging.getLogger(__name__)
+
+app = FastAPI()
+
+# === Globals ===
+llm = None
+embeddings = None
+vectorstore = None
+retriever = None
+quiz_chain = None
+grade_chain = None
 
 class QuizRequest(BaseModel):
-    search_query: str
+    topic: str
 
 class GradeRequest(BaseModel):
-    answers: str
-
-def get_llm():
-    return ChatGroq(
-        model="meta-llama/llama-4-scout-17b-16e-instruct",
-        temperature=0,
-        max_tokens=1024,
-        api_key=GROQ_API_KEY
-    )
+    answers: str  # string of Q/A pairs
 
-def wikipedia_query(search_query: str):
+@app.on_event("startup")
+def load_components():
+    global llm, embeddings, vectorstore, retriever, quiz_chain, grade_chain
     try:
-        docs = WikipediaLoader(query=search_query, load_max_docs=2).load()
-        return docs
+        api_key = os.getenv("API_KEY")
+        if not api_key:
+            logger.error("API_KEY environment variable is not set or empty.")
+            raise RuntimeError("API_KEY environment variable is not set or empty.")
+        logger.info("API_KEY is set.")
+
+        # 1) Init LLM & Embeddings
+        llm = ChatGroq(
+            model="meta-llama/llama-4-scout-17b-16e-instruct",
+            temperature=0,
+            max_tokens=1024,
+            api_key=api_key,
+        )
+        embeddings = HuggingFaceEmbeddings(
+            model_name="intfloat/multilingual-e5-large",
+            model_kwargs={"device": "cpu"},
+            encode_kwargs={"normalize_embeddings": True},
+        )
+
+        # 2) Load FAISS indexes
+        for zip_name, dir_name in [("faiss_index.zip", "faiss_index"), ("faiss_index(1).zip", "faiss_index_extra")]:
+            if not os.path.exists(dir_name):
+                with zipfile.ZipFile(zip_name, 'r') as z:
+                    z.extractall(dir_name)
+                logger.info(f"Unzipped {zip_name} to {dir_name}.")
+            else:
+                logger.info(f"Directory {dir_name} already exists.")
+
+        vs1 = FAISS.load_local("faiss_index", embeddings, allow_dangerous_deserialization=True)
+        logger.info("FAISS index 1 loaded.")
+        vs2 = FAISS.load_local("faiss_index_extra", embeddings, allow_dangerous_deserialization=True)
+        logger.info("FAISS index 2 loaded.")
+
+        vs1.merge_from(vs2)
+        vectorstore = vs1
+        logger.info("Merged FAISS indexes into a single vectorstore.")
+
+        retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
+
+        # Quiz generation chain
+        quiz_prompt = PromptTemplate(
+            template="""
+Generate a quiz on the topic "{topic}" using **only** the information in the "Retrieved context".
+Include clear questions and multiple-choice options (A, B, C, D).
+If context is insufficient, reply with "I don't know".
+
+Retrieved context:
+{context}
+
+Quiz topic:
+{topic}
+
+Quiz:
+""",
+            input_variables=["context", "topic"],
+        )
+        quiz_chain = RetrievalQA.from_chain_type(
+            llm=llm,
+            chain_type="stuff",
+            retriever=retriever,
+            return_source_documents=False,
+            chain_type_kwargs={"prompt": quiz_prompt},
+        )
+        logger.info("Quiz chain ready.")
+
+        # Grade quiz chain
+        grade_prompt = PromptTemplate(
+            template="""
+Grade the following quiz answers based on the "Retrieved context".
+Provide a score and brief feedback for each question.
+If context is insufficient to grade, say "I don't know" for that question.
+
+Retrieved context:
+{context}
+
+User answers:
+{answers}
+
+Grading:
+""",
+            input_variables=["context", "answers"],
+        )
+        grade_chain = RetrievalQA.from_chain_type(
+            llm=llm,
+            chain_type="stuff",
+            retriever=retriever,
+            return_source_documents=False,
+            chain_type_kwargs={"prompt": grade_prompt},
+        )
+        logger.info("Grading chain ready.")
+
     except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Wikipedia query failed: {e}")
+        logger.error("Error loading components", exc_info=True)
+        raise
 
 @app.get("/")
-async def root():
-    return {
-        "message": "Welcome to the Quiz Generator API!  \n"
-                   "POST /generate_quiz/ to create a new quiz.  \n"
-                   "POST /grade_quiz/ to grade your answers against the last quiz."
-    }
-
-@app.post("/generate_quiz/")
-async def generate_quiz(request: QuizRequest):
-    # fetch & store context
-
-    llm = get_llm()
-    search_query=llm.invoke(f'write one word which should be search for this question {request.search_query}')
-    
-    context_docs = wikipedia_query(search_query.content)
-    context_text = str(context_docs)
-    STORE["context"] = context_text
-
-    # generate quiz
-    
-    prompt = f"""
-You are a quiz generator assistant. Create a quiz for the given context.
-
-Instructions:
-- Do not write answers in the quiz.
-- Quiz should be based on the following context:
-
-context:
-{context_text}
-
-question:
-generate quiz on {request.search_query}
-
-Your response:
-"""
-    result = llm.invoke(prompt)
-    STORE["quiz"] = result.content
-
-    return {
-        "quiz": result.content
-    }
-
-@app.post("/grade_quiz/")
-async def grade_quiz(request: GradeRequest):
-    # ensure we have a quiz to grade
-    if STORE["quiz"] is None or STORE["context"] is None:
-        raise HTTPException(status_code=400, detail="No quiz available. Call /generate_quiz/ first.")
-
-    llm = get_llm()
-    prompt = f"""
-Check the quiz answers and give marks and also provide the correct answers. Use the following context to check the quiz. Return only the total mark and the correct answers.
-Do not write anything else expect the marks and correct answers do not generate new quiz.
-quiz:
-{STORE['quiz']}
-
-answers:
-{request.answers}
-
-context:
-{STORE['context']}
-"""
-    result = llm.invoke(prompt)
-    return {
-        "grade": result.content
-    }
+def root():
+    return {"message": "API is up and running!"}
+
+@app.post("/quiz")
+def create_quiz(request: QuizRequest):
+    try:
+        logger.info("Generating quiz for topic: %s", request.topic)
+        result = quiz_chain.invoke({"query": request.topic})
+        logger.info("Quiz generated successfully.")
+        return {"quiz": result.get("result")}
+    except Exception as e:
+        logger.error("Error generating quiz", exc_info=True)
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/grade")
+def grade_quiz(request: GradeRequest):
+    try:
+        logger.info("Grading quiz with provided answers.")
+        result = grade_chain.invoke({"query": request.answers})
+        logger.info("Quiz graded successfully.")
+        return {"grading": result.get("result")}
+    except Exception as e:
+        logger.error("Error grading quiz", exc_info=True)
+        raise HTTPException(status_code=500, detail=str(e))