client.py

"""ClaimSense Pro client + 12 verb builder helpers.

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

    from client import ClaimsEnvProClient, open_claim, verify_purchase, approve

    async with ClaimsEnvProClient("https://akhiilll-claims-env-pro.hf.space") as env:
        obs = await env.reset(seed=2)
        obs = await env.step(open_claim(1))
        obs = await env.step(verify_purchase())
        obs = await env.step(approve(8500))
"""

from __future__ import annotations

from typing import Dict

from openenv.core import EnvClient
from openenv.core.client_types import StepResult
from openenv.core.env_server.types import State

try:
    from .models import ClaimsAction, ClaimsObservation
except ImportError:
    from models import ClaimsAction, ClaimsObservation  # type: ignore[no-redef]


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


class ClaimsEnvProClient(EnvClient[ClaimsAction, ClaimsObservation, State]):
    """Typed OpenEnv client for the 12-verb adjudication gym."""

    def _step_payload(self, action: ClaimsAction) -> Dict:
        return {"message": action.message}

    def _parse_result(self, payload: Dict) -> StepResult[ClaimsObservation]:
        body = payload.get("observation", payload)
        observation = ClaimsObservation(
            dashboard=body.get("dashboard", ""),
            step_number=body.get("step_number", 0),
            total_steps=body.get("total_steps", 50),
            claims_in_queue=body.get("claims_in_queue", 0),
            claims_processed=body.get("claims_processed", 0),
            claims_appealed=body.get("claims_appealed", 0),
            active_claim_id=body.get("active_claim_id", -1),
            correct_decisions=body.get("correct_decisions", 0),
            wrong_decisions=body.get("wrong_decisions", 0),
            fraud_caught=body.get("fraud_caught", 0),
            fraud_missed=body.get("fraud_missed", 0),
            accuracy_score=body.get("accuracy_score", 0.0),
            fraud_score=body.get("fraud_score", 0.0),
            payout_accuracy=body.get("payout_accuracy", 0.0),
            task_name=body.get("task_name", ""),
            shift_pressure=body.get("shift_pressure", "normal"),
            dense_step_reward=body.get("dense_step_reward", 0.0),
            final_score=body.get("final_score", 0.0),
            revealed_info=body.get("revealed_info", {}),
            done=payload.get("done", False),
            reward=payload.get("reward"),
            metadata=body.get("metadata", {}),
        )
        return StepResult(
            observation=observation,
            reward=payload.get("reward"),
            done=payload.get("done", False),
        )

    def _parse_state(self, payload: Dict) -> State:
        return State(
            episode_id=payload.get("episode_id"),
            step_count=payload.get("step_count", 0),
        )


# ---------------------------------------------------------------------------
# Verb builders — one helper per verb
# ---------------------------------------------------------------------------


def view_queue() -> ClaimsAction:
    return ClaimsAction(message="VIEW_QUEUE")


def open_claim(claim_id: int) -> ClaimsAction:
    return ClaimsAction(message=f"OPEN_CLAIM {claim_id}")


def review_documents() -> ClaimsAction:
    return ClaimsAction(message="REVIEW_DOCUMENTS")


def check_policy() -> ClaimsAction:
    return ClaimsAction(message="CHECK_POLICY")


def investigate_fraud() -> ClaimsAction:
    return ClaimsAction(message="INVESTIGATE_FRAUD")


def verify_purchase() -> ClaimsAction:
    """NEW in claims-env-pro — Plaid bank-feed audit."""
    return ClaimsAction(message="VERIFY_PURCHASE")


def request_info(doc_type: str) -> ClaimsAction:
    return ClaimsAction(message=f"REQUEST_INFO {doc_type}")


def approve(amount: float) -> ClaimsAction:
    return ClaimsAction(message=f"APPROVE {amount:.2f}")


def deny(reason: str = "fraud_detected") -> ClaimsAction:
    return ClaimsAction(message=f"DENY {reason}")


def escalate() -> ClaimsAction:
    return ClaimsAction(message="ESCALATE")


def handle_appeal(claim_id: int) -> ClaimsAction:
    return ClaimsAction(message=f"HANDLE_APPEAL {claim_id}")


def end_shift() -> ClaimsAction:
    return ClaimsAction(message="END_SHIFT")


def combo(*messages: str) -> ClaimsAction:
    """Compose up to 3 commands into a single semicolon-separated turn."""
    return ClaimsAction(message="; ".join(messages[:3]))


__all__ = [
    "ClaimsEnvProClient",
    "view_queue", "open_claim", "review_documents", "check_policy",
    "investigate_fraud", "verify_purchase", "request_info",
    "approve", "deny", "escalate", "handle_appeal", "end_shift", "combo",
]