Add provider-agnostic model agent (vLLM/OpenRouter/Bedrock)

- providers.py: ChatProvider abstraction; OpenAICompatibleProvider
(covers vLLM + OpenRouter + OpenAI by base_url), BedrockProvider stub
with precise NotImplementedError. Pure ProviderConfig selection.
- agent.py: ModelAgent — Training-compatible text briefing, OpenAI
function-calling tool schema (move_units/attack_unit/observe) filtered
by scenario tools, tolerant tool-call->Command parsing with aliases,
optional minimap PNG multimodal (graceful text-only fallback),
bounded chat history with stale-image stripping. Exposes agent_fn for
eval_core.
- tests/test_agent.py: offline FakeProvider drives the full
ModelAgent->eval_core->live-Rust loop; schema/briefing/parsing units;
opt-in OpenRouter live smoke (skips without OPENROUTER_API_KEY).

22 passed, 1 skipped.

openra_bench/__init__.py

@@ -4,4 +4,4 @@ planning, on the Rust OpenRA environment, reusing OpenRA-RL-Training components.
 See EVAL_STACK_PLAN.md for architecture and phasing.
 """
 
-__all__ = ["rust_adapter", "eval_core"]
+__all__ = ["rust_adapter", "eval_core", "agent", "providers", "scenarios"]

openra_bench/agent.py

@@ -0,0 +1,242 @@
+"""Provider-agnostic model agent.
+
+Turns a `RustObsAdapter.render_state()` into a Training-compatible text
+briefing (+ optional minimap image), calls a `ChatProvider`, and parses
+tool calls back into `openra_train.Command` objects. Exposes an
+`agent_fn` matching `eval_core`'s `(render_state, Command) -> [Command]`
+contract.
+
+Tool contract mirrors OpenRA-RL-Training so models trained there behave
+consistently: `move_units(unit_ids, target_x, target_y)`,
+`attack_unit(unit_ids, target_id)`, `observe()`. The scenario's `tools`
+list filters which are offered.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+from .providers import ChatProvider, ProviderConfig, make_provider
+
+logger = logging.getLogger(__name__)
+
+SYSTEM_PROMPT = (
+    "You are commanding units in Command & Conquer: Red Alert.\n"
+    "Each turn you receive a BRIEFING (and, when available, a MINIMAP image: "
+    "bright=visible, dim=explored, black=unknown fog).\n"
+    "Units are listed as `<id> <type> @(x,y)` (with `-> (tx,ty)` if moving).\n"
+    "Pass numeric unit IDs to tools, e.g. unit_ids=[1004,1005].\n"
+    "Every turn MUST include at least one tool call. Think briefly, then act."
+)
+
+_TOOL_SCHEMAS: dict[str, dict] = {
+    "move_units": {
+        "type": "function",
+        "function": {
+            "name": "move_units",
+            "description": "Move the given units to a map cell. Units auto-fire "
+            "opportunistically en route. Use to position/scout/retreat.",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "unit_ids": {"type": "array", "items": {"type": "integer"}},
+                    "target_x": {"type": "integer"},
+                    "target_y": {"type": "integer"},
+                },
+                "required": ["unit_ids", "target_x", "target_y"],
+            },
+        },
+    },
+    "attack_unit": {
+        "type": "function",
+        "function": {
+            "name": "attack_unit",
+            "description": "Order the given units to pathfind to and focus-fire "
+            "a specific enemy actor id until it dies.",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "unit_ids": {"type": "array", "items": {"type": "integer"}},
+                    "target_id": {"type": "integer"},
+                },
+                "required": ["unit_ids", "target_id"],
+            },
+        },
+    },
+    "observe": {
+        "type": "function",
+        "function": {
+            "name": "observe",
+            "description": "Take no action; advance the game and re-observe.",
+            "parameters": {"type": "object", "properties": {}},
+        },
+    },
+}
+# Aliases tolerated from models trained on slightly different names.
+_TOOL_ALIASES = {"attack_target": "attack_unit", "stop_units": "observe"}
+
+
+def _tool_schemas(allowed: list[str] | None) -> list[dict]:
+    names = list(_TOOL_SCHEMAS) if not allowed else allowed
+    out = [_TOOL_SCHEMAS[n] for n in names if n in _TOOL_SCHEMAS]
+    if "observe" not in {t["function"]["name"] for t in out}:
+        out.append(_TOOL_SCHEMAS["observe"])  # always allow a no-op
+    return out
+
+
+def build_briefing(render_state: dict, objective: str = "") -> str:
+    """Training-style text state. Self-contained (no engine handles)."""
+    lines: list[str] = []
+    if objective:
+        lines.append(f"OBJECTIVE: {objective}")
+    lines.append(
+        f"tick={render_state.get('game_tick', 0)} "
+        f"explored={render_state.get('explored_percent', 0.0):.1f}%"
+    )
+    own = render_state.get("units_summary", []) or []
+    lines.append(f"\nYOUR UNITS ({len(own)}):")
+    for u in own:
+        act = u.get("activity")
+        suffix = f", {act}" if act and act != "idle" else ""
+        lines.append(
+            f"  {u['id']} {u.get('type') or 'unit'} @({u['cell_x']},{u['cell_y']}){suffix}"
+        )
+    enemy = render_state.get("enemy_summary", []) or []
+    if enemy:
+        lines.append(f"\nVISIBLE ENEMIES ({len(enemy)}):")
+        for e in enemy:
+            kind = "building" if e.get("is_building") else (e.get("type") or "unit")
+            lines.append(f"  {e['id']} {kind} @({e['cell_x']},{e['cell_y']})")
+    else:
+        lines.append("\nVISIBLE ENEMIES: none (scout the fog)")
+    return "\n".join(lines)
+
+
+def _render_minimap_b64(render_state: dict) -> str | None:
+    """Best-effort minimap PNG. Returns None (text-only fallback) when
+    terrain isn't resolvable — vision degrades gracefully in Phase 0."""
+    try:
+        from openra_rl_training.training.minimap_renderer import render_minimap
+
+        return render_minimap(
+            terrain_png=render_state.get("terrain_png"),  # None in Phase 0 -> None
+            map_width=render_state.get("map_width", 64),
+            map_height=render_state.get("map_height", 64),
+            bounds_x=render_state.get("bounds_x", 0),
+            bounds_y=render_state.get("bounds_y", 0),
+            own_units=render_state.get("units_summary", []) or [],
+            enemy_units=render_state.get("enemy_summary", []) or [],
+            ascii_minimap=render_state.get("minimap", ""),
+        )
+    except Exception as e:  # noqa: BLE001 — vision is optional
+        logger.debug("minimap render skipped: %s", e)
+        return None
+
+
+def _to_commands(tool_calls: list[dict], Command: Any) -> list:
+    cmds = []
+    for call in tool_calls:
+        name = _TOOL_ALIASES.get(call.get("name", ""), call.get("name", ""))
+        args = call.get("arguments") or {}
+        try:
+            if name == "move_units":
+                ids = [str(i) for i in args["unit_ids"]]
+                cmds.append(
+                    Command.move_units(ids, int(args["target_x"]), int(args["target_y"]))
+                )
+            elif name == "attack_unit":
+                ids = [str(i) for i in args["unit_ids"]]
+                cmds.append(Command.attack_unit(ids, str(args["target_id"])))
+            elif name == "observe":
+                cmds.append(Command.observe())
+        except (KeyError, TypeError, ValueError) as e:
+            logger.debug("dropping malformed tool call %s: %s", call, e)
+    return cmds
+
+
+class ModelAgent:
+    """One instance per episode (keeps bounded chat history).
+
+    Usage:
+        agent = ModelAgent(cfg, allowed_tools=compiled.scenario.tools,
+                            objective=compiled.scenario.description)
+        result = run_level(compiled, agent.agent_fn, seed=...)
+    """
+
+    def __init__(
+        self,
+        cfg: ProviderConfig,
+        allowed_tools: list[str] | None = None,
+        objective: str = "",
+        provider: ChatProvider | None = None,
+    ):
+        self.cfg = cfg
+        self.objective = objective
+        self.tools = _tool_schemas(allowed_tools)
+        self.provider = provider or make_provider(cfg)
+        self.history: list[dict] = [{"role": "system", "content": SYSTEM_PROMPT}]
+        self.stats = {"turns": 0, "tool_calls": 0, "empty_replies": 0}
+
+    def _user_message(self, render_state: dict) -> dict:
+        text = build_briefing(render_state, self.objective)
+        if self.cfg.vision:
+            b64 = _render_minimap_b64(render_state)
+            if b64:
+                return {
+                    "role": "user",
+                    "content": [
+                        {"type": "text", "text": text},
+                        {
+                            "type": "image_url",
+                            "image_url": {"url": f"data:image/png;base64,{b64}"},
+                        },
+                    ],
+                }
+        return {"role": "user", "content": text}
+
+    @staticmethod
+    def _strip_old_images(history: list[dict]) -> None:
+        """Keep only the latest image to bound ViT token cost (mirrors
+        Training's _strip_historical_images)."""
+        seen = False
+        for msg in reversed(history):
+            c = msg.get("content")
+            if isinstance(c, list):
+                if not seen:
+                    seen = True
+                    continue
+                msg["content"] = " ".join(
+                    p.get("text", "") for p in c if p.get("type") == "text"
+                )
+
+    def agent_fn(self, render_state: dict, Command: Any) -> list:
+        self.stats["turns"] += 1
+        self.history.append(self._user_message(render_state))
+        self._strip_old_images(self.history)
+        reply = self.provider.complete(self.history, self.tools)
+        self.history.append(
+            {
+                "role": "assistant",
+                "content": reply.text or "",
+                "tool_calls": [
+                    {
+                        "id": f"c{i}",
+                        "type": "function",
+                        "function": {"name": c["name"], "arguments": c["arguments"]},
+                    }
+                    for i, c in enumerate(reply.tool_calls)
+                ],
+            }
+        )
+        cmds = _to_commands(reply.tool_calls, Command)
+        self.stats["tool_calls"] += len(cmds)
+        if not cmds:
+            self.stats["empty_replies"] += 1
+            cmds = [Command.observe()]
+        # Satisfy the OpenAI contract: every tool_call needs a tool result.
+        for i in range(len(reply.tool_calls)):
+            self.history.append(
+                {"role": "tool", "tool_call_id": f"c{i}", "content": "ok"}
+            )
+        return cmds

