synth-ai 0.2.0__py3-none-any.whl → 0.2.1.dev0__py3-none-any.whl

synth_ai/environments/reproducibility/tree.py

@@ -0,0 +1,364 @@
+"""
+NOTE - first pass was o3-generated. Mostly bc idrk what I 'want' from this yet ...
+trajectory_tree_store.py
+~~~~~~~~~~~~~~~~~~~~~~~~
+A minimal search-tree wrapper that pairs
+
+  • an *in-memory* NetworkX DiGraph (parent ⇢ children edges)
+  • a *content-addressable* FilesystemSnapshotStore (heavy blobs)
+
+so you can implement things like LATS / MCTS without bringing in the
+big “backend.production” code-base.
+"""
+
+from __future__ import annotations
+
+import json
+import sqlite3
+import gzip
+import pickle
+import logging
+from pathlib import Path
+from typing import Any, Dict, Optional, Tuple, Iterable
+
+import networkx as nx
+
+# from filesystem_snapshot_store import FilesystemSnapshotStore  # ← your re-impl
+
+log = logging.getLogger(__name__)
+
+
+# --------------------------------------------------------------------------- #
+# lightweight metadata record                                                 #
+# --------------------------------------------------------------------------- #
+import os
+import hashlib
+import logging
+from typing import Union
+
+log = logging.getLogger(__name__)
+
+# Default directory for storing snapshots relative to some base path
+# This could be configured via environment variables or settings later.
+DEFAULT_SNAPSHOT_DIR = Path(os.getenv("SNAPSHOT_STORE_PATH", "/tmp/agent_snapshots"))
+
+
+class FilesystemSnapshotStore:
+    """
+    Stores and retrieves environment state snapshots on the filesystem.
+
+    Uses content-addressable storage: the key (ID) for a snapshot
+    is the SHA-256 hash of its compressed content.
+    """
+
+    def __init__(self, base_dir: Union[str, Path] = DEFAULT_SNAPSHOT_DIR):
+        self.base_dir = Path(base_dir)
+        try:
+            self.base_dir.mkdir(parents=True, exist_ok=True)
+            log.info(f"Initialized snapshot store at: {self.base_dir}")
+        except OSError as e:
+            log.error(
+                f"Failed to create snapshot directory {self.base_dir}: {e}",
+                exc_info=True,
+            )
+            raise
+
+    def _get_path(self, key: str) -> Path:
+        """Constructs the full path for a given snapshot key."""
+        # Maybe add subdirectories later for large numbers of files, e.g., key[:2]/key[2:]
+        filename = f"{key}.snapshot.gz"
+        return self.base_dir / filename
+
+    def write(self, blob: Union[bytes, Dict[str, Any]]) -> str:
+        """
+        Stores a snapshot blob (bytes or dict) and returns its SHA-256 key.
+
+        • Dicts → pickle → gzip
+        • Bytes already gzip-compressed (magic 0x1f 0x8b) are stored as-is
+          to avoid double compression.
+        """
+        try:
+            if isinstance(blob, dict):
+                compressed_blob = gzip.compress(pickle.dumps(blob))
+            elif isinstance(blob, bytes):
+                # Skip re-compression if data is already gzipped
+                compressed_blob = blob if blob[:2] == b"\x1f\x8b" else gzip.compress(blob)
+            else:
+                raise TypeError(f"Unsupported blob type for snapshot store: {type(blob)}")
+
+            key = hashlib.sha256(compressed_blob).hexdigest()
+            path = self._get_path(key)
+            if not path.exists():
+                path.write_bytes(compressed_blob)
+            return key
+        except Exception as e:
+            log.error(f"Failed to write snapshot: {e}", exc_info=True)
+            raise
+
+    def read(self, key: str) -> Optional[bytes]:
+        """
+        Retrieves the raw *compressed* snapshot bytes for a given key.
+
+        Returns None if the key is not found.
+        Deserialization (decompression, unpickling) is the responsibility
+        of the caller (e.g., ReproducibleResource.from_snapshot).
+        """
+        filepath = self._get_path(key)
+        if not filepath.exists():
+            log.warning(f"Snapshot key not found: {key}")
+            return None
+        try:
+            with open(filepath, "rb") as f:
+                compressed_blob = f.read()
+            return compressed_blob
+        except OSError as e:
+            log.error(f"Failed to read snapshot {key} from {filepath}: {e}", exc_info=True)
+            return None  # Or re-raise? Returning None might be safer.
+
+    def exists(self, key: str) -> bool:
+        """Checks if a snapshot with the given key exists."""
+        return self._get_path(key).exists()
+
+
+# Global instance (optional, could use dependency injection)
+# snapshot_store = FilesystemSnapshotStore()
+
+
+class TrajectorySnapshot:
+    """
+    A *metadata* header for one node in the search tree.
+    The heavy serialized-state bytes live only in the snapshot store.
+    """
+
+    __slots__ = (
+        "snap_id",
+        "parent_id",
+        "depth",
+        "action",
+        "reward",
+        "terminated",
+        "info",
+    )
+
+    def __init__(
+        self,
+        snap_id: str,
+        parent_id: Optional[str],
+        depth: int,
+        action: Optional[Any],
+        reward: float = 0.0,
+        terminated: bool = False,
+        info: Optional[Dict[str, Any]] = None,
+    ):
+        self.snap_id = snap_id
+        self.parent_id = parent_id
+        self.depth = depth
+        self.action = action
+        self.reward = reward
+        self.terminated = bool(terminated)
+        self.info = info or {}
+
+    # helpful for printing / debugging
+    def __repr__(self) -> str:  # pragma: no cover
+        a = json.dumps(self.action) if self.action is not None else "∅"
+        return (
+            f"TrajSnap(id={self.snap_id[:7]}…, depth={self.depth}, "
+            f"action={a}, reward={self.reward}, term={self.terminated})"
+        )
+
+
+# --------------------------------------------------------------------------- #
+# tree manager                                                                #
+# --------------------------------------------------------------------------- #
+
+
+class TrajectoryTreeStore:
+    """
+    ❑  Adds snapshots (root / children) and keeps the DAG in-memory
+    ❑  Optionally mirrors headers to a tiny SQLite DB (so you can kill +
+       resume a long search)
+    ❑  Hands out raw snapshot **bytes**; decoding is up to the caller.
+    """
+
+    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - #
+    # construction                                                            #
+    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - #
+
+    def __init__(
+        self,
+        snapshot_store: Optional[FilesystemSnapshotStore] = None,
+        *,
+        db_path: Optional[Path | str] = None,
+    ):
+        self.snap_store = snapshot_store or FilesystemSnapshotStore()
+        self.graph: nx.DiGraph = nx.DiGraph()
+        self.db_path = Path(db_path).expanduser() if db_path else None
+        if self.db_path:
+            self._init_sqlite()
+
+    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - #
+    # public API                                                              #
+    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - #
+
+    # insertion -------------------------------------------------------------
+
+    def add_root(self, snapshot_blob: bytes, *, info: Dict[str, Any] | None = None) -> str:
+        """Insert the very first node and return its content-hash key."""
+        snap_id = self.snap_store.write(snapshot_blob)
+        self._add_node(TrajectorySnapshot(snap_id, None, 0, None, 0.0, False, info))
+        return snap_id
+
+    def add_child(
+        self,
+        parent_id: str,
+        snapshot_blob: bytes,
+        *,
+        action: Any,
+        reward: float,
+        terminated: bool = False,
+        info: Dict[str, Any] | None = None,
+    ) -> str:
+        """Attach `snapshot_blob` as a child reached by `action` from *parent_id*."""
+        if parent_id not in self.graph:
+            raise KeyError(f"Parent snapshot {parent_id[:8]}… not in tree")
+        depth = self.graph.nodes[parent_id]["meta"].depth + 1  # type: ignore[index]
+        snap_id = self.snap_store.write(snapshot_blob)
+        meta = TrajectorySnapshot(snap_id, parent_id, depth, action, reward, terminated, info)
+        self._add_node(meta)  # records node + (maybe) SQLite
+        self.graph.add_edge(parent_id, snap_id, action=action, reward=reward)  # NX edge attrs
+        return snap_id
+
+    # read-side helpers -----------------------------------------------------
+
+    def get_children(self, snap_id: str) -> Tuple[str, ...]:
+        return tuple(self.graph.successors(snap_id))
+
+    def get_parent(self, snap_id: str) -> Optional[str]:
+        preds = tuple(self.graph.predecessors(snap_id))
+        return preds[0] if preds else None
+
+    def is_leaf(self, snap_id: str) -> bool:
+        return self.graph.out_degree(snap_id) == 0
+
+    # simple enumerations useful for MCTS / LATS ---------------------------
+
+    def iter_leaves(self) -> Iterable[str]:
+        """Yield snapshot-ids that currently have no children."""
+        return (n for n in self.graph.nodes if self.is_leaf(n))
+
+    def path_to_root(self, snap_id: str) -> Tuple[str, ...]:
+        """Return (snap_id, …, root_id)"""
+        path = [snap_id]
+        while (p := self.get_parent(path[-1])) is not None:
+            path.append(p)
+        return tuple(path)
+
+    def reconstruct_actions(self, snap_id: str) -> Tuple[Any, ...]:
+        """Return the sequence of *actions* from the root → `snap_id`."""
+        actions = []
+        for child, parent in zip(self.path_to_root(snap_id)[:-1], self.path_to_root(snap_id)[1:]):
+            actions.append(self.graph.edges[parent, child]["action"])
+        return tuple(reversed(actions))
+
+    # snapshot access -------------------------------------------------------
+
+    def load_snapshot_blob(self, snap_id: str) -> bytes:
+        blob = self.snap_store.read(snap_id)
+        if blob is None:
+            raise FileNotFoundError(f"Snapshot {snap_id[:8]}… missing on disk")
+        return blob
+
+    def load_pickled_payload(self, snap_id: str) -> Any:
+        """Decompress + unpickle whatever you stored under this id."""
+        return pickle.loads(gzip.decompress(self.load_snapshot_blob(snap_id)))
+
+    # mutation --------------------------------------------------------------
+
+    def prune_subtree(self, root_id: str) -> None:
+        """Remove *root_id* and all its descendants from the in-mem graph and DB."""
+        doomed = list(nx.dfs_preorder_nodes(self.graph, root_id))
+        self.graph.remove_nodes_from(doomed)
+        if self.db_path:
+            with sqlite3.connect(self.db_path) as conn:
+                conn.executemany("DELETE FROM nodes WHERE snap_id = ?;", ((n,) for n in doomed))
+                conn.executemany(
+                    "DELETE FROM edges WHERE parent_id = ? OR child_id = ?;",
+                    ((n, n) for n in doomed),
+                )
+                conn.commit()
+
+    def wipe(self) -> None:
+        """Clear the *entire* tree (does **not** delete snapshot files)."""
+        self.graph.clear()
+        if self.db_path:
+            with sqlite3.connect(self.db_path) as conn:
+                conn.executescript("DELETE FROM nodes; DELETE FROM edges;")
+                conn.commit()
+
+    # ------------------------------------------------------------------- #
+    # internal helpers                                                     #
+    # ------------------------------------------------------------------- #
+
+    def _add_node(self, meta: TrajectorySnapshot) -> None:
+        self.graph.add_node(meta.snap_id, meta=meta)
+        if self.db_path:
+            self._sqlite_insert(meta)
+
+    # ------------------------------------------------------------------- #
+    # tiny SQLite backing store (optional)                                 #
+    # ------------------------------------------------------------------- #
+
+    def _init_sqlite(self) -> None:
+        self.db_path.parent.mkdir(parents=True, exist_ok=True)
+        with sqlite3.connect(self.db_path) as conn:
+            conn.executescript(
+                """
+                CREATE TABLE IF NOT EXISTS nodes(
+                    snap_id   TEXT PRIMARY KEY,
+                    parent_id TEXT,
+                    depth     INTEGER,
+                    action    TEXT,
+                    reward    REAL,
+                    terminated INTEGER,
+                    info      TEXT
+                );
+                CREATE TABLE IF NOT EXISTS edges(
+                    parent_id TEXT,
+                    child_id  TEXT,
+                    action    TEXT,
+                    reward    REAL,
+                    PRIMARY KEY(parent_id, child_id)
+                );
+                """
+            )
+            conn.commit()
+
+    def _sqlite_insert(self, meta: TrajectorySnapshot) -> None:
+        with sqlite3.connect(self.db_path) as conn:
+            conn.execute(
+                """INSERT OR IGNORE INTO nodes
+                   (snap_id, parent_id, depth, action, reward, terminated, info)
+                   VALUES (?,?,?,?,?,?,?)""",
+                (
+                    meta.snap_id,
+                    meta.parent_id,
+                    meta.depth,
+                    json.dumps(meta.action),
+                    meta.reward,
+                    int(meta.terminated),
+                    json.dumps(meta.info),
+                ),
+            )
+            if meta.parent_id:
+                conn.execute(
+                    """INSERT OR IGNORE INTO edges
+                       (parent_id, child_id, action, reward)
+                       VALUES (?,?,?,?)""",
+                    (
+                        meta.parent_id,
+                        meta.snap_id,
+                        json.dumps(meta.action),
+                        meta.reward,
+                    ),
+                )
+            conn.commit()

