routers/agent.py

from fastapi import APIRouter, HTTPException, UploadFile, File, Form
from pydantic import BaseModel
from typing import List, Optional
import json
import io
from pypdf import PdfReader
import docx
from agent.graph import app as agent_app

router = APIRouter(
    prefix="/agent",
    tags=["agent"],
    responses={404: {"description": "Not found"}},
)

class GenerateQuestionRequest(BaseModel):
    description: str
    existing_questions: List[str]

class GenerateQuestionResponse(BaseModel):
    question_text: str
    options: List[str]
    correct_option_index: int

@router.post("/generate-question", response_model=GenerateQuestionResponse)
async def generate_question(
    description: str = Form(...),
    existing_questions: str = Form(default="[]"),
    file: Optional[UploadFile] = File(None)
):
    try:
        existing_questions_list = json.loads(existing_questions)
        
        document_content = ""
        if file:
            content = await file.read()
            filename = file.filename.lower()
            
            if filename.endswith(".pdf"):
                pdf = PdfReader(io.BytesIO(content))
                text = ""
                for page in pdf.pages:
                    text += page.extract_text() + "\n"
                document_content = text
            elif filename.endswith(".docx"):
                doc = docx.Document(io.BytesIO(content))
                text = ""
                for para in doc.paragraphs:
                    text += para.text + "\n"
                document_content = text
            elif filename.endswith(".txt") or filename.endswith(".md"):
                document_content = content.decode("utf-8")

        inputs = {
            "qcm_description": description,
            "existing_questions": existing_questions_list,
            "document_content": document_content,
            "generated_question": None
        }
        
        result = await agent_app.ainvoke(inputs)
        
        generated = result.get("generated_question")
        if not generated:
            raise HTTPException(status_code=500, detail="Failed to generate question")
            
        return generated
    except Exception as e:
        print(f"Error in generate_question: {e}")
        raise HTTPException(status_code=500, detail=str(e))