Spaces:
Running
Running
| """ | |
| API Models and Schemas for HBV AI Assistant | |
| """ | |
| from pydantic import BaseModel, Field, validator | |
| 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") | |
| # ============================================================================ | |
| # HBV PATIENT ASSESSMENT MODELS | |
| # ============================================================================ | |
| class HBVPatientInput(BaseModel): | |
| """Input model for HBV patient assessment""" | |
| sex: str = Field(..., description="Patient sex: Male / Female") | |
| age: int = Field(..., description="Patient age in years", ge=0, le=120) | |
| pregnancy_status: str = Field(..., description="Pregnancy status: Not pregnant / Pregnant") | |
| hbsag_status: str = Field(..., description="HBsAg status: Positive / Negative") | |
| duration_hbsag_months: int = Field(..., description="Duration of HBsAg positivity in months", ge=0) | |
| hbv_dna_level: float = Field(..., description="HBV DNA level in IU/mL", ge=0) | |
| hbeag_status: str = Field(..., description="HBeAg status: Positive / Negative") | |
| alt_level: float = Field(..., description="ALT level in U/L", ge=0) | |
| fibrosis_stage: str = Field(..., description="Fibrosis/Cirrhosis stage: F0-F1 / F2-F3 / F4") | |
| necroinflammatory_activity: str = Field(..., description="Necroinflammatory activity: A0 / A1 / A2 / A3") | |
| extrahepatic_manifestations: bool = Field(..., description="Presence of extrahepatic manifestations") | |
| immunosuppression_status: Optional[str] = Field(None, description="Immunosuppression status: None / Chemotherapy / Other") | |
| coinfections: List[str] = Field(default_factory=list, description="Coinfections: HIV, HCV, HDV") | |
| family_history_cirrhosis_hcc: bool = Field(..., description="Family history of Cirrhosis or HCC (first-degree relative)") | |
| other_comorbidities: Optional[List[str]] = Field(None, description="Other comorbidities") | |
| def validate_sex(cls, v): | |
| if v not in ['Male', 'Female']: | |
| raise ValueError('Sex must be either Male or Female') | |
| return v | |
| def validate_pregnancy(cls, v): | |
| if v not in ['Not pregnant', 'Pregnant']: | |
| raise ValueError('Pregnancy status must be either "Not pregnant" or "Pregnant"') | |
| return v | |
| def validate_hbsag(cls, v): | |
| if v not in ['Positive', 'Negative']: | |
| raise ValueError('HBsAg status must be either Positive or Negative') | |
| return v | |
| def validate_hbeag(cls, v): | |
| if v not in ['Positive', 'Negative']: | |
| raise ValueError('HBeAg status must be either Positive or Negative') | |
| return v | |
| def validate_fibrosis(cls, v): | |
| if v not in ['F0-F1', 'F2-F3', 'F4']: | |
| raise ValueError('Fibrosis stage must be F0-F1, F2-F3, or F4') | |
| return v | |
| def validate_necroinflammatory(cls, v): | |
| if v not in ['A0', 'A1', 'A2', 'A3']: | |
| raise ValueError('Necroinflammatory activity must be A0, A1, A2, or A3') | |
| return v | |
| class HBVAssessmentResponse(BaseModel): | |
| """Response model for HBV patient assessment""" | |
| eligible: bool = Field(..., description="Whether patient is eligible for treatment") | |
| recommendations: str = Field(..., description="Detailed recommendations with citations from SASLT 2021 guidelines including page numbers") | |
| class TextAssessmentInput(BaseModel): | |
| """Input model for text-based HBV patient assessment""" | |
| text_input: str = Field(..., description="Free-form text containing patient data", min_length=10) | |
| class Config: | |
| json_schema_extra = { | |
| "example": { | |
| "text_input": "45-year-old male patient\\nHBsAg: Positive for 12 months\\nHBV DNA: 5000 IU/mL\\nHBeAg: Positive\\nALT: 80 U/L\\nFibrosis stage: F2-F3\\nNecroinflammatory activity: A2\\nNo extrahepatic manifestations\\nNo immunosuppression\\nNo coinfections (HIV, HCV, HDV)\\nNo family history of cirrhosis or HCC" | |
| } | |
| } | |
| class ChatRequest(BaseModel): | |
| """Request model for chat interactions""" | |
| query: str = Field(..., description="Doctor's question about HBV guidelines", min_length=1) | |
| session_id: str = Field(default="default", description="Session identifier for conversation continuity") | |
| patient_context: Optional[HBVPatientInput] = Field(None, description="Optional patient context from assessment") | |
| assessment_result: Optional[HBVAssessmentResponse] = Field(None, description="Optional assessment result for context") | |
| class ChatResponse(BaseModel): | |
| """Response model for chat interactions""" | |
| response: str = Field(..., description="AI response to the doctor's question") | |
| session_id: str = Field(..., description="Session identifier") | |