Update main.py

main.py

@@ -2,24 +2,28 @@
 
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
-from typing import List, Optional
+from typing import Optional
 from langchain.document_loaders import WikipediaLoader
 from langchain_groq import ChatGroq
+import os
 
-app = FastAPI()
+app = FastAPI(title="Quiz Generator API")
 
-# Replace with your actual Groq API key
-GROQ_API_KEY = "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"
+# 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.genenv('api_key')
 class QuizRequest(BaseModel):
     search_query: str
     complexity: str = "Hard"
     quiz_type: str = "MCQs"
 
 class GradeRequest(BaseModel):
-    quiz: str
     answers: str
-    context: Optional[str] = None
 
 def get_llm():
     return ChatGroq(
@@ -29,56 +33,75 @@ def get_llm():
         api_key=GROQ_API_KEY
     )
 
-def wikipedia_query(search_query):
+def wikipedia_query(search_query: str):
     try:
-        data = WikipediaLoader(query=search_query, load_max_docs=2).load()
-        return data
+        docs = WikipediaLoader(query=search_query, load_max_docs=2).load()
+        return docs
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Wikipedia query failed: {e}")
 
+@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):
-    context = wikipedia_query(request.search_query)
-    llm = get_llm()
+    # fetch & store context
+    context_docs = wikipedia_query(request.search_query)
+    context_text = str(context_docs)
+    STORE["context"] = context_text
 
+    # generate quiz
+    llm = get_llm()
     prompt = f"""
-    You are a quiz generator assistant. Create a quiz for the given context.
+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:
+Instructions:
+- Do not write answers in the quiz.
+- Quiz should be based on the following context:
 
-    context:
-    {context}
+context:
+{context_text}
 
-    quiz complexity:
-    {request.complexity}
+quiz complexity:
+{request.complexity}
 
-    quiz type:
-    {request.quiz_type}
-
-    Your response:
-    """
+quiz type:
+{request.quiz_type}
 
+Your response:
+"""
     result = llm.invoke(prompt)
-    return {"quiz": result.content, "context": str(context)}
+    STORE["quiz"] = result.content
+
+    return {
+        "quiz": result.content
+    }
 
 @app.post("/grade_quiz/")
 async def grade_quiz(request: GradeRequest):
-    llm = get_llm()
+    # 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. Write only the marks.
-
-    quiz:
-    {request.quiz}
+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.
 
-    answers:
-    {request.answers}
+quiz:
+{STORE['quiz']}
 
-    context:
-    {request.context}
-    """
+answers:
+{request.answers}
 
+context:
+{STORE['context']}
+"""
     result = llm.invoke(prompt)
-    return {"grade": result.content}
+    return {
+        "grade": result.content
+    }