Spaces:
Running
Running
| """The OpenEnv environment: one berth-planning episode per session, three MCP tools. | |
| get_situation() the quay, notices, ships already alongside, closures and the ships to berth | |
| check_plan(plan) rule violations and cost of the agent's own plan (budgeted; never a score or a reference) | |
| submit_plan(plan) the final answer: ends the episode; the rubric grades it once | |
| `reset(task_id=...)` or `reset(split=, index=)` (or `seed=` for a random task) picks a task. Tasks are served through | |
| OpenEnv's Task API (`list_splits`, `num_tasks`, `get_task`, `get_task_range`). The reward is `BerthPlanRubric` | |
| applied to the `submit_plan` action; every other step has reward 0.0. An episode that runs out of tool calls ends | |
| with reward 0.0. | |
| """ | |
| from __future__ import annotations | |
| import json | |
| import os | |
| import random | |
| import threading | |
| import time | |
| import uuid | |
| from collections import OrderedDict | |
| from typing import Annotated, Any | |
| from fastmcp import FastMCP | |
| from openenv.core.env_server.mcp_environment import MCPEnvironment | |
| from openenv.core.env_server.mcp_types import CallToolAction | |
| from openenv.core.env_server.types import Action, EnvironmentMetadata, Observation, State | |
| from pydantic import BaseModel, Field | |
| from berth_core import PROMPT_VERSION, Task, TaskPack, evaluate, load_pack, parse_plan, plan_to_list, rules, situation | |
| from berth_core.model import PlanError | |
| try: | |
| from .rubric import SUBMIT_TOOL, BerthPlanRubric | |
| except ImportError: # flat layout | |
| from rubric import SUBMIT_TOOL, BerthPlanRubric | |
| MAX_CHECKS = int(os.environ.get("BERTH_MAX_CHECKS", "10")) | |
| MAX_TOOL_CALLS = int(os.environ.get("BERTH_MAX_TOOL_CALLS", "24")) | |
| class PlanEntry(BaseModel): | |
| ship: int = Field(description="ship id from the table") | |
| berth_hour: int = Field(description="hour the ship berths (>= its arrival hour)") | |
| section: int = Field(description="first (lowest-numbered) section the ship occupies") | |
| cranes: int | None = Field(default=None, description="quay cranes working the ship (tasks with crane rules; " | |
| "defaults to the ship's planned cranes)") | |
| PlanArg = Annotated[list[PlanEntry] | str, Field( | |
| description='one entry per ship: [{"ship": 0, "berth_hour": 36, "section": 22}, ...]')] | |
| class BerthState(State): | |
| task_id: str | None = None | |
| split: str | None = None | |
| checks_used: int = 0 | |
| tool_calls: int = 0 | |
| done: bool = False | |
| reward: float | None = None | |
| grade: dict | None = None | |
| class Episode: | |
| """Everything one episode did, kept for the viewer's live page (bounded store, newest first).""" | |
| def __init__(self, task: Task, episode_id: str): | |
| self.task = task | |
| self.episode_id = episode_id | |
| self.started = time.time() | |
| self.steps: list[dict] = [] | |
| self.done = False | |
| self.grade: dict | None = None | |
| self.end_reason: str | None = None | |
| def summary(self) -> dict: | |
| return {"episode_id": self.episode_id, "task_id": self.task.task_id, "started": self.started, | |
| "done": self.done, "end_reason": self.end_reason, "steps": self.steps, "grade": self.grade} | |
| class _Store: | |
| def __init__(self, cap: int = 200): | |
| self.cap = cap | |
| self._d: OrderedDict[str, Episode] = OrderedDict() | |
| self._lock = threading.Lock() | |
| def add(self, ep: Episode): | |
| with self._lock: | |
| self._d[ep.episode_id] = ep | |
| while len(self._d) > self.cap: | |
| self._d.popitem(last=False) | |
| def get(self, eid: str) -> Episode | None: | |
| return self._d.get(eid) | |
| def recent(self, n: int = 50) -> list[Episode]: | |
| return list(reversed(self._d.values()))[:n] | |
| STORE = _Store() | |
| def _json(obj: Any) -> str: | |
| return json.dumps(obj, ensure_ascii=False, separators=(",", ":")) | |
| def _plan_result(task: Task, plan) -> dict: | |
| r = evaluate(task, plan) | |
| return { | |
| "feasible": r.feasible, | |
| "violations": r.violations, | |
| "cost": r.cost, "delay_cost": r.delay_cost, "moves": r.moves, | |
| "ships": [{"ship": s.ship, "departure": s.departure, "delay_h": s.delay_h, "moved": s.moved, "cost": s.cost} | |
| | ({"cranes": s.cranes, "hours_alongside": s.departure - s.berth_hour} if s.cranes is not None else {}) | |
| for s in r.ships if s.berth_hour is not None], | |
| } | |
| class BerthPlanningEnvironment(MCPEnvironment): | |
| """Re-plan a real week of container berthings at a Barcelona quay after a disruption.""" | |
| SUPPORTS_CONCURRENT_SESSIONS = True | |
| def get_metadata(self) -> EnvironmentMetadata: | |
| """What OpenEnv's /metadata reports (the web UI shows it too).""" | |
| return EnvironmentMetadata( | |
| name="PortSimEnv v1", | |
| description="Re-plan a broken week of container-ship dockings at a Port of Barcelona quay, built from the " | |
| "port's real 2024 calls. Tools: get_situation, check_plan (10 per episode), submit_plan " | |
| "(graded once, deterministically, against the CP-SAT optimum).", | |
| version="1.0.0", | |
| author="FineEnvs", | |
| documentation_url="https://github.com/adithya-s-k/FineEnvs/tree/main/07-simulation-environments/portsim-v1", | |
| ) | |
| def __init__(self, pack: TaskPack | None = None): | |
| self.pack = pack or load_pack() | |
| self._episode: Episode | None = None | |
| self._state = BerthState() | |
| mcp = FastMCP("berth_planning") | |
| def get_situation() -> str: | |
| """The quay, notices, ships already alongside, closed sections and the table of ships to berth.""" | |
| ep = self._require() | |
| return situation(ep.task) | |
| def check_plan(plan: PlanArg) -> str: | |
| """Check your plan: rule violations per ship, departure and delay per ship, and the plan's cost. | |
| Uses one of your checks. It does not grade the plan.""" | |
| ep = self._require(open_only=True) | |
| if self._state.checks_used >= MAX_CHECKS: | |
| return _json({"error": f"no checks left (you had {MAX_CHECKS}); submit_plan when ready"}) | |
| self._state.checks_used += 1 | |
| try: | |
| parsed, problems = parse_plan(ep.task, _raw(plan)) | |
| except PlanError as e: | |
| out = {"error": str(e), "checks_left": MAX_CHECKS - self._state.checks_used} | |
| ep.steps.append({"turn": self._state.tool_calls, "tool": "check_plan", "plan": None, "result": out}) | |
| return _json(out) | |
| out = _plan_result(ep.task, parsed) | |
| if problems: | |
| out["entry_problems"] = problems | |
| out["feasible"] = False | |
| out["checks_left"] = MAX_CHECKS - self._state.checks_used | |
| ep.steps.append({"turn": self._state.tool_calls, "tool": "check_plan", "plan": plan_to_list(parsed), | |
| "result": out}) | |
| return _json(out) | |
| def submit_plan(plan: PlanArg) -> str: | |
| """Submit your final berth plan. Ends the episode; the plan is graded once.""" | |
| ep = self._require(open_only=True) | |
| try: | |
| parsed, problems = parse_plan(ep.task, _raw(plan)) | |
| except PlanError as e: | |
| parsed, problems = {}, [str(e)] | |
| # The rubric grades the same arguments again in step(); this result is what the agent reads. | |
| g = self.rubric.graded(CallToolAction(tool_name=SUBMIT_TOOL, arguments={"plan": _raw(plan)})) | |
| ep.done = True | |
| ep.end_reason = "submitted" | |
| ep.grade = g.as_dict() if g else None | |
| ep.steps.append({"turn": self._state.tool_calls, "tool": "submit_plan", "plan": plan_to_list(parsed), | |
| "result": {"grade": ep.grade}}) | |
| out = {"submitted": True, "feasible": g.feasible if g else False, "reward": g.reward if g else 0.0} | |
| if g and g.feasible: | |
| out.update(cost=g.cost, delay_cost=g.delay_cost, moves=g.moves) | |
| elif g: | |
| out.update(violations=g.violations[:20], entry_problems=g.parse_problems) | |
| return _json(out) | |
| super().__init__(mcp) | |
| self.rubric = BerthPlanRubric() | |
| # ------------------------------------------------------------ Task API | |
| def list_splits(self): | |
| return [{"name": s, "type": s, "num_tasks": self.pack.count(s)} for s in self.pack.splits()] | |
| def num_tasks(self, split): | |
| return self.pack.count(split) | |
| def get_task(self, split, index): | |
| return self.pack.public(self.pack.at(split, int(index))) | |
| def get_task_range(self, split, start=None, stop=None): | |
| n = self.pack.count(split) | |
| start, stop = int(start or 0), int(n if stop is None else stop) | |
| return [self.pack.public(self.pack.at(split, i)) for i in range(max(0, start), min(n, stop))] | |
| def list_tasks(self, split): | |
| return self.get_task_range(split) | |
| # ------------------------------------------------------------ lifecycle | |
| def reset(self, seed: int | None = None, episode_id: str | None = None, split: str | None = None, | |
| index: int | None = None, task_id: str | None = None, **kwargs: Any) -> Observation: | |
| if task_id is not None: | |
| task = self.pack.get(task_id) | |
| else: | |
| split = split or ("train" if "train" in self.pack.splits() else self.pack.splits()[0]) | |
| n = self.pack.count(split) | |
| index = random.Random(seed).randrange(n) if index is None else int(index) | |
| task = self.pack.at(split, index) | |
| ep = Episode(task, str(episode_id)[:64] if episode_id else uuid.uuid4().hex[:12]) | |
| STORE.add(ep) | |
| self._episode = ep | |
| self._state = BerthState(episode_id=ep.episode_id, step_count=0, task_id=task.task_id, split=task.split) | |
| self._reset_rubric() | |
| self.rubric.set_task(task) | |
| return Observation(done=False, reward=None, metadata={ | |
| "episode_id": ep.episode_id, | |
| "task_id": task.task_id, | |
| "split": task.split, | |
| "difficulty": task.difficulty, | |
| "instructions": rules(task, MAX_CHECKS), | |
| "situation": situation(task), | |
| "max_checks": MAX_CHECKS, | |
| "max_tool_calls": MAX_TOOL_CALLS, | |
| "tools": ["get_situation", "check_plan", "submit_plan"], | |
| "prompt_version": PROMPT_VERSION, | |
| "viewer_url": f"/viewer/#/live/{ep.episode_id}", | |
| }) | |
| def _require(self, open_only: bool = False) -> Episode: | |
| if self._episode is None: | |
| raise ValueError("no episode: call reset first") | |
| if open_only and self._episode.done: | |
| raise ValueError("the episode is over: the plan was already submitted") | |
| return self._episode | |
| def state(self) -> BerthState: | |
| return self._state | |
| # ------------------------------------------------------------ stepping | |
| def _after_step(self, action: Action, obs: Observation) -> Observation: | |
| ep = self._episode | |
| if ep is None or not isinstance(action, CallToolAction): | |
| return obs | |
| self._state.tool_calls += 1 | |
| self._state.step_count += 1 | |
| obs.reward = self._apply_rubric(action, obs) # 0.0 unless this was submit_plan | |
| if ep.done: | |
| g = self.rubric.last_grade | |
| obs.done = True | |
| obs.metadata = {**(obs.metadata or {}), "grade": g.as_dict() if g else ep.grade, | |
| "rubric": {name: r.last_score for name, r in self.rubric.named_rubrics()}, | |
| "episode_id": ep.episode_id} | |
| self._state.done, self._state.reward = True, obs.reward | |
| self._state.grade = obs.metadata["grade"] | |
| elif self._state.tool_calls >= MAX_TOOL_CALLS: | |
| ep.done, ep.end_reason = True, "tool_call_limit" | |
| obs.done, obs.reward = True, 0.0 | |
| obs.metadata = {**(obs.metadata or {}), "end_reason": "tool_call_limit", "episode_id": ep.episode_id} | |
| self._state.done, self._state.reward = True, 0.0 | |
| return obs | |
| def step(self, action: Action, timeout_s: float | None = None, **kwargs: Any) -> Observation: | |
| return self._after_step(action, super().step(action, timeout_s=timeout_s, **kwargs)) | |
| async def step_async(self, action: Action, timeout_s: float | None = None, **kwargs: Any) -> Observation: | |
| return self._after_step(action, await super().step_async(action, timeout_s=timeout_s, **kwargs)) | |
| def _step_impl(self, action: Action, timeout_s: float | None = None, **kwargs) -> Observation: | |
| return Observation(done=False, reward=0.0, metadata={"error": "use the MCP tools: get_situation, check_plan, submit_plan"}) | |
| def _raw(plan): | |
| """Tool arguments arrive as PlanEntry models (or a JSON string); the core parser takes plain data.""" | |
| if isinstance(plan, list): | |
| return [p.model_dump(exclude_none=True) if isinstance(p, BaseModel) else p for p in plan] | |
| return plan | |