evd_agent/models.py

from __future__ import annotations

from datetime import datetime, timezone
from enum import Enum
from typing import Dict, List, Literal, Optional, Set

from pydantic import BaseModel, Field


class CaseClassification(str, Enum):
    NOT_SUSPECTED = "No Case"
    SUSPECTED = "Suspected Case"
    PROBABLE = "Probable Case"


class InterviewStatus(str, Enum):
    IN_PROGRESS = "in_progress"
    COMPLETE = "complete"


class ChatTurn(BaseModel):
    role: Literal["assistant", "clinician"]
    content: str
    timestamp: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))


class PatientFacts(BaseModel):
    temperature_c: Optional[float] = None
    fever_reported: Optional[bool] = None
    sudden_onset_fever: Optional[bool] = None

    headache: Optional[bool] = None
    lethargy: Optional[bool] = None
    loss_of_appetite: Optional[bool] = None
    muscle_pain: Optional[bool] = None
    joint_pain: Optional[bool] = None
    stomach_pain: Optional[bool] = None
    difficulty_swallowing: Optional[bool] = None
    vomiting: Optional[bool] = None
    difficulty_breathing: Optional[bool] = None
    diarrhea: Optional[bool] = None
    hiccups: Optional[bool] = None

    unexplained_bleeding: Optional[bool] = None
    sudden_unexplained_death: Optional[bool] = None
    patient_deceased: Optional[bool] = None

    exposure_known_case_21d: Optional[bool] = None
    exposure_outbreak_area_21d: Optional[bool] = None
    travel_outbreak_area_21d: Optional[bool] = None
    attended_funeral_21d: Optional[bool] = None
    healthcare_worker_exposure_21d: Optional[bool] = None
    epidemiological_link_known_case: Optional[bool] = None

    lab_confirmation_available: Optional[bool] = None
    clinician_assessed_consistent: Optional[bool] = None
    failed_treatment: Optional[bool] = None

    location: Optional[str] = None


class CountyRisk(BaseModel):
    """Kenya county-level risk intelligence for adaptive questioning."""

    county_name: str
    risk_tier: Literal["very_high", "high", "medium", "low"] = "low"
    is_border_county: bool = False
    corridor_flags: List[str] = Field(default_factory=list)
    relevant_poes: List[str] = Field(default_factory=list)
    high_risk_profiles: List[str] = Field(default_factory=list)
    key_risk_factors: List[str] = Field(default_factory=list)
    source_week: Optional[str] = None
    last_updated: Optional[str] = None


class EpidemiologicalContext(BaseModel):
    country: str = "Unknown"
    district: str = "Unknown"
    active_outbreak_districts: List[str] = Field(default_factory=list)
    neighboring_outbreak_districts: List[str] = Field(default_factory=list)
    cross_border_alerts: List[str] = Field(default_factory=list)
    recent_confirmed_cases: int = 0
    community_transmission: bool = False
    health_facility_alerts: List[str] = Field(default_factory=list)
    county_risks: Dict[str, CountyRisk] = Field(default_factory=dict)
    last_updated: Optional[str] = None


class DecisionOutput(BaseModel):
    classification: CaseClassification = CaseClassification.NOT_SUSPECTED
    triggered_rule: str = "No case definition currently met."
    evidence: List[str] = Field(default_factory=list)
    recommended_action: str = "Continue routine triage and monitor for evolving symptoms."
    confidence: float = 0.0
    should_stop_interview: bool = False


class RiskProfile(BaseModel):
    internal_score: float = 0.0
    risk_level: Literal["LOW", "MODERATE", "HIGH", "CRITICAL"] = "LOW"
    dominant_factors: List[str] = Field(default_factory=list)


class LLMInterviewPlan(BaseModel):
    summary_known: str = ""
    missing_evidence: List[str] = Field(default_factory=list)
    fact_updates: PatientFacts = Field(default_factory=PatientFacts)
    evidence_statements: List[str] = Field(default_factory=list)
    reasoning: str = ""
    classification: Optional[CaseClassification] = None
    criteria_matched: List[str] = Field(default_factory=list)
    criteria_not_met: List[str] = Field(default_factory=list)
    triggered_rule: str = ""
    recommended_action: str = ""
    next_question: Optional[str] = None
    should_stop_interview: bool = False
    confidence: float = 0.0


class InterviewState(BaseModel):
    session_id: str
    status: InterviewStatus = InterviewStatus.IN_PROGRESS
    facts: PatientFacts = Field(default_factory=PatientFacts)
    context: EpidemiologicalContext = Field(default_factory=EpidemiologicalContext)
    history: List[ChatTurn] = Field(default_factory=list)
    asked_questions: Set[str] = Field(default_factory=set)
    followup_question_count: int = 0
    pending_question_key: Optional[str] = None
    pending_question_text: Optional[str] = None
    decision: DecisionOutput = Field(default_factory=DecisionOutput)
    risk_profile: RiskProfile = Field(default_factory=RiskProfile)
    llm_summary: Optional[str] = None
    missing_evidence: List[str] = Field(default_factory=list)
    rationale_log: List[str] = Field(default_factory=list)


class TurnResult(BaseModel):
    assistant_message: str
    decision: DecisionOutput
    risk_profile: RiskProfile
    next_question_key: Optional[str] = None
    llm_summary: Optional[str] = None
    state_updates: Dict[str, str] = Field(default_factory=dict)