Akshaya0810/invoice-processing-env

1"""2invoice_processing_environment.py3 4Wraps InvoiceProcessingEnv in the official openenv.core.Environment base class.5This gives us /health, /metadata, /schema, /mcp, /reset, /step, /state all 6generated automatically by create_app().7"""8from __future__ import annotations9 10from typing import Any, Optional11 12from openenv.core.env_server.interfaces import Environment13from openenv.core.env_server.types import (14    Action as OpenEnvAction,15    Observation as OpenEnvObservation,16    State as OpenEnvState,17    EnvironmentMetadata,18)19from pydantic import Field20 21# ── Typed models that satisfy OpenEnv base classes ──────────────────────────────22 23class InvoiceAction(OpenEnvAction):24    """One agent action in the invoice processing workflow."""25    action_type: str = Field(26        description=(27            "EXTRACT_FIELDS | MATCH_TO_PO | VERIFY_TAX_COMPLIANCE | "28            "APPROVE_PAYMENT | FLAG_DISCREPANCY | REQUEST_MORE_INFO | REJECT"29        )30    )31    invoice_id: str = Field(description="Target invoice ID")32    payload: dict[str, Any] = Field(33        default_factory=dict,34        description=(35            "Optional payload. FLAG_DISCREPANCY: {flag_reason: str}. "36            "MATCH_TO_PO: {po_id: str}."37        ),38    )39    task_id: str = Field(40        default="easy",41        description="Task to run: easy | medium | hard. Only used on reset.",42    )43 44class InvoiceObservation(OpenEnvObservation):45    """What the agent sees after each step."""46    task_id: str = ""47    inbox: list[dict[str, Any]] = Field(default_factory=list)48    current_invoice_id: Optional[str] = None49    extracted_fields: Optional[dict[str, Any]] = None50    matched_po: Optional[dict[str, Any]] = None51    compliance_flags: list[str] = Field(default_factory=list)52    current_balance: float = 0.053    step_number: int = 054    invoices_resolved: int = 055    invoices_remaining: int = 056    message: str = ""57    last_action_error: Optional[str] = None58    reward: float = 0.059    done: bool = False60 61class InvoiceState(OpenEnvState):62    """Full internal state including grader score."""63    task_id: str = ""64    step_number: int = 065    max_steps: int = 1066    all_resolved: bool = False67    current_balance: float = 0.068    cumulative_reward: float = 0.069    grader_score: float = 0.000170    invoices: dict[str, Any] = Field(default_factory=dict)71    decisions: dict[str, str] = Field(default_factory=dict)72    flags_raised: dict[str, list[str]] = Field(default_factory=dict)73    tasks: list[dict[str, Any]] = Field(default_factory=list)74 75    model_config = {"extra": "allow"}76 77# ── Environment implementation ──────────────────────────────────────────────────78 79class InvoiceProcessingEnvironment(Environment[InvoiceAction, InvoiceObservation, InvoiceState]):80    """81    Accounts Payable automation environment — OpenEnv-compliant wrapper.82 83    Wraps the core InvoiceProcessingEnv state machine and exposes it84    through the official openenv.core Environment interface so that85    create_app() can auto-generate all required HTTP endpoints.86    """87 88    SUPPORTS_CONCURRENT_SESSIONS = False89 90    def __init__(self, **kwargs: Any) -> None:91        super().__init__(**kwargs)92        # Import here to avoid circular imports93        from app.env import InvoiceProcessingEnv94        self._env = InvoiceProcessingEnv()95        self._current_task_id = "easy"96        self._last_reward = 0.097        self._done = False98 99    def get_metadata(self) -> EnvironmentMetadata:100        return EnvironmentMetadata(101            name="invoice-processing-env",102            description=(103                "Accounts Payable automation environment. An AI agent processes "104                "batches of vendor invoices: extracting fields, matching purchase "105                "orders, verifying tax compliance, detecting fraud, and routing "106                "payment decisions. Covers the full SAP Ariba / Coupa AP workflow."107            ),108            version="1.0.0",109        )110 111    def reset(112        self,113        seed: Optional[int] = None,114        episode_id: Optional[str] = None,115        **kwargs: Any,116    ) -> InvoiceObservation:117        task_id = kwargs.get("task_id", self._current_task_id) or "easy"118        self._current_task_id = task_id119        self._done = False120        self._last_reward = 0.0121 122        obs = self._env.reset(task_id=task_id)123        return self._wrap_observation(obs, reward=0.0, done=False)124 125    def step(126        self,127        action: InvoiceAction,128        timeout_s: Optional[float] = None,129        **kwargs: Any,130    ) -> InvoiceObservation:131        from app.actions import Action, ActionType132 133        # create_app creates a fresh env instance per request — auto-reset if needed134        if not self._env._initialised:135            self._env.reset(task_id=self._current_task_id)136 137        # Handle task switch via action138        if action.task_id and action.task_id != self._current_task_id:139            self._current_task_id = action.task_id140            self._env.reset(task_id=self._current_task_id)141 142        try:143            core_action = Action(144                action_type=ActionType(action.action_type),145                invoice_id=action.invoice_id,146                payload=action.payload or {},147            )148        except Exception as exc:149            obs = self._env._build_observation(f"Invalid action: {exc}")150            return self._wrap_observation(obs, reward=-0.05, done=False)151 152        result = self._env.step(core_action)153        self._last_reward = result.reward.value154        self._done = result.done155        return self._wrap_observation(result.observation, result.reward.value, result.done)156 157    @property158    def state(self) -> InvoiceState:159        # Auto-reset if fresh instance (create_app creates new env per request)160        if not self._env._initialised:161            self._env.reset(task_id=self._current_task_id)162 163        s = self._env.state()164 165        # Build tasks list with grader scores — this is what the validator reads166        tasks = [167            {168                "id": "easy",169                "description": "Single clean invoice with a matching PO",170                "max_steps": 10,171                "grader_score": s["grader_score"] if s.get("task_id") == "easy" else 0.0001,172            },173            {174                "id": "medium",175                "description": "Three invoices: tax mismatch and duplicate",176                "max_steps": 20,177                "grader_score": s["grader_score"] if s.get("task_id") == "medium" else 0.0001,178            },179            {180                "id": "hard",181                "description": "Eight-invoice batch with fraud and FX edge cases",182                "max_steps": 40,183                "grader_score": s["grader_score"] if s.get("task_id") == "hard" else 0.0001,184            },185        ]186 187        return InvoiceState(188            task_id=s.get("task_id", ""),189            step_number=s.get("step_number", 0),190            max_steps=s.get("max_steps", 10),191            all_resolved=s.get("all_resolved", False),192            current_balance=s.get("current_balance", 0.0),193            cumulative_reward=s.get("cumulative_reward", 0.0),194            grader_score=float(s.get("grader_score", 0.0)),195            invoices=s.get("invoices", {}),196            decisions=s.get("decisions", {}),197            flags_raised=s.get("flags_raised", {}),198            tasks=tasks,199        )200 201    def _wrap_observation(202        self,203        obs: Any,204        reward: float,205        done: bool,206    ) -> InvoiceObservation:207        """Convert internal Observation to OpenEnv InvoiceObservation."""208        def _ser(x: Any) -> Any:209            if x is None:210                return None211            if hasattr(x, "model_dump"):212                return x.model_dump(mode="json")213            return x214 215        inbox = [_ser(inv) for inv in (obs.inbox if hasattr(obs, "inbox") else [])]216        flags = [217            f.value if hasattr(f, "value") else str(f)218            for f in (obs.compliance_flags if hasattr(obs, "compliance_flags") else [])219        ]220 221        return InvoiceObservation(222            task_id=getattr(obs, "task_id", ""),223            inbox=inbox,224            current_invoice_id=getattr(obs, "current_invoice_id", None),225            extracted_fields=_ser(getattr(obs, "extracted_fields", None)),226            matched_po=_ser(getattr(obs, "matched_po", None)),227            compliance_flags=flags,228            current_balance=getattr(obs, "current_balance", 0.0),229            step_number=getattr(obs, "step_number", 0),230            invoices_resolved=getattr(obs, "invoices_resolved", 0),231            invoices_remaining=getattr(obs, "invoices_remaining", 0),232            message=getattr(obs, "message", ""),233            last_action_error=getattr(obs, "last_action_error", None),234            reward=reward,235            done=done,236        )237