src/api/models.py

"""
Pydantic request / response models for the FastAPI backend.

Endpoints and their schemas:
    GET  /api/health                       → HealthResponse
    POST /api/session/create               → CreateSessionResponse
    POST /api/session/{id}/message         → MessageRequest → MessageResponse
    POST /api/session/{id}/upload          → (multipart UploadFile) → UploadResponse
    POST /api/session/{id}/validate-entities → ValidateEntitiesRequest → MessageResponse
    GET  /api/session/{id}/history         → HistoryResponse
    POST /api/classify                     → ClassifyRequest → ClassifyResponse (debug)
    POST /api/extract                      → ExtractRequest  → ExtractResponse  (debug)
"""

from pydantic import BaseModel


# ---------------------------------------------------------------------------
# Shared sub-model
# ---------------------------------------------------------------------------

class EntityOut(BaseModel):
    """Entity span in the same schema as EvidenceNER's Entity dataclass."""
    text: str
    label: str
    start: int
    end: int
    confidence: float


# ---------------------------------------------------------------------------
# Health
# ---------------------------------------------------------------------------

class HealthResponse(BaseModel):
    status: str          # "ok" | "degraded"
    components: dict     # component_name → "ok" | "keyword_fallback" | "rule_fallback" | "error: …"


# ---------------------------------------------------------------------------
# Session lifecycle
# ---------------------------------------------------------------------------

class CreateSessionResponse(BaseModel):
    session_id: str


# ---------------------------------------------------------------------------
# Message
# ---------------------------------------------------------------------------

class MessageRequest(BaseModel):
    text: str


class RedactionSpanOut(BaseModel):
    """One redacted PII span, returned to the user's own browser for the
    side-by-side reveal. Never forwarded to any third-party API."""
    entity_type: str
    original: str
    placeholder: str
    start: int
    end: int


class MessageResponse(BaseModel):
    reply: str
    pii_redacted: bool
    pii_types_found: list[str]
    # Side-by-side reveal payload (Feature 1). Defaults keep older callers /
    # the validate-entities and escalation handlers (no redaction) valid.
    original_text: str = ""
    redacted_text: str = ""
    redactions: list[RedactionSpanOut] = []


# ---------------------------------------------------------------------------
# Upload
# ---------------------------------------------------------------------------

class UploadResponse(BaseModel):
    filename: str
    raw_text: str
    entities: list[EntityOut]


# ---------------------------------------------------------------------------
# HITL validate-entities
# ---------------------------------------------------------------------------

class ValidateEntitiesRequest(BaseModel):
    entities: dict   # {entity_type: corrected_value}, e.g. {"ORG": "HDFC Bank"}


# ---------------------------------------------------------------------------
# History
# ---------------------------------------------------------------------------

class HistoryMessage(BaseModel):
    role: str        # "user" | "assistant"
    content: str


class HistoryResponse(BaseModel):
    session_id: str
    history: list[HistoryMessage]


# ---------------------------------------------------------------------------
# Privacy audit trail
# ---------------------------------------------------------------------------

class AuditEntry(BaseModel):
    """One outbound/processing event in the privacy audit trail."""
    timestamp: str          # ISO-8601 UTC
    event: str              # "outbound_to_anthropic" | "document_local"
    description: str        # human-readable summary
    transmitted_text: str   # exactly what was sent externally (already redacted)
    pii_types_found: list[str]
    pii_count: int
    leak_check: str         # "passed" | "failed" | "n/a"


class AuditResponse(BaseModel):
    session_id: str
    entries: list[AuditEntry]


# ---------------------------------------------------------------------------
# Debug endpoints
# ---------------------------------------------------------------------------

class ClassifyRequest(BaseModel):
    text: str


class ClassifyResponse(BaseModel):
    domain: str
    confidence: float
    all_probs: dict


class ExtractRequest(BaseModel):
    text: str


class ExtractResponse(BaseModel):
    entities: list[EntityOut]