support_ticket_env/models.py

from __future__ import annotations

from typing import Annotated, Any, Literal, TypeAlias

from pydantic import BaseModel, ConfigDict, Field, TypeAdapter


class SearchKBAction(BaseModel):
    model_config = ConfigDict(extra="forbid")

    action_type: Literal["search_kb"]
    query: str = Field(min_length=1)


class LookupAccountAction(BaseModel):
    model_config = ConfigDict(extra="forbid")

    action_type: Literal["lookup_account"]
    customer_id: str = Field(min_length=1)


class SendReplyAction(BaseModel):
    model_config = ConfigDict(extra="forbid")

    action_type: Literal["send_reply"]
    message: str = Field(min_length=1)


class IssueRefundAction(BaseModel):
    model_config = ConfigDict(extra="forbid")

    action_type: Literal["issue_refund"]
    amount_cents: int = Field(gt=0)
    reason_code: Literal["duplicate_charge"]


class ResolveTicketAction(BaseModel):
    model_config = ConfigDict(extra="forbid")

    action_type: Literal["resolve_ticket"]
    resolution_code: Literal["password_reset_guidance", "billing_refund_processed"]


class EscalateTicketAction(BaseModel):
    model_config = ConfigDict(extra="forbid")

    action_type: Literal["escalate_ticket"]
    queue: Literal["support_lead", "legal_data_incident"]
    priority: Literal["P2", "P0"]
    summary: str = Field(min_length=1)


SupportTicketAction: TypeAlias = Annotated[
    SearchKBAction
    | LookupAccountAction
    | SendReplyAction
    | IssueRefundAction
    | ResolveTicketAction
    | EscalateTicketAction,
    Field(discriminator="action_type"),
]

ACTION_ADAPTER = TypeAdapter(SupportTicketAction)

ACTION_TYPE_NAMES = [
    "search_kb",
    "lookup_account",
    "send_reply",
    "issue_refund",
    "resolve_ticket",
    "escalate_ticket",
]


def parse_action(value: SupportTicketAction | dict[str, Any]) -> SupportTicketAction:
    return ACTION_ADAPTER.validate_python(value)


class ConversationTurn(BaseModel):
    model_config = ConfigDict(extra="forbid")

    role: Literal["customer", "agent"]
    message: str
    step_index: int = Field(ge=0)


class KBSearchResult(BaseModel):
    model_config = ConfigDict(extra="forbid")

    tool_name: Literal["search_kb"]
    success: bool
    query: str
    article_ids: list[str] = Field(default_factory=list)
    snippets: list[str] = Field(default_factory=list)
    message: str | None = None


class AccountLookupResult(BaseModel):
    model_config = ConfigDict(extra="forbid")

    tool_name: Literal["lookup_account"]
    success: bool
    customer_id: str
    account_summary: dict[str, Any] = Field(default_factory=dict)
    message: str | None = None


class ReplyResult(BaseModel):
    model_config = ConfigDict(extra="forbid")

    tool_name: Literal["send_reply"]
    success: bool
    message_preview: str
    message: str | None = None


class RefundResult(BaseModel):
    model_config = ConfigDict(extra="forbid")

    tool_name: Literal["issue_refund"]
    success: bool
    refunded: bool
    amount_cents: int
    reason_code: str
    message: str | None = None


class ResolveResult(BaseModel):
    model_config = ConfigDict(extra="forbid")

    tool_name: Literal["resolve_ticket"]
    success: bool
    resolution_code: str
    ticket_status: Literal["resolved"]
    message: str | None = None


class EscalationResult(BaseModel):
    model_config = ConfigDict(extra="forbid")

    tool_name: Literal["escalate_ticket"]
    success: bool
    queue: str
    priority: str
    summary: str
    ticket_status: Literal["escalated"]
    message: str | None = None


class ErrorToolResult(BaseModel):
    model_config = ConfigDict(extra="forbid")

    tool_name: Literal["error"]
    success: Literal[False]
    error_code: str
    message: str


ToolResult: TypeAlias = Annotated[
    KBSearchResult
    | AccountLookupResult
    | ReplyResult
    | RefundResult
    | ResolveResult
    | EscalationResult
    | ErrorToolResult,
    Field(discriminator="tool_name"),
]


class ScoreCriterion(BaseModel):
    model_config = ConfigDict(extra="forbid")

    criterion_id: str
    label: str
    weight: float = Field(ge=0.0, le=1.0)
    earned: bool
    contribution: float = Field(ge=0.0, le=1.0)


class TaskScorecard(BaseModel):
    model_config = ConfigDict(extra="forbid")

    task_id: str
    score: float = Field(ge=0.0, le=1.0)
    criteria: list[ScoreCriterion]


class SupportTicketObservation(BaseModel):
    model_config = ConfigDict(extra="forbid")

    task_id: str
    ticket_id: str
    ticket_status: Literal["open", "resolved", "escalated"]
    customer_id: str
    organization_name: str
    subject: str
    customer_message: str
    conversation_history: list[ConversationTurn]
    last_tool_result: ToolResult | None = None
    steps_taken: int = Field(ge=0)
    steps_remaining: int = Field(ge=0)
    available_action_types: list[str]
    last_action_error: str | None = None
    known_facts: dict[str, Any] = Field(default_factory=dict)


class SupportTicketStepResult(BaseModel):
    model_config = ConfigDict(extra="forbid")

    observation: SupportTicketObservation
    reward: float
    done: bool
    info: dict[str, Any] = Field(default_factory=dict)