ankigen_core/agents/schemas.py

"""
Pydantic schemas for structured outputs from agents.
These schemas ensure type safety and eliminate JSON parsing errors.
"""

from typing import List, Dict, Any, Optional
from pydantic import BaseModel, Field
from enum import Enum


class DifficultyLevel(str, Enum):
    """Difficulty levels for flashcards"""

    BEGINNER = "beginner"
    INTERMEDIATE = "intermediate"
    ADVANCED = "advanced"


class CardType(str, Enum):
    """Types of flashcards"""

    BASIC = "basic"
    CLOZE = "cloze"


class CardFrontSchema(BaseModel):
    """Schema for the front of a flashcard"""

    question: str = Field(..., description="The question or prompt for the flashcard")


class CardBackSchema(BaseModel):
    """Schema for the back of a flashcard"""

    answer: str = Field(..., description="The main answer to the question")
    explanation: str = Field(..., description="Detailed explanation of the answer")
    example: str = Field(..., description="A concrete example illustrating the concept")


class CardMetadataSchema(BaseModel):
    """Schema for flashcard metadata"""

    topic: str = Field(..., description="The main topic of the card")
    subject: str = Field(..., description="The subject area (e.g., Biology, History)")
    difficulty: DifficultyLevel = Field(..., description="The difficulty level")
    tags: Optional[List[str]] = Field(
        None, description="Relevant tags for categorization"
    )
    learning_outcomes: Optional[List[str]] = Field(
        None, description="What the learner should achieve"
    )
    prerequisites: Optional[List[str]] = Field(
        None, description="Required prior knowledge"
    )
    related_concepts: Optional[List[str]] = Field(
        None, description="Related concepts to explore"
    )
    estimated_time: Optional[str] = Field(None, description="Estimated time to learn")
    common_mistakes: Optional[List[str]] = Field(
        None, description="Common mistakes to avoid"
    )
    memory_aids: Optional[List[str]] = Field(
        None, description="Memory aids or mnemonics"
    )
    real_world_applications: Optional[List[str]] = Field(
        None, description="Real-world applications"
    )


class CardSchema(BaseModel):
    """Complete schema for a flashcard"""

    card_type: CardType = Field(..., description="The type of flashcard")
    front: CardFrontSchema = Field(..., description="The front of the card")
    back: CardBackSchema = Field(..., description="The back of the card")
    metadata: CardMetadataSchema = Field(..., description="Metadata about the card")
    enhancement_notes: Optional[str] = Field(
        None, description="Notes about enhancements made"
    )


class CardsGenerationSchema(BaseModel):
    """Schema for multiple cards generation"""

    cards: List[CardSchema] = Field(..., description="List of generated flashcards")


class JudgeDecisionSchema(BaseModel):
    """Schema for judge decisions"""

    approved: bool = Field(..., description="Whether the card is approved")
    score: float = Field(
        ..., ge=0.0, le=1.0, description="Quality score between 0 and 1"
    )
    feedback: str = Field(..., description="Detailed feedback about the card")
    improvements: Optional[List[str]] = Field(
        None, description="Suggested improvements"
    )
    reasoning: str = Field(..., description="Detailed reasoning for the decision")
    confidence: float = Field(
        ..., ge=0.0, le=1.0, description="Confidence in the decision"
    )
    metadata: Optional[Dict[str, Any]] = Field(None, description="Additional metadata")


class EnhancementSchema(BaseModel):
    """Schema for card enhancements"""

    enhanced_card: CardSchema = Field(..., description="The enhanced flashcard")
    enhancement_summary: str = Field(..., description="Summary of what was enhanced")
    enhancement_details: Optional[Dict[str, Any]] = Field(
        None, description="Detailed enhancement information"
    )


class GenerationRequestSchema(BaseModel):
    """Schema for generation requests"""

    topic: str = Field(..., description="The topic to generate cards for")
    subject: str = Field(..., description="The subject area")
    num_cards: int = Field(..., ge=1, le=20, description="Number of cards to generate")
    difficulty: DifficultyLevel = Field(..., description="Target difficulty level")
    context: Optional[Dict[str, Any]] = Field(None, description="Additional context")
    preferences: Optional[Dict[str, Any]] = Field(None, description="User preferences")


class TokenUsageSchema(BaseModel):
    """Schema for token usage tracking"""

    prompt_tokens: int = Field(..., ge=0, description="Number of tokens in the prompt")
    completion_tokens: int = Field(
        ..., ge=0, description="Number of tokens in the completion"
    )
    total_tokens: int = Field(..., ge=0, description="Total tokens used")
    estimated_cost: float = Field(..., ge=0.0, description="Estimated cost in USD")
    model: str = Field(..., description="Model used for the request")


class AutoConfigSchema(BaseModel):
    """Schema for auto-configuration based on subject analysis"""

    # What to search for in Context7
    library_search_term: str = Field(
        ...,
        description="Library name to search for in Context7 (e.g., 'pandas', 'react', 'tensorflow')",
    )

    # Specific topic within the library (optional)
    documentation_focus: Optional[str] = Field(
        None,
        description="Specific topic/area within the library documentation to focus on",
    )

    # Suggested settings based on subject analysis
    topic_number: int = Field(
        ..., ge=2, le=20, description="Number of topics to generate (2-20)"
    )
    cards_per_topic: int = Field(
        ..., ge=2, le=30, description="Number of cards per topic (2-30)"
    )
    learning_preferences: str = Field(
        ..., description="Learning preferences and focus areas for card generation"
    )
    generate_cloze: bool = Field(
        ...,
        description="Whether to generate cloze cards (true for syntax/code, false for concepts)",
    )
    model_choice: str = Field(
        ...,
        description="Recommended model: 'gpt-4.1' for complex topics, 'gpt-4.1-nano' for simpler topics",
    )

    # Analysis metadata
    subject_type: str = Field(
        ...,
        description="Type of subject: 'concepts', 'syntax', 'api', 'theory', 'practical'",
    )
    scope: str = Field(
        ..., description="Scope of the subject: 'narrow', 'medium', 'broad'"
    )
    rationale: str = Field(
        ..., description="Brief explanation of why these settings were chosen"
    )