app/backend/models.py

"""
Pydantic models for request/response validation
"""

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


class ModerateRequest(BaseModel):
    text: str = Field(..., description="Text to moderate")
    model: str = Field(
        default="lionguard-2.1",
        description="Model to use: lionguard-2, lionguard-2.1, or lionguard-2-lite"
    )


class CategoryScore(BaseModel):
    name: str
    emoji: str
    max_score: float


class ModerateResponse(BaseModel):
    binary_score: float
    binary_verdict: str  # "pass", "warn", "fail"
    binary_percentage: int
    categories: List[CategoryScore]
    text_id: str
    model_used: str


class FeedbackRequest(BaseModel):
    text_id: str = Field(..., description="ID of the text being voted on")
    agree: bool = Field(..., description="True for thumbs up, False for thumbs down")


class FeedbackResponse(BaseModel):
    success: bool
    message: str


class ChatMessage(BaseModel):
    role: str
    content: str


class ChatHistories(BaseModel):
    no_moderation: List[ChatMessage] = Field(default_factory=list)
    openai_moderation: List[ChatMessage] = Field(default_factory=list)
    lionguard: List[ChatMessage] = Field(default_factory=list)


class ChatRequest(BaseModel):
    message: str = Field(..., description="Message to send to all guardrails")
    model: str = Field(
        default="lionguard-2.1",
        description="LionGuard model variant to use"
    )
    histories: ChatHistories = Field(default_factory=ChatHistories)


class ChatResponse(BaseModel):
    histories: ChatHistories
    lionguard_score: Optional[float] = None