"""Tool-calling agent loop against a berth-planning server, for three provider APIs. anthropic:[@low] Anthropic Messages API (streamed; thinking blocks kept in the history) openai:[@low|@high] OpenAI Responses API (reasoning models only accept function tools here) hf: 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 {}) @property 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 @dataclass class Call: id: str name: str arguments: dict error: str | None = None raw: str = "" @dataclass 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 "" in text: think, text = text.split("", 1) think = think.replace("", "") 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