gatewaysdk 0.3.2__py3-none-any.whl

gatewaysdk/environment/_world.py

@@ -0,0 +1,346 @@
+"""Worlds — deterministic instances of a mock-tools environment's seeded state.
+
+A world is one sqlite file, branched (file-copied) from the environment's
+prepared template db, plus the seams that let an agent's tool calls execute
+against it: ``execute``, ``as_registry`` (the exact replay ``ToolStubRegistry``
+mapping contract), ``wrap_tools``, and grading (``verify`` / ``state``).
+"""
+
+from __future__ import annotations
+
+import shutil
+import sqlite3
+import time
+import uuid
+from contextlib import nullcontext
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Mapping, Optional, Union
+
+try:
+    from ..replay import StubRegistryMapping
+except ImportError:
+    # ``replay`` pulls in ``requests`` + the sessions stack. A lean consumer
+    # that only needs worlds + verifiers (e.g. the benchmark runner's
+    # mock_tools engine, which vendors just this ``environment`` subpackage —
+    # no ``replay``) still imports ``World``/``verify``; only ``as_registry``
+    # (the replay registry contract) is unavailable there, and it is never
+    # called on that path. The full SDK always has ``replay``, so real
+    # behaviour is unchanged.
+    StubRegistryMapping = dict  # type: ignore[assignment,misc]
+from . import runtime
+from .errors import ToolRouteError, VerifierError
+
+if TYPE_CHECKING:  # pragma: no cover
+    from .core import Environment
+
+
+@dataclass
+class VerifierResult:
+    """One verifier outcome. ``score=None`` stays ``None`` (house rule: null
+    scores are their own category, never coerced to 0) and yields
+    ``passed=None``; otherwise ``passed = score >= 1.0``."""
+
+    name: str
+    passed: Optional[bool]
+    score: Optional[float]
+    detail: str = ""
+    rewards: Dict[str, float] = field(default_factory=dict)
+
+
+def split_sql_statements(script: str) -> List[str]:
+    """Split a SQL script into complete statements (via
+    ``sqlite3.complete_statement`` accumulation, so semicolons inside string
+    literals don't split)."""
+    statements: List[str] = []
+    buffer = ""
+    for line in script.splitlines(keepends=True):
+        buffer += line
+        if sqlite3.complete_statement(buffer):
+            stripped = buffer.strip()
+            if stripped:
+                statements.append(stripped)
+            buffer = ""
+    if buffer.strip():
+        statements.append(buffer.strip())
+    return statements
+
+
+_NO_INPUT = object()
+
+
+def _stub_input(tool_input: Any, kwargs: Dict[str, Any]) -> Any:
+    """One rule for callable seams: an explicit positional input (even None)
+    is the input; keyword arguments are the input; nothing means omitted."""
+    if tool_input is not _NO_INPUT:
+        return tool_input
+    return kwargs if kwargs else _NO_INPUT
+
+
+class WorldRegistry(StubRegistryMapping):
+    """``world.as_registry()`` — the replay registry contract (mapping with
+    ``__missing__``, ``lookup``, ``call_log``/``called_names``) backed by the
+    live env world instead of recordings."""
+
+    def __init__(self, world: "World") -> None:
+        self._world = world
+
+        def make(name: str) -> Callable[..., Any]:
+            def stub(tool_input: Any = _NO_INPUT, **kwargs: Any) -> Any:
+                return world._registry_call(name, _stub_input(tool_input, kwargs))
+
+            stub.__name__ = name if name.isidentifier() else "stub"
+            stub.__qualname__ = f"gateway_env.{name}"
+            return stub
+
+        super().__init__(
+            {name: make(name) for name in world.env._handlers}, make
+        )
+
+    @property
+    def call_log(self) -> List[Dict[str, Any]]:
+        return self._world.call_log
+
+    def lookup(self, tool_name: str, tool_input: Any = _NO_INPUT) -> Any:
+        return self._world._registry_call(tool_name, tool_input)
+
+    def called_names(self) -> List[str]:
+        return [entry["tool"] for entry in self._world.call_log]
+
+
+class World:
+    """One live, isolated instance of the environment's seeded state."""
+
+    def __init__(self, env: "Environment", *, world_id: Optional[str] = None, snapshot: Optional[Path] = None) -> None:
+        self.env = env
+        self.id = world_id or uuid.uuid4().hex[:8]
+        self.db_path: Path = env._worlds_dir / f"world-{self.id}.db"
+        self.call_log: List[Dict[str, Any]] = []
+        from .schema import bundle_mode, check_bundle
+
+        self.contract = check_bundle(env.bundle_path) if bundle_mode(env.bundle_path) is not None else None
+        self.store = None
+        shutil.copyfile(snapshot if snapshot is not None else env._template_db, self.db_path)
+        try:
+            self._conn = self._connect()
+        except BaseException:
+            self.db_path.unlink(missing_ok=True)
+            raise
+
+    def _connect(self) -> sqlite3.Connection:
+        # autocommit + cross-thread: agents call tools from worker threads, and
+        # file copies (branch/snapshot/reset) must always see committed state.
+        if self.contract is not None:
+            from .schema import SQLiteStore
+
+            self.store = SQLiteStore.open(self.db_path, self.contract)
+            return self.store._conn  # Trusted host only; handlers receive ReadConnection.
+        conn = sqlite3.connect(self.db_path, isolation_level=None, check_same_thread=False)
+        conn.row_factory = sqlite3.Row
+        return conn
+
+    # -- execution seams -------------------------------------------------
+
+    def execute(self, tool_name: str, tool_input: Any = _NO_INPUT) -> Any:
+        """Execute atomically in schema mode, validating both input and output."""
+        from .schema.store import ReadConnection
+        from .schema.validation import validate
+
+        handler = self.env._handlers.get(tool_name)
+        if handler is None or (self.contract is not None and tool_name not in self.contract.tools):
+            raise ToolRouteError(f"tool '{tool_name}' has no declared env handler")
+        supplied = ({} if self.contract else None) if tool_input is _NO_INPUT else tool_input
+        started = time.perf_counter()
+        with self.store.transaction() if self.store is not None else nullcontext():
+            if self.contract is not None:
+                validate(
+                    self.contract.tools[tool_name]["input_schema"],
+                    supplied,
+                    location=f"world.json#/x-gateway/tools/{tool_name}/input",
+                )
+            token = runtime._active_connection.set(ReadConnection(self.store) if self.store is not None else self._conn)
+            store_token = runtime._active_store.set(self.store)
+            try:
+                if isinstance(supplied, Mapping):
+                    result = handler(**supplied)
+                elif supplied is None:
+                    result = handler()
+                else:
+                    result = handler(supplied)
+                if self.contract is not None:
+                    validate(
+                        self.contract.tools[tool_name]["output_schema"],
+                        result,
+                        location=f"world.json#/x-gateway/tools/{tool_name}/output",
+                    )
+            finally:
+                runtime._active_connection.reset(token)
+                runtime._active_store.reset(store_token)
+        duration_ms = (time.perf_counter() - started) * 1000.0
+        self.call_log.append({"tool": tool_name, "match": "env", "input": supplied})
+        self.env._record_route(tool_name, "env", duration_ms, self.id)
+        return result
+
+    def _registry_call(
+        self, tool_name: str, tool_input: Any = _NO_INPUT
+    ) -> Any:
+        if tool_name in self.env._handlers and self.env.resolve_route(tool_name) == "env":
+            return self.execute(tool_name, tool_input)
+        route = self.env.resolve_route(tool_name)
+        logged = None if tool_input is _NO_INPUT else tool_input
+        self.call_log.append({"tool": tool_name, "match": "miss", "input": logged})
+        if route == "prod":
+            raise ToolRouteError(
+                f"tool '{tool_name}' routes to prod, but the registry seam has no "
+                "live implementation to call — pass the real callable via "
+                "world.wrap_tools()/your own registry, or rely on the ambient "
+                "define_tool wrapper (with env:), which passes prod-routed calls "
+                "through to the original function."
+            )
+        raise ToolRouteError(
+            f"tool '{tool_name}' has no env handler "
+            f"(served: {sorted(self.env._handlers)}) and no prod route."
+        )
+
+    def as_registry(self) -> WorldRegistry:
+        """Tool-name → callable mapping, drop-in for agent tool layers —
+        identical contract to ``ToolStubRegistry.as_registry()`` in replay."""
+        return WorldRegistry(self)
+
+    def wrap_tools(
+        self, tools: Union[Mapping[str, Any], List[Any]]
+    ) -> Union[Dict[str, Any], List[Any]]:
+        """Substitute env-served callables by name in a list/dict of tools;
+        everything not served (or not env-routed) passes through unchanged."""
+
+        def routed_to_env(name: Optional[str]) -> bool:
+            return bool(name) and name in self.env._handlers and self.env.resolve_route(name) == "env"
+
+        def stub_for(name: str) -> Callable[..., Any]:
+            def stub(tool_input: Any = _NO_INPUT, **kwargs: Any) -> Any:
+                return self.execute(name, _stub_input(tool_input, kwargs))
+
+            stub.__name__ = name if name.isidentifier() else "stub"
+            return stub
+
+        if isinstance(tools, Mapping):
+            return {
+                name: (stub_for(name) if routed_to_env(name) else value)
+                for name, value in tools.items()
+            }
+        wrapped: List[Any] = []
+        for item in tools:
+            name = getattr(item, "name", None)
+            if not isinstance(name, str):
+                name = getattr(item, "__name__", None)
+            wrapped.append(stub_for(name) if routed_to_env(name) else item)
+        return wrapped
+
+    # -- state lifecycle -------------------------------------------------
+
+    def reset(self) -> None:
+        """Back to the seed, in place (same world instance)."""
+        if self.contract is not None:
+            from .schema import check_bundle
+
+            if check_bundle(self.env.bundle_path).lock != self.contract.lock:
+                from .schema import ContractError
+
+                raise ContractError("state_contract", "Bundle contract changed; create a new world explicitly")
+        if self.store is not None:
+            self.store.close()
+        else:
+            self._conn.close()
+        shutil.copyfile(self.env._template_db, self.db_path)
+        self._conn = self._connect()
+
+    def restore(self, source: Union[str, Path]) -> None:
+        """Replace this world's state with a file :meth:`snapshot` wrote, in place.
+        The file is opened the way a fresh copy of the seed is, so a schema-backed
+        world audits it against its contract."""
+        if self.store is not None:
+            self.store.close()
+        else:
+            self._conn.close()
+        shutil.copyfile(source, self.db_path)
+        self._conn = self._connect()
+
+    def snapshot(self, destination: Union[str, Path]) -> Path:
+        """Pin this world's current state to a file (for debugging fixtures)."""
+        if self.store is not None:
+            return self.store.snapshot(destination)
+        dest = Path(destination)
+        dest.parent.mkdir(parents=True, exist_ok=True)
+        shutil.copyfile(self.db_path, dest)
+        return dest
+
+    def close(self) -> None:
+        try:
+            if self.store is not None:
+                self.store.close()
+            else:
+                self._conn.close()
+        except sqlite3.Error:  # pragma: no cover — close is best-effort
+            pass
+
+    # -- grading ---------------------------------------------------------
+
+    def state(self, sql: str, *params: Any) -> List[Dict[str, Any]]:
+        """Read-only SQL against the world's db (SELECT/WITH only)."""
+        if self.store is not None:
+            return [dict(row) for row in self.store.query(sql, params)]
+        head = sql.lstrip().split(None, 1)[0].upper() if sql.strip() else ""
+        if head not in ("SELECT", "WITH"):
+            raise ValueError(
+                "world.state() is read-only — only SELECT/WITH queries are "
+                "allowed (use tool handlers or verifiers for writes)."
+            )
+        rows = self._conn.execute(sql, params).fetchall()
+        return [dict(row) for row in rows]
+
+    def verify(self, name: str) -> VerifierResult:
+        """Run ``verifiers/<name>.sql`` against this world.
+
+        Statements run in order; the score is the first column of the first
+        row of the LAST row-returning statement. NULL / no rows → score None
+        (passed None) — never coerced to 0.
+        """
+        verifier_name = name[:-4] if name.endswith(".sql") else name
+        path = self.env.bundle_path / "verifiers" / f"{verifier_name}.sql"
+        if not path.is_file():
+            available = sorted(p.stem for p in (self.env.bundle_path / "verifiers").glob("*.sql"))
+            raise VerifierError(f"unknown verifier '{name}' (available: {available})")
+
+        script = path.read_text(encoding="utf-8")
+        score: Optional[float] = None
+        saw_rows = False
+        last_sql = ""
+        for statement in split_sql_statements(script):
+            try:
+                if self.store is not None:
+                    rows = self.store.query(statement, private=True)
+                    saw_rows = True
+                    last_sql = statement
+                    value = rows[0][0] if rows else None
+                    score = float(value) if value is not None else None
+                    continue
+                cursor = self._conn.execute(statement)
+            except sqlite3.Error as error:
+                raise VerifierError(
+                    f"verifier '{verifier_name}' failed: {error}\n  SQL: {statement}"
+                ) from error
+            if cursor.description is not None:  # row-returning statement
+                saw_rows = True
+                last_sql = statement
+                row = cursor.fetchone()
+                value = row[0] if row is not None else None
+                score = float(value) if value is not None else None
+
+        if not saw_rows:
+            detail = "no row-returning statement — score is null"
+        elif score is None:
+            detail = f"null score from: {last_sql}"
+        else:
+            detail = f"score {score} from: {last_sql}"
+        passed = None if score is None else score >= 1.0
+        return VerifierResult(name=verifier_name, passed=passed, score=score, detail=detail)