agent-learning-kit 0.1.0__py3-none-any.whl

fi/alk/harness/world/stores/sqlite.py

@@ -0,0 +1,245 @@
+"""Records in a SQLite file, or in memory.
+
+The default, because it needs nothing installed and nothing standing up. Saving to disk is a
+copy of the database rather than a dump, so loading it back is exact and fast, which matters
+because every scenario and every probe starts from that copy.
+"""
+
+from __future__ import annotations
+
+import json
+import shutil
+import sqlite3
+from pathlib import Path
+from typing import Any, Mapping, Sequence
+
+from . import Records, Snapshot, StoreError
+
+# What SQLite hands out that is not itself a record. Only present once a table is declared
+# AUTOINCREMENT, which is why its absence is normal rather than a gap.
+COUNTERS = "sqlite_sequence"
+
+
+def _sqlite_value(value: Any) -> Any:
+    """Encode structured seed values deterministically for SQLite TEXT columns."""
+    if isinstance(value, (dict, list)):
+        return json.dumps(value, sort_keys=True, separators=(",", ":"))
+    return value
+
+
+class SqliteStore(Records):
+    engine = "sqlite"
+    key = "sqlite"
+    FILE = "world.sqlite"
+
+    def __init__(self, database: str | Path = ":memory:", **_ignored: Any) -> None:
+        self.database = str(database)
+        self.connection = sqlite3.connect(self.database, check_same_thread=False)
+        self.connection.execute("PRAGMA foreign_keys = ON")
+
+    # -- lifecycle -------------------------------------------------------------------
+
+    def start(self) -> None:
+        """Already up. Connecting is what ``__init__`` did, and there is no server to wait for."""
+
+    def stop(self) -> None:
+        self.close()
+
+    def dsn(self) -> str:
+        return f"sqlite:///{self.database}"
+
+    # -- statements ------------------------------------------------------------------
+
+    def execute(self, statement: str, params: Sequence[Any] = ()) -> int:
+        cursor = self.connection.execute(
+            statement, tuple(_sqlite_value(value) for value in params)
+        )
+        self.connection.commit()
+        return cursor.rowcount
+
+    def apply(self, script: str) -> None:
+        """Several statements at once, which is how a schema or a seed arrives."""
+        if not script.strip():
+            return
+        self.connection.executescript(script)
+        self.connection.commit()
+
+    # The name this had before a store was asked to speak more than SQL.
+    script = apply
+
+    def query(self, statement: str, params: Sequence[Any] = ()) -> list[dict[str, Any]]:
+        cursor = self.connection.execute(statement, tuple(params))
+        columns = [column[0] for column in (cursor.description or [])]
+        return [dict(zip(columns, row)) for row in cursor.fetchall()]
+
+    # -- records ---------------------------------------------------------------------
+
+    def collections(self) -> list[str]:
+        found = self.query(
+            "SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%'"
+        )
+        return [row["name"] for row in found]
+
+    def holds(self, collection: str) -> bool:
+        return bool(
+            self.query(
+                "SELECT name FROM sqlite_master WHERE type='table' AND name = ?",
+                [collection],
+            )
+        )
+
+    def records(self, collection: str) -> list[dict[str, Any]]:
+        return self.query(f'SELECT * FROM "{collection}"')
+
+    def add(self, collection: str, record: Mapping[str, Any]) -> int:
+        columns = ", ".join(f'"{name}"' for name in record)
+        marks = ", ".join("?" for _ in record)
+        return self.execute(
+            f'INSERT INTO "{collection}" ({columns}) VALUES ({marks})',
+            list(record.values()),
+        )
+
+    def amend(
+        self, collection: str, key: str, changes: Mapping[str, Any], *, by: str = ""
+    ) -> int:
+        if not by:
+            raise KeyError(
+                f"{collection} is a table, so changing a record needs the column it is keyed on"
+            )
+        sets = ", ".join(f'"{name}" = ?' for name in changes)
+        return self.execute(
+            f'UPDATE "{collection}" SET {sets} WHERE "{by}" = ?',
+            [*changes.values(), key],
+        )
+
+    def remove(self, collection: str, key: str = "", *, by: str = "") -> int:
+        if key and not by:
+            raise KeyError(
+                f"{collection} is a table, so removing one record needs the column it is keyed on"
+            )
+        sql = f'DELETE FROM "{collection}"' + (f' WHERE "{by}" = ?' if key else "")
+        return self.execute(sql, [key] if key else [])
+
+    # -- going back, in memory -------------------------------------------------------
+
+    def freeze(self) -> Snapshot:
+        return Snapshot(rows=self.state(), counters=self._counters())
+
+    def restore(self, snapshot: Snapshot) -> None:
+        """Put the rows back, and the counters behind them.
+
+        Foreign keys are dropped for the duration rather than the tables being sorted into
+        dependency order: any order is wrong for some schema, and a restore that fails on a
+        schema the agent really has is worse than one that trusts the snapshot it took itself.
+        """
+        self.connection.execute("PRAGMA foreign_keys = OFF")
+        try:
+            for name in self.collections():
+                self.connection.execute(f'DELETE FROM "{name}"')
+            for name, rows in snapshot.rows.items():
+                for row in rows:
+                    if not row:
+                        continue
+                    columns = ", ".join(f'"{column}"' for column in row)
+                    marks = ", ".join("?" for _ in row)
+                    self.connection.execute(
+                        f'INSERT INTO "{name}" ({columns}) VALUES ({marks})',
+                        [_sqlite_value(value) for value in row.values()],
+                    )
+            self._reinstate(snapshot.counters)
+            self.connection.commit()
+        finally:
+            self.connection.execute("PRAGMA foreign_keys = ON")
+
+    def _counters(self) -> dict[str, int]:
+        if not self.holds(COUNTERS):
+            return {}
+        return {
+            row["name"]: row["seq"]
+            for row in self.query(f"SELECT name, seq FROM {COUNTERS}")
+        }
+
+    def _reinstate(self, counters: Mapping[str, int]) -> None:
+        if not self.holds(COUNTERS):
+            return
+        self.connection.execute(f"DELETE FROM {COUNTERS}")
+        for name, seq in counters.items():
+            self.connection.execute(
+                f"INSERT INTO {COUNTERS} (name, seq) VALUES (?, ?)", (name, seq)
+            )
+
+    def clear(self) -> None:
+        """Empty every table, whatever references what.
+
+        Foreign keys are suspended for the duration rather than the tables being sorted into
+        dependency order. Deleting them one at a time in the wrong order fails on the referenced
+        ones, and a caller that swallows those failures is left believing it emptied a store that
+        still holds most of its data.
+        """
+        self.connection.execute("PRAGMA foreign_keys = OFF")
+        try:
+            for name in self.collections():
+                self.connection.execute(f'DELETE FROM "{name}"')
+            self.connection.commit()
+        finally:
+            self.connection.execute("PRAGMA foreign_keys = ON")
+
+    def take(self, held: str | Path) -> None:
+        """Become a copy of another SQLite database: the agent's own.
+
+        The whole file, schema and data together, rather than rows read out and written back. An
+        agent's real store carries things a reconstruction loses: its exact types, its indexes,
+        its keys, and every oddity in the data that its queries were actually written against.
+        """
+        origin = sqlite3.connect(f"file:{Path(held)}?mode=ro", uri=True)
+        try:
+            with self.connection:
+                origin.backup(self.connection)
+        finally:
+            origin.close()
+
+    # -- going back, on disk ---------------------------------------------------------
+
+    def save_to(self, path: str | Path) -> None:
+        root = Path(path)
+        root.mkdir(parents=True, exist_ok=True)
+        # Settled first: an open read transaction makes the copy fail, and a handler that only
+        # read leaves one behind.
+        self.connection.commit()
+        held = root / self.FILE
+        # A world whose live database already is the saved file has nothing to copy, and copying
+        # it would be a backup onto its own file. SQLite retries a locked destination rather than
+        # refusing, so that does not fail: it hangs, with no error and no timeout, and the build
+        # stops dead somewhere nobody is looking.
+        if self._same_file(held):
+            return
+        copy = sqlite3.connect(held)
+        with copy:
+            self.connection.backup(copy)
+        copy.close()
+
+    def _same_file(self, held: Path) -> bool:
+        if self.database == ":memory:":
+            return False
+        live = Path(self.database)
+        if not live.exists() or not held.exists():
+            return str(live) == str(held)
+        return live.samefile(held)
+
+    def load_from(self, path: str | Path) -> None:
+        held = Path(path) / self.FILE
+        if not held.exists():
+            raise StoreError(f"no saved store at {held}")
+        if self.database == ":memory:":
+            origin = sqlite3.connect(held)
+            with self.connection:
+                origin.backup(self.connection)
+            origin.close()
+            return
+        self.connection.close()
+        shutil.copyfile(held, self.database)
+        self.connection = sqlite3.connect(self.database, check_same_thread=False)
+        self.connection.execute("PRAGMA foreign_keys = ON")
+
+    def close(self) -> None:
+        self.connection.close()

