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

examples/rl/run_eval.py

@@ -5,27 +5,30 @@ from __future__ import annotations
 
 import argparse
 import asyncio
+import contextlib
 import json
 import os
-from pathlib import Path
-from typing import Any, Dict, List, Optional, Tuple
+import tomllib
+from typing import Any
 
 import httpx
-import tomllib
 
 
 class TaskAppClient:
     """Minimal async client for math single-step task app."""
 
-    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 = {"X-API-Key": self.api_key} if self.api_key else {}
         self._client = httpx.AsyncClient(
-            base_url=self.base_url, headers=headers, timeout=httpx.Timeout(120.0), follow_redirects=True
+            base_url=self.base_url,
+            headers=headers,
+            timeout=httpx.Timeout(120.0),
+            follow_redirects=True,
         )
         return self
 
@@ -39,36 +42,37 @@ class TaskAppClient:
         if self._client is None:
             headers = {"X-API-Key": self.api_key} if self.api_key else {}
             self._client = httpx.AsyncClient(
-                base_url=self.base_url, headers=headers, timeout=httpx.Timeout(120.0), follow_redirects=True
+                base_url=self.base_url,
+                headers=headers,
+                timeout=httpx.Timeout(120.0),
+                follow_redirects=True,
             )
         return self._client
 
-    async def initialize(self, split: str, seed: int | None) -> Dict[str, Any]:
-        payload: Dict[str, Any] = {"config": {"split": split}}
+    async def initialize(self, split: str, seed: int | None) -> dict[str, Any]:
+        payload: dict[str, Any] = {"config": {"split": split}}
         if seed is not None:
             payload["seed"] = seed
         resp = await self.client.post("/env/math/initialize", json=payload)
         resp.raise_for_status()
         return resp.json()
 
-    async def step(self, env_id: str, tool_calls: List[Dict[str, Any]]) -> Dict[str, Any]:
+    async def step(self, env_id: str, tool_calls: list[dict[str, Any]]) -> dict[str, Any]:
         payload = {"env_id": env_id, "action": {"tool_calls": tool_calls}}
         resp = await self.client.post("/env/math/step", json=payload)
         resp.raise_for_status()
         return resp.json()
 
     async def terminate(self, env_id: str) -> None:
-        try:
+        with contextlib.suppress(Exception):
             await self.client.post("/env/math/terminate", json={"env_id": env_id})
-        except Exception:
-            pass
 
-    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 rollout(self, payload: Dict[str, Any]) -> Dict[str, Any]:
+    async def rollout(self, payload: dict[str, Any]) -> dict[str, Any]:
         resp = await self.client.post("/rollout", json=payload)
         resp.raise_for_status()
         return resp.json()
@@ -76,10 +80,10 @@ class TaskAppClient:
     async def post_inference(
         self,
         url: str,
-        payload: Dict[str, Any],
+        payload: dict[str, Any],
         *,
-        headers: Dict[str, str] | None = None,
-    ) -> Dict[str, Any]:
+        headers: dict[str, str] | None = None,
+    ) -> dict[str, Any]:
         async with httpx.AsyncClient(timeout=httpx.Timeout(60.0)) as c:
             resp = await c.post(url, json=payload, headers=headers)
         resp.raise_for_status()
@@ -90,7 +94,7 @@ TOOL_NAME = "math_submit"
 DEFAULT_SPLIT = os.getenv("MATH_EVAL_DEFAULT_SPLIT", "validation")
 
 
-def _math_tool_schema() -> List[Dict[str, Any]]:
+def _math_tool_schema() -> list[dict[str, Any]]:
     return [
         {
             "type": "function",
@@ -103,8 +107,7 @@ def _math_tool_schema() -> List[Dict[str, Any]]:
                         "answer": {
                             "type": "string",
                             "description": "Final answer in simplest form",
-                        }
-                        ,
+                        },
                         "explanation": {
                             "type": "string",
                             "description": "Optional explanation of reasoning",
@@ -118,7 +121,7 @@ def _math_tool_schema() -> List[Dict[str, Any]]:
     ]
 
 
-def _build_messages(problem: str) -> List[Dict[str, Any]]:
+def _build_messages(problem: str) -> list[dict[str, Any]]:
     return [
         {
             "role": "system",
@@ -134,18 +137,18 @@ def _build_messages(problem: str) -> List[Dict[str, Any]]:
     ]
 
 
-def _parse_tool_calls(data: Dict[str, Any]) -> List[Dict[str, Any]]:
+def _parse_tool_calls(data: dict[str, Any]) -> list[dict[str, Any]]:
     choices = data.get("choices") or []
     if not choices:
         return []
     message = choices[0].get("message") or {}
     raw_calls = message.get("tool_calls") or []
-    tool_calls: List[Dict[str, Any]] = []
+    tool_calls: list[dict[str, Any]] = []
     for call in raw_calls:
         function = call.get("function") or {}
         name = function.get("name")
         arguments = function.get("arguments")
-        parsed_args: Dict[str, Any]
+        parsed_args: dict[str, Any]
         if isinstance(arguments, str):
             try:
                 parsed_args = json.loads(arguments)
@@ -159,7 +162,7 @@ def _parse_tool_calls(data: Dict[str, Any]) -> List[Dict[str, Any]]:
     return tool_calls
 
 
-def _detect_provider(model: str, hint: Optional[str]) -> str:
+def _detect_provider(model: str, hint: str | None) -> str:
     if hint:
         return hint.lower()
     lowered = (model or "").lower()
@@ -188,10 +191,10 @@ async def _choose_actions(
     provider: str,
     model: str,
     problem: str,
-    policy_cfg: Dict[str, Any],
-) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]:
+    policy_cfg: dict[str, Any],
+) -> tuple[list[dict[str, Any]], dict[str, Any]]:
     messages = _build_messages(problem)
