src/models/response_models.py

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

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):
    cpt_codes: list
    cpt_explanation: str
    icd_codes: list
    icd_explanation: str


# PII Detection Detail Model
class PIIDetail(BaseModel):
    """Details of detected PII entity"""
    entity_type: str = Field(description="Type of PII detected (e.g., PERSON, PHONE_NUMBER)")
    start: int = Field(description="Start position in original text")
    end: int = Field(description="End position in original text")
    score: float = Field(description="Confidence score")


# Updated File Upload Response with PII info
class FileUploadResponse(BaseModel):
    """Response model for file upload endpoint with PII removal info"""
    success: bool = Field(description="Whether file processing was successful")
    filename: str = Field(description="Name of uploaded file")
    extracted_text_length: int = Field(description="Length of extracted text (after PII removal)")
    pii_removed: bool = Field(description="Whether PII was detected and removed")
    pii_count: int = Field(description="Number of PII entities removed")
    cpt_codes: list = Field(description="List of CPT codes")
    cpt_explanation: str = Field(description="Explanation of CPT codes")
    icd_codes: list = Field(description="List of ICD codes")
    icd_explanation: str = Field(description="Explanation of ICD codes")
    
    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"
            }
        }