fi/alk/harness/world/stores/written.py

@@ -0,0 +1,182 @@
+"""A store the build stage wrote, for an engine nobody taught the harness in advance.
+
+This is what keeps "whatever the agent uses" from meaning "whatever we got around to
+shipping". When the build stage finds an agent on an engine with no store in the tree, it
+writes one: the image to run, the port it listens on, how to build a connection string for it,
+and five functions saying how to talk to it.
+
+None of that is trusted. A reset written by a model for an engine nobody reviewed is exactly
+the thing that fails silently -- rows go back, a counter does not, and every scenario after the
+first is measured against a world that drifted. So a written store is registered, not accepted:
+whether it works is decided by ``prove_store``, which runs the same change twice from the same
+starting point and compares. Nothing here has to be right for that gate to be meaningful, which
+is the only reason writing it at build time is safe at all.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Callable
+
+from . import Snapshot, StoreError, register_store
+from .container import ContainerStore
+
+# The functions a written store defines. Fewer would not be enough for an arbitrary engine, and
+# more would be us guessing at what engines have in common. The last three are what a scenario's
+# own setup lands on: without them the environment can be stood up and read, but nothing can
+# change a little of it, so every scenario would run against the same base.
+REQUIRED = ("connect", "apply", "state", "freeze", "restore", "add", "amend", "remove")
+
+API = (
+    "Your code defines exactly these functions:\n"
+    "    def connect(dsn)                      -> a live client, already connected\n"
+    "    def apply(db, script)                 -> run statements: migrations, or a seed\n"
+    "    def state(db)                         -> {group: [row, ...]} for everything it holds\n"
+    "    def freeze(db)                        -> (rows, counters)\n"
+    "    def restore(db, rows, counters)       -> put both back exactly\n"
+    "    def add(db, group, record)            -> insert one record, return how many landed\n"
+    "    def amend(db, group, key, changes, by)-> update records where `by` equals `key`\n"
+    "    def remove(db, group, key, by)        -> delete those records; no key means all of them\n"
+    "Import whatever driver this engine needs at the top of the file; if it is not installed "
+    "you will be told which one is missing. `counters` is anything that keeps counting after "
+    "the rows are gone -- a sequence, an auto-increment. Restoring rows without it gives the "
+    "next scenario ids continuing from the last one. Engines that hand out nothing of the sort "
+    "return {}. Read state in a stable order, or a check comparing the first row is reading a "
+    "coin toss. `add`, `amend` and `remove` are what a scenario's setup calls, so they are the "
+    "difference between a suite of scenarios and one base world tested many times."
+)
+
+
+def _compile(code: str, engine: str) -> dict[str, Callable[..., Any]]:
+    """Turn the written code into its five functions, or say precisely what is missing."""
+    namespace: dict[str, Any] = {}
+    try:
+        exec(compile(code, f"<store:{engine}>", "exec"), namespace)  # nosec B102
+    except ImportError as exc:
+        raise StoreError(
+            f"{engine} ops import something that is not installed: {exc}. Install the driver "
+            f"this engine needs, or use an engine whose driver is already present."
+        ) from exc
+    except SyntaxError as exc:
+        raise StoreError(f"{engine} ops do not parse: {exc}") from exc
+
+    missing = [name for name in REQUIRED if not callable(namespace.get(name))]
+    if missing:
+        raise StoreError(
+            f"{engine} ops define {', '.join(sorted(n for n in REQUIRED if n not in missing)) or 'nothing'}"
+            f" but not {', '.join(missing)}.\n{API}"
+        )
+    return {name: namespace[name] for name in REQUIRED}
+
+
+def register_written(
+    *,
+    engine: str,
+    image: str,
+    container_port: int,
+    code: str,
+    boot_env: dict[str, str] | None = None,
+    dsn_template: str = "",
+) -> type[ContainerStore]:
+    """Teach the harness an engine from code written at build time.
+
+    Registered rather than accepted: this makes the engine available to ``declare_engine``, and
+    says nothing at all about whether its reset is correct. That is ``prove_store``'s to decide.
+    """
+    if not engine.strip():
+        raise StoreError("an engine needs a name")
+    if not image.strip():
+        raise StoreError(f"{engine} needs an image to run")
+    if not container_port:
+        raise StoreError(f"{engine} needs the port it listens on")
+    ops = _compile(code, engine)
+
+    class WrittenStore(ContainerStore):
+        pass
+
+    WrittenStore.engine = engine
+    WrittenStore.image = image
+    WrittenStore.container_port = container_port
+    WrittenStore.boot_env = dict(boot_env or {})
+    WrittenStore._ops = ops  # type: ignore[attr-defined]
+    WrittenStore._dsn_template = (
+        dsn_template or "{engine}://{user}:{password}@{host}:{port}/{database}"
+    )
+
+    def dsn(self: ContainerStore) -> str:
+        host, port = self.address()
+        return type(self)._dsn_template.format(  # type: ignore[attr-defined]
+            engine=type(self).engine,
+            user=self.user,
+            password=self.password,
+            host=host,
+            port=port,
+            database=self.database,
+        )
+
+    def _client(self: ContainerStore) -> Any:
+        """A fresh client per operation.
+
+        Never held open, for the same reason the Postgres store does not hold one: an idle
+        transaction of ours blocks the reset, and a reset that hangs on the harness's own
+        connection is a very expensive thing to debug.
+        """
+        return type(self)._ops["connect"](self.dsn())  # type: ignore[attr-defined]
+
+    def _with(self: ContainerStore, name: str, *args: Any) -> Any:
+        db = self._client()  # type: ignore[attr-defined]
+        try:
+            return type(self)._ops[name](db, *args)  # type: ignore[attr-defined]
+        finally:
+            closer = getattr(db, "close", None)
+            if callable(closer):
+                closer()
+
+    def probe(self: ContainerStore) -> None:
+        db = self._client()  # type: ignore[attr-defined]
+        closer = getattr(db, "close", None)
+        if callable(closer):
+            closer()
+
+    def apply(self: ContainerStore, script: str) -> None:
+        if not script.strip():
+            return
+        self._with("apply", script)  # type: ignore[attr-defined]
+
+    def state(self: ContainerStore) -> dict[str, list[dict[str, Any]]]:
+        return self._with("state")  # type: ignore[attr-defined]
+
+    def freeze(self: ContainerStore) -> Snapshot:
+        rows, counters = self._with("freeze")  # type: ignore[attr-defined]
+        return Snapshot(rows=rows, counters=counters or {})
+
+    def restore(self: ContainerStore, snapshot: Snapshot) -> None:
+        self._with("restore", snapshot.rows, snapshot.counters)  # type: ignore[attr-defined]
+
+    def add(self: ContainerStore, collection: str, record: Any) -> int:
+        return int(self._with("add", collection, dict(record)) or 0)  # type: ignore[attr-defined]
+
+    def amend(
+        self: ContainerStore, collection: str, key: str, changes: Any, *, by: str = ""
+    ) -> int:
+        return int(self._with("amend", collection, key, dict(changes), by) or 0)  # type: ignore[attr-defined]
+
+    def remove(
+        self: ContainerStore, collection: str, key: str = "", *, by: str = ""
+    ) -> int:
+        return int(self._with("remove", collection, key, by) or 0)  # type: ignore[attr-defined]
+
+    WrittenStore.add = add  # type: ignore[assignment]
+    WrittenStore.amend = amend  # type: ignore[assignment]
+    WrittenStore.remove = remove  # type: ignore[assignment]
+    WrittenStore.dsn = dsn  # type: ignore[assignment]
+    WrittenStore._client = _client  # type: ignore[attr-defined]
+    WrittenStore._with = _with  # type: ignore[attr-defined]
+    WrittenStore.probe = probe  # type: ignore[assignment]
+    WrittenStore.apply = apply  # type: ignore[assignment]
+    WrittenStore.state = state  # type: ignore[assignment]
+    WrittenStore.freeze = freeze  # type: ignore[assignment]
+    WrittenStore.restore = restore  # type: ignore[assignment]
+    WrittenStore.__name__ = f"{engine.title().replace('_', '')}Store"
+
+    register_store(engine, WrittenStore)
+    return WrittenStore