eval-protocol 0.0.3__py3-none-any.whl

eval_protocol/agent/resources/sql_resource.py

@@ -0,0 +1,271 @@
+"""
+SQLResource: A ForkableResource for managing SQL database states, initially focusing on SQLite.
+"""
+
+import os
+import shutil
+import sqlite3
+import uuid
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+from ..resource_abc import ForkableResource
+
+
+class SQLResource(ForkableResource):
+    """
+    A ForkableResource for managing SQL database states, primarily SQLite.
+
+    Manages a SQLite database file, allowing it to be initialized with a schema
+    and seed data, forked (by copying the DB file), checkpointed (by copying),
+    and restored.
+
+    Attributes:
+        _config (Dict[str, Any]): Configuration for the resource.
+        _db_path (Optional[Path]): Path to the current SQLite database file.
+        _base_db_path (Optional[Path]): Path to the initially set up database, used for forking.
+        _temp_dir (Path): Directory to store database files.
+    """
+
+    def __init__(self) -> None:
+        self._config: Dict[str, Any] = {}
+        self._db_path: Optional[Path] = None
+        self._base_db_path: Optional[Path] = None
+        # Consider making temp_dir configurable or using a more robust temp solution
+        self._temp_dir = Path("./.rk_temp_dbs").resolve()  # Ensure absolute path
+        self._temp_dir.mkdir(parents=True, exist_ok=True)
+
+    def _get_db_connection(self) -> sqlite3.Connection:
+        if not self._db_path:
+            raise ConnectionError("Database path not set. Call setup() or fork() first.")
+        # Set timeout to prevent indefinite hangs
+        return sqlite3.connect(str(self._db_path), timeout=10)
+
+    async def setup(self, config: Dict[str, Any]) -> None:
+        """
+        Initializes the SQLite database.
+
+        Args:
+            config: Configuration dictionary. Expected keys:
+                - 'db_type' (str): Must be 'sqlite'.
+                - 'db_name' (Optional[str]): Name for the database file. Defaults to a UUID.
+                - 'schema_file' (Optional[str]): Path to an SQL file to execute for schema setup.
+                - 'seed_data_file' (Optional[str]): Path to an SQL file for initial data seeding.
+                - 'schema_sql' (Optional[str]): SQL string for schema setup.
+                - 'seed_sql' (Optional[str]): SQL string for initial data seeding.
+        """
+        self._config = config.copy()
+        db_type = self._config.get("db_type", "sqlite")
+        if db_type != "sqlite":
+            raise ValueError("SQLResource currently only supports 'sqlite'.")
+
+        db_name = self._config.get("db_name", f"db_{uuid.uuid4().hex}.sqlite")
+        self._base_db_path = self._temp_dir / db_name
+        self._db_path = self._base_db_path  # Initially, the current DB is the base DB
+
+        # Ensure a fresh start if the base DB file already exists from a previous run
+        if self._base_db_path is not None and self._base_db_path.exists():
+            self._base_db_path.unlink()
+
+        conn = self._get_db_connection()
+        try:
+            with conn:
+                # Apply schema
+                schema_file = self._config.get("schema_file")
+                if schema_file and Path(schema_file).exists():
+                    with open(schema_file, "r") as f:
+                        conn.executescript(f.read())
+
+                schema_sql = self._config.get("schema_sql")
+                if schema_sql:
+                    conn.executescript(schema_sql)
+
+                # Apply seed data
+                seed_data_file = self._config.get("seed_data_file")
+                if seed_data_file and Path(seed_data_file).exists():
+                    with open(seed_data_file, "r") as f:
+                        conn.executescript(f.read())
+
+                seed_sql = self._config.get("seed_sql")
+                if seed_sql:
+                    conn.executescript(seed_sql)
+        finally:
+            conn.close()
+
+    async def fork(self) -> "SQLResource":
+        """
+        Creates a new SQLResource instance with a copy of the base database state.
+        If called on an already forked resource, it forks from its current state.
+        """
+        if not self._db_path or not self._db_path.exists():
+            raise RuntimeError("Cannot fork: original database does not exist or setup was not called.")
+
+        forked_resource = SQLResource()
+        forked_resource._config = self._config.copy()
+        forked_resource._temp_dir = self._temp_dir  # Share the same temp dir base
+
+        # The new fork's base is the current state of this resource
+        forked_resource._base_db_path = self._db_path
+
+        # Create a new unique DB file for this fork
+        forked_db_name = f"fork_{uuid.uuid4().hex}.sqlite"
+        forked_resource._db_path = self._temp_dir / forked_db_name
+
+        shutil.copyfile(str(self._db_path), str(forked_resource._db_path))
+        return forked_resource
+
+    async def checkpoint(self) -> Dict[str, Any]:
+        """
+        Returns a serializable representation of the resource's current state.
+        For SQLite, this involves copying the database file to a checkpoint location
+        and returning the path.
+        """
+        if not self._db_path or not self._db_path.exists():
+            raise RuntimeError("Cannot checkpoint: database does not exist.")
+
+        checkpoint_name = f"checkpoint_{self._db_path.stem}_{uuid.uuid4().hex}.sqlite"
+        checkpoint_path = self._temp_dir / checkpoint_name
+        shutil.copyfile(str(self._db_path), str(checkpoint_path))
+        return {"db_type": "sqlite", "checkpoint_path": str(checkpoint_path)}
+
+    async def restore(self, state_data: Dict[str, Any]) -> None:
+        """
+        Restores the resource's state from a previously checkpointed state.
+        For SQLite, this means copying the checkpointed DB file to become the current DB.
+        """
+        db_type = state_data.get("db_type")
+        checkpoint_path_str = state_data.get("checkpoint_path")
+
+        if db_type != "sqlite" or not checkpoint_path_str:
+            raise ValueError("Invalid state_data for SQLite restore.")
+
+        checkpoint_path = Path(checkpoint_path_str)
+        if not checkpoint_path.exists():
+            raise FileNotFoundError(f"Checkpoint file not found: {checkpoint_path}")
+
+        # If current db_path is not set (e.g. fresh resource), assign one
+        if not self._db_path:
+            self._db_path = self._temp_dir / f"restored_{uuid.uuid4().hex}.sqlite"
+
+        shutil.copyfile(str(checkpoint_path), str(self._db_path))
+        self._base_db_path = self._db_path  # The restored state becomes the new base for future forks
+
+    async def step(self, action_name: str, action_params: Dict[str, Any]) -> Any:
+        """
+        Executes a SQL query on the database.
+
+        Args:
+            action_name: Should be 'execute_sql'.
+            action_params: Dictionary containing:
+                - 'query' (str): The SQL query to execute.
+                - 'parameters' (Optional[Dict | List]): Parameters for the query.
+                - 'fetch_mode' (Optional[str]): 'one', 'all', or 'val'. If None, no fetch.
+
+        Returns:
+            Query result based on fetch_mode, or rowcount for DML.
+        """
+        if action_name != "execute_sql":
+            raise NotImplementedError(f"Action '{action_name}' not supported by SQLResource.")
+
+        query = action_params.get("query")
+        if not query:
+            raise ValueError("Missing 'query' in action_params for 'execute_sql'.")
+
+        params = action_params.get("parameters", [])
+        fetch_mode = action_params.get("fetch_mode")  # 'one', 'all', 'val'
+
+        conn = self._get_db_connection()
+        try:
+            with conn:
+                cursor = conn.cursor()
+                cursor.execute(query, params)
+
+                if fetch_mode == "one":
+                    columns = [desc[0] for desc in cursor.description]
+                    row = cursor.fetchone()
+                    return dict(zip(columns, row)) if row else None
+                elif fetch_mode == "all":
+                    columns = [desc[0] for desc in cursor.description]
+                    rows = cursor.fetchall()
+                    return [dict(zip(columns, row)) for row in rows]
+                elif fetch_mode == "val":
+                    row = cursor.fetchone()
+                    return row[0] if row else None
+                else:  # DML or no fetch needed
+                    return {"rowcount": cursor.rowcount}
+        finally:
+            conn.close()
+
+    async def get_observation(self) -> Dict[str, Any]:
+        """
+        Returns the current observable state of the resource.
+        For SQLResource, this could be the path to the DB or a status message.
+        """
+        return {
+            "db_type": "sqlite",
+            "db_path": str(self._db_path) if self._db_path else None,
+            "status": ("ready" if self._db_path and self._db_path.exists() else "uninitialized"),
+        }
+
+    async def get_tools_spec(self) -> List[Dict[str, Any]]:
+        """
+        Returns tool specifications for interacting with the SQL database.
+        """
+        return [
+            {
+                "type": "function",
+                "function": {
+                    "name": "execute_sql",
+                    "description": "Executes a SQL query against the database. "
+                    "Use 'fetch_mode' to control return value: "
+                    "'one' for a single row, "
+                    "'all' for all rows, "
+                    "'val' for a single value from the first row. "
+                    "If 'fetch_mode' is not provided, returns rowcount for DML statements.",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {
+                            "query": {
+                                "type": "string",
+                                "description": "The SQL query to execute.",
+                            },
+                            "parameters": {
+                                "type": "array",  # Or object for named parameters, sqlite3 supports both
+                                "description": "Parameters for the SQL query (optional).",
+                                "items": {"type": "any"},
+                            },
+                            "fetch_mode": {
+                                "type": "string",
+                                "enum": ["one", "all", "val"],
+                                "description": "Specifies how to fetch results (optional).",
+                            },
+                        },
+                        "required": ["query"],
+                    },
+                },
+            }
+        ]
+
+    async def close(self) -> None:
+        """
+        Cleans up by deleting the created database file(s).
+        More robust cleanup of the _temp_dir might be needed if it's shared or persistent.
+        """
+        if self._db_path and self._db_path.exists():
+            try:
+                self._db_path.unlink()
+            except OSError as e:
+                print(f"Error deleting database file {self._db_path}: {e}")
+
+        # Potentially clean up base_db_path if it's different and also temporary
+        # if self._base_db_path and self._base_db_path.exists() and self._base_db_path != self._db_path:
+        #     try:
+        #         self._base_db_path.unlink()
+        #     except OSError:
+        #         pass # ignore if it was already deleted or moved
+
+        # For now, we don't delete the _temp_dir itself, as it might contain checkpoints
+        # or other DBs from concurrent operations. A more sophisticated cleanup strategy
+        # for _temp_dir might be needed for long-running processes.
+        self._db_path = None
+        self._base_db_path = None