ahagestedt/quickbooks-qa

1"""2Base QA Environment — OpenEnv-compatible environment for QA evaluation.3 4All software-specific environments inherit from BaseQAEnvironment.5This provides the core reset/step logic with mock QA evaluation.6 7Conforms to openenv.core Environment[ActT, ObsT, StateT] API (v0.2.x):8  - reset(seed, episode_id, **kwargs) -> Observation9  - step(action, timeout_s, **kwargs) -> Observation10  - state (property) -> State11"""12 13from __future__ import annotations14 15import logging16import uuid17from pathlib import Path18from typing import Any, Optional19 20from openenv.core.env_server.interfaces import Environment21 22from rl_qa_base.models import QAAction, QAObservation, QAState23from rl_qa_base.data_loader import load_environments24from rl_qa_base.qa_engine import (25    QAEngine,26    generate_mock_ticket,27    generate_mock_response,28    generate_rubric_prompt,29)30 31logger = logging.getLogger(__name__)32 33 34class BaseQAEnvironment(Environment[QAAction, QAObservation, QAState]):35    """36    Parametric QA evaluation environment.37 38    Each instance represents one software product with N workflow tasks.39    Tasks are selected via reset_with_task(task_index=N).40 41    Lifecycle:42        1. reset() -> QAObservation (initial observation with task details)43        2. step(QAAction) -> QAObservation (evaluation result with reward)44 45    Episodes are single-step: one step() call evaluates and returns done=True.46    """47 48    def __init__(49        self,50        software_name: str,51        data_path: str | Path,52        seed: Optional[int] = None,53        transform=None,54    ) -> None:55        super().__init__(transform=transform)56        self.software_name = software_name57        self.data_path = Path(data_path)58        self.tasks = load_environments(self.data_path)59        self.engine = QAEngine(seed=seed)60 61        # Current episode state62        self._current_task: dict[str, Any] = {}63        self._state = QAState()64        self._done = True65        self._current_task_index = 066 67        logger.info(68            "Initialized %s environment with %d tasks from %s",69            software_name, len(self.tasks), data_path,70        )71 72    @property73    def task_count(self) -> int:74        return len(self.tasks)75 76    def get_task_list(self) -> list[dict[str, str]]:77        """Return summary of all available tasks (for /api/tasks endpoint)."""78        return [79            {80                "index": i,81                "env_id": t["env_id"],82                "workflow": t["workflow"],83                "api_endpoint": t["api_endpoint"],84            }85            for i, t in enumerate(self.tasks)86        ]87 88    def reset(89        self,90        seed: Optional[int] = None,91        episode_id: Optional[str] = None,92        **kwargs: Any,93    ) -> QAObservation:94        """Reset the environment and return initial observation.95 96        Args:97            seed: Optional random seed for reproducibility.98            episode_id: Optional custom episode identifier.99            **kwargs: Additional parameters (task_index supported).100        """101        task_index = kwargs.get("task_index", self._current_task_index)102        return self.reset_with_task(task_index=task_index, seed=seed, episode_id=episode_id)103 104    def reset_with_task(105        self,106        task_index: int = 0,107        seed: Optional[int] = None,108        episode_id: Optional[str] = None,109    ) -> QAObservation:110        """Reset with a specific task index.111 112        This is the internal method called by the custom /reset endpoint.113        """114        if seed is not None:115            self.engine = QAEngine(seed=seed)116 117        if task_index < 0 or task_index >= len(self.tasks):118            task_index = 0119 120        self._current_task_index = task_index121        self._current_task = self.tasks[task_index]122        self._done = False123 124        # Build state125        ep_id = episode_id or f"EP-{uuid.uuid4().hex[:8].upper()}"126        self._state = QAState(127            episode_id=ep_id,128            step_count=0,129            environment_def=self._current_task,130            task_index=task_index,131            is_done=False,132        )133 134        # Build observation135        rubric = generate_rubric_prompt(self._current_task)136        observation = QAObservation(137            done=False,138            reward=None,139            task_id=self._current_task["env_id"],140            task_index=task_index,141            software=self._current_task["software"],142            workflow=self._current_task["workflow"],143            api_endpoint=self._current_task.get("api_endpoint", ""),144            expected_outcome=self._current_task.get("expected_outcome", ""),145            instruction=rubric,146            step_number=0,147            max_steps=1,148        )149 150        return self._apply_transform(observation)151 152    def step(153        self,154        action: QAAction,155        timeout_s: Optional[float] = None,156        **kwargs: Any,157    ) -> QAObservation:158        """Execute one evaluation step.159 160        Args:161            action: QAAction with optional ticket_text, response_text.162            timeout_s: Optional timeout (unused, for interface compatibility).163 164        Returns:165            QAObservation with evaluation results and reward.166        """167        if self._done:168            return QAObservation(169                done=True,170                reward=0.0,171                task_id=self._current_task.get("env_id", ""),172                software=self._current_task.get("software", ""),173                workflow=self._current_task.get("workflow", ""),174                step_number=1,175                max_steps=1,176            )177 178        # Handle finish action179        if action.action_type == "finish":180            self._done = True181            self._state.is_done = True182            return QAObservation(183                done=True,184                reward=0.0,185                task_id=self._current_task["env_id"],186                software=self._current_task["software"],187                workflow=self._current_task["workflow"],188                step_number=1,189                max_steps=1,190            )191 192        # Run QA evaluation193        result = self.engine.evaluate_mock(194            env_def=self._current_task,195            ticket_text=action.ticket_text,196            response_text=action.response_text,197        )198 199        # Compute reward and gating200        reward = self.engine.compute_reward(result)201        gating = self.engine.compute_gating(result.overall_confidence)202 203        # Update state204        self._state.step_count += 1205        self._state.ticket_text = action.ticket_text or generate_mock_ticket(self._current_task)206        self._state.response_text = action.response_text or generate_mock_response(self._current_task)207        self._state.customer_tier = action.customer_tier or "standard"208        self._state.reasoning_steps = [209            {210                "step_name": s.step_name,211                "reasoning": s.reasoning,212                "confidence": s.confidence,213            }214            for s in result.reasoning_steps215        ]216        self._state.grade = {217            "overall_score": result.grade.overall_score,218            "safety_violation": result.grade.safety_violation,219            "summary": result.grade.summary,220        }221        self._state.overall_confidence = result.overall_confidence222        self._state.gating_decision = gating223        self._state.reward_value = reward224        self._state.is_done = True225        self._done = True226 227        # Build observation with evaluation results228        observation = QAObservation(229            done=True,230            reward=reward,231            task_id=self._current_task["env_id"],232            task_index=self._state.task_index,233            software=self._current_task["software"],234            workflow=self._current_task["workflow"],235            api_endpoint=self._current_task.get("api_endpoint", ""),236            expected_outcome=self._current_task.get("expected_outcome", ""),237            instruction=generate_rubric_prompt(self._current_task),238            ticket=self._state.ticket_text,239            annotator_response=self._state.response_text,240            evaluation={241                "reasoning_steps": self._state.reasoning_steps,242                "grade": self._state.grade,243                "overall_confidence": result.overall_confidence,244                "model_name": result.model_name,245            },246            gating_decision=gating,247            step_number=1,248            max_steps=1,249        )250 251        return self._apply_transform(observation)252 253    @property254    def state(self) -> QAState:255        """Get the current environment state."""256        return self._state257