schema.py

"""
schema.py
---------
Pydantic data contracts for the PayParse NLU pipeline.

Defines:
- Supported intents (`IntentType`)
- Extracted transaction entities (`TransactionEntities`)
- The strict NLU output schema enforced on Gemini (`NLUResult`)
- The multi-turn dialogue state (`DialogueState`)
- FastAPI request/response models (`ChatRequest`, `ChatResponse`)
- The slot-filling contract (`REQUIRED_SLOTS`)
"""

from __future__ import annotations

import re
from enum import Enum
from typing import Optional

from pydantic import BaseModel, Field, field_validator


# ---------------------------------------------------------------------------
# Intents
# ---------------------------------------------------------------------------
class IntentType(str, Enum):
    """Supported user intents (GoPay-style financial actions)."""

    TRANSFER_UANG = "transfer_uang"   # e.g. "tf goceng ke budi"
    BELI_PULSA = "beli_pulsa"         # e.g. "beliin pulsa 25rb ke 0812..."
    BAYAR_PLN = "bayar_pln"           # e.g. "bayar listrik id 4512..."
    PESAN_GOJEK = "pesan_gojek"       # e.g. "gojek ke stasiun"
    PESAN_GOFOOD = "pesan_gofood"     # e.g. "gofood nasi goreng"
    UNKNOWN = "unknown"               # out-of-scope / unintelligible


# Slot-filling contract: which entity fields MUST be present per intent
# before the transaction payload can be confirmed.
REQUIRED_SLOTS: dict[IntentType, list[str]] = {
    IntentType.TRANSFER_UANG: ["recipient", "recipient_phone", "amount"],
    IntentType.BELI_PULSA: ["phone_number", "amount"],
    IntentType.BAYAR_PLN: ["customer_id"],
    IntentType.PESAN_GOJEK: ["tujuan"],
    IntentType.PESAN_GOFOOD: ["makanan"],
    IntentType.UNKNOWN: [],
}


# ---------------------------------------------------------------------------
# Entities
# ---------------------------------------------------------------------------
class TransactionEntities(BaseModel):
    """Entities extracted from a user utterance. All fields optional —
    the state machine decides what is still missing."""

    recipient: Optional[str] = Field(
        default=None,
        description="Name of the transfer recipient, e.g. 'Budi', 'mama'.",
    )
    recipient_phone: Optional[str] = Field(
        default=None,
        description=(
            "Phone number of the transfer recipient, digits only. "
            "Filled by the state manager after resolving target_kontak, "
            "or directly if the user dictates digits."
        ),
    )
    amount: Optional[int] = Field(
        default=None,
        description=(
            "Amount in Indonesian Rupiah as an integer. "
            "Slang must be normalized: goceng=5000, ceban=10000, "
            "gocap=50000, cepek=100000, seceng=1000, '50rb'=50000, "
            "'sejuta'=1000000."
        ),
    )
    phone_number: Optional[str] = Field(
        default=None,
        description="Destination phone number for pulsa top-up, digits only.",
    )
    target_kontak: Optional[str] = Field(
        default=None,
        description=(
            "Contact name or pronoun used INSTEAD of raw digits when the user "
            "refers to a phone number indirectly, e.g. 'Anton', 'nomorku', "
            "'nomor ini'. Resolved to an actual number by the state manager."
        ),
    )
    customer_id: Optional[str] = Field(
        default=None,
        description="PLN customer ID / meter number, digits only.",
    )
    provider: Optional[str] = Field(
        default=None,
        description="Telco provider if mentioned, e.g. 'Telkomsel', 'XL'.",
    )
    asal: Optional[str] = Field(
        default=None,
        description=(
            "Origin location for Gojek ride booking. "
            "Defaults to 'Bogor' if not specified by the user."
        ),
    )
    tujuan: Optional[str] = Field(
        default=None,
        description=(
            "Destination location for Gojek ride booking, "
            "e.g. 'stasiun', 'bandara', 'mall botani'."
        ),
    )
    makanan: Optional[str] = Field(
        default=None,
        description=(
            "Food item ordered via GoFood, "
            "e.g. 'nasi goreng', 'ayam geprek', 'mie ayam'."
        ),
    )

    @field_validator("amount")
    @classmethod
    def amount_must_be_positive(cls, v: Optional[int]) -> Optional[int]:
        if v is not None and v <= 0:
            return None  # discard nonsense amounts instead of erroring
        return v

    @field_validator("phone_number", "customer_id", "recipient_phone")
    @classmethod
    def digits_only(cls, v: Optional[str]) -> Optional[str]:
        """Normalize '0812-3456 7890' -> '081234567890'; drop empty results."""
        if v is None:
            return None
        digits = re.sub(r"\D", "", v)
        return digits or None

    @field_validator("target_kontak")
    @classmethod
    def clean_contact(cls, v: Optional[str]) -> Optional[str]:
        """Trim whitespace; drop empty strings so lookups stay predictable."""
        return v.strip() or None if v else None

    def filled_fields(self) -> set[str]:
        """Names of fields that currently hold a value."""
        return {k for k, v in self.model_dump().items() if v is not None}

    def merge(self, new: "TransactionEntities") -> "TransactionEntities":
        """Non-destructive merge: new values fill gaps or override,
        but never erase an existing value with None."""
        data = self.model_dump()
        for key, value in new.model_dump().items():
            if value is not None:
                data[key] = value
        return TransactionEntities(**data)


