client.py

"""Python client for the ClaimSense adjudication gym.

Wraps OpenEnv's HTTP client so notebooks can talk to a remote Space
without crafting JSON manually::

    async with AdjudicatorClient("https://your-space.hf.space") as env:
        obs = await env.reset()
        result = await env.step(lookup_policy())

Convenience builders at the bottom (``lookup_policy``, ``risk_score``,
``settle`` …) save one import per call site. Their action_type strings
match the gym's vocabulary exactly.
"""

from __future__ import annotations

from typing import Any, Optional

from openenv.core import EnvClient
from openenv.core.env_client import StepResult

from .models import AdjudicatorAction, AdjudicatorObservation, AdjudicatorState


# ---------------------------------------------------------------------------
# Client
# ---------------------------------------------------------------------------


class AdjudicatorClient(
    EnvClient[AdjudicatorAction, AdjudicatorObservation, AdjudicatorState]
):
    """Thin OpenEnv client with typed payloads for the adjudication gym."""

    # OpenEnv asks subclasses how to serialise/deserialise.

    def _step_payload(self, action: AdjudicatorAction) -> dict[str, Any]:
        return {
            "action_type": action.action_type,
            "claim_id": action.claim_id,
            "parameters": action.parameters,
        }

    def _parse_result(
        self, payload: dict[str, Any]
    ) -> StepResult[AdjudicatorObservation]:
        body = payload.get("observation", payload)
        observation = AdjudicatorObservation(
            claim_id=body.get("claim_id", ""),
            claim_type=body.get("claim_type", ""),
            claim_amount_requested=body.get("claim_amount_requested", 0.0),
            claimant_name=body.get("claimant_name", ""),
            incident_date=body.get("incident_date", ""),
            description=body.get("description", ""),
            system_response=body.get("system_response", ""),
            action_success=body.get("action_success", True),
            revealed_info=body.get("revealed_info", {}),
            available_actions=body.get("available_actions", []),
            time_elapsed_minutes=body.get("time_elapsed_minutes", 0),
            queries_made=body.get("queries_made", 0),
            is_terminal=body.get("is_terminal", False),
            terminal_reason=body.get("terminal_reason", ""),
        )
        return StepResult(
            observation=observation,
            reward=payload.get("reward", 0.0),
            done=observation.is_terminal,
        )

    def _parse_state(self, payload: dict[str, Any]) -> AdjudicatorState:
        return AdjudicatorState(
            episode_id=payload.get("episode_id", ""),
            claim_id=payload.get("claim_id", ""),
            claim_type=payload.get("claim_type", ""),
            claim_amount_requested=payload.get("claim_amount_requested", 0.0),
            actions_taken=payload.get("actions_taken", 0),
            queries_made=payload.get("queries_made", 0),
            time_elapsed_minutes=payload.get("time_elapsed_minutes", 0),
            total_reward=payload.get("total_reward", 0.0),
        )


# ---------------------------------------------------------------------------
# Action builders
# ---------------------------------------------------------------------------


def _build(
    verb: str,
    *,
    claim_id: str = "",
    parameters: Optional[dict[str, Any]] = None,
) -> AdjudicatorAction:
    """Internal helper: tighten the boilerplate of constructing actions."""
    return AdjudicatorAction(
        action_type=verb,
        claim_id=claim_id,
        parameters=parameters or {},
    )


def lookup_policy(claim_id: str = "") -> AdjudicatorAction:
    """Ask the policy registry for coverage details."""
    return _build("query_policy", claim_id=claim_id)


def pull_history(claim_id: str = "") -> AdjudicatorAction:
    """Pull the claimant's historical claim record."""
    return _build("query_claim_history", claim_id=claim_id)


def risk_score(claim_id: str = "") -> AdjudicatorAction:
    """Run the fraud-scoring engine."""
    return _build("check_fraud", claim_id=claim_id)


def request_evidence(
    doc_types: list[str], claim_id: str = ""
) -> AdjudicatorAction:
    """Request supporting documents (photos, reports, …)."""
    return _build(
        "request_documents",
        claim_id=claim_id,
        parameters={"doc_types": list(doc_types)},
    )


def confirm_coverage(damage_type: str, claim_id: str = "") -> AdjudicatorAction:
    """Verify whether a particular damage type is covered."""
    return _build(
        "verify_coverage",
        claim_id=claim_id,
        parameters={"damage_type": damage_type},
    )


def audit_transactions(claim_id: str = "") -> AdjudicatorAction:
    """Cross-reference the claim with bank-feed (Plaid) transactions."""
    return _build("verify_purchase", claim_id=claim_id)


def compute_settlement(amount: float, claim_id: str = "") -> AdjudicatorAction:
    """Apply deductible and limit to compute the canonical payout."""
    return _build(
        "calculate_payout",
        claim_id=claim_id,
        parameters={"amount": amount},
    )


def settle(
    payout: float, reason: str = "Claim approved", claim_id: str = ""
) -> AdjudicatorAction:
    """Terminal: approve the claim with the supplied payout."""
    return _build(
        "approve",
        claim_id=claim_id,
        parameters={"payout": payout, "reason": reason},
    )


def reject(reason: str = "Claim denied", claim_id: str = "") -> AdjudicatorAction:
    """Terminal: deny the claim with a reason."""
    return _build("deny", claim_id=claim_id, parameters={"reason": reason})


def route_to_supervisor(
    reason: str = "Requires senior review", claim_id: str = ""
) -> AdjudicatorAction:
    """Terminal: hand the claim to a senior adjuster."""
    return _build("escalate", claim_id=claim_id, parameters={"reason": reason})


# ---------------------------------------------------------------------------
# Backwards-compatible aliases (legacy names from the original release)
# ---------------------------------------------------------------------------

ClaimsEnv = AdjudicatorClient
query_policy = lookup_policy
query_claim_history = pull_history
check_fraud = risk_score
request_documents = request_evidence
verify_coverage = confirm_coverage
verify_purchase = audit_transactions
calculate_payout = compute_settlement
approve = settle
deny = reject
escalate = route_to_supervisor


__all__ = [
    "AdjudicatorClient",
    "ClaimsEnv",
    "lookup_policy",
    "pull_history",
    "risk_score",
    "request_evidence",
    "confirm_coverage",
    "audit_transactions",
    "compute_settlement",
    "settle",
    "reject",
    "route_to_supervisor",
    # legacy
    "query_policy",
    "query_claim_history",
    "check_fraud",
    "request_documents",
    "verify_coverage",
    "verify_purchase",
    "calculate_payout",
    "approve",
    "deny",
    "escalate",
]