models.py

"""
Data models for LedgerShield.

Defines the core dataclasses and Pydantic models used throughout the
LedgerShield benchmark, including:

- **Type aliases**: Domain-specific Literal types for actions, decisions,
  and task families.
- **LedgerShieldReward**: Pydantic model for structured reward payloads.
- **ToolResult**: Result of a single tool invocation.
- **CaseDecision**: Agent's final decision submission.
- **LedgerShieldAction**: Gymnasium-style action (action_type + payload).
- **LedgerShieldObservation**: Full observation at each step.
- **LedgerShieldState**: Internal episode state (not visible to agent).

TypedDict Internal Returns (Phase 4.5):
    StepResultDict and ScoreBreakdownDict provide formalized typing
    for internal return values.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Literal, TypedDict

from openenv_compat import Action, Observation, State
from pydantic import BaseModel, Field

# ── Type aliases ─────────────────────────────────────────────────────────────

InvestigationActionType = Literal[
    "zoom",
    "get_doc_crop",
    "ocr",
    "lookup_vendor",
    "lookup_vendor_history",
    "lookup_policy",
    "lookup_po",
    "lookup_receipt",
    "search_ledger",
    "inspect_email_thread",
    "compare_bank_account",
]

InterventionActionType = Literal[
    "request_callback_verification",
    "freeze_vendor_profile",
    "request_bank_change_approval_chain",
    "request_po_reconciliation",
    "request_additional_receipt_evidence",
    "route_to_procurement",
    "route_to_security",
    "flag_duplicate_cluster_review",
    "create_human_handoff",
]

ActionType = Literal[
    "zoom",
    "get_doc_crop",
    "ocr",
    "lookup_vendor",
    "lookup_vendor_history",
    "lookup_policy",
    "lookup_po",
    "lookup_receipt",
    "search_ledger",
    "inspect_email_thread",
    "compare_bank_account",
    "request_callback_verification",
    "freeze_vendor_profile",
    "request_bank_change_approval_chain",
    "request_po_reconciliation",
    "request_additional_receipt_evidence",
    "route_to_procurement",
    "route_to_security",
    "flag_duplicate_cluster_review",
    "create_human_handoff",
    "submit_decision",
]

DecisionType = Literal["PAY", "HOLD", "NEEDS_REVIEW", "ESCALATE_FRAUD"]
TaskType = Literal["task_a", "task_b", "task_c", "task_d", "task_e"]


# ── TypedDict for internal returns (Phase 4.5) ──────────────────────────────

class StepResultDict(TypedDict, total=False):
    """Typed dictionary for step() return payloads."""
    observation: dict[str, Any]
    reward: float
    done: bool
    truncated: bool
    terminated: bool
    info: dict[str, Any]


class ScoreBreakdownDict(TypedDict, total=False):
    """Typed dictionary for score_submission() breakdown."""
    field_score: float
    line_item_score: float
    evidence_score: float
    decision_score: float
    discrepancy_score: float
    duplicate_score: float
    fraud_score: float
    reason_score: float
    policy_score: float
    counterfactual_score: float
    investigation_score: float
    intervention_score: float
    resolution_state_score: float
    calibration_score: float
    efficiency_score: float
    outcome_score: float
    pressure_event_score: float
    callback_interpretation_score: float
    cross_invoice_link_score: float
    campaign_detection_score: float
    compliance_score: float
    compliance_adjustment: float
    compliance_penalty: float
    currency_validation_score: float
    currency_adjustment: float
    cross_invoice_link_matches: float
    counterfactual_doc_refs: float
    degenerate_penalty: float
    error: float


# ── Pydantic models ─────────────────────────────────────────────────────────

class LedgerShieldReward(BaseModel):
    """Structured reward payload returned at each step.

    Attributes:
        value: Scalar reward value (may be shaped or terminal).
        terminal: Whether this is the episode-ending reward.
        components: Breakdown by reward source (shaping, cost, etc.).
        metadata: Additional context (action_type, terminal_reason, etc.).
    """
    value: float
    terminal: bool = False
    components: dict[str, float] = Field(default_factory=dict)
    metadata: dict[str, Any] = Field(default_factory=dict)


# ── Dataclasses ──────────────────────────────────────────────────────────────

@dataclass
class ToolResult:
    """Result of a single tool invocation.

    Attributes:
        tool_name: Name of the tool that was called.
        success: Whether the call succeeded.
        payload: Returned data from the tool.
        cost: Budget cost of the call.
        message: Human-readable result message.
        novel_signal_count: New risk signals discovered by this call.
        revealed_artifact_ids: Artifact IDs revealed by this call.
    """
    tool_name: str
    success: bool
    payload: dict[str, Any] = field(default_factory=dict)
    cost: float = 0.0
    message: str = ""
    novel_signal_count: int = 0
    revealed_artifact_ids: list[str] = field(default_factory=list)


@dataclass
class CaseDecision:
    """Agent's final decision submission for a case.

    Contains all the structured outputs the agent must produce,
    including the decision, supporting evidence, risk assessments,
    policy checks, and counterfactual reasoning.

    Attributes:
        case_id: The case being decided.
        decision: One of PAY, HOLD, NEEDS_REVIEW, ESCALATE_FRAUD.
        risk_score: Agent's assessed risk level (0.0–1.0).
        confidence: Agent's confidence in its decision (0.0–1.0).
        extracted_fields: Key-value pairs extracted from documents.
        line_items: Itemized list of invoice line items.
        discrepancies: List of identified discrepancies.
        duplicate_links: IDs of potential duplicate invoices.
        fraud_flags: Identified fraud indicator types.
        reason_codes: Canonical reason codes for the decision.
        policy_checks: Policy verification results.
        evidence_map: Evidence references keyed by claim type.
        counterfactual: Hypothetical alternative scenario analysis.
        notes: Free-text investigation notes.
        recommended_next_action: Suggested follow-up action.
        handoff_packet: Structured data for human handoff.
        intervention_log: Record of intervention actions taken.
    """
    case_id: str
    decision: DecisionType
    risk_score: float = 0.0
    confidence: float = 0.5
    extracted_fields: dict[str, Any] = field(default_factory=dict)
    line_items: list[dict[str, Any]] = field(default_factory=list)
    discrepancies: list[str] = field(default_factory=list)
    duplicate_links: list[str] = field(default_factory=list)
    fraud_flags: list[str] = field(default_factory=list)
    reason_codes: list[str] = field(default_factory=list)
    policy_checks: dict[str, str] = field(default_factory=dict)
    evidence_map: dict[str, Any] = field(default_factory=dict)
    counterfactual: str = ""
    notes: str = ""
    recommended_next_action: str = ""
    handoff_packet: dict[str, Any] = field(default_factory=dict)
    intervention_log: list[dict[str, Any]] = field(default_factory=list)


@dataclass
class LedgerShieldAction(Action):
    """Agent action consisting of an action type and payload.

    Attributes:
        action_type: Which tool/intervention/decision to invoke.
        payload: Tool-specific parameters.
    """
    action_type: ActionType
    payload: dict[str, Any] = field(default_factory=dict)


@dataclass
class LedgerShieldObservation(Observation):
    """Full observation available to the agent at each step.

    Contains everything the agent can see: documents, artifacts,
    budget status, risk signals, and the last tool result.

    Attributes:
        case_id: Current case identifier.
        task_type: Task family (task_a through task_e).
        instruction: Natural language task instruction.
        visible_documents: Catalog of visible document metadata.
        revealed_artifacts: List of investigation artifacts.
        pending_events: Async events waiting to resolve.
        budget_remaining: Remaining investigation budget.
        budget_total: Total budget for the episode.
        step_count: Current step number.
        max_steps: Maximum allowed steps.
        case_clock: Logical clock for the case.
        risk_snapshot: Current risk signal summary.
        investigation_status: Investigation progress metrics.
        last_tool_result: Result from the most recent action.
        messages: System messages for the agent.
        allowed_actions: List of valid action types.
        available_interventions: List of intervention action types.
        case_metadata: Additional case context (due date, labels).
        portfolio_context: Cross-case portfolio information.
    """
    case_id: str = ""
    task_type: str = ""
    instruction: str = ""
    visible_documents: list[dict[str, Any]] = field(default_factory=list)
    revealed_artifacts: list[dict[str, Any]] = field(default_factory=list)
    pending_events: list[dict[str, Any]] = field(default_factory=list)
    budget_remaining: float = 0.0
    budget_total: float = 0.0
    step_count: int = 0
    max_steps: int = 0
    case_clock: int = 0
    risk_snapshot: dict[str, Any] = field(default_factory=dict)
    investigation_status: dict[str, Any] = field(default_factory=dict)
    last_tool_result: dict[str, Any] = field(default_factory=dict)
    messages: list[str] = field(default_factory=list)
    allowed_actions: list[str] = field(default_factory=list)
    available_interventions: list[str] = field(default_factory=list)
    case_metadata: dict[str, Any] = field(default_factory=dict)
    portfolio_context: dict[str, Any] = field(default_factory=dict)


@dataclass
class LedgerShieldState(State):
    """Internal episode state (not directly visible to agent).

    Tracks everything the environment needs to manage the episode,
    including hidden risk signals, trajectory, and scoring metadata.

    Attributes:
        episode_id: Unique ID for this episode.
        case_id: The loaded case ID.
        task_type: Task family.
        budget_total: Total investigation budget.
        budget_remaining: Remaining budget.
        max_steps: Maximum allowed steps.
        step_count: Current step number.
        case_clock: Logical case clock.
        submitted: Whether a decision has been submitted.
        final_score: Final graded score (set at submission).
        unsafe_outcome: Whether the outcome was unsafe.
        visible_doc_ids: IDs of documents the agent can see.
        revealed_artifact_ids: IDs of revealed investigation artifacts.
        tool_trace: Full trace of all tool calls and results.
        trajectory: Simplified trajectory for grading.
        interventions_taken: List of intervention records.
        observed_risk_signals: Risk signals the agent has discovered.
        hidden_risk_signals: All risk signals (including undiscovered).
        final_outcome: Simulated outcome dict (set at submission).
        handoff_packet: Agent's handoff data for human review.
        pending_event_ids: IDs of pending async events.
        portfolio_metrics: Cross-case portfolio metrics.
        decision_readiness: Computed readiness score (0–1).
        difficulty: Case difficulty level.
        terminal_reason: Why the episode ended.
        pressure_events_seen: IDs of pressure events encountered.
        pressure_resistance_score: Score for resisting adversarial pressure.
        contrastive_pair_id: ID linking contrastive pair cases.
    """
    episode_id: str = ""
    case_id: str = ""
    task_type: str = ""
    budget_total: float = 15.0
    budget_remaining: float = 15.0
    max_steps: int = 20
    step_count: int = 0
    case_clock: int = 0
    submitted: bool = False
    final_score: float = 0.0
    unsafe_outcome: bool = False
    visible_doc_ids: list[str] = field(default_factory=list)
    revealed_artifact_ids: list[str] = field(default_factory=list)
    tool_trace: list[dict[str, Any]] = field(default_factory=list)
    trajectory: list[dict[str, Any]] = field(default_factory=list)
    interventions_taken: list[dict[str, Any]] = field(default_factory=list)
    observed_risk_signals: list[str] = field(default_factory=list)
    hidden_risk_signals: list[str] = field(default_factory=list)
    final_outcome: dict[str, Any] = field(default_factory=dict)
    handoff_packet: dict[str, Any] = field(default_factory=dict)
    pending_event_ids: list[str] = field(default_factory=list)
    portfolio_metrics: dict[str, Any] = field(default_factory=dict)
    decision_readiness: float = 0.0
    difficulty: str = "medium"
    terminal_reason: str = ""
    pressure_events_seen: list[str] = field(default_factory=list)
    pressure_resistance_score: float = 0.0
    contrastive_pair_id: str = ""