# ---------------------------------------------------------------------------
# NLU output (strict schema enforced on Gemini structured output)
# ---------------------------------------------------------------------------
class NLUResult(BaseModel):
    """Single-turn NLU extraction result returned by the LLM."""

    intent: IntentType = Field(description="The classified user intent.")
    entities: TransactionEntities = Field(
        default_factory=TransactionEntities,
        description="Entities explicitly stated in the utterance. "
        "Never invent values that were not said.",
    )
    confidence: float = Field(
        default=0.0,
        ge=0.0,
        le=1.0,
        description="Model self-estimated confidence in the intent (0-1).",
    )
    normalized_text: str = Field(
        default="",
        description="The utterance rewritten in clean formal Indonesian "
        "with slang and typos resolved.",
    )


# ---------------------------------------------------------------------------
# Dialogue state (multi-turn)
# ---------------------------------------------------------------------------
class DialogueStatus(str, Enum):
    COLLECTING = "collecting"   # waiting for missing slots
    CONFIRMING = "confirming"   # all slots filled, awaiting user's ya/tidak
    COMPLETE = "complete"       # confirmed, payload released
    UNKNOWN = "unknown"         # out-of-scope query


class DialogueState(BaseModel):
    """Persistent state of one conversation session."""

    session_id: str
    intent: IntentType = IntentType.UNKNOWN
    entities: TransactionEntities = Field(default_factory=TransactionEntities)
    missing_slots: list[str] = Field(default_factory=list)
    status: DialogueStatus = DialogueStatus.UNKNOWN
    turn_count: int = 0
    history: list[str] = Field(
        default_factory=list,
        description="Raw user utterances in order, for LLM context.",
    )


# ---------------------------------------------------------------------------
# API models
# ---------------------------------------------------------------------------
class ChatRequest(BaseModel):
    session_id: str = Field(min_length=1, description="Client-chosen session key.")
    message: str = Field(min_length=1, description="Raw user utterance.")
    resolve_contacts_locally: bool = Field(
        default=False,
        description=(
            "When true (mobile client), the server will NOT look up its own "
            "contact directory and instead signal the client to resolve "
            "target_kontak into a phone number on-device."
        ),
    )


class ChatResponse(BaseModel):
    session_id: str
    reply: str = Field(description="Assistant reply (follow-up question or confirmation).")
    intent: IntentType
    is_complete: bool
    missing_slots: list[str]
    payload: Optional[dict] = Field(
        default=None,
        description="Final transaction payload, present only when complete.",
    )
    needs_contact_resolution: bool = Field(
        default=False,
        description=(
            "True when the server asks the client to resolve target_kontak "
            "into a phone number on-device (mobile flow). The client should "
            "open a contact picker, find the number, then send it back as "
            "the next message."
        ),
    )
    target_kontak: Optional[str] = Field(
        default=None,
        description=(
            "The contact name/pronoun the client should resolve, present "
            "only when needs_contact_resolution is true."
        ),
    )