api/models.py

"""
API Models and Schemas for Medical RAG AI Advisor
"""
from pydantic import BaseModel, Field
from typing import Optional, List, Dict, Any
from enum import Enum


class QueryType(str, Enum):
    """Types of medical queries supported"""
    GENERAL = "general"
    DRUG_INTERACTION = "drug_interaction"
    SIDE_EFFECTS = "side_effects"
    GUIDELINES = "guidelines"
    COMPARISON = "comparison"


class QueryRequest(BaseModel):
    """Request model for medical queries"""
    query: str = Field(..., description="Medical question or query", min_length=1)
    query_type: Optional[QueryType] = Field(None, description="Type of medical query")
    context: Optional[str] = Field(None, description="Additional context for the query")
    patient_info: Optional[Dict[str, Any]] = Field(None, description="Patient information if relevant")


class QueryResponse(BaseModel):
    """Response model for medical queries"""
    response: str = Field(..., description="AI-generated medical response")
    sources: Optional[List[str]] = Field(None, description="Sources used for the response")
    confidence: Optional[float] = Field(None, description="Confidence score of the response")
    query_type: Optional[QueryType] = Field(None, description="Detected or specified query type")
    processing_time: Optional[float] = Field(None, description="Time taken to process the query")


class StreamChunk(BaseModel):
    """Model for streaming response chunks"""
    chunk: str = Field(..., description="Chunk of the streaming response")
    is_final: bool = Field(False, description="Whether this is the final chunk")


class SideEffectReport(BaseModel):
    """Model for side effect reporting"""
    drug_name: str = Field(..., description="Name of the drug")
    side_effects: str = Field(..., description="Reported side effects")
    patient_age: Optional[int] = Field(None, description="Patient age")
    patient_gender: Optional[str] = Field(None, description="Patient gender")
    dosage: Optional[str] = Field(None, description="Drug dosage")
    duration: Optional[str] = Field(None, description="Duration of treatment")
    severity: Optional[str] = Field(None, description="Severity of side effects")
    outcome: Optional[str] = Field(None, description="Outcome of the side effects")
    additional_details: Optional[str] = Field(None, description="Additional clinical details")


class SideEffectResponse(BaseModel):
    """Response model for side effect reporting"""
    report_id: str = Field(..., description="Unique identifier for the report")
    status: str = Field(..., description="Status of the report submission")
    message: str = Field(..., description="Confirmation message")
    recommendations: Optional[List[str]] = Field(None, description="Clinical recommendations")


class ComparisonRequest(BaseModel):
    """Request model for provider/treatment comparisons"""
    providers: List[str] = Field(..., description="List of providers to compare", min_items=2)
    criteria: Optional[List[str]] = Field(None, description="Specific criteria for comparison")


class ComparisonResponse(BaseModel):
    """Response model for provider/treatment comparisons"""
    comparison: str = Field(..., description="Detailed comparison analysis")
    summary: Dict[str, Any] = Field(..., description="Summary of key differences")
    recommendations: Optional[str] = Field(None, description="Recommendations based on comparison")


class InitializationStatus(BaseModel):
    """Initialization status response model"""
    is_complete: bool = Field(..., description="Whether initialization is complete")
    status_message: str = Field(..., description="Current initialization status")
    is_successful: bool = Field(..., description="Whether initialization was successful")
    error: Optional[str] = Field(None, description="Initialization error if any")


class HealthStatus(BaseModel):
    """Health check response model"""
    status: str = Field(..., description="API health status")
    version: str = Field(..., description="API version")
    timestamp: str = Field(..., description="Current timestamp")
    components: Dict[str, str] = Field(..., description="Status of system components")
    initialization: Optional[InitializationStatus] = Field(None, description="Background initialization status")


class ErrorResponse(BaseModel):
    """Error response model"""
    error: str = Field(..., description="Error type")
    message: str = Field(..., description="Error message")
    details: Optional[Dict[str, Any]] = Field(None, description="Additional error details")
    timestamp: str = Field(..., description="Error timestamp")