openra_bench/providers.py

@@ -0,0 +1,158 @@
+"""Provider abstraction for the eval agent.
+
+One small interface, `ChatProvider.complete(messages, tools) -> ChatReply`.
+Adapters:
+
+* `OpenAICompatibleProvider` — OpenAI Chat Completions wire format. Covers
+  local **vLLM** (matches Training's rollout path) and **OpenRouter**
+  (the Phase-0 test target) by base_url alone.
+* `BedrockProvider` — AWS Bedrock Converse. Stubbed with a precise
+  NotImplementedError so the wiring exists before the dependency does.
+
+Selection is pure config (`ProviderConfig`); no provider-specific code
+leaks into the agent.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+from dataclasses import dataclass, field
+from typing import Any, Literal
+
+import httpx
+
+ProviderName = Literal["openai", "vllm", "openrouter", "bedrock"]
+
+# Convenience presets; base_url/api_key_env still overridable in config.
+_PRESETS: dict[str, dict[str, str]] = {
+    "openrouter": {
+        "base_url": "https://openrouter.ai/api/v1",
+        "api_key_env": "OPENROUTER_API_KEY",
+    },
+    "vllm": {
+        "base_url": "http://localhost:8100/v1",
+        "api_key_env": "VLLM_API_KEY",  # vLLM ignores the value
+    },
+    "openai": {
+        "base_url": "https://api.openai.com/v1",
+        "api_key_env": "OPENAI_API_KEY",
+    },
+}
+
+
+@dataclass
+class ProviderConfig:
+    provider: ProviderName = "openrouter"
+    model: str = "anthropic/claude-3.5-sonnet"
+    base_url: str | None = None
+    api_key_env: str | None = None
+    temperature: float = 0.7
+    max_tokens: int = 1024
+    timeout_s: float = 120.0
+    vision: bool = True
+    extra_headers: dict[str, str] = field(default_factory=dict)
+
+    def resolved_base_url(self) -> str:
+        if self.base_url:
+            return self.base_url
+        preset = _PRESETS.get(self.provider)
+        if not preset:
+            raise ValueError(
+                f"no base_url and no preset for provider {self.provider!r}"
+            )
+        return preset["base_url"]
+
+    def resolved_api_key(self) -> str:
+        env = self.api_key_env or _PRESETS.get(self.provider, {}).get("api_key_env")
+        if not env:
+            raise ValueError(f"no api_key_env for provider {self.provider!r}")
+        key = os.environ.get(env, "")
+        if not key and self.provider != "vllm":
+            raise RuntimeError(
+                f"{env} not set — required for provider {self.provider!r}"
+            )
+        return key or "not-needed"
+
+
+@dataclass
+class ChatReply:
+    """Normalized model reply."""
+
+    text: str
+    tool_calls: list[dict]  # [{"name": str, "arguments": dict}]
+    raw: dict = field(default_factory=dict)
+
+
+class ChatProvider:
+    def complete(self, messages: list[dict], tools: list[dict]) -> ChatReply:
+        raise NotImplementedError
+
+
+class OpenAICompatibleProvider(ChatProvider):
+    """OpenAI /chat/completions with `tools`. vLLM + OpenRouter + OpenAI."""
+
+    def __init__(self, cfg: ProviderConfig):
+        self.cfg = cfg
+        self._client = httpx.Client(timeout=cfg.timeout_s)
+
+    def complete(self, messages: list[dict], tools: list[dict]) -> ChatReply:
+        cfg = self.cfg
+        headers = {
+            "Authorization": f"Bearer {cfg.resolved_api_key()}",
+            "Content-Type": "application/json",
+            **cfg.extra_headers,
+        }
+        body: dict[str, Any] = {
+            "model": cfg.model,
+            "messages": messages,
+            "temperature": cfg.temperature,
+            "max_tokens": cfg.max_tokens,
+        }
+        if tools:
+            body["tools"] = tools
+            body["tool_choice"] = "auto"
+        resp = self._client.post(
+            f"{cfg.resolved_base_url()}/chat/completions",
+            headers=headers,
+            json=body,
+        )
+        resp.raise_for_status()
+        data = resp.json()
+        msg = data["choices"][0]["message"]
+        calls: list[dict] = []
+        for tc in msg.get("tool_calls") or []:
+            fn = tc.get("function", {})
+            args = fn.get("arguments", {})
+            if isinstance(args, str):
+                try:
+                    args = json.loads(args or "{}")
+                except json.JSONDecodeError:
+                    args = {}
+            calls.append({"name": fn.get("name", ""), "arguments": args})
+        return ChatReply(text=msg.get("content") or "", tool_calls=calls, raw=data)
+
+    def close(self) -> None:
+        self._client.close()
+
+
+class BedrockProvider(ChatProvider):
+    """AWS Bedrock Converse. Wired but not yet implemented."""
+
+    def __init__(self, cfg: ProviderConfig):
+        self.cfg = cfg
+
+    def complete(self, messages: list[dict], tools: list[dict]) -> ChatReply:
+        raise NotImplementedError(
+            "BedrockProvider not implemented yet. Use provider='openrouter' "
+            "or 'vllm' for Phase 0; Bedrock Converse adapter is a tracked "
+            "follow-up (needs boto3 + message/tool shape translation)."
+        )
+
+
+def make_provider(cfg: ProviderConfig) -> ChatProvider:
+    if cfg.provider == "bedrock":
+        return BedrockProvider(cfg)
+    if cfg.provider in ("openai", "vllm", "openrouter"):
+        return OpenAICompatibleProvider(cfg)
+    raise ValueError(f"unknown provider {cfg.provider!r}")

tests/test_agent.py

@@ -0,0 +1,149 @@
+"""Agent + provider tests.
+
+Offline by default: a FakeProvider returns scripted tool calls so the
+briefing build, tool-schema filtering, tool-call parsing, and the full
+ModelAgent -> eval_core -> live-Rust loop are all asserted without a
+network. The OpenRouter live test is opt-in (skips without an API key).
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import pytest
+
+ot = pytest.importorskip("openra_train", reason="Rust env wheel not installed")
+
+from openra_bench.agent import ModelAgent, build_briefing, _to_commands, _tool_schemas
+from openra_bench.providers import ChatProvider, ChatReply, ProviderConfig
+
+TRAIN = Path("/Users/berta/Projects/OpenRA-RL-Training")
+PACKS = Path(__file__).parent.parent / "openra_bench" / "scenarios" / "packs"
+
+
+class FakeProvider(ChatProvider):
+    """Returns a fixed sequence of tool-call sets, then observe()."""
+
+    def __init__(self, script: list[list[dict]]):
+        self.script = script
+        self.i = 0
+        self.seen_messages: list[list[dict]] = []
+
+    def complete(self, messages, tools):
+        self.seen_messages.append(messages)
+        calls = self.script[self.i] if self.i < len(self.script) else [
+            {"name": "observe", "arguments": {}}
+        ]
+        self.i += 1
+        return ChatReply(text="", tool_calls=calls)
+
+
+def test_tool_schema_filtering():
+    only_move = _tool_schemas(["move_units"])
+    names = {t["function"]["name"] for t in only_move}
+    assert "move_units" in names
+    assert "attack_unit" not in names
+    assert "observe" in names, "a no-op must always be offered"
+    assert {t["function"]["name"] for t in _tool_schemas(None)} == {
+        "move_units",
+        "attack_unit",
+        "observe",
+    }
+
+
+def test_build_briefing_format():
+    rs = {
+        "game_tick": 120,
+        "explored_percent": 12.5,
+        "units_summary": [
+            {"id": "1001", "type": "jeep", "cell_x": 5, "cell_y": 6, "activity": "idle"}
+        ],
+        "enemy_summary": [],
+    }
+    b = build_briefing(rs, objective="find the base")
+    assert "OBJECTIVE: find the base" in b
+    assert "1001 jeep @(5,6)" in b
+    assert "tick=120" in b and "explored=12.5%" in b
+    assert "none (scout the fog)" in b
+
+
+def test_tool_call_parsing_and_aliases():
+    cmds = _to_commands(
+        [
+            {"name": "move_units", "arguments": {"unit_ids": [1, 2], "target_x": 9, "target_y": 4}},
+            {"name": "attack_target", "arguments": {"unit_ids": [3], "target_id": 77}},  # alias
+            {"name": "observe", "arguments": {}},
+            {"name": "garbage", "arguments": {}},  # dropped
+            {"name": "move_units", "arguments": {"unit_ids": [1]}},  # malformed -> dropped
+        ],
+        ot.Command,
+    )
+    # 3 valid (move, attack-alias, observe); 2 dropped
+    assert len(cmds) == 3
+
+
+def test_model_agent_drives_live_rust_with_fake_provider():
+    from openra_bench.eval_core import run_level
+    from openra_bench.scenarios import load_pack
+    from openra_bench.scenarios.loader import compile_level
+
+    pack = load_pack(PACKS / "perception-frontier-reading.yaml")
+    compiled = compile_level(pack, "easy")
+
+    # Scripted "scout east" behaviour, then observe forever.
+    fake = FakeProvider(
+        [[{"name": "move_units", "arguments": {"unit_ids": [1001], "target_x": 100, "target_y": 20}}]]
+        * 8
+    )
+    agent = ModelAgent(
+        ProviderConfig(vision=False),
+        allowed_tools=compiled.scenario.tools,
+        objective=compiled.scenario.description,
+        provider=fake,
+    )
+    res = run_level(compiled, agent.agent_fn, seed=1)
+
+    assert res.outcome in {"win", "draw", "loss"}
+    assert res.turns >= 1 and len(res.trace) == res.turns
+    assert agent.stats["turns"] == res.turns
+    # Provider actually saw a system prompt + a user briefing.
+    first = fake.seen_messages[0]
+    assert first[0]["role"] == "system"
+    assert any(m["role"] == "user" for m in first)
+
+
+def test_history_strips_stale_images():
+    hist = [
+        {"role": "user", "content": [{"type": "text", "text": "t1"}, {"type": "image_url", "image_url": {}}]},
+        {"role": "user", "content": [{"type": "text", "text": "t2"}, {"type": "image_url", "image_url": {}}]},
+    ]
+    ModelAgent._strip_old_images(hist)
+    assert isinstance(hist[0]["content"], str) and hist[0]["content"] == "t1"
+    assert isinstance(hist[1]["content"], list)  # newest image kept
+
+
+@pytest.mark.skipif(
+    not os.environ.get("OPENROUTER_API_KEY"),
+    reason="set OPENROUTER_API_KEY to run the live provider test",
+)
+def test_openrouter_live_smoke():
+    from openra_bench.eval_core import run_level
+    from openra_bench.scenarios import load_pack
+    from openra_bench.scenarios.loader import compile_level
+
+    pack = load_pack(PACKS / "perception-frontier-reading.yaml")
+    compiled = compile_level(pack, "easy")
+    agent = ModelAgent(
+        ProviderConfig(
+            provider="openrouter",
+            model=os.environ.get("OPENROUTER_MODEL", "anthropic/claude-3.5-sonnet"),
+            vision=False,
+            max_tokens=512,
+        ),
+        allowed_tools=compiled.scenario.tools,
+        objective=compiled.scenario.description,
+    )
+    res = run_level(compiled, agent.agent_fn, seed=1)
+    assert res.outcome in {"win", "draw", "loss"}
+    assert agent.stats["tool_calls"] >= 1, "model issued no usable tool calls"