"""AP Commander — Core AP Clerk Environment (OpenEnv interface: reset / step / state).""" from __future__ import annotations from typing import Optional, Tuple, Dict, Any, List from .models import APObservation, APAction, APReward, DecisionType from .tasks import TASKS, grade_action # OpenEnv base class — use real package if available, else local stub try: from openenv.core.env_server import Environment as _OpenEnvBase except ImportError: class _OpenEnvBase: # type: ignore[no-redef] """Stub matching the OpenEnv Environment interface.""" def reset(self, seed=None, episode_id=None, **kwargs): raise NotImplementedError def step(self, action, timeout_s=None, **kwargs): raise NotImplementedError @property def state(self): raise NotImplementedError # Intermediate actions: do not end the episode, reveal context instead _INTERMEDIATE = frozenset({ DecisionType.QUERY_VENDOR, DecisionType.ESCALATE, DecisionType.HOLD, DecisionType.HYPOTHETICAL, }) class APClerkEnvironment(_OpenEnvBase): """ AP Clerk environment. Intermediate actions (QUERY_VENDOR / ESCALATE / HOLD / HYPOTHETICAL) reveal context without ending the episode. Terminal actions grade and close it. Fixed seed → reproducible episode; None → fresh random each call. """ def __init__(self) -> None: self._task_id: Optional[str] = None self._observation: Optional[APObservation] = None self._step_count: int = 0 self._done: bool = False self._episode_score: float = 0.01 self._max_steps: int = 1 self._context_store: List[str] = [] def reset(self, task_id: str, seed: Optional[int] = None) -> APObservation: if task_id not in TASKS: raise ValueError( f"Unknown task_id {task_id!r}. " f"Valid options: {list(TASKS.keys())}" ) spec = TASKS[task_id] self._task_id = task_id self._step_count = 0 self._done = False self._episode_score = 0.01 self._max_steps = spec.max_steps obs = spec.generator(seed=seed) obs.step_count = 0 obs.max_steps = self._max_steps # Context notes pre-generated by the generator are hidden until revealed # by an intermediate action (ESCALATE / QUERY_VENDOR / HOLD). self._context_store = list(obs.context_notes) obs.context_notes = [] obs.action_history = [] self._observation = obs return obs def step(self, action: APAction) -> Tuple[APObservation, APReward, bool, Dict[str, Any]]: if self._observation is None: raise RuntimeError("Call reset(task_id) before step().") if self._done: raise RuntimeError("Episode already finished. Call reset() to start a new one.") self._step_count += 1 is_intermediate = action.decision in _INTERMEDIATE at_limit = self._step_count >= self._max_steps # ── HYPOTHETICAL action (training self-play, no context reveal) ───────── if action.decision == DecisionType.HYPOTHETICAL and not at_limit: self._observation.action_history.append({ "step": self._step_count, "decision": "HYPOTHETICAL", "reason_code": action.reason_code.value, "explanation": action.explanation, }) self._observation.step_count = self._step_count # Simulate a brief outcome hint without revealing graded context hint = ( f"[HYPOTHETICAL] If you chose {action.reason_code.value}: " f"outcome would depend on whether the invoice data supports it. " "Review the invoice total, PO amounts, and GRN quantities carefully before committing." ) self._observation.context_notes.append(hint) reward = APReward( score=0.01, breakdown={"hypothetical_step": action.reason_code.value}, feedback="Hypothetical path explored. Now commit to a real decision.", done=False, ) info: Dict[str, Any] = { "task_id": self._task_id, "step_count": self._step_count, "episode_score": 0.01, "hypothetical": True, } return self._observation, reward, False, info # ── Intermediate action (episode continues) ─────────────────────────── if is_intermediate and not at_limit: if self._context_store: prefix_map = { DecisionType.ESCALATE: "[MANAGER]", DecisionType.QUERY_VENDOR: "[VENDOR]", DecisionType.HOLD: "[COMPLIANCE]", } preferred_prefix = prefix_map.get(action.decision, "") # Prefer context note matching the action type; fall back to FIFO matched_idx = next( (i for i, n in enumerate(self._context_store) if preferred_prefix and n.startswith(preferred_prefix)), 0, ) note = self._context_store.pop(matched_idx) self._observation.context_notes.append(note) # History lets graders award the process bonus for correct sequences self._observation.action_history.append({ "step": self._step_count, "decision": action.decision.value, "reason_code": action.reason_code.value, "explanation": action.explanation, }) self._observation.step_count = self._step_count steps_left = self._max_steps - self._step_count context_revealed = self._observation.context_notes[-1] \ if self._observation.context_notes else None reward = APReward( score=0.01, breakdown={ "intermediate_step": action.decision.value, "steps_remaining": steps_left, }, feedback=( f"Intermediate step recorded: {action.decision.value}. " + (f"Context revealed — {context_revealed}" if context_revealed else "No additional context available. Proceed to final decision.") ), done=False, ) info = { "task_id": self._task_id, "step_count": self._step_count, "episode_score": 0.01, "intermediate": True, } return self._observation, reward, False, info # ── Terminal action (or forced terminal at step limit) ──────────────── self._observation.action_history.append({ "step": self._step_count, "decision": action.decision.value, "reason_code": action.reason_code.value, "explanation": action.explanation, }) reward = grade_action(self._task_id, self._observation, action) # Clamp to open interval (0, 1) as required by the evaluator clamped_score = max(0.01, min(0.99, reward.score)) reward = APReward( score=clamped_score, breakdown=reward.breakdown, feedback=reward.feedback, done=reward.done, ) self._done = True self._episode_score = reward.score self._observation.step_count = self._step_count info = { "task_id": self._task_id, "step_count": self._step_count, "episode_score": self._episode_score, } return self._observation, reward, self._done, info @property def state(self) -> Dict[str, Any]: return { "task_id": self._task_id, "step_count": self._step_count, "done": self._done, "episode_score": self._episode_score, "current_observation": self._observation, } @staticmethod def list_tasks() -> Dict[str, Dict[str, str]]: return { tid: { "name": spec.name, "difficulty": spec.difficulty, "description": spec.description, } for tid, spec in TASKS.items() }