Spaces:
Running
Running
| """Tool-calling agent loop against a berth-planning server, for three provider APIs. | |
| anthropic:<model>[@low] Anthropic Messages API (streamed; thinking blocks kept in the history) | |
| openai:<model>[@low|@high] OpenAI Responses API (reasoning models only accept function tools here) | |
| hf:<org/model[:provider]> Hugging Face router, OpenAI-compatible chat completions (streamed) | |
| Every tool call goes through the env server over OpenEnv's WebSocket session, so the reward recorded is the server's | |
| rubric. One episode: the system prompt is the task's rules, the first user message its situation; the model gets | |
| `max_turns` turns (a turn may hold several tool calls); each tool result says how many turns are left; the last | |
| turn asks it to submit. Two turns in a row without a tool call end the episode. An episode without `submit_plan` | |
| scores 0. | |
| `run_episode` returns the rollout record the viewer reads (messages with reasoning, every plan the agent sent with the | |
| server's answer, the final grade, token usage). | |
| """ | |
| from __future__ import annotations | |
| import json | |
| import os | |
| import time | |
| import uuid | |
| from dataclasses import dataclass, field | |
| from typing import Any | |
| from openenv.core.env_server.mcp_types import CallToolAction | |
| from openenv.core.mcp_client import MCPToolClient | |
| REQUEST_TIMEOUT = 900.0 | |
| NOTES_CHARS = 8000 | |
| OPENING = ("{situation}\n\nMake the new berth plan. Use check_plan to test drafts, then call submit_plan with your final " | |
| "plan.") | |
| # ================================================================ env session | |
| class EnvSession: | |
| """One WebSocket session on a berth-planning server (one episode at a time).""" | |
| def __init__(self, base_url: str): | |
| self.client = MCPToolClient(base_url.rstrip("/"), message_timeout_s=300.0).sync() | |
| self._tools = None | |
| def reset(self, task_id: str, episode_id: str | None = None) -> dict: | |
| kw = {"task_id": task_id} | |
| if episode_id: | |
| kw["episode_id"] = episode_id | |
| return dict(self.client.reset(**kw).observation.metadata or {}) | |
| def tools(self): | |
| if self._tools is None: | |
| self._tools = self.client.list_tools() | |
| return self._tools | |
| def call(self, name: str, arguments: dict) -> tuple[str, bool, float | None, dict]: | |
| step = self.client.step(CallToolAction(tool_name=name, arguments=arguments)) | |
| obs = step.observation | |
| err = getattr(obs, "error", None) | |
| res = getattr(obs, "result", None) | |
| text = None | |
| if isinstance(res, dict): | |
| blocks = res.get("content") or [] | |
| text = "\n".join(b.get("text", "") for b in blocks if b.get("type") == "text") or None | |
| if text is None and res.get("data") is not None: | |
| text = res["data"] if isinstance(res["data"], str) else json.dumps(res["data"]) | |
| elif res is not None: | |
| data = getattr(res, "data", None) | |
| text = data if isinstance(data, str) else json.dumps(data) if data is not None else str(res) | |
| if err: | |
| msg = err.get("message") if isinstance(err, dict) else getattr(err, "message", str(err)) | |
| text = json.dumps({"error": str(msg)}) | |
| return text or "", bool(step.done or obs.done), step.reward if step.reward is not None else obs.reward, \ | |
| dict(obs.metadata or {}) | |
| def close(self): | |
| try: | |
| self.client.close() | |
| except Exception: | |
| pass | |
| # ================================================================ model adapters | |
| class Call: | |
| id: str | |
| name: str | |
| arguments: dict | |
| error: str | None = None | |
| raw: str = "" | |
| class Turn: | |
| text: str | |
| reasoning: str | |
| calls: list[Call] | |
| stop: str | |
| usage: dict = field(default_factory=dict) | |
| def plan_value(plan: Any) -> Any: | |
| """A plan as sent (a list, or a JSON string of one) -> the list, for the record the viewer reads.""" | |
| if isinstance(plan, str): | |
| try: | |
| plan = json.loads(plan) | |
| except json.JSONDecodeError: | |
| return plan | |
| if isinstance(plan, dict) and isinstance(plan.get("plan"), list): | |
| plan = plan["plan"] | |
| return plan | |
| def _args(raw: str) -> tuple[dict, str | None]: | |
| try: | |
| v = json.loads(raw or "{}") | |
| except json.JSONDecodeError as e: | |
| return {}, f"arguments are not valid JSON: {e.msg}" | |
| return (v, None) if isinstance(v, dict) else ({}, "arguments must be a JSON object") | |
| class AnthropicAgent: | |
| def __init__(self, model: str, tools, system: str, max_tokens: int, effort: str = "default"): | |
| import anthropic | |
| self.client = anthropic.Anthropic(max_retries=4, timeout=REQUEST_TIMEOUT) | |
| self.model, self.system, self.max_tokens = model, system, max_tokens | |
| self.extra = {"extra_body": {"output_config": {"effort": effort}}} if effort in ("low", "medium", "high") else {} | |
| self.tools = [{"name": t.name, "description": t.description or "", | |
| "input_schema": t.input_schema or {"type": "object", "properties": {}}} for t in tools] | |
| self.messages: list[dict] = [] | |
| self._pending: list[dict] = [] | |
| def user(self, text: str): | |
| self._pending.append({"type": "text", "text": text}) | |
| def notes(self, text: str): | |
| self._pending.append({"type": "text", "text": f"(Your notes from the cut-off turn)\n{text}"}) | |
| def results(self, pairs: list[tuple[Call, str]]): | |
| blocks = [{"type": "tool_result", "tool_use_id": c.id, "content": [{"type": "text", "text": out}]} | |
| for c, out in pairs] | |
| self._pending = blocks + self._pending | |
| def step(self) -> Turn: | |
| if self._pending: | |
| self.messages.append({"role": "user", "content": self._pending}) | |
| self._pending = [] | |
| with self.client.messages.stream(model=self.model, system=self.system, messages=self.messages, | |
| tools=self.tools, max_tokens=self.max_tokens, **self.extra) as s: | |
| msg = s.get_final_message() | |
| self.messages.append({"role": "assistant", "content": [b.model_dump(exclude_none=True) for b in msg.content]}) | |
| text, reasoning, calls = "", "", [] | |
| for b in msg.content: | |
| if b.type == "text": | |
| text += b.text | |
| elif b.type == "thinking": | |
| reasoning += getattr(b, "thinking", "") or "" | |
| elif b.type == "tool_use": | |
| args = b.input if isinstance(b.input, dict) else {} | |
| calls.append(Call(b.id, b.name, args, None, json.dumps(b.input))) | |
| u = msg.usage | |
| usage = {"input": (u.input_tokens or 0) + (getattr(u, "cache_read_input_tokens", 0) or 0) | |
| + (getattr(u, "cache_creation_input_tokens", 0) or 0), "output": u.output_tokens or 0} | |
| return Turn(text, reasoning, calls, msg.stop_reason or "", usage) | |
| class OpenAIResponsesAgent: | |
| def __init__(self, model: str, tools, system: str, max_tokens: int, effort: str = "medium"): | |
| from openai import OpenAI | |
| self.client = OpenAI(max_retries=3, timeout=REQUEST_TIMEOUT) | |
| self.model, self.system, self.max_tokens, self.effort = model, system, max_tokens, effort | |
| self.tools = [{"type": "function", "name": t.name, "description": t.description or "", | |
| "parameters": t.input_schema or {"type": "object", "properties": {}}} for t in tools] | |
| self.input: list = [] | |
| self._texts: list[str] = [] | |
| def user(self, text: str): | |
| self._texts.append(text) | |
| def notes(self, text: str): | |
| self._texts.append(f"(Your notes from the cut-off turn)\n{text}") | |
| def results(self, pairs): | |
| for c, out in pairs: | |
| self.input.append({"type": "function_call_output", "call_id": c.id, "output": out}) | |
| def step(self) -> Turn: | |
| if self._texts: | |
| self.input.append({"role": "user", "content": [{"type": "input_text", "text": "\n\n".join(self._texts)}]}) | |
| self._texts = [] | |
| resp = self.client.responses.create(model=self.model, instructions=self.system, input=self.input, | |
| tools=self.tools, max_output_tokens=self.max_tokens, store=False, | |
| include=["reasoning.encrypted_content"], | |
| reasoning={"effort": self.effort, "summary": "auto"}) | |
| self.input.extend(resp.output) | |
| text, reasoning, calls = "", "", [] | |
| for item in resp.output: | |
| if item.type == "message": | |
| text += "".join(getattr(c, "text", "") for c in item.content) | |
| elif item.type == "reasoning": | |
| reasoning += "\n".join(getattr(s, "text", "") for s in (getattr(item, "summary", None) or [])) | |
| elif item.type == "function_call": | |
| args, err = _args(item.arguments) | |
| calls.append(Call(item.call_id, item.name, args, err, item.arguments or "")) | |
| u = resp.usage | |
| return Turn(text, reasoning, calls, resp.status or "", {"input": getattr(u, "input_tokens", 0) or 0, | |
| "output": getattr(u, "output_tokens", 0) or 0}) | |
| class ChatAgent: | |
| """OpenAI-compatible chat completions, streamed (long generations stall behind the HF router otherwise).""" | |
| def __init__(self, model: str, tools, system: str, max_tokens: int, base_url: str, api_key: str, | |
| effort: str = "default"): | |
| from openai import OpenAI | |
| self.client = OpenAI(base_url=base_url, api_key=api_key, max_retries=3, timeout=REQUEST_TIMEOUT) | |
| # "org/model:prov1|prov2" -> try prov1, fall back to prov2 on a rejected request | |
| name, _, provs = model.rpartition(":") if ":" in model else (model, "", "") | |
| chain = [f"{name}:{p}" for p in provs.split("|")] if provs and "|" in provs else [model] | |
| self.model, self.fallbacks, self.provider_switches = chain[0], chain[1:], [] | |
| self.max_tokens = max_tokens | |
| self.extra = {"reasoning_effort": effort} if effort in ("low", "medium", "high") else {} | |
| self.tools = [{"type": "function", "function": {"name": t.name, "description": t.description or "", | |
| "parameters": t.input_schema or {"type": "object", "properties": {}}}} | |
| for t in tools] | |
| self.messages: list[dict] = [{"role": "system", "content": system}] | |
| self._texts: list[str] = [] | |
| def _create(self): | |
| """Stream a completion; on a request the provider rejects or fails (context too long, 4xx, 5xx - not 429), move on | |
| to the next provider in the spec (`org/model:prov1|prov2`) for the rest of the episode.""" | |
| import openai | |
| while True: | |
| try: | |
| return self.client.chat.completions.create( | |
| model=self.model, messages=self.messages, tools=self.tools, max_tokens=self.max_tokens, | |
| stream=True, stream_options={"include_usage": True}, **self.extra) | |
| except (openai.BadRequestError, openai.UnprocessableEntityError, openai.APIStatusError) as e: | |
| status = getattr(e, "status_code", 0) or 0 | |
| if not self.fallbacks or status == 429: | |
| raise | |
| self.provider_switches.append(f"{self.model} -> {self.fallbacks[0]} ({status}: {str(e)[:120]})") | |
| self.model = self.fallbacks.pop(0) | |
| def user(self, text: str): | |
| self._texts.append(text) | |
| def notes(self, text: str): | |
| # The cut-off turn left an empty assistant message; its reasoning becomes that message's content. | |
| if self.messages and self.messages[-1]["role"] == "assistant" and not self.messages[-1].get("content"): | |
| self.messages[-1]["content"] = f"(notes so far, cut off by the output limit)\n{text}" | |
| else: | |
| self._texts.append(f"(Your notes from the cut-off turn)\n{text}") | |
| def results(self, pairs): | |
| for c, out in pairs: | |
| self.messages.append({"role": "tool", "tool_call_id": c.id, "content": out}) | |
| def step(self) -> Turn: | |
| if self._texts: | |
| self.messages.append({"role": "user", "content": "\n\n".join(self._texts)}) | |
| self._texts = [] | |
| stream = self._create() | |
| content, reasoning, calls, stop, usage = [], [], {}, "", {"input": 0, "output": 0} | |
| # Some providers ignore max_tokens (Novita served 65-75k-token turns for Qwen3.8); cap it here by characters. | |
| budget, seen = self.max_tokens * 5, 0 | |
| for chunk in stream: | |
| if seen > budget: | |
| stop = "length (harness cap)" | |
| stream.close() | |
| break | |
| if getattr(chunk, "usage", None): | |
| usage = {"input": chunk.usage.prompt_tokens or 0, "output": chunk.usage.completion_tokens or 0} | |
| if not chunk.choices: | |
| continue | |
| ch = chunk.choices[0] | |
| d = ch.delta | |
| if ch.finish_reason: | |
| stop = ch.finish_reason | |
| if d is None: | |
| continue | |
| if d.content: | |
| content.append(d.content) | |
| seen += len(d.content) | |
| extra = getattr(d, "model_extra", None) or {} | |
| r = getattr(d, "reasoning_content", None) or extra.get("reasoning_content") or extra.get("reasoning") | |
| if isinstance(r, str): | |
| reasoning.append(r) | |
| seen += len(r) | |
| for tc in d.tool_calls or []: | |
| e = calls.setdefault(tc.index or 0, {"id": "", "name": "", "args": ""}) | |
| if tc.id: | |
| e["id"] = tc.id | |
| if tc.function and tc.function.name: | |
| e["name"] = tc.function.name if not e["name"] else e["name"] | |
| if tc.function and tc.function.arguments: | |
| e["args"] += tc.function.arguments | |
| text = "".join(content) | |
| think = "".join(reasoning) | |
| if not think and "</think>" in text: | |
| think, text = text.split("</think>", 1) | |
| think = think.replace("<think>", "") | |
| out = [] | |
| for i in sorted(calls): | |
| e = calls[i] | |
| args, err = _args(e["args"]) | |
| out.append(Call(e["id"] or f"call_{uuid.uuid4().hex[:8]}", e["name"], args, err, e["args"])) | |
| msg: dict[str, Any] = {"role": "assistant", "content": text} | |
| if out: | |
| msg["tool_calls"] = [{"id": c.id, "type": "function", "function": {"name": c.name, "arguments": c.raw or "{}"}} | |
| for c in out] | |
| self.messages.append(msg) | |
| return Turn(text, think, out, stop, usage) | |
| def make_agent(spec: str, tools, system: str, max_tokens: int): | |
| base, _, effort = spec.partition("@") | |
| provider, _, model = base.partition(":") | |
| effort = effort or "default" | |
| if provider == "anthropic": | |
| return AnthropicAgent(model, tools, system, max_tokens, effort) | |
| if provider == "openai": | |
| return OpenAIResponsesAgent(model, tools, system, max_tokens, "medium" if effort == "default" else effort) | |
| if provider == "hf": | |
| return ChatAgent(model, tools, system, max_tokens, os.environ.get("HF_ROUTER_URL", "https://router.huggingface.co/v1"), | |
| os.environ["HF_TOKEN"], effort) | |
| if provider == "vllm": | |
| return ChatAgent(model, tools, system, max_tokens, os.environ["VLLM_BASE_URL"], | |
| os.environ.get("VLLM_API_KEY", "none"), effort) | |
| raise ValueError(f"unknown provider in {spec!r}: use anthropic:, openai:, hf: or vllm:") | |
| # ================================================================ episode | |
| def run_episode(spec: str, base_url: str, task_id: str, max_turns: int = 12, max_tokens: int = 24000, | |
| wall_clock_s: float = 1800.0) -> dict: | |
| env = EnvSession(base_url) | |
| started = time.time() | |
| episode_id = uuid.uuid4().hex[:12] | |
| record: dict[str, Any] = {"model": spec, "task_id": task_id, "episode_id": episode_id, | |
| "started": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime(started)), | |
| "messages": [], "steps": [], "final": {"submitted": False, "plan": None, "grade": None}, | |
| "reward": 0.0, "usage": {"input_tokens": 0, "output_tokens": 0}, "turns": 0, | |
| "end_reason": None, "errors": []} | |
| try: | |
| meta = env.reset(task_id, episode_id) | |
| system = meta["instructions"] | |
| opening = OPENING.format(situation=meta["situation"]) | |
| record["messages"] += [{"role": "system", "content": system}, {"role": "user", "content": opening}] | |
| agent = make_agent(spec, env.tools, system, max_tokens) | |
| agent.user(opening) | |
| empty = 0 | |
| for turn in range(max_turns): | |
| if time.time() - started > wall_clock_s: | |
| record["end_reason"] = "wall_clock" | |
| break | |
| last = turn == max_turns - 1 | |
| if last: | |
| note = "This is your last turn: call submit_plan now with your best plan." | |
| agent.user(note) | |
| record["messages"].append({"role": "user", "content": note}) | |
| t = agent.step() | |
| record["turns"] += 1 | |
| record["usage"]["input_tokens"] += t.usage.get("input", 0) | |
| record["usage"]["output_tokens"] += t.usage.get("output", 0) | |
| record["messages"].append({"role": "assistant", "content": t.text, "reasoning": t.reasoning, "stop": t.stop, | |
| "tool_calls": [{"id": c.id, "name": c.name, "arguments": c.raw or json.dumps(c.arguments)} | |
| for c in t.calls]}) | |
| if not t.calls: | |
| empty += 1 | |
| if empty >= 2 or last: | |
| record["end_reason"] = "no_tool_call" | |
| break | |
| if t.stop.startswith("length") and t.reasoning: | |
| # A turn cut off mid-thought loses its reasoning; hand the tail back as notes so the next turn | |
| # continues instead of starting over. | |
| agent.notes(t.reasoning[-NOTES_CHARS:]) | |
| nudge = ("Your last turn ran out of output tokens before calling a tool. Your notes from it are above. " | |
| "Keep reasoning short now and call check_plan with your current draft, or submit_plan.") | |
| else: | |
| nudge = ("No tool was called. Plans only count through the tools: call check_plan to test a draft " | |
| "or submit_plan with your final plan.") | |
| agent.user(nudge) | |
| record["messages"].append({"role": "user", "content": nudge}) | |
| continue | |
| empty = 0 | |
| pairs, done = [], False | |
| for c in t.calls: | |
| if done: | |
| out = json.dumps({"error": "the episode is over"}) | |
| elif c.error: | |
| out = json.dumps({"error": c.error}) | |
| else: | |
| out, done, reward, meta = env.call(c.name, c.arguments) | |
| if c.name in ("check_plan", "submit_plan"): | |
| try: | |
| res = json.loads(out) | |
| except json.JSONDecodeError: | |
| res = {"text": out} | |
| record["steps"].append({"turn": turn + 1, "tool": c.name, "plan": plan_value(c.arguments.get("plan")), | |
| "result": res}) | |
| if done: | |
| record["reward"] = float(reward or 0.0) | |
| if c.name == "submit_plan": | |
| record["final"] = {"submitted": True, "plan": plan_value(c.arguments.get("plan")), | |
| "grade": meta.get("grade"), "rubric": meta.get("rubric")} | |
| record["end_reason"] = "submitted" | |
| else: | |
| record["end_reason"] = meta.get("end_reason", "server_ended") | |
| left = max_turns - turn - 1 | |
| text = out if done else f"{out}\n(turns left: {left})" | |
| pairs.append((c, text)) | |
| record["messages"].append({"role": "tool", "tool_call_id": c.id, "name": c.name, "content": out}) | |
| agent.results(pairs) | |
| if done: | |
| break | |
| else: | |
| record["end_reason"] = record["end_reason"] or "turn_limit" | |
| record["end_reason"] = record["end_reason"] or "turn_limit" | |
| except Exception as exc: # provider or transport failure: recorded, scored as it stands | |
| record["errors"].append(f"{type(exc).__name__}: {str(exc)[:800]}") | |
| record["end_reason"] = record["end_reason"] or "error" | |
| finally: | |
| env.close() | |
| switches = getattr(locals().get("agent"), "provider_switches", None) | |
| if switches: | |
| record["provider_switches"] = switches | |
| record["seconds"] = round(time.time() - started, 1) | |
| return record | |