synth_ai/environments/service/app.py

@@ -0,0 +1,78 @@
+import sys
+import os  # Added to ensure os is available before use
+
+# Ensure local 'src' directory is on PYTHONPATH for dev installs
+# Current file: <repo>/src/synth_env/service/app.py
+# We want to add <repo>/src to sys.path (two levels up)
+_src_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
+if _src_dir not in sys.path:
+    sys.path.insert(0, _src_dir)
+
+print(f"SYS.PATH IN APP.PY: {sys.path}")
+import logging
+
+from fastapi import FastAPI
+from synth_ai.environments.service.registry import list_supported_env_types, register_environment
+from synth_ai.environments.service.core_routes import api_router
+from synth_ai.environments.service.external_registry import (
+    ExternalRegistryConfig,
+    load_external_environments,
+)
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Register built-in environments at import time
+import synth_ai.environments.examples.sokoban.environment as sok
+
+register_environment("Sokoban", sok.SokobanEnvironment)
+import synth_ai.environments.examples.crafter_classic.environment as cc
+
+register_environment("CrafterClassic", cc.CrafterClassicEnvironment)
+import synth_ai.environments.examples.verilog.environment as ve
+
+register_environment("Verilog", ve.VerilogEnvironment)
+import synth_ai.environments.examples.tictactoe.environment as ttt
+
+register_environment("TicTacToe", ttt.TicTacToeEnvironment)
+import synth_ai.environments.examples.nethack.environment as nh
+
+register_environment("NetHack", nh.NetHackEnvironment)
+# AlgoTune excluded from package due to size/complexity
+# import synth_ai.environments.examples.algotune.environment as at
+# register_environment("AlgoTune", at.AlgoTuneEnvironment)
+import synth_ai.environments.examples.minigrid.environment as mg
+
+register_environment("MiniGrid", mg.MiniGridEnvironment)
+import synth_ai.environments.examples.enron.environment as enron
+
+register_environment("Enron", enron.EnronEnvironment)
+
+app = FastAPI(title="Environment Service")
+
+
+@app.on_event("startup")
+async def startup_event():
+    """Load external environments on startup."""
+    # Support configuration-based loading for external environments
+    # You can set EXTERNAL_ENVIRONMENTS env var with JSON config
+    external_config = os.getenv("EXTERNAL_ENVIRONMENTS")
+    if external_config:
+        try:
+            import json
+
+            config_data = json.loads(external_config)
+            config = ExternalRegistryConfig(
+                external_environments=config_data.get("external_environments", [])
+            )
+            load_external_environments(config)
+        except Exception as e:
+            logger.error(f"Failed to load external environment config: {e}")
+
+    # Log all registered environments
+    logger.info(f"Registered environments: {list_supported_env_types()}")
+
+
+# Mount the main API router
+app.include_router(api_router, tags=["environments"])