Update main.py

main.py

@@ -1,81 +1,63 @@
 import os
-import json
-from typing import List
+import shutil
 from fastapi import FastAPI, UploadFile, File, HTTPException
-from pydantic import BaseModel
-from langchain_groq import ChatGroq
-from langchain.document_loaders import PyPDFLoader
+from fastapi.responses import JSONResponse
+from pdf_processor import extract_answers_from_pdf, evaluate_student
 
-# Securely load your Groq API key from environment variables
-API_KEY = os.getenv("GROQ_API_KEY")
-if not API_KEY:
-    raise ValueError("GROQ_API_KEY environment variable not set.")
+app = FastAPI()
 
-app = FastAPI(title="PDF Question Extractor", version="1.0")
+# Directory to temporarily store uploaded files.
+UPLOAD_DIR = "uploads"
+if not os.path.exists(UPLOAD_DIR):
+    os.makedirs(UPLOAD_DIR)
 
-# Define the expected JSON response schema
-class ExtractionResult(BaseModel):
-    answers: List[str]
-
-# Initialize the language model (LLM)
-def get_llm():
-    return ChatGroq(
-        model="llama-3.3-70b-versatile",
-        temperature=0,
-        max_tokens=1024,
-        api_key=API_KEY
-    )
-
-llm = get_llm()
-
-# Root endpoint: Provides a welcome message and instructions
-@app.get("/")
-async def root():
-    return {
-        "message": "Welcome to the PDF Question Extractor API.",
-        "usage": "POST your PDF to /extract-answers/ to extract answers."
-    }
-
-# PDF extraction endpoint: Processes a PDF file upload
-@app.post("/extract-answers/")
-async def extract_answers(file: UploadFile = File(...)):
+@app.post("/extract/")
+async def extract_pdf(file: UploadFile = File(...)):
+    """
+    Endpoint to extract answers from a PDF file.
+    """
     try:
-        # Save the uploaded file temporarily
-        file_path = f"./temp_{file.filename}"
-        with open(file_path, "wb") as buffer:
-            buffer.write(file.file.read())
-
-        # Load and split the PDF into pages
-        loader = PyPDFLoader(file_path)
-        pages = loader.load_and_split()
-        all_page_content = "\n".join(page.page_content for page in pages)
-
-        # Generate the JSON schema from the Pydantic model
-        schema_dict = ExtractionResult.model_json_schema()
-        schema = json.dumps(schema_dict, indent=2)
-
-        # Build the prompt with system and user messages
-        system_message = (
-            "You are a document analysis tool that extracts the options and correct answers "
-            "from the provided document content. The output must be a JSON object that strictly follows the schema: " 
-            + schema
-        )
-        user_message = (
-            "Please extract the correct answers and options (A, B, C, D, E) from the following document content:\n\n"
-            + all_page_content
-        )
-        prompt = system_message + "\n\n" + user_message
-
-        # Invoke the LLM and request a JSON response
-        response = llm.invoke(prompt, response_format={"type": "json_object"})
-
-        # Validate and parse the JSON response using Pydantic
-        result = ExtractionResult.model_validate_json(response.content)
-
-        # Cleanup the temporary file
-        os.remove(file_path)
-
-        return result.model_dump()
-
+        file_location = os.path.join(UPLOAD_DIR, file.filename)
+        with open(file_location, "wb") as f:
+            shutil.copyfileobj(file.file, f)
+        
+        result = extract_answers_from_pdf(file_location)
+        return JSONResponse(content=result.model_dump())
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+    finally:
+        if os.path.exists(file_location):
+            os.remove(file_location)
+
+@app.post("/evaluate/")
+async def evaluate_pdfs(answer_key_file: UploadFile = File(...), student_file: UploadFile = File(...)):
+    """
+    Endpoint to evaluate student answers by comparing the answer key and student's answer PDFs.
+    """
+    try:
+        answer_key_path = os.path.join(UPLOAD_DIR, answer_key_file.filename)
+        student_path = os.path.join(UPLOAD_DIR, student_file.filename)
+        
+        with open(answer_key_path, "wb") as f:
+            shutil.copyfileobj(answer_key_file.file, f)
+        with open(student_path, "wb") as f:
+            shutil.copyfileobj(student_file.file, f)
+        
+        # Extract answers from both PDFs.
+        answer_key_result = extract_answers_from_pdf(answer_key_path)
+        student_result = extract_answers_from_pdf(student_path)
+        
+        # Evaluate the student answers.
+        evaluation = evaluate_student(answer_key_result, student_result)
+        return JSONResponse(content=evaluation.model_dump())
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
+    finally:
+        if os.path.exists(answer_key_path):
+            os.remove(answer_key_path)
+        if os.path.exists(student_path):
+            os.remove(student_path)
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=True)