gatewaysdk 0.3.2__py3-none-any.whl

gatewaysdk/environment/schema/tools_world.py

@@ -0,0 +1,322 @@
+"""A mock-tools world on the World Host.
+
+A mock-tools bundle (``gateway-env.toml`` at its root, ``tools/*.py`` handlers,
+``db/`` and ``verifiers/*.sql``) has no HTTP surface: an agent reaches it
+through the ``tools`` surface, the platform's ``call``/``grade``/``reset``/
+``seed`` rendezvous. On the pod engine that rendezvous is a runner sidecar and
+``deploy/benchmark-runner/engines/schema_session.py``; on the host it is this
+module, driven by :class:`~gatewaysdk.environment.schema.host.WorldHost`.
+
+The session contract is the runner's: ``call`` executes a handler through the
+schema contract, ``grade`` runs the task's verifiers, ``reset`` copies the
+template back and re-arms the task context, ``seed`` lands rows through the
+contract. The handlers are world-authored code, so a tools version always
+lives in a version worker; the supervisor only detects the kind.
+"""
+
+from __future__ import annotations
+
+import json
+import shutil
+import uuid
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+from .. import _bundle
+from .._materialize import materialize_template
+from .._world import World
+from ..errors import WorldStartError
+from . import snapshot as snapshots
+from .compiler import ContractError, bundle_mode, check_bundle, fail
+from .data import seed_batch
+
+Handler = Callable[..., Any]
+
+#: The reward key the scalar goes under; a verifier may not claim it.
+REWARD = "reward"
+#: The opt-in table a bundle declares to read the active task from SQL.
+TASK_CONTEXT_TABLE = "benchmark_runtime_context"
+
+
+@dataclass(frozen=True)
+class ToolsTask:
+    """One task of a mock-tools bundle: the prompt and the checks that grade it.
+    Mirrors the runner's ``MockTask`` minus the eval-loop fields (user
+    simulation, dependency worlds), which a session does not run."""
+
+    id: str
+    prompt: str
+    verifier: str | None = None
+    verifiers: tuple[str, ...] = ()
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+    def selected(self, all_verifiers: list[str]) -> list[str]:
+        """The verifiers this task grades with: its own list, else the
+        primary one, else every verifier the bundle ships."""
+        if self.verifiers:
+            return list(self.verifiers)
+        if self.verifier is not None:
+            return [self.verifier]
+        return list(all_verifiers)
+
+
+@dataclass
+class ToolsVersion:
+    """A mock-tools version as the host caches it: the bundle, its tasks and
+    verifiers, the materialized template every session copies, and (in a
+    worker) the loaded handlers."""
+
+    bundle: Path
+    manifest: _bundle.EnvManifest
+    verifiers: list[str]
+    tasks: dict[str, ToolsTask]
+    template: Path
+    handlers: dict[str, Handler] | None = None
+
+
+def is_tools_bundle(bundle: Path) -> bool:
+    return (bundle / "gateway-env.toml").is_file() and not (bundle / "gateway-api-world.toml").is_file()
+
+
+def verifier_names(bundle: Path) -> list[str]:
+    verifiers_dir = bundle / "verifiers"
+    if not verifiers_dir.is_dir():
+        return []
+    return sorted(p.stem for p in verifiers_dir.glob("*.sql"))
+
+
+def _task(raw: dict, index: int) -> ToolsTask | None:
+    prompt = raw.get("request") or raw.get("prompt") or raw.get("instruction") or raw.get("goal")
+    if not isinstance(prompt, str) or not prompt.strip():
+        return None
+    task_id = str(raw.get("id") or raw.get("name") or f"task-{index}")
+    verifier = raw.get("verifier")
+    listed = raw.get("verifiers")
+    metadata = raw.get("metadata")
+    return ToolsTask(
+        id=task_id,
+        prompt=prompt.strip(),
+        verifier=verifier if isinstance(verifier, str) and verifier.strip() else None,
+        verifiers=(
+            tuple(dict.fromkeys(str(v).strip() for v in listed if str(v).strip())) if isinstance(listed, list) else ()
+        ),
+        metadata=metadata if isinstance(metadata, dict) else {},
+    )
+
+
+def _tasks_from_rows(rows: list) -> list[ToolsTask]:
+    out = []
+    for index, raw in enumerate(rows):
+        task = _task(raw, index) if isinstance(raw, dict) else None
+        if task is not None:
+            out.append(task)
+    return out
+
+
+def load_tasks(bundle: Path, manifest: _bundle.EnvManifest, verifiers: list[str]) -> dict[str, ToolsTask]:
+    """The bundle's tasks, in the runner's precedence: ``[[tasks]]`` in the
+    manifest, else ``tasks.jsonl`` / ``tasks.json``, else one task per
+    verifier. A bundle with neither has nothing to grade and is refused."""
+    declared = manifest.raw.get("tasks")
+    tasks = _tasks_from_rows(declared) if isinstance(declared, list) else []
+    if not tasks:
+        for name in ("tasks.jsonl", "tasks.json"):
+            path = bundle / name
+            if not path.is_file():
+                continue
+            text = path.read_text(encoding="utf-8")
+            if name.endswith(".jsonl"):
+                rows = [json.loads(line) for line in text.splitlines() if line.strip()]
+            else:
+                data = json.loads(text)
+                rows = data if isinstance(data, list) else data.get("tasks", [])
+            tasks = _tasks_from_rows(rows)
+            if tasks:
+                break
+    if not tasks:
+        tasks = [
+            ToolsTask(
+                id=name,
+                prompt=f"Drive the environment into the state checked by the verifier '{name}'.",
+                verifier=name,
+            )
+            for name in verifiers
+        ]
+    if not tasks:
+        fail(
+            "host",
+            "mock-tools bundle declares no tasks and ships no verifiers/*.sql: nothing to grade",
+        )
+    by_id: dict[str, ToolsTask] = {}
+    for task in tasks:
+        if task.id in by_id:
+            fail("host", f"task id '{task.id}' is declared twice")
+        by_id[task.id] = task
+    return by_id
+
+
+def load_tools_version(
+    bundle: Path,
+    template: Path,
+    snapshot: dict[str, str] | None,
+    *,
+    run_handlers: bool,
+) -> ToolsVersion:
+    """Cache a mock-tools version: verify and pin its snapshot, materialize the
+    template once, read its tasks, and in a worker import its handlers."""
+    if not is_tools_bundle(bundle):
+        fail("host", f"{bundle} is not a mock-tools bundle (no gateway-env.toml)")
+    manifest = _bundle.parse_manifest(bundle)
+    if manifest.raw.get("dependencies"):
+        fail("host", "the World Host does not open task sets with dependency worlds yet; use the pod engine")
+    if bundle_mode(bundle) is None:
+        fail(
+            "host",
+            "the World Host opens schema-backed mock-tools worlds only ([schema] language in "
+            "gateway-env.toml); a db/schema.sql world runs on the pod engine",
+        )
+    contract = check_bundle(bundle)
+    if snapshot:
+        snapshots.fetch(bundle, contract, snapshot["url"], snapshot["sha256"], snapshot["digest"])
+    if not template.is_file():
+        materialize_template(bundle, template, require_sql=True)
+    verifiers = verifier_names(bundle)
+    if REWARD in verifiers:
+        fail("host", f"verifier name '{REWARD}' is reserved for the scalar grade")
+    tasks = load_tasks(bundle, manifest, verifiers)
+    for task in tasks.values():
+        unknown = set(task.selected(verifiers)) - set(verifiers)
+        if unknown:
+            fail("host", f"task '{task.id}' names verifiers the bundle lacks: {', '.join(sorted(unknown))}")
+    handlers = None
+    if run_handlers:
+        try:
+            handlers = _bundle.load_handlers(bundle, manifest)
+        except WorldStartError as error:
+            fail("host", str(error))
+    return ToolsVersion(
+        bundle=bundle,
+        manifest=manifest,
+        verifiers=verifiers,
+        tasks=tasks,
+        template=template,
+        handlers=handlers,
+    )
+
+
+class _EnvHost:
+    """What :class:`World` reads of an ``Environment``: the template, a scratch
+    dir, the handlers, the bundle path, and always-``env`` routing."""
+
+    def __init__(self, version: ToolsVersion, worlds_dir: Path) -> None:
+        self.bundle_path = version.bundle
+        self._worlds_dir = worlds_dir
+        self._template_db = version.template
+        self._handlers = version.handlers or {}
+
+    def resolve_route(self, tool_name: str) -> str:
+        return "env"
+
+    def _record_route(self, *args: Any, **kwargs: Any) -> None:
+        return None
+
+
+class ToolsWorld:
+    """One session of a mock-tools version: a copy of the template in the
+    session's own directory, executed through the bundle's handlers."""
+
+    def __init__(self, version: ToolsVersion, session_dir: Path, task: str) -> None:
+        if version.handlers is None:
+            fail("host", "a mock-tools session runs world-authored handlers; only a version worker opens one")
+        selected = version.tasks.get(task)
+        if selected is None:
+            fail("host", f"no task '{task}' in this world; it has {', '.join(sorted(version.tasks))}")
+            raise AssertionError  # fail() never returns
+        self.version = version
+        self.task = selected
+        self._dir = Path(session_dir)
+        self._dir.mkdir(parents=True, exist_ok=True)
+        self.world = World(_EnvHost(version, self._dir), world_id=uuid.uuid4().hex)  # type: ignore[arg-type]
+        self._arm()
+
+    def _arm(self) -> None:
+        """Write the active task into the bundle's opt-in context table, so its
+        SQL verifiers can key off the scenario without a tool that sets it."""
+        conn = self.world._conn
+        declared = conn.execute(
+            "SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = ?", (TASK_CONTEXT_TABLE,)
+        ).fetchall()
+        if not declared:
+            return
+        scenario_id = str(self.task.metadata.get("scenario_id") or self.task.id)
+        conn.execute(f"DELETE FROM {TASK_CONTEXT_TABLE}")
+        conn.execute(
+            f"INSERT INTO {TASK_CONTEXT_TABLE} (context_key, scenario_id) VALUES ('active_task', ?)",
+            (scenario_id,),
+        )
+
+    def call(self, tool: str, args: dict) -> dict:
+        if tool not in (self.version.handlers or {}):
+            known = ", ".join(sorted(self.version.handlers or {})) or "(none)"
+            fail("host", f"unknown tool '{tool}'; this world serves: {known}")
+        try:
+            result = self.world.execute(tool, args)
+        except ContractError:
+            raise
+        except Exception as error:  # noqa: BLE001 - the handler's failure is the call's answer
+            return {"status": 500, "body": {"error": str(error)[:2000]}}
+        return {"status": 200, "body": result}
+
+    def grade(self) -> dict:
+        rewards: dict[str, float | None] = {}
+        for name in self.task.selected(self.version.verifiers):
+            rewards[name] = self.world.verify(name).score
+        if self.task.verifier is not None and self.task.verifier in rewards:
+            reward = rewards[self.task.verifier]
+        else:
+            numeric = [s for s in rewards.values() if isinstance(s, (int, float))]
+            reward = sum(numeric) / len(numeric) if numeric else None
+        return {REWARD: reward, "rewards": rewards}
+
+    def reset(self) -> dict:
+        self.world.reset()
+        self._arm()
+        return {"reset": True}
+
+    def export(self, destination: Path) -> Path:
+        """The session's state as one sqlite file, taken consistently."""
+        return self.world.snapshot(destination)
+
+    def restore(self, source: Path) -> None:
+        """Continue from a file :meth:`export` wrote."""
+        self.world.restore(source)
+        self._arm()
+
+    def seed(self, args: dict) -> dict:
+        rows = args.get("rows")
+        mode = args.get("mode", "append")
+        if not isinstance(rows, dict):
+            fail("host", "seed needs args.rows as {entity: [row, ...]}")
+        store, contract = self.world.store, self.world.contract
+        if store is None or contract is None:
+            fail("host", "seed needs a schema-backed world (gateway-world/1); this bundle has none")
+            raise AssertionError  # fail() never returns
+        batch = seed_batch(store, contract, rows, mode, batch_id=f"host-seed-{uuid.uuid4().hex}")
+        applied = store.apply_batch(batch)
+        return {
+            "seeded": applied.get("applied", sum(len(v) for v in rows.values())),
+            "mode": mode,
+            "entity_counts": {entity: len(store.list(entity)) for entity in rows if entity in contract.entities},
+        }
+
+    def state(self) -> dict:
+        store, contract = self.world.store, self.world.contract
+        if store is None or contract is None:
+            return {"entities": {}}
+        return {"entities": {entity: store.list(entity) for entity in contract.entities}}
+
+    def close(self) -> None:
+        self.world.close()
+        shutil.rmtree(self._dir, ignore_errors=True)