gatewaysdk 0.3.2__py3-none-any.whl

gatewaysdk/worlds.py

@@ -0,0 +1,639 @@
+"""World-over-tools: run an agent's real tool surface against a synced world.
+
+The dev-loop story in one file:
+
+    from gatewaysdk import worlds
+
+    world = worlds.load("google-drive-sim@wip-charles", variant="alice")
+
+    @world.over(prod_list_files)          # same name, same signature
+    def list_files(folder=None): ...
+
+    # wrapped calls now read/write alice's SEEDED database instead of prod;
+    # with no world loaded the decorator is inert and prod impls run.
+
+    result = world.run(my_agent)          # grade the end state, file the
+                                          # rollout back as a version-pinned
+                                          # run (branch tag included)
+
+A world is a pulled container bundle (branch-aware — `slug@branch`), one task
+VARIANT of it (its own database, materialized fresh from the task's seed), and
+the bundle's own tool handlers bound to that database. Wrapping maps a prod
+tool to the bundle handler of the SAME name and fails loudly on a tool the
+world does not serve — a silent passthrough would let a "world run" quietly
+hit production, which is the exact accident this module exists to prevent.
+
+Conventions a bundle opts into (see samples/environments/google-drive-sim):
+  * task dirs carry ``[metadata] variant = "<name>"`` in task.toml;
+  * ``environment/`` holds ``schema.sql`` + ``seed.sql`` + ``tools/`` with a
+    module exposing ``TOOLS = {name: callable}`` and optionally
+    ``DB_ENV_VAR`` naming the env var its tools read the db path from
+    (default ``GATEWAY_WORLD_DB``);
+  * ``environment/grade.py`` reads ``DRIVE_DB``-style env + ``GRADE_OUT`` and
+    writes ``reward.json`` — which makes SDK grading identical to the harbor
+    container's own verifier.
+
+Stdlib only, like the rest of benchmark_hub.
+"""
+
+from __future__ import annotations
+
+import functools
+import importlib.util
+import json
+import os
+import re
+import sqlite3
+import subprocess
+import sys
+import tempfile
+import time
+from pathlib import Path
+from typing import Any, Callable, Dict, List, Optional
+
+from .benchmark_hub.client import (
+    BenchmarkPushError,
+    GatewayBenchmarkClient,
+    local_tree_state,
+)
+from .benchmark_hub.evals import GatewayEvalsClient
+
+try:  # py311+
+    import tomllib
+except ImportError:  # pragma: no cover
+    tomllib = None  # type: ignore[assignment]
+
+
+class WorldError(RuntimeError):
+    pass
+
+
+def _task_variant(task_dir: Path) -> Optional[str]:
+    toml_path = task_dir / "task.toml"
+    if not toml_path.is_file():
+        return None
+    if tomllib is not None:
+        try:
+            data = tomllib.loads(toml_path.read_text(encoding="utf-8"))
+            v = (data.get("metadata") or {}).get("variant")
+            return str(v) if v else None
+        except Exception:
+            pass
+    m = re.search(r'^\s*variant\s*=\s*"([^"]+)"', toml_path.read_text(encoding="utf-8"), re.M)
+    return m.group(1) if m else None
+
+
+def _pick_task(bundle: Path, variant: Optional[str]) -> Path:
+    tasks = sorted(p for p in (bundle / "tasks").iterdir() if (p / "task.toml").is_file()) if (bundle / "tasks").is_dir() else []
+    if not tasks:
+        raise WorldError(f"{bundle} has no tasks/*/task.toml — not a harbor world")
+    if variant is None:
+        return tasks[0]
+    for t in tasks:
+        if _task_variant(t) == variant:
+            return t
+    for t in tasks:  # dir-name fallback, per the accepted identity rule
+        if variant in t.name:
+            return t
+    known = [f"{t.name} (variant={_task_variant(t) or '—'})" for t in tasks]
+    raise WorldError(f"no task matches variant {variant!r}; have: {', '.join(known)}")
+
+
+def tasks(bundle_dir: str | Path) -> List[Dict[str, str]]:
+    """Rich task listing for a bundle: one dict per task with its name,
+    variant, and instruction — what a suite runner or a notebook needs to
+    work with tasks as data instead of shelling out to the CLI.
+
+        first = worlds.load("google-drive-sim@main")
+        for task in worlds.tasks(first.bundle_dir):
+            print(task["variant"], "—", task["instruction"][:60])
+    """
+    bundle = Path(bundle_dir)
+    listing: List[Dict[str, str]] = []
+    for name in variants(bundle):
+        task_dir = _pick_task(bundle, name)
+        instruction_file = task_dir / "instruction.md"
+        listing.append(
+            {
+                "name": task_dir.name,
+                "variant": name,
+                "instruction": (
+                    instruction_file.read_text(encoding="utf-8")
+                    if instruction_file.is_file()
+                    else ""
+                ),
+            }
+        )
+    return listing
+
+
+def variants(bundle_dir: str | Path) -> List[str]:
+    """Variant names a bundle serves, in task order — one entry per
+    ``tasks/*/task.toml``, named by the task's own variant when it declares
+    one and by the directory otherwise. This is the suite-runner's loop:
+
+        first = worlds.load("google-drive-sim@main")
+        for name in worlds.variants(first.bundle_dir):
+            world = worlds.load_path(first.bundle_dir, variant=name)
+    """
+    bundle = Path(bundle_dir)
+    tasks = (
+        sorted(p for p in (bundle / "tasks").iterdir() if (p / "task.toml").is_file())
+        if (bundle / "tasks").is_dir()
+        else []
+    )
+    if not tasks:
+        raise WorldError(f"{bundle} has no tasks/*/task.toml — not a harbor world")
+    return [_task_variant(t) or t.name for t in tasks]
+
+
+def _load_tools_module(task_dir: Path):
+    tools_dir = task_dir / "environment" / "tools"
+    candidates = sorted(tools_dir.glob("*.py")) if tools_dir.is_dir() else []
+    for path in candidates:
+        if path.name.startswith("_"):
+            continue
+        spec = importlib.util.spec_from_file_location(f"gateway_world_tools_{path.stem}_{id(path)}", path)
+        if spec is None or spec.loader is None:
+            continue
+        module = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(module)
+        if hasattr(module, "TOOLS") and isinstance(module.TOOLS, dict):
+            return module
+    raise WorldError(
+        f"{task_dir.name} serves no tools module (environment/tools/*.py exposing TOOLS)"
+    )
+
+
+class World:
+    def __init__(
+        self,
+        *,
+        bundle_dir: Path,
+        task_dir: Path,
+        variant: Optional[str],
+        slug: Optional[str],
+        branch: Optional[str],
+        version_id: Optional[str],
+        file_block: Optional[str] = None,
+    ) -> None:
+        from .environment.schema.compiler import require_legacy_bundle
+
+        require_legacy_bundle(bundle_dir, "Harbor world sessions")
+        self.bundle_dir = bundle_dir
+        self.task_dir = task_dir
+        self.variant = variant
+        self.slug = slug
+        self.branch = branch
+        self.version_id = version_id
+        # Why a run of THIS world cannot be filed (dirty checkout, no hub
+        # link). None whenever version_id is set.
+        self.file_block = file_block
+        self.call_log: List[Dict[str, Any]] = []
+
+        env_dir = task_dir / "environment"
+        schema = env_dir / "schema.sql"
+        if not schema.is_file():
+            raise WorldError(f"{task_dir.name} has no environment/schema.sql to materialize")
+        self.db_path = str(Path(tempfile.mkdtemp(prefix="gateway-world-")) / "world.db")
+        conn = sqlite3.connect(self.db_path)
+        try:
+            conn.executescript(schema.read_text(encoding="utf-8"))
+            seed = env_dir / "seed.sql"
+            if seed.is_file():
+                conn.executescript(seed.read_text(encoding="utf-8"))
+            conn.commit()
+        finally:
+            conn.close()
+
+        self._module = _load_tools_module(task_dir)
+        self._db_env_var = getattr(self._module, "DB_ENV_VAR", "GATEWAY_WORLD_DB")
+        os.environ[self._db_env_var] = self.db_path
+
+        # Handlers bound to THIS world, each recording into the call log so an
+        # SDK run files a real trajectory, not just a score.
+        self.tools: Dict[str, Callable[..., Any]] = {
+            name: self._record(name, fn) for name, fn in self._module.TOOLS.items()
+        }
+
+    def _record(self, name: str, fn: Callable[..., Any]) -> Callable[..., Any]:
+        @functools.wraps(fn)
+        def recorded(*args: Any, **kwargs: Any) -> Any:
+            entry: Dict[str, Any] = {
+                "tool": name,
+                "arguments": json.dumps({"args": args, "kwargs": kwargs}, default=str)[:2000],
+            }
+            os.environ[self._db_env_var] = self.db_path
+            try:
+                result = fn(*args, **kwargs)
+                entry["result_preview"] = json.dumps(result, default=str)[:2000]
+                return result
+            except Exception as exc:
+                entry["error"] = str(exc)[:500]
+                raise
+            finally:
+                self.call_log.append(entry)
+
+        return recorded
+
+    def over(
+        self, prod_fn: Callable[..., Any], *, name: Optional[str] = None
+    ) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
+        """Decorator: route calls to the world's handler of the same name.
+
+        ``@world.over(prod_list_files)`` on a function keeps the prod
+        implementation as the fallback ONLY when this world is inert (never
+        constructed / explicitly disabled) — a loaded world that does not
+        serve the tool raises instead of silently touching production.
+        """
+        def decorate(fn: Callable[..., Any]) -> Callable[..., Any]:
+            # The DECORATED function's name is the agent's tool name; the prod
+            # fn is only the implementation being displaced.
+            tool_name = name or fn.__name__
+            if tool_name not in self.tools:
+                raise WorldError(
+                    f"world {self.slug or self.bundle_dir.name!r} does not serve "
+                    f"tool {tool_name!r}; it serves: {', '.join(sorted(self.tools))}"
+                )
+            world_fn = self.tools[tool_name]
+
+            @functools.wraps(prod_fn)
+            def routed(*args: Any, **kwargs: Any) -> Any:
+                return world_fn(*args, **kwargs)
+
+            return routed
+
+        return decorate
+
+    def wrap(self, prod_tools: Dict[str, Callable[..., Any]]) -> Dict[str, Callable[..., Any]]:
+        """Wholesale mapping: {name: prod_fn} → {name: world_fn}, fail-loud on
+        any name the world does not serve."""
+        missing = sorted(set(prod_tools) - set(self.tools))
+        if missing:
+            raise WorldError(
+                f"world does not serve: {', '.join(missing)}; it serves: {', '.join(sorted(self.tools))}"
+            )
+        return {name: self.tools[name] for name in prod_tools}
+
+    def override(
+        self, name: str, fn: Callable[..., Any]
+    ) -> Callable[..., Any]:
+        """The INVERSE of wrap(): keep the world's seeded database, tasks and
+        graders — swap ONE tool's implementation with your LOCAL one. This is
+        the schema-stable dev loop for tool logic changes: the local fn reads
+        the same database the world materialized (the module convention — the
+        env var named by ``DB_ENV_VAR`` already points at it), the task's own
+        grader judges the end state, and calls are recorded into the filed
+        call log like any world tool. Unknown names fail loudly: overriding a
+        tool the world does not serve would silently test nothing."""
+        if name not in self.tools:
+            raise WorldError(
+                f"world does not serve a tool named {name!r}; it serves: "
+                f"{', '.join(sorted(self.tools))}"
+            )
+        recorded = self._record(name, fn)
+        self.tools[name] = recorded
+        return recorded
+
+    def grade(self) -> Dict[str, Any]:
+        """Run the task's OWN grader against the world's end state — the same
+        code the harbor container runs, pointed at the materialized db."""
+        grader = self.task_dir / "environment" / "grade.py"
+        if not grader.is_file():
+            raise WorldError(f"{self.task_dir.name} ships no environment/grade.py")
+        out_dir = tempfile.mkdtemp(prefix="gateway-grade-")
+        env = dict(os.environ)
+        env[self._db_env_var] = self.db_path
+        env["DRIVE_DB"] = env.get("DRIVE_DB", self.db_path)
+        env["GRADE_OUT"] = out_dir
+        proc = subprocess.run(
+            [sys.executable, str(grader)], env=env, capture_output=True, text=True, timeout=120
+        )
+        reward_path = Path(out_dir) / "reward.json"
+        if proc.returncode != 0 or not reward_path.is_file():
+            raise WorldError(
+                f"grader failed (exit {proc.returncode}): {proc.stderr.strip()[:400]}"
+            )
+        return json.loads(reward_path.read_text(encoding="utf-8"))
+
+    def run(
+        self,
+        agent_fn: Callable[["World"], Any],
+        *,
+        model: str = "sdk-dev-loop",
+        file: bool = True,
+        host: Optional[str] = None,
+        session_id: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        """Execute ``agent_fn(world)``, grade the end state, and (by default)
+        file the result back as a version-pinned run — landing on the same
+        Runs/Rollouts surfaces as hosted runs, branch tag included (the
+        server stamps metadata.ref from the pinned version).
+
+        Pass ``session_id`` when the agent traces itself (the id given to
+        ``tracing.session(...)``): the platform reads ``info.sessionId`` off
+        the sample to link the run to the traced session — that link is the
+        trajectory. The filed sample carries the task name, the instruction
+        as the prompt, and ``agent_fn``'s return value as the completion, so
+        the run page shows the episode, not an empty row."""
+        started = time.time()
+        agent_error: Optional[str] = None
+        final: Any = None
+        try:
+            final = agent_fn(self)
+        except Exception as exc:
+            agent_error = str(exc)[:500]
+        verdict = self.grade()
+        reward = verdict.get("reward")
+        subs = {k: v for k, v in verdict.items() if k != "reward"}
+        result: Dict[str, Any] = {
+            "reward": reward,
+            "rewards": subs,
+            "agentError": agent_error,
+            "calls": len(self.call_log),
+            "durationSeconds": round(time.time() - started, 3),
+            **({"sessionId": session_id} if session_id else {}),
+        }
+        if not file:
+            return result
+        if not self.version_id:
+            raise WorldError(
+                self.file_block
+                or "this world was loaded from a local path — pass file=False, "
+                "or load it from the hub so runs can pin a version"
+            )
+        instruction_file = self.task_dir / "instruction.md"
+        instruction = (
+            instruction_file.read_text(encoding="utf-8")
+            if instruction_file.is_file()
+            else ""
+        )
+        evals = GatewayEvalsClient.from_env(host)
+        evaluation_id = evals.create_evaluation(
+            self.version_id,
+            model=model,
+            rollouts_per_example=1,
+            metadata={
+                "source": "sdk-world-run",
+                "variant": self.variant,
+                **({"sessionId": session_id} if session_id else {}),
+            },
+        )
+        evals.push_samples(
+            evaluation_id,
+            [
+                {
+                    "exampleId": self.task_dir.name,
+                    "task": self.task_dir.name,
+                    "rolloutIndex": 0,
+                    "prompt": instruction[:20000] or None,
+                    "completion": (
+                        str(final)[:20000] if final is not None else None
+                    ),
+                    "reward": reward,
+                    "score": reward,
+                    "correct": (reward >= 1.0) if isinstance(reward, (int, float)) else None,
+                    "rewards": subs,
+                    "info": {
+                        "variant": self.variant,
+                        "source": "sdk-world-run",
+                        "agentError": agent_error,
+                        "callLog": self.call_log[-50:],
+                        # The platform's session link + token-usage join read
+                        # exactly this key.
+                        **({"sessionId": session_id} if session_id else {}),
+                    },
+                }
+            ],
+        )
+        evals.finalize_evaluation(
+            evaluation_id,
+            status="FAILED" if agent_error else "COMPLETED",
+            metrics={
+                # avg_reward is what the Runs list renders as the score —
+                # a finalize without it files a run that LOOKS unscored.
+                **({"avg_reward": reward} if isinstance(reward, (int, float)) else {}),
+                "reward": reward,
+                **{f"rewards.{k}": v for k, v in subs.items()},
+            },
+        )
+        result["evaluationId"] = evaluation_id
+        return result
+
+
+def load_path(bundle_dir: str | Path, *, variant: Optional[str] = None) -> World:
+    """Load a world from a LOCAL directory. A dir that is a hub checkout
+    (`gateway bench pull`/`checkout`/`push` wrote `.gateway/`) whose tree
+    still matches its recorded base loads with that base's slug/branch/
+    version — runs file exactly like `load()`. A drifted checkout or an
+    ad-hoc dir still runs and grades, but filing is refused with the reason
+    (pinning drifted code to a version would fabricate provenance)."""
+    bundle = Path(bundle_dir)
+    task = _pick_task(bundle, variant)
+    state = local_tree_state(bundle)
+    meta = state["meta"] or {}
+    clean = bool(state["clean"]) and bool(meta.get("versionId"))
+    if clean:
+        slug, branch, version_id, file_block = (
+            meta.get("slug"),
+            meta.get("branch") or meta.get("ref"),
+            meta.get("versionId"),
+            None,
+        )
+    elif meta.get("versionId"):
+        dirty: List[str] = state["dirty"]
+        shown = ", ".join(dirty[:5]) + (", …" if len(dirty) > 5 else "")
+        slug, branch, version_id = meta.get("slug"), None, None
+        file_block = (
+            f"this checkout differs from "
+            f"{meta.get('branch') or meta.get('ref') or 'its base'}@"
+            f"{str(meta.get('contentHash') or '')[:8]} in {len(dirty)} "
+            f"file(s) ({shown}) — push first (`gateway bench push`) so the "
+            f"run pins the code it actually executed, or pass file=False"
+        )
+    else:
+        slug, branch, version_id = None, None, None
+        file_block = (
+            "this directory has no hub link (.gateway/.env-metadata.json) — "
+            "push it (`gateway bench push`) so runs can pin a version, or "
+            "pass file=False"
+        )
+    return World(
+        bundle_dir=bundle,
+        task_dir=task,
+        variant=variant or _task_variant(task) or task.name,
+        slug=slug,
+        branch=branch,
+        version_id=version_id,
+        file_block=file_block,
+    )
+
+
+def load(
+    slug_ref: str,
+    *,
+    variant: Optional[str] = None,
+    host: Optional[str] = None,
+) -> World:
+    """Pull ``slug`` / ``slug@branch|tag|semver|hash`` from the hub and load
+    one variant of it. The pulled branch travels with the world, so filed
+    runs carry it."""
+    client = GatewayBenchmarkClient.from_env(host)
+    target = tempfile.mkdtemp(prefix="gateway-world-pull-")
+    pulled = client.pull(slug_ref, target=target)
+    bundle = Path(pulled["target"])
+    task = _pick_task(bundle, variant)
+    return World(
+        bundle_dir=bundle,
+        task_dir=task,
+        variant=variant or _task_variant(task) or task.name,
+        slug=pulled.get("slug"),
+        branch=pulled.get("branch"),
+        version_id=pulled.get("versionId"),
+    )
+
+# ── Object model: WorldSet / Task / TaskEnv ──────────────────────────────────
+# `worlds.open()` is the handle FOR THE WORLD (one pull, one version pin);
+# `Task` is the scenario as data; `TaskEnv` (== the classic `World`) is one
+# task's materialized instance. `worlds.load()` remains sugar over this and
+# behaves exactly as before.
+
+TaskEnv = World
+"""The materialized per-task instance — seeded db, recorded tools, grader.
+The class historically named `World`; both names stay valid."""
+
+
+class Task:
+    """One scenario of a WorldSet, as data first: inspect `variant`,
+    `instruction`, and `name` without paying for any seeding. `start()`
+    materializes the environment; `run()` is start → agent → grade → file."""
+
+    def __init__(self, world_set: "WorldSet", variant: str) -> None:
+        self.world_set = world_set
+        self.variant = variant
+        task_dir = _pick_task(world_set.bundle_dir, variant)
+        self.name = task_dir.name
+        instruction_file = task_dir / "instruction.md"
+        self.instruction = (
+            instruction_file.read_text(encoding="utf-8")
+            if instruction_file.is_file()
+            else ""
+        )
+
+    def start(self) -> TaskEnv:
+        return load_path(self.world_set.bundle_dir, variant=self.variant)
+
+    def run(
+        self,
+        agent_fn: Callable[[TaskEnv], Any],
+        *,
+        model: Optional[str] = None,
+        file: bool = True,
+        session_id: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        env = self.start()
+        return env.run(
+            agent_fn,
+            file=file,
+            session_id=session_id,
+            # Only override the env's default label when a model was named.
+            **({"model": model} if model else {}),
+        )
+
+    def __repr__(self) -> str:  # notebook-friendly
+        return f"Task({self.variant!r}, name={self.name!r})"
+
+
+class WorldSet:
+    """THE world: one pull, one pinned version, every task. Iterable and
+    indexable by variant; `run_all` is the whole suite as one primitive."""
+
+    def __init__(self, bundle_dir: Path, *, slug: Optional[str],
+                 branch: Optional[str], version_id: Optional[str]) -> None:
+        self.bundle_dir = Path(bundle_dir)
+        self.slug = slug
+        self.branch = branch
+        self.version_id = version_id
+        self.tasks: List[Task] = [Task(self, v) for v in variants(self.bundle_dir)]
+
+    def task(self, variant: str) -> Task:
+        for task in self.tasks:
+            if task.variant == variant or variant in task.name:
+                return task
+        known = ", ".join(t.variant for t in self.tasks)
+        raise WorldError(f"no task matches {variant!r}; have: {known}")
+
+    def __iter__(self):
+        return iter(self.tasks)
+
+    def __len__(self) -> int:
+        return len(self.tasks)
+
+    def run_all(
+        self,
+        agent_fn: Callable[[TaskEnv], Any],
+        *,
+        model: Optional[str] = None,
+        tasks: Optional[List[str]] = None,
+        fail_under: Optional[float] = None,
+        file: bool = True,
+    ) -> Dict[str, Any]:
+        """Run the agent across every task (or the named subset), grading and
+        filing each. Returns {results: {variant: result}, mean, passed} —
+        `passed` is None without a bar, else mean >= fail_under."""
+        chosen = (
+            [self.task(v) for v in tasks] if tasks is not None else self.tasks
+        )
+        results: Dict[str, Dict[str, Any]] = {}
+        for task in chosen:
+            results[task.variant] = task.run(agent_fn, model=model, file=file)
+        rewards = [float(r.get("reward") or 0.0) for r in results.values()]
+        mean = sum(rewards) / len(rewards) if rewards else 0.0
+        return {
+            "results": results,
+            "mean": mean,
+            "passed": None if fail_under is None else mean >= fail_under,
+        }
+
+    def __repr__(self) -> str:
+        ref = f"@{self.branch}" if self.branch else ""
+        return f"WorldSet({(self.slug or str(self.bundle_dir)) + ref!r}, tasks={len(self.tasks)})"
+
+
+def open(  # noqa: A001 - deliberate, mirrors pathlib-style opens
+    slug_ref: str,
+    *,
+    host: Optional[str] = None,
+) -> WorldSet:
+    """Open a world from the hub: ONE pull shared by every task.
+
+        world = worlds.open("google-drive-sim@main")
+        world.task("alice").run(my_agent, model="openai/gpt-5.6-terra")
+        world.run_all(my_agent, fail_under=0.6)
+    """
+    client = GatewayBenchmarkClient.from_env(host)
+    target = tempfile.mkdtemp(prefix="gateway-world-pull-")
+    pulled = client.pull(slug_ref, target=target)
+    bundle = Path(pulled["target"])
+    state = local_tree_state(bundle)
+    meta = state["meta"] or {}
+    return WorldSet(
+        bundle,
+        slug=meta.get("slug"),
+        branch=meta.get("branch") or meta.get("ref"),
+        version_id=meta.get("versionId"),
+    )
+
+
+def open_path(bundle_dir: str | Path) -> WorldSet:
+    """Open a LOCAL checkout as a WorldSet (same filing rules as load_path)."""
+    bundle = Path(bundle_dir)
+    state = local_tree_state(bundle)
+    meta = state["meta"] or {}
+    clean = bool(state["clean"]) and bool(meta.get("versionId"))
+    return WorldSet(
+        bundle,
+        slug=meta.get("slug"),
+        branch=(meta.get("branch") or meta.get("ref")) if clean else None,
+        version_id=meta.get("versionId") if clean else None,
+    )