src/models.py

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

ReportSection = Literal[
    "FINDINGS",
    "IMAGE_CHARACTERISTICS",
    "CLINICAL_INTERPRETATION",
    "MANAGEMENT_CONSIDERATIONS",
    "LIMITATIONS",
    "REFERENCES",
]

RootCause = Literal[
    "CHUNK_ERROR",
    "RETRIEVAL_MISS",
    "LLM_HALLUCINATION",
    "UNKNOWN",
]


class ChunkRef(BaseModel):
    chroma_id: str
    content: str
    source: str
    page: int
    similarity_score: float


class CorrectionRequest(BaseModel):
    session_id: str
    original_report: str
    retrieved_context: List[ChunkRef]
    features: dict
    section: ReportSection
    correction_text: str
    corrected_snippet: Optional[str] = None


class CorrectionResponse(BaseModel):
    correction_id: str
    root_cause: RootCause
    action_taken: str
    affected_chunk_ids: List[str]
    verification_query: Optional[str] = None


class CorrectionRecord(BaseModel):
    correction_id: str
    session_id: str
    timestamp: str
    section: str
    correction_text: str
    corrected_snippet: Optional[str]
    root_cause: str
    action_taken: str
    affected_chunk_ids: str  # JSON-encoded list
    original_report: str
    features_json: str
    retrieved_context_json: str
    verified: bool = False
    verification_result: Optional[str] = None


class VerifyRequest(BaseModel):
    correction_id: str
    features: dict