| """Typed models for the RecallTrace OpenEnv environment.""" |
|
|
| from __future__ import annotations |
|
|
| from enum import Enum |
| from typing import Any, Dict, List, Optional |
|
|
| from pydantic import BaseModel, ConfigDict, Field |
|
|
|
|
| class ActionType(str, Enum): |
| INSPECT_NODE = "inspect_node" |
| TRACE_LOT = "trace_lot" |
| QUARANTINE = "quarantine" |
| NOTIFY = "notify" |
| FINALIZE = "finalize" |
|
|
|
|
| class RecallAction(BaseModel): |
| """Action submitted by an agent.""" |
|
|
| model_config = ConfigDict(extra="forbid") |
|
|
| type: ActionType |
| node_id: Optional[str] = None |
| lot_id: Optional[str] = None |
| quantity: Optional[int] = Field(default=None, ge=1) |
| rationale: Optional[str] = None |
|
|
|
|
| class RewardSignal(BaseModel): |
| """Typed reward payload.""" |
|
|
| model_config = ConfigDict(extra="forbid") |
|
|
| value: float = Field(ge=-1.0, le=1.0) |
| reason: str |
| components: Dict[str, float] = Field(default_factory=dict) |
|
|
|
|
| class InspectionEvidence(BaseModel): |
| """Evidence revealed after inspecting a node.""" |
|
|
| model_config = ConfigDict(extra="allow") |
|
|
| status: str |
| unsafe_quantity: int = Field(ge=0) |
| evidence: str |
| safe_quantity: Optional[int] = Field(default=None, ge=0) |
|
|
|
|
| class TaskDefinition(BaseModel): |
| """Static task descriptor.""" |
|
|
| model_config = ConfigDict(extra="forbid") |
|
|
| task_id: str |
| name: str |
| difficulty: str |
| objective: str |
| max_steps: int = Field(ge=1) |
|
|
|
|
| class RecallObservation(BaseModel): |
| """Observable state exposed to the agent.""" |
|
|
| model_config = ConfigDict(extra="forbid") |
|
|
| task_id: str |
| phase: int |
| recall_notice: str |
| available_actions: List[str] |
| inventory: Dict[str, Dict[str, int]] |
| discovered_shipments: Dict[str, List[str]] |
| inspected_nodes: List[str] |
| inspection_results: Dict[str, Dict[str, InspectionEvidence]] |
| trace_results: Dict[str, Dict[str, Any]] |
| notified_nodes: List[str] |
| quarantined_inventory: Dict[str, Dict[str, int]] |
| history: List[str] |
| steps_taken: int = Field(ge=0) |
| remaining_step_budget: int = Field(ge=0) |
|
|
|
|
| class StepInfo(BaseModel): |
| """Structured info payload returned after each step.""" |
|
|
| model_config = ConfigDict(extra="allow") |
|
|
| message: str |
| action_type: str |
| score: Optional[float] = Field(default=None, ge=0.0, le=1.0) |
| reward_breakdown: Dict[str, float] = Field(default_factory=dict) |
|
|
|
|
| class EnvironmentState(BaseModel): |
| """Full internal state for debugging and grading.""" |
|
|
| model_config = ConfigDict(extra="forbid") |
|
|
| done: bool |
| task: TaskDefinition |
| steps_taken: int = Field(ge=0) |
| state_data: Dict[str, Any] |
| ground_truth: Dict[str, Any] |
|
|
|
|
| class TaskGrade(BaseModel): |
| """Deterministic grader output.""" |
|
|
| model_config = ConfigDict(extra="forbid") |
|
|
| task_id: str |
| score: float = Field(ge=0.0, le=1.0) |
| success: bool |
| steps_taken: int = Field(ge=0) |
| max_steps: int = Field(ge=1) |
| reward_total: float |
| final_info: Dict[str, Any] |
|
|
