synth-ai 0.2.9.dev4__py3-none-any.whl → 0.2.9.dev6__py3-none-any.whl

examples/warming_up_to_rl/run_eval.py

@@ -5,17 +5,22 @@ Baseline evaluation script (public-friendly skeleton)
 - Uses a TaskAppClient interface (to be implemented in synth-ai SDK)
 - Keeps structure aligned with research/testing/crafter eval harness
 """
+
 from __future__ import annotations
-import os
+
+import argparse
+import asyncio
+import contextlib
 import json
+import os
 import re
-from typing import Any, Dict, List, Optional
-from collections import Counter
-import asyncio
-import httpx
-import argparse
 import tomllib
+from collections import Counter
 from pathlib import Path
+from typing import Any
+
+import httpx
+
 
 class TaskAppClient:
     """Minimal async client for the task app initialize/step/terminate routes.
@@ -23,12 +28,12 @@ class TaskAppClient:
     This is a public-friendly shim for examples, pending SDK surface consolidation.
     """
 
-    def __init__(self, base_url: str, api_key: Optional[str] = None) -> None:
+    def __init__(self, base_url: str, api_key: str | None = None) -> None:
         self.base_url = base_url.rstrip("/")
         self.api_key = api_key
-        self._client: Optional[httpx.AsyncClient] = None
+        self._client: httpx.AsyncClient | None = None
 
-    async def __aenter__(self) -> "TaskAppClient":
+    async def __aenter__(self) -> TaskAppClient:
         headers = {}
         if self.api_key:
             headers["X-API-Key"] = self.api_key
@@ -54,9 +59,9 @@ class TaskAppClient:
             )
         return self._client
 
-    async def initialize(self, env_name: str, config: Dict[str, Any]) -> Dict[str, Any]:
+    async def initialize(self, env_name: str, config: dict[str, Any]) -> dict[str, Any]:
         """POST /env/{env_name}/initialize (compat route supported in task app)."""
-        payload: Dict[str, Any] = {
+        payload: dict[str, Any] = {
             "seed": config.get("seed"),
         }
         # Allow both world_config and config inputs; env routes will normalize difficulty
@@ -68,29 +73,31 @@ class TaskAppClient:
         resp.raise_for_status()
         return resp.json()
 
-    async def step(self, env_name: str, env_id: str, tool_calls: List[Dict[str, Any]]) -> Dict[str, Any]:
+    async def step(
+        self, env_name: str, env_id: str, tool_calls: list[dict[str, Any]]
+    ) -> dict[str, Any]:
         """POST /env/{env_name}/step with wrapped tool_calls in action."""
         payload = {"env_id": env_id, "action": {"tool_calls": tool_calls}}
         resp = await self.client.post(f"/env/{env_name}/step", json=payload)
         resp.raise_for_status()
         return resp.json()
 
-    async def terminate(self, env_name: str, env_id: str) -> Dict[str, Any]:
+    async def terminate(self, env_name: str, env_id: str) -> dict[str, Any]:
         resp = await self.client.post(f"/env/{env_name}/terminate", json={"env_id": env_id})
         resp.raise_for_status()
         return resp.json()
 
-    async def get_info(self) -> Dict[str, Any]:
+    async def get_info(self) -> dict[str, Any]:
         resp = await self.client.get("/info")
         resp.raise_for_status()
         return resp.json()
 
-    async def proxy_groq_chat(self, payload: Dict[str, Any]) -> Dict[str, Any]:
+    async def proxy_groq_chat(self, payload: dict[str, Any]) -> dict[str, Any]:
         resp = await self.client.post("/proxy/groq/v1/chat/completions", json=payload)
         resp.raise_for_status()
         return resp.json()
 
-    async def vllm_chat(self, vllm_base_url: str, payload: Dict[str, Any]) -> Dict[str, Any]:
+    async def vllm_chat(self, vllm_base_url: str, payload: dict[str, Any]) -> dict[str, Any]:
         async with httpx.AsyncClient(base_url=vllm_base_url.rstrip("/"), timeout=60.0) as c:
             resp = await c.post("/v1/chat/completions", json=payload)
             # Do not raise for status to surface body in errors
@@ -102,11 +109,21 @@ class TaskAppClient:
                 return {"error": data}
             return data
 
-    async def rollout(self, *, run_id: str, env_name: str, seed: int, difficulty: str, policy_name: str, policy_config: Dict[str, Any], max_turns: int) -> Dict[str, Any]:
-        ops: List[str] = []
+    async def rollout(
+        self,
+        *,
+        run_id: str,
+        env_name: str,
+        seed: int,
+        difficulty: str,
+        policy_name: str,
+        policy_config: dict[str, Any],
+        max_turns: int,
+    ) -> dict[str, Any]:
+        ops: list[str] = []
         for _ in range(max_turns):
             ops.extend(["agent", "env"])
-        payload: Dict[str, Any] = {
+        payload: dict[str, Any] = {
             "run_id": run_id,
             "env": {
                 "env_name": env_name,
@@ -128,35 +145,41 @@ class TaskAppClient:
         resp.raise_for_status()
         return resp.json()
 
+
 TASK_APP_URL = os.getenv("TASK_APP_URL", "https://YOUR-TASK-APP.modal.run").rstrip("/")
 MODEL = os.getenv("EVAL_MODEL", "qwen/qwen3-32b")
 NUM_EPISODES = int(os.getenv("NUM_EPISODES", "3"))
 MAX_TURNS = int(os.getenv("MAX_TURNS", "10"))
 CONCURRENCY = int(os.getenv("CONCURRENCY", "1"))
 
-def _interact_tool_schema() -> List[Dict[str, Any]]:
-    return [{
-        "type": "function",
-        "function": {
-            "name": "interact",
-            "description": "Perform actions in the Crafter environment.",
-            "parameters": {
-                "type": "object",
-                "properties": {
-                    "actions": {"type": "array", "items": {"type": "string"}},
-                    "reasoning": {"type": "string"},
+
+def _interact_tool_schema() -> list[dict[str, Any]]:
+    return [
+        {
+            "type": "function",
+            "function": {
+                "name": "interact",
+                "description": "Perform actions in the Crafter environment.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "actions": {"type": "array", "items": {"type": "string"}},
+                        "reasoning": {"type": "string"},
+                    },
+                    "required": ["actions", "reasoning"],
                 },
-                "required": ["actions", "reasoning"],
             },
-        },
-    }]
+        }
+    ]
+
 
-def _build_messages_from_observation(observation: Dict[str, Any], history: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+def _build_messages_from_observation(
+    observation: dict[str, Any], history: list[dict[str, Any]]
+) -> list[dict[str, Any]]:
     inv = observation.get("inventory") or {}
     pos = observation.get("player_position") or []
     ach = observation.get("achievements_status") or {}
-    turns_taken = observation.get("num_steps_taken") or 0
-    user_lines: List[str] = []
+    user_lines: list[str] = []
     user_lines.append("Environment: CrafterClassic")
     user_lines.append(f"Player position: {pos}")
     user_lines.append(f"Inventory: {json.dumps(inv, ensure_ascii=False)}")
@@ -171,7 +194,8 @@ def _build_messages_from_observation(observation: Dict[str, Any], history: List[
     content = "\n".join(user_lines)
     return [{"role": "user", "content": content}]
 
-def _parse_tool_calls_from_openai_response(data: Dict[str, Any]) -> List[str]:
+
+def _parse_tool_calls_from_openai_response(data: dict[str, Any]) -> list[str]:
     try:
         choices = data.get("choices")
         if isinstance(choices, list) and choices:
@@ -203,7 +227,11 @@ def _parse_tool_calls_from_openai_response(data: Dict[str, Any]) -> List[str]:
         if isinstance(content, str):
             text = content
         elif isinstance(content, list):
-            text = "\n".join(str(part.get("text")) for part in content if isinstance(part, dict) and part.get("text"))
+            text = "\n".join(
+                str(part.get("text"))
+                for part in content
+                if isinstance(part, dict) and part.get("text")
+            )
         for raw in re.findall(r"\{[\s\S]*\}", text or ""):
             try:
                 obj = json.loads(raw)
@@ -217,9 +245,16 @@ def _parse_tool_calls_from_openai_response(data: Dict[str, Any]) -> List[str]:
         pass
     return []
 
-async def _choose_actions_via_llm(client: TaskAppClient, provider: str, model: str, observation: Dict[str, Any], history: List[Dict[str, Any]]) -> List[str]:
+
+async def _choose_actions_via_llm(
+    client: TaskAppClient,
+    provider: str,
+    model: str,
+    observation: dict[str, Any],
+    history: list[dict[str, Any]],
+) -> list[str]:
     messages = _build_messages_from_observation(observation, history)
-    payload: Dict[str, Any] = {
+    payload: dict[str, Any] = {
         "model": model,
         "messages": messages,
         "tools": _interact_tool_schema(),
@@ -245,33 +280,40 @@ async def _choose_actions_via_llm(client: TaskAppClient, provider: str, model: s
     actions = _parse_tool_calls_from_openai_response(data)
     return actions or []
 
-def _expand_actions_to_tool_calls(actions: List[str]) -> List[Dict[str, Any]]:
-    out: List[Dict[str, Any]] = []
+
+def _expand_actions_to_tool_calls(actions: list[str]) -> list[dict[str, Any]]:
+    out: list[dict[str, Any]] = []
     for a in actions[:5]:
         out.append({"tool": "interact", "args": {"action": a}})
     return out
 
+
 def _detect_provider(model: str) -> str:
     m = (model or "").lower()
     if "qwen/qwen3-32b" in m or "qwen-2.5-" in m or m.startswith("groq:"):
         return "groq"
     return "vllm"
 
-def _rollout_inference_url_from_cfg(cfg: Dict[str, Any], default_vllm: Optional[str]) -> Optional[str]:
+
+def _rollout_inference_url_from_cfg(cfg: dict[str, Any], default_vllm: str | None) -> str | None:
     # Prefer explicit inference_url in TOML; else fall back to discovered vLLM base
     url = cfg.get("inference_url")
     if isinstance(url, str) and url:
         return url
     return default_vllm
 
-async def eval_episode(client: TaskAppClient, seed: int) -> Dict[str, Any]:
+
+async def eval_episode(client: TaskAppClient, seed: int) -> dict[str, Any]:
     env_name = "CrafterClassic"
-    history: List[Dict[str, Any]] = []
+    history: list[dict[str, Any]] = []
     achievements: set[str] = set()
     turns = 0
 
     # Initialize environment
-    init_cfg: Dict[str, Any] = {"seed": seed, "world_config": {"difficulty": os.getenv("DIFFICULTY", "easy")}}
+    init_cfg: dict[str, Any] = {
+        "seed": seed,
+        "world_config": {"difficulty": os.getenv("DIFFICULTY", "easy")},
+    }
     created = await client.initialize(env_name, init_cfg)
     env_id = created.get("env_id")
     if not isinstance(env_id, str) or not env_id:
@@ -285,7 +327,9 @@ async def eval_episode(client: TaskAppClient, seed: int) -> Dict[str, Any]:
     try:
         while turns < MAX_TURNS and not done:
             # Ask LLM for actions; fallback to a simple exploratory pair
-            chosen_actions = await _choose_actions_via_llm(client, provider, MODEL, observation, history)
+            chosen_actions = await _choose_actions_via_llm(
+                client, provider, MODEL, observation, history
+            )
             if not chosen_actions:
                 chosen_actions = ["move_up", "do"]
             tool_calls = _expand_actions_to_tool_calls(chosen_actions)
@@ -299,13 +343,12 @@ async def eval_episode(client: TaskAppClient, seed: int) -> Dict[str, Any]:
                 if isinstance(nxt, dict):
                     observation = nxt
     finally:
-        try:
+        with contextlib.suppress(Exception):
             await client.terminate(env_name, env_id)
-        except Exception:
-            pass
 
     return {"seed": seed, "turns": turns, "achievements": sorted(achievements)}
 
+
 async def main() -> None:
     # Best-effort load local .env if present (ensures ENVIRONMENT_API_KEY for rollout)
     try:
@@ -322,13 +365,17 @@ async def main() -> None:
     except Exception:
         pass
 
-    parser = argparse.ArgumentParser(description="Baseline eval against task app with optional TOML config")
+    parser = argparse.ArgumentParser(
+        description="Baseline eval against task app with optional TOML config"
+    )
     parser.add_argument("--toml", help="Path to TOML config file", default=None)
-    parser.add_argument("--use-rollout", action="store_true", help="Use server-side rollout endpoint for eval")
+    parser.add_argument(
+        "--use-rollout", action="store_true", help="Use server-side rollout endpoint for eval"
+    )
     args = parser.parse_args()
 
     global TASK_APP_URL, MODEL, NUM_EPISODES, MAX_TURNS, CONCURRENCY
-    cfg: Dict[str, Any] = {}
+    cfg: dict[str, Any] = {}
     if args.toml:
         with open(args.toml, "rb") as f:
             cfg = tomllib.load(f)
@@ -346,10 +393,14 @@ async def main() -> None:
             if env_url:
                 TASK_APP_URL = env_url.rstrip("/")
             else:
-                raise RuntimeError("TASK_APP_URL is a placeholder. Set task_app_url in TOML or export TASK_APP_URL.")
+                raise RuntimeError(
+                    "TASK_APP_URL is a placeholder. Set task_app_url in TOML or export TASK_APP_URL."
+                )
 
     print(f"Task App: {TASK_APP_URL}")
-    print(f"Model: {MODEL} Episodes: {NUM_EPISODES} Max turns: {MAX_TURNS} Concurrency: {CONCURRENCY}")
+    print(
+        f"Model: {MODEL} Episodes: {NUM_EPISODES} Max turns: {MAX_TURNS} Concurrency: {CONCURRENCY}"
+    )
     sem = asyncio.Semaphore(max(CONCURRENCY, 1))
     async with TaskAppClient(TASK_APP_URL, api_key=os.getenv("ENVIRONMENT_API_KEY")) as client:
         if args.use_rollout:
@@ -359,16 +410,24 @@ async def main() -> None:
             inf_url = _rollout_inference_url_from_cfg(cfg, default_vllm)
             if not inf_url:
                 raise RuntimeError("Could not resolve inference URL for rollout")
+
             async def _run(seed: int):
                 async with sem:
                     try:
                         run_id = f"eval-{seed}"
                         # Build policy config from TOML (explicit control; no server-side guessing)
-                        policy_cfg: Dict[str, Any] = {
+                        policy_cfg: dict[str, Any] = {
                             "model": cfg.get("model", MODEL),
                             "inference_url": inf_url,
                         }
-                        for k in ("max_tokens", "temperature", "top_p", "thinking_mode", "thinking_budget", "use_tools"):
+                        for k in (
+                            "max_tokens",
+                            "temperature",
+                            "top_p",
+                            "thinking_mode",
+                            "thinking_budget",
+                            "use_tools",
+                        ):
                             if k in cfg and cfg.get(k) is not None:
                                 policy_cfg[k] = cfg.get(k)
 
@@ -385,8 +444,16 @@ async def main() -> None:
                         ach = []
                         try:
                             trajs = r.get("trajectories") or []
-                            final_obs = (trajs[0].get("final") or {}).get("observation") if trajs and isinstance(trajs[0], dict) else None
-                            ach_map = (final_obs or {}).get("achievements_status") if isinstance(final_obs, dict) else None
+                            final_obs = (
+                                (trajs[0].get("final") or {}).get("observation")
+                                if trajs and isinstance(trajs[0], dict)
+                                else None
+                            )
+                            ach_map = (
+                                (final_obs or {}).get("achievements_status")
+                                if isinstance(final_obs, dict)
+                                else None
+                            )
                             if isinstance(ach_map, dict):
                                 ach = sorted([k for k, v in ach_map.items() if v])
                         except Exception:
@@ -401,7 +468,11 @@ async def main() -> None:
                         return {"seed": seed, "turns": length, "achievements": ach}
                     except Exception as e:
                         return {"seed": seed, "turns": 0, "achievements": [], "error": str(e)}
-            results = await asyncio.gather(*[asyncio.create_task(_run(i)) for i in range(1, NUM_EPISODES + 1)], return_exceptions=False)
+
+            results = await asyncio.gather(
+                *[asyncio.create_task(_run(i)) for i in range(1, NUM_EPISODES + 1)],
+                return_exceptions=False,
+            )
             # Aggregate summary
             counts = [len(r.get("achievements") or []) for r in results if isinstance(r, dict)]
             turns = [int(r.get("turns") or 0) for r in results if isinstance(r, dict)]
@@ -424,11 +495,16 @@ async def main() -> None:
             }
             print(json.dumps(summary, indent=2))
         else:
+
             async def _run(seed: int):
                 async with sem:
                     return await eval_episode(client, seed)
-            results = await asyncio.gather(*[asyncio.create_task(_run(i)) for i in range(1, NUM_EPISODES + 1)])
+
+            results = await asyncio.gather(
+                *[asyncio.create_task(_run(i)) for i in range(1, NUM_EPISODES + 1)]
+            )
             print(json.dumps({"episodes": results}, indent=2))
 
+
 if __name__ == "__main__":
     asyncio.run(main())