ahagestedt/quickbooks-qa

1"""2QuickBooks QA Environment — OpenEnv WebSocket client.3 4Connects to a running QuickBooks QA environment server via WebSocket5and provides a Python API for running episodes.6 7Usage:8    from client import QuickbooksQAClient9 10    with QuickbooksQAClient("http://localhost:8000") as client:11        result = client.reset()12        print(result.observation)13 14        from rl_qa_base.models import QAAction15        result = client.step(QAAction(action_type="evaluate"))16        print(f"Reward: {result.reward}, Done: {result.done}")17"""18 19from __future__ import annotations20 21from typing import Any, Dict22 23from openenv.core import EnvClient24from openenv.core.client_types import StepResult25 26from rl_qa_base.models import QAAction, QAObservation, QAState27 28 29class QuickbooksQAClient(EnvClient[QAAction, QAObservation, QAState]):30    """WebSocket client for the QuickBooks QA environment."""31 32    def _step_payload(self, action: QAAction) -> dict:33        """Convert QAAction to the JSON payload expected by the server."""34        d = action.model_dump(exclude_none=True)35        return d36 37    def _parse_result(self, payload: dict) -> StepResult[QAObservation]:38        """Parse server response into a StepResult[QAObservation]."""39        obs_data = payload.get("observation", payload)40        reward = payload.get("reward", None)41        done = payload.get("done", False)42 43        observation = QAObservation(**obs_data)44        return StepResult(45            observation=observation,46            reward=reward,47            done=done,48        )49 50    def _parse_state(self, payload: dict) -> QAState:51        """Parse server state response into a QAState."""52        return QAState(**payload)53 54 55# ---------------------------------------------------------------------------56# CLI demo57# ---------------------------------------------------------------------------58 59def _demo():60    """Run a quick demo episode."""61    url = "http://localhost:8000"62    print(f"Connecting to {url}...")63 64    with QuickbooksQAClient(url) as client:65        # Reset with first task66        result = client.reset()67        obs = result.observation68        print(f"\n--- Task: {obs.workflow} ({obs.task_id}) ---")69        print(f"Software: {obs.software}")70        print(f"API Endpoint: {obs.api_endpoint}")71 72        # Run evaluation73        action = QAAction(action_type="evaluate")74        result = client.step(action)75        obs = result.observation76        print(f"\n--- Evaluation Result ---")77        print(f"Reward: {result.reward}")78        print(f"Done: {result.done}")79        if obs.evaluation:80            print(f"Score: {obs.evaluation.get('grade', {}).get('overall_score')}/5")81            print(f"Confidence: {obs.evaluation.get('overall_confidence', 0):.3f}")82        print(f"Gating: {obs.gating_decision}")83 84 85if __name__ == "__main__":86    _demo()87