models/generate.py

"""Pydantic models for text generation"""

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

class GenerateRequest(BaseModel):
    """Request model for argument generation"""
    model_config = ConfigDict(
        json_schema_extra={
            "example": {
                "topic": "Assisted suicide should be a criminal offence",
                "position": "positive"  # "positive" or "negative"
            }
        }
    )
    
    topic: str = Field(..., min_length=5, max_length=1000, 
                       description="The debate topic or statement")
    position: str = Field(..., min_length=5, max_length=50,
                         description="The stance to take")


class GenerateResponse(BaseModel):
    """Response model for argument generation"""
    model_config = ConfigDict(
        json_schema_extra={
            "example": {
                "topic": "Assisted suicide should be a criminal offence",
                "position": "positive",  # "positive" or "negative"
                "argument": "People have the right to choose how they end their lives",
                "timestamp": "2024-11-15T10:30:00"
            }
        }
    )
    
    topic: str  
    position: str
    argument: str
    timestamp: str
    timestamp: str


class BatchGenerateRequest(BaseModel):
    """Request model for batch argument generation"""
    items: List[GenerateRequest]


class BatchGenerateResponse(BaseModel):
    """Response model for batch argument generation"""
    results: List[GenerateResponse]
    model_info: Optional[str] = "KPA T5 Generation Model"
    timestamp: str