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 dataclasses import dataclass

# ── 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
    proper_score: float
    brier_score: float
    log_score: float
    penalized_brier_score: float
    causal_score: float
    causal_association_score: float
    causal_intervention_score: float
    d_separation_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
    certificate_score: float
    certificate_validity_score: float
    certificate_support_score: float
    certificate_stability_score: float
    certificate_minimality_score: float
    certificate_unsupported_claim_rate: float
    certificate_adjustment: float
    institutional_loss_score: float
    degenerate_penalty: float
    error: float


# ── Reward model (lightweight dataclass fallback) ──────────────────────────

@dataclass
class LedgerShieldReward:
    """Structured reward payload returned at each step.

    This lightweight dataclass replaces the previous Pydantic model when
    Pydantic is not available in the runtime. It provides the same fields
    and default factories used throughout the codebase.
    """
    value: float = 0.0
    terminal: bool = False
    components: dict[str, float] = field(default_factory=dict)
    metadata: dict[str, Any] = field(default_factory=dict)
    
    # Provide a Pydantic-compatible serialization shim used in some runtime
    # paths. The project previously relied on BaseModel.model_dump(), so
    # offer a lightweight equivalent to avoid runtime errors when Pydantic
    # is not installed.
    def model_dump(self) -> dict[str, Any]:
        return {
            "value": float(self.value),
            "terminal": bool(self.terminal),
            "components": dict(self.components or {}),
            "metadata": dict(self.metadata or {}),
        }


# ── 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.
        predicted_probabilities: Probability distribution over latent hypotheses.
        counterfactual: Hypothetical alternative scenario analysis.
        decision_certificate: Machine-checkable proof graph tying evidence,
            policies, hypotheses, interventions, and counterfactuals to the
            payment decision.
        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)
    predicted_probabilities: dict[str, float] = field(default_factory=dict)
    counterfactual: str = ""
    decision_certificate: dict[str, Any] = field(default_factory=dict)
    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.
        sprt_state: Public ASHTG SPRT state for the active case.
        tool_rankings: VoI ranking over currently available tools.
        reward_machine: Reward-machine progress snapshot.
        institutional_memory: Public persistent portfolio memory and loss state.
    """
    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)
    sprt_state: dict[str, Any] = field(default_factory=dict)
    tool_rankings: dict[str, Any] = field(default_factory=dict)
    reward_machine: dict[str, Any] = field(default_factory=dict)
    institutional_memory: 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.
        sprt_state: Serialized SPRT state for sequential hypothesis testing.
        tool_rankings: Latest VoI ranking over available tools.
        reward_machine_state: Serialized reward-machine state.
        calibration_running_average: Running calibration proxy across the episode.
        institutional_metrics: Persistent portfolio metrics after this episode.
        decision_certificate_report: Verifier report for the submitted decision
            certificate.
    """
    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 = ""
    sprt_state: dict[str, Any] = field(default_factory=dict)
    tool_rankings: dict[str, Any] = field(default_factory=dict)
    reward_machine_state: dict[str, Any] = field(default_factory=dict)
    calibration_running_average: float = 0.0
    institutional_metrics: dict[str, Any] = field(default_factory=dict)
    decision_certificate_report: dict[str, Any] = field(default_factory=dict)