src/models/response_models.py

from pydantic import BaseModel
from typing import List, Optional

class ICDCode(BaseModel):
    code: str
    description: str
    explanation: str

class CPTCode(BaseModel):
    code: str
    description: str
    explanation: str

class ProviderNotesRequest(BaseModel):
    provider_notes: str = Field(
        ..., 
        description="The medical provider notes to analyze",
        min_length=10,
        example="Patient presents with acute bronchitis. Performed comprehensive examination and prescribed antibiotics."
    )
    
    class Config:
        json_schema_extra = {
            "example": {
                "provider_notes": "Patient presents with acute bronchitis. Cough for 5 days, productive with yellow sputum. Lung exam reveals diffuse wheezing. Prescribed azithromycin 500mg."
            }
        }

class ProviderNote(BaseModel):
    note: str

class CodingResponse(BaseModel):
    icd_codes: List[ICDCode]
    cpt_codes: List[CPTCode]
    overall_summary: str

    class Config:
        json_schema_extra = {
            "example": {
                "icd_codes": [
                    {
                        "code": "J20.9",
                        "description": "Acute bronchitis, unspecified",
                        "explanation": "Patient presents with acute bronchitis as documented in provider notes"
                    }
                ],
                "cpt_codes": [
                    {
                        "code": "99213",
                        "description": "Office visit, established patient",
                        "explanation": "Comprehensive examination performed as documented"
                    }
                ],
                "overall_summary": "Patient encounter for acute bronchitis with examination and treatment"
            }
        }

class FileUploadResponse(BaseModel):
    success: bool
    filename: str
    extracted_text_length: int
    pii_removed: bool
    pii_count: int
    icd_codes: List[ICDCode]
    cpt_codes: List[CPTCode]
    overall_summary: str
    
    class Config:
        json_schema_extra = {
            "example": {
                "success": True,
                "filename": "provider_notes.txt",
                "extracted_text_length": 450,
                "pii_removed": True,
                "pii_count": 3,
                "cpt_codes": ["99213", "93000"],
                "cpt_explanation": "Office visit and EKG",
                "icd_codes": ["I20.0", "R07.9"],
                "icd_explanation": "Unstable angina and chest pain"
            }
        }