agents/states.py

from typing import Literal, Optional, List
from pydantic import BaseModel, Field


class QuizTask(BaseModel):
    """
    Schema for a single quiz task that the agent should generate or return.
    """

    task_id: int = Field(
        ...,
        description=(
            "Unique integer identifier of the task within a quiz. "
            "Start from 1 and increment by 1 for each new task."
        ),
    )

    task: str = Field(
        ...,
        description=(
            "The question text shown to the student. "
            "Must be a complete, clear instruction or question in plain language."
        ),
    )

    task_type: Literal["fill_gap", "multiple_choice", "type_in"] = Field(
        ...,
        description=(
            "The interaction type for this task:\n"
            "- 'fill_gap': a sentence or formula with a missing part the student must fill in.\n"
            "- 'multiple_choice': student chooses exactly one option from 'answer_options'.\n"
            "- 'type_in': open-ended question where student types a short free-text answer."
        ),
    )

    answer_options: Optional[List[str]] = Field(
        None,
        description=(
            "List of answer options, in the exact order they should be shown to the student. "
            "Required when task_type = 'multiple_choice', 'fill_gap'. "
            "Must be None for 'type_in'."
        ),
    )

    correct_answer: Optional[str] = Field(
        None,
        description=(
            "The correct answer in normalized string form.\n"
            "- For 'fill_gap' and 'multiple_choice', this MUST exactly match one element of 'answer_options'.\n"
            "- For 'type_in', this is the canonical correct answer "
            "(e.g. '42', 'O(n log n)'); you may later implement fuzzy matching if needed."
        ),
    )


class Quiz(BaseModel):
    """Collection of quiz tasks generated by the Examiner agent."""
    tasks: List[QuizTask] = Field(
        ...,
        description=(
            "A collection of quiz tasks that the agent should generate for the given document.\n"
        )
    )


class WorkflowState(BaseModel):
    """State for the LangGraph workflow."""
    pdf_path: str = Field(default="", description="Path to the uploaded PDF")
    summary: str = Field(default="", description="Summary of the document")
    quiz: Optional[Quiz] = Field(default=None, description="Generated quiz")
    user_answers: Optional[List[str]] = Field(default=None, description="User's quiz answers")
    feedback: str = Field(default="", description="Supervisor feedback on quiz results")
    messages: List[str] = Field(default_factory=list, description="Conversation history")