-    payload: Dict[str, Any] = {
+    payload: dict[str, Any] = {
         "model": model,
         "messages": messages,
         "tools": _math_tool_schema(),
@@ -203,9 +206,7 @@ async def _choose_actions(
 
     if provider == "groq":
         # Task app proxies Groq requests; reuse existing headers on the client
-        response = await client.client.post(
-            "/proxy/groq/v1/chat/completions", json=payload
-        )
+        response = await client.client.post("/proxy/groq/v1/chat/completions", json=payload)
         response.raise_for_status()
         body = response.json()
     else:
@@ -223,23 +224,20 @@ async def _choose_actions(
                 headers=headers or None,
             )
         except httpx.ReadTimeout as exc:
-            raise RuntimeError(
-                "Inference request timed out. Check the inference service." ) from exc
+            raise RuntimeError("Inference request timed out. Check the inference service.") from exc
         try:
             body = response.json()
         except Exception:
             body = {"raw": response.text[:800]}
         if response.status_code >= 500:
-            raise RuntimeError(
-                f"Inference server error {response.status_code}: {body}")
+            raise RuntimeError(f"Inference server error {response.status_code}: {body}")
         if response.status_code >= 400:
-            raise RuntimeError(
-                f"Inference request invalid ({response.status_code}): {body}")
+            raise RuntimeError(f"Inference request invalid ({response.status_code}): {body}")
     tool_calls = _parse_tool_calls(body)
     return tool_calls, body
 
 
-def _tool_to_answer(tool_calls: List[Dict[str, Any]]) -> str:
+def _tool_to_answer(tool_calls: list[dict[str, Any]]) -> str:
     if not tool_calls:
         return ""
     args = tool_calls[0].get("args") or {}
@@ -251,11 +249,11 @@ async def eval_episode(
     client: TaskAppClient,
     *,
     split: str,
-    seed: Optional[int],
+    seed: int | None,
     model: str,
     provider: str,
-    policy_cfg: Dict[str, Any],
-) -> Dict[str, Any]:
+    policy_cfg: dict[str, Any],
+) -> dict[str, Any]:
     created = await client.initialize(split, seed)
     env_id = created["env_id"]
     observation = created.get("observation") or {}
@@ -288,10 +286,10 @@ async def eval_via_rollout(
     *,
     run_id: str,
     split: str,
-    seed: Optional[int],
+    seed: int | None,
     model: str,
-    policy_cfg: Dict[str, Any],
-) -> Dict[str, Any]:
+    policy_cfg: dict[str, Any],
+) -> dict[str, Any]:
     payload = {
         "run_id": run_id,
         "env": {
@@ -314,6 +312,7 @@ async def eval_via_rollout(
     steps = traj.get("steps") or []
     step = steps[0] if steps else {}
     info = step.get("info") or {}
+    observation = step.get("obs") or {}
     return {
         "seed": seed,
         "split": split,
@@ -328,14 +327,14 @@ async def eval_via_rollout(
     }
 
 
-def _load_config(path: Optional[str]) -> Dict[str, Any]:
+def _load_config(path: str | None) -> dict[str, Any]:
     if not path:
         return {}
     with open(path, "rb") as fh:
         return tomllib.load(fh)
 
 
-def _default_policy_cfg(cfg: Dict[str, Any]) -> Dict[str, Any]:
+def _default_policy_cfg(cfg: dict[str, Any]) -> dict[str, Any]:
     policy = dict(cfg.get("policy") or {})
     if "inference_url" not in policy:
         env_url = os.getenv("INFERENCE_URL")
@@ -371,15 +370,17 @@ async def main() -> None:
     seed_start = int(cfg.get("seed_start") or 0)
 
     policy_cfg = _default_policy_cfg(cfg)
-    provider_hint = cfg.get("provider") or cfg.get("policy", {}).get("provider") or policy_cfg.get("provider")
+    provider_hint = (
+        cfg.get("provider") or cfg.get("policy", {}).get("provider") or policy_cfg.get("provider")
+    )
     provider = _detect_provider(model, provider_hint)
     policy_cfg.pop("provider", None)
 
     api_key = os.getenv("ENVIRONMENT_API_KEY")
 
     successes = 0
-    failures: Dict[str, int] = {}
-    results: List[Dict[str, Any]] = []
+    failures: dict[str, int] = {}
+    results: list[dict[str, Any]] = []
 
     async with TaskAppClient(task_app_url, api_key=api_key) as client:
         for episode in range(episodes):
@@ -412,7 +413,7 @@ async def main() -> None:
             problem = data.get("problem")
             tool_calls = data.get("tool_calls") or []
             print(
-                f"Episode {episode+1}/{episodes} seed={seed} status={status} reward={data.get('reward')}\n"
+                f"Episode {episode + 1}/{episodes} seed={seed} status={status} reward={data.get('reward')}\n"
                 f"  problem: {problem!r}\n"
                 f"  tool   : {tool_calls!r}\n"
                 f"  answer : {answer!r}\n  expected: {expected!r}",

examples/rl/run_rl_and_save.py

@@ -7,14 +7,14 @@ import argparse
 import json
 import os
 import sys
+import tomllib
 from pathlib import Path
-from typing import Any, Dict
+from typing import Any
 
 import requests
-import tomllib
 
 
-def _load_toml(path: Path) -> Dict[str, Any]:
+def _load_toml(path: Path) -> dict[str, Any]:
     if not path.exists():
         print(f"config not found: {path}", file=sys.stderr)
         sys.exit(2)
@@ -24,10 +24,18 @@ def _load_toml(path: Path) -> Dict[str, Any]:
 
 def main() -> None:
     parser = argparse.ArgumentParser(description="Create math RL job via backend RL endpoint")
-    parser.add_argument("--backend", default=os.getenv("BACKEND_BASE_URL", "http://localhost:8000/api"))
+    parser.add_argument(
+        "--backend", default=os.getenv("BACKEND_BASE_URL", "http://localhost:8000/api")
+    )
     parser.add_argument("--config", required=True, help="Path to RL TOML config")
-    parser.add_argument("--task-url", default=os.getenv("TASK_APP_URL", ""), help="Override task service URL")
-    parser.add_argument("--idempotency", default=os.getenv("RL_IDEMPOTENCY_KEY", ""), help="Optional Idempotency-Key header")
+    parser.add_argument(
+        "--task-url", default=os.getenv("TASK_APP_URL", ""), help="Override task service URL"
+    )
+    parser.add_argument(
+        "--idempotency",
+        default=os.getenv("RL_IDEMPOTENCY_KEY", ""),
+        help="Optional Idempotency-Key header",
+    )
     args = parser.parse_args()
 
     cfg_path = Path(args.config).expanduser()
@@ -35,19 +43,29 @@ def main() -> None:
 
     services = cfg.get("services") if isinstance(cfg.get("services"), dict) else {}
 
-    task_url = (args.task_url or "").strip() or (os.getenv("TASK_APP_URL") or "").strip() or (services.get("task_url") or "").strip()
+    task_url = (
+        (args.task_url or "").strip()
+        or (os.getenv("TASK_APP_URL") or "").strip()
+        or (services.get("task_url") or "").strip()
+    )
     if not task_url:
-        print("Missing task service URL. Provide --task-url or set TASK_APP_URL or services.task_url in TOML", file=sys.stderr)
+        print(
+            "Missing task service URL. Provide --task-url or set TASK_APP_URL or services.task_url in TOML",
+            file=sys.stderr,
+        )
         sys.exit(2)
 
     model_cfg = cfg.get("model") if isinstance(cfg.get("model"), dict) else {}
     has_source = bool((model_cfg.get("source") or "").strip())
     has_base = bool((model_cfg.get("base") or "").strip())
     if has_source == has_base:
-        print("Model section must specify exactly one of [model].source or [model].base", file=sys.stderr)
+        print(
+            "Model section must specify exactly one of [model].source or [model].base",
+            file=sys.stderr,
+        )
         sys.exit(2)
 
-    payload: Dict[str, Any] = {
+    payload: dict[str, Any] = {
         "job_type": "rl",
         "compute": cfg.get("compute", {}),
         "data": {
@@ -59,7 +77,7 @@ def main() -> None:
 
     backend = str(args.backend).rstrip("/")
     url = f"{backend}/rl/jobs"
-    api_key = (os.getenv("SYNTH_API_KEY") or os.getenv("synth_key") or "").strip()
+    api_key = (os.getenv("SYNTH_API_KEY") or os.getenv("SYNTH_KEY") or "").strip()
     if not api_key:
         print("Missing SYNTH_API_KEY in env", file=sys.stderr)
         sys.exit(2)
@@ -91,4 +109,3 @@ def main() -> None:
 
 if __name__ == "__main__":
     main()
-