api/schemas.py

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


class ChatRequest(BaseModel):
    """Chat request schema"""
    message: str = Field(..., min_length=1, max_length=1000)
    session_id: str = Field(default="default")
    domain: Optional[str] = Field(default="general_medical")
    enable_reasoning: bool = Field(default=True)
    use_real_time: bool = Field(default=True)
    max_papers: int = Field(default=15, ge=1, le=50)

    class Config:
        json_schema_extra = {
            "example": {
                "message": "Compare deep learning vs traditional ML in medical imaging",
                "session_id": "user_123",
                "domain": "medical_imaging",
                "enable_reasoning": True,
                "use_real_time": True,
                "max_papers": 15
            }
        }


class ChatResponse(BaseModel):
    """Chat response schema"""
    success: bool
    message: str
    session_id: str
    processing_time: Optional[float] = None
    confidence_score: Optional[float] = None
    papers_used: Optional[int] = None
    reasoning_used: Optional[bool] = None
    error: Optional[str] = None
    raw_response: Optional[Dict[str, Any]] = None

    class Config:
        json_schema_extra = {
            "example": {
                "success": True,
                "message": "Here's the comparison...",
                "session_id": "user_123",
                "processing_time": 23.86,
                "confidence_score": 86.7,
                "papers_used": 13,
                "reasoning_used": True
            }
        }


class SessionCreate(BaseModel):
    """Session creation schema"""
    user_id: Optional[str] = None
    initial_domain: Optional[str] = "general_medical"


class DomainInfo(BaseModel):
    """Domain information"""
    id: str
    name: str
    icon: str
    description: Optional[str] = None


class WebSocketMessage(BaseModel):
    """WebSocket message schema"""
    type: str  # "message", "typing", "error", "init_session"
    message: Optional[str] = None
    session_id: Optional[str] = None
    domain: Optional[str] = None
    enable_reasoning: Optional[bool] = None
    use_real_time: Optional[bool] = None
    max_papers: Optional[int] = None