interview_schema.py

"""Typed contract between the interview model personas and the existing intake schema.

ExtractedIntake is the narrow, validated surface the Intake-Extractor persona may
fill; extracted_to_intake() bridges it into the deterministic pipeline's
PatientProfile + StructuredIntake. Extractor strings are MODEL OUTPUT and must
clear model_text_is_safe before they may impersonate patient-reported intake.
"""
from __future__ import annotations

from typing import Any, Literal

from pydantic import BaseModel, ConfigDict, Field

from schema import PatientProfile, StructuredIntake, model_text_is_safe


# Mirrors app.CHECK_MAP values / StructuredIntake booleans. Kept as a module
# constant (not imported from app) to avoid an import cycle with the Gradio app;
# tests/test_interview_extract.py asserts equality with CHECK_MAP.
INTAKE_BOOL_FIELDS: tuple[str, ...] = (
    "biting_pain",
    "hot_cold_sensitivity",
    "pain_prevents_sleep",
    "swelling",
    "rapidly_spreading_swelling",
    "fever_or_unwell",
    "breathing_or_swallowing_issue",
    "limited_opening_or_locked_jaw",
    "loose_crown_or_bridge",
    "trauma_or_sudden_bite_change",
    "numbness_or_neuro_symptoms",
    "chest_pain_or_jaw_pain_with_exertion",
    "jaw_pain_with_chewing_relieved_by_rest",
    "vision_scalp_or_new_headache",
    "gum_pimple_or_drainage",
    "bruising_or_burning_after_root_canal",
)

# Extractor-authored free-text fields that flow into StructuredIntake / profile.
_GUARDED_TEXT_FIELDS = (
    "name",
    "chief_concern",
    "tooth_or_area",
    "recent_dental_work",
    "symptom_duration",
    "meds",
    "allergies",
    "goals",
)


ODIPARAAxis = Literal[
    "onset",
    "duration",
    "intensity",
    "progression",
    "aggravating",
    "relieving",
    "associated",
]

DentalDetail = Literal[
    "location",
    "radiation",
    "character",
    "treatment_or_trauma_context",
    "constant_or_episodic",
    "episode_or_lingering_duration",
    "day_or_night_pattern",
    "current_and_worst_intensity",
    "functional_impact",
    "change_over_time",
    "spontaneous_or_provoked",
    "thermal_trigger",
    "bite_or_jaw_trigger",
    "relief_measures",
    "local_associated_symptoms",
    "regional_or_jaw_symptoms",
]

InterviewAxis = Literal[
    "",
    "chief_concern",
    "character_radiation",
    "onset",
    "duration",
    "intensity",
    "progression",
    "aggravating",
    "relieving",
    "associated",
    "dental_history",
    "medical_history",
    "red_flag_infection",
    "red_flag_airway",
    "goals",
]


class NextQuestion(BaseModel):
    """One History-Taker turn: a single intake question, never a conclusion."""

    model_config = ConfigDict(extra="ignore")

    question: str = Field(default="", max_length=300)
    axis: InterviewAxis = ""
    covered_axes: list[ODIPARAAxis] = Field(default_factory=list)
    covered_details: list[DentalDetail] = Field(default_factory=list)


class ExtractedIntake(BaseModel):
    """Everything the Intake-Extractor persona may report from the transcript."""

    model_config = ConfigDict(extra="ignore")

    name: str = ""
    chief_concern: str = ""
    tooth_or_area: str = ""
    recent_dental_work: str = ""
    symptom_duration: str = ""
    pain_score: int = Field(default=0, ge=0, le=10)

    biting_pain: bool = False
    hot_cold_sensitivity: bool = False
    pain_prevents_sleep: bool = False
    swelling: bool = False
    rapidly_spreading_swelling: bool = False
    fever_or_unwell: bool = False
    breathing_or_swallowing_issue: bool = False
    limited_opening_or_locked_jaw: bool = False
    loose_crown_or_bridge: bool = False
    trauma_or_sudden_bite_change: bool = False
    numbness_or_neuro_symptoms: bool = False
    chest_pain_or_jaw_pain_with_exertion: bool = False
    jaw_pain_with_chewing_relieved_by_rest: bool = False
    vision_scalp_or_new_headache: bool = False
    gum_pimple_or_drainage: bool = False
    bruising_or_burning_after_root_canal: bool = False

    age: int | None = Field(default=None, ge=0, le=120)
    language: Literal["English", "Arabic", "Bilingual"] = "English"
    meds: str = ""
    allergies: str = ""
    goals: str = ""


def _strict_schema(model_cls: type[BaseModel]) -> dict[str, Any]:
    """JSON schema for vLLM/xgrammar structured outputs: closed object."""
    schema = model_cls.model_json_schema()
    schema["additionalProperties"] = False
    return schema


def intake_json_schema() -> dict[str, Any]:
    return _strict_schema(ExtractedIntake)


def next_question_json_schema() -> dict[str, Any]:
    return _strict_schema(NextQuestion)


def _guarded(value: str) -> str:
    """Blank extractor text that reads as diagnosis/treatment, keep the rest.

    Blanking (not erroring) is fail-safe here: every guarded field has a
    deterministic downstream fallback, and the raw patient transcript — which is
    allowed to contain anything — still reaches the rules engine as the story.
    """
    cleaned = (value or "").strip()
    if not cleaned:
        return ""
    return cleaned if model_text_is_safe(cleaned) else ""


def extracted_to_intake(
    extracted: ExtractedIntake,
) -> tuple[PatientProfile, StructuredIntake]:
    """Bridge validated extractor output into the deterministic pipeline types.

    Booleans pass through unguarded: they only feed evaluate_red_flags, where a
    spurious True can over-escalate (fail-safe direction) but never suppress a
    flag, because story-text matching runs independently of the booleans.
    """
    guarded = {field: _guarded(getattr(extracted, field)) for field in _GUARDED_TEXT_FIELDS}

    profile = PatientProfile(
        name=guarded["name"],
        age=extracted.age,
        language=extracted.language,
        meds=guarded["meds"],
        allergies=guarded["allergies"],
        goals=guarded["goals"],
    )
    intake_values: dict[str, Any] = {
        "chief_concern": guarded["chief_concern"],
        "tooth_or_area": guarded["tooth_or_area"],
        "recent_dental_work": guarded["recent_dental_work"],
        "symptom_duration": guarded["symptom_duration"],
        "pain_score": extracted.pain_score,
    }
    for field in INTAKE_BOOL_FIELDS:
        intake_values[field] = getattr(extracted, field)
    return profile, StructuredIntake(**intake_values)