yamlgraph 0.1.1__py3-none-any.whl

yamlgraph/schema_loader.py

@@ -0,0 +1,160 @@
+"""Dynamic Pydantic model generation from YAML schema definitions.
+
+This module enables defining output schemas in YAML prompt files,
+making prompts fully self-contained with their expected output structure.
+
+Example YAML schema:
+    schema:
+      name: MyOutputModel
+      fields:
+        title:
+          type: str
+          description: "The output title"
+        confidence:
+          type: float
+          constraints: {ge: 0.0, le: 1.0}
+"""
+
+import re
+from pathlib import Path
+from typing import Any
+
+import yaml
+from pydantic import Field, create_model
+
+# =============================================================================
+# Type Resolution
+# =============================================================================
+
+# Mapping from type strings to Python types
+TYPE_MAP: dict[str, type] = {
+    "str": str,
+    "int": int,
+    "float": float,
+    "bool": bool,
+    "Any": Any,
+}
+
+
+def resolve_type(type_str: str, field_name: str | None = None) -> type:
+    """Resolve a type string to a Python type.
+
+    Supports:
+        - Basic types: str, int, float, bool, Any
+        - Generic types: list[str], list[int], dict[str, str]
+
+    Args:
+        type_str: Type string like "str", "list[str]", "dict[str, Any]"
+        field_name: Optional field name for better error messages
+
+    Returns:
+        Python type annotation
+
+    Raises:
+        ValueError: If type string is not recognized
+    """
+    # Check basic types first
+    if type_str in TYPE_MAP:
+        return TYPE_MAP[type_str]
+
+    # Handle list[T] pattern
+    list_match = re.match(r"list\[(\w+)\]", type_str)
+    if list_match:
+        inner_type = resolve_type(list_match.group(1), field_name)
+        return list[inner_type]
+
+    # Handle dict[K, V] pattern
+    dict_match = re.match(r"dict\[(\w+),\s*(\w+)\]", type_str)
+    if dict_match:
+        key_type = resolve_type(dict_match.group(1), field_name)
+        value_type = resolve_type(dict_match.group(2), field_name)
+        return dict[key_type, value_type]
+
+    # Provide helpful error with supported types
+    supported = ", ".join(TYPE_MAP.keys())
+    context = f" for field '{field_name}'" if field_name else ""
+    raise ValueError(
+        f"Unknown type: '{type_str}'{context}. "
+        f"Supported types: {supported}, list[T], dict[K, V]"
+    )
+
+
+# =============================================================================
+# Model Building
+# =============================================================================
+
+
+def build_pydantic_model(schema: dict) -> type:
+    """Build a Pydantic model dynamically from a schema dict.
+
+    Args:
+        schema: Schema definition with 'name' and 'fields' keys
+            Example:
+                {
+                    "name": "MyOutputModel",
+                    "fields": {
+                        "title": {"type": "str", "description": "..."},
+                        "score": {"type": "float", "constraints": {"ge": 0}},
+                    }
+                }
+
+    Returns:
+        Dynamically created Pydantic model class
+    """
+    model_name = schema["name"]
+    field_definitions = {}
+
+    for field_name, field_def in schema["fields"].items():
+        # Resolve the type - pass field_name for better error messages
+        field_type = resolve_type(field_def["type"], field_name)
+
+        # Handle optional fields
+        is_optional = field_def.get("optional", False)
+        if is_optional:
+            field_type = field_type | None
+
+        # Build Field kwargs
+        field_kwargs: dict[str, Any] = {}
+
+        if "description" in field_def:
+            field_kwargs["description"] = field_def["description"]
+
+        if "default" in field_def:
+            field_kwargs["default"] = field_def["default"]
+        elif is_optional:
+            field_kwargs["default"] = None
+
+        # Add constraints (ge, le, min_length, max_length, etc.)
+        if constraints := field_def.get("constraints"):
+            field_kwargs.update(constraints)
+
+        # Create field tuple: (type, Field(...))
+        if field_kwargs:
+            field_definitions[field_name] = (field_type, Field(**field_kwargs))
+        else:
+            field_definitions[field_name] = (field_type, ...)
+
+    return create_model(model_name, **field_definitions)
+
+
+# =============================================================================
+# YAML Loading
+# =============================================================================
+
+
+def load_schema_from_yaml(yaml_path: str | Path) -> type | None:
+    """Load a Pydantic model from a prompt YAML file's schema block.
+
+    Args:
+        yaml_path: Path to the YAML prompt file
+
+    Returns:
+        Dynamically created Pydantic model, or None if no schema defined
+    """
+    with open(yaml_path) as f:
+        config = yaml.safe_load(f)
+
+    if "schema" not in config:
+        return None
+
+    return build_pydantic_model(config["schema"])

yamlgraph/storage/__init__.py

@@ -0,0 +1,17 @@
+"""Storage utilities for persistence and export."""
+
+from yamlgraph.storage.database import YamlGraphDB
+from yamlgraph.storage.export import (
+    export_state,
+    export_summary,
+    list_exports,
+    load_export,
+)
+
+__all__ = [
+    "YamlGraphDB",
+    "export_state",
+    "export_summary",
+    "list_exports",
+    "load_export",
+]

yamlgraph/storage/checkpointer.py

@@ -0,0 +1,72 @@
+"""LangGraph native checkpointer integration.
+
+Provides SQLite-based checkpointing for graph state persistence,
+enabling time travel, replay, and resume from any checkpoint.
+"""
+
+import sqlite3
+from pathlib import Path
+from typing import Any
+
+from langgraph.checkpoint.sqlite import SqliteSaver
+from langgraph.graph.state import CompiledStateGraph
+
+from yamlgraph.config import DATABASE_PATH
+
+
+def get_checkpointer(db_path: str | Path | None = None) -> SqliteSaver:
+    """Get a SQLite checkpointer for graph compilation.
+
+    The checkpointer enables:
+    - Automatic state persistence after each node
+    - Time travel via get_state_history()
+    - Resume from any checkpoint
+    - Fault tolerance with pending writes
+
+    Args:
+        db_path: Path to SQLite database file.
+                 Defaults to outputs/yamlgraph.db
+
+    Returns:
+        SqliteSaver instance for use with graph.compile()
+
+    Example:
+        >>> checkpointer = get_checkpointer()
+        >>> graph = workflow.compile(checkpointer=checkpointer)
+        >>> result = graph.invoke(input, {"configurable": {"thread_id": "abc"}})
+    """
+    if db_path is None:
+        db_path = DATABASE_PATH
+
+    path = Path(db_path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+
+    conn = sqlite3.connect(str(path), check_same_thread=False)
+    return SqliteSaver(conn)
+
+
+def get_state_history(
+    graph: CompiledStateGraph,
+    thread_id: str,
+) -> list[Any]:
+    """Get checkpoint history for a thread.
+
+    Returns checkpoints in reverse chronological order (most recent first).
+
+    Args:
+        graph: Compiled graph with checkpointer
+        thread_id: Thread identifier to query
+
+    Returns:
+        List of StateSnapshot objects, or empty list if thread doesn't exist
+
+    Example:
+        >>> history = get_state_history(graph, "my-thread")
+        >>> for snapshot in history:
+        ...     print(f"Step {snapshot.metadata.get('step')}: {snapshot.values}")
+    """
+    config = {"configurable": {"thread_id": thread_id}}
+    try:
+        return list(graph.get_state_history(config))
+    except Exception:
+        return []

yamlgraph/storage/database.py

@@ -0,0 +1,320 @@
+"""SQLite Storage - Simple persistence for pipeline state.
+
+Provides a lightweight wrapper around SQLite for storing
+and retrieving pipeline execution state.
+
+Supports optional connection pooling for high-throughput scenarios.
+"""
+
+import json
+import sqlite3
+import threading
+from contextlib import contextmanager
+from datetime import datetime
+from pathlib import Path
+from queue import Empty, Queue
+from typing import Iterator
+
+from pydantic import BaseModel
+
+from yamlgraph.config import DATABASE_PATH
+
+
+class ConnectionPool:
+    """Thread-safe SQLite connection pool.
+
+    Maintains a pool of reusable connections for high-throughput scenarios.
+    Connections are returned to the pool after use instead of being closed.
+    """
+
+    def __init__(self, db_path: Path, pool_size: int = 5):
+        """Initialize connection pool.
+
+        Args:
+            db_path: Path to SQLite database
+            pool_size: Maximum number of connections to maintain
+        """
+        self._db_path = db_path
+        self._pool_size = pool_size
+        self._pool: Queue[sqlite3.Connection] = Queue(maxsize=pool_size)
+        self._lock = threading.Lock()
+        self._total_connections = 0
+
+    def _create_connection(self) -> sqlite3.Connection:
+        """Create a new database connection."""
+        conn = sqlite3.connect(self._db_path, check_same_thread=False)
+        conn.row_factory = sqlite3.Row
+        return conn
+
+    @contextmanager
+    def get_connection(self) -> Iterator[sqlite3.Connection]:
+        """Get a connection from the pool.
+
+        Creates a new connection if pool is empty and under limit.
+
+        Yields:
+            Database connection (returned to pool on exit)
+        """
+        conn = None
+        try:
+            # Try to get from pool
+            try:
+                conn = self._pool.get_nowait()
+            except Empty:
+                # Pool empty - create new connection if under limit
+                with self._lock:
+                    if self._total_connections < self._pool_size:
+                        conn = self._create_connection()
+                        self._total_connections += 1
+                    else:
+                        # At limit - block waiting for connection
+                        pass
+
+                if conn is None:
+                    conn = self._pool.get()  # Blocking wait
+
+            yield conn
+
+        finally:
+            # Return connection to pool
+            if conn is not None:
+                try:
+                    self._pool.put_nowait(conn)
+                except Exception:
+                    # Pool full, close connection
+                    conn.close()
+                    with self._lock:
+                        self._total_connections -= 1
+
+    def close_all(self) -> None:
+        """Close all connections in the pool."""
+        while True:
+            try:
+                conn = self._pool.get_nowait()
+                conn.close()
+            except Empty:
+                break
+        with self._lock:
+            self._total_connections = 0
+
+
+class YamlGraphDB:
+    """SQLite wrapper for yamlgraph state persistence.
+
+    Supports two connection modes:
+    - Default: Creates new connection per operation (simple, safe)
+    - Pooled: Reuses connections from pool (high-throughput)
+
+    Example:
+        # Default mode (simple)
+        db = YamlGraphDB()
+
+        # Pooled mode (high-throughput)
+        db = YamlGraphDB(use_pool=True, pool_size=10)
+    """
+
+    def __init__(
+        self,
+        db_path: str | Path | None = None,
+        use_pool: bool = False,
+        pool_size: int = 5,
+    ):
+        """Initialize database connection.
+
+        Args:
+            db_path: Path to SQLite database file (default: outputs/yamlgraph.db)
+            use_pool: Enable connection pooling for high-throughput scenarios
+            pool_size: Maximum connections in pool (only used if use_pool=True)
+        """
+        if db_path is None:
+            db_path = DATABASE_PATH
+        self.db_path = Path(db_path)
+        self.db_path.parent.mkdir(parents=True, exist_ok=True)
+
+        self._use_pool = use_pool
+        self._pool: ConnectionPool | None = None
+        if use_pool:
+            self._pool = ConnectionPool(self.db_path, pool_size)
+
+        self._init_db()
+
+    @contextmanager
+    def _get_connection(self) -> Iterator[sqlite3.Connection]:
+        """Get a database connection.
+
+        Uses pool if enabled, otherwise creates new connection.
+
+        Yields:
+            Database connection
+        """
+        if self._pool is not None:
+            with self._pool.get_connection() as conn:
+                yield conn
+        else:
+            conn = sqlite3.connect(self.db_path)
+            conn.row_factory = sqlite3.Row
+            try:
+                yield conn
+            finally:
+                conn.close()
+
+    def close(self) -> None:
+        """Close database connections.
+
+        For pooled mode, closes all connections in pool.
+        """
+        if self._pool is not None:
+            self._pool.close_all()
+
+    def _init_db(self):
+        """Initialize database tables."""
+        with self._get_connection() as conn:
+            conn.execute("""
+                CREATE TABLE IF NOT EXISTS pipeline_runs (
+                    id INTEGER PRIMARY KEY AUTOINCREMENT,
+                    thread_id TEXT NOT NULL,
+                    created_at TEXT NOT NULL,
+                    updated_at TEXT NOT NULL,
+                    status TEXT NOT NULL DEFAULT 'running',
+                    state_json TEXT NOT NULL
+                )
+            """)
+            conn.execute("""
+                CREATE INDEX IF NOT EXISTS idx_thread_id
+                ON pipeline_runs(thread_id)
+            """)
+            conn.commit()
+
+    def save_state(self, thread_id: str, state: dict, status: str = "running") -> int:
+        """Save pipeline state.
+
+        Args:
+            thread_id: Unique identifier for this run
+            state: State dictionary to persist
+            status: Current status (running, completed, failed)
+
+        Returns:
+            Row ID of the saved state
+        """
+        now = datetime.now().isoformat()
+        state_json = json.dumps(self._serialize_state(state), default=str)
+
+        with self._get_connection() as conn:
+            # Check if thread exists
+            existing = conn.execute(
+                "SELECT id FROM pipeline_runs WHERE thread_id = ?", (thread_id,)
+            ).fetchone()
+
+            if existing:
+                conn.execute(
+                    """UPDATE pipeline_runs
+                       SET updated_at = ?, status = ?, state_json = ?
+                       WHERE thread_id = ?""",
+                    (now, status, state_json, thread_id),
+                )
+                conn.commit()
+                return existing["id"]
+            else:
+                cursor = conn.execute(
+                    """INSERT INTO pipeline_runs
+                       (thread_id, created_at, updated_at, status, state_json)
+                       VALUES (?, ?, ?, ?, ?)""",
+                    (thread_id, now, now, status, state_json),
+                )
+                conn.commit()
+                return cursor.lastrowid
+
+    def load_state(self, thread_id: str) -> dict | None:
+        """Load pipeline state by thread ID.
+
+        Args:
+            thread_id: Unique identifier for the run
+
+        Returns:
+            State dictionary or None if not found
+        """
+        with self._get_connection() as conn:
+            row = conn.execute(
+                "SELECT state_json FROM pipeline_runs WHERE thread_id = ?", (thread_id,)
+            ).fetchone()
+
+            if row:
+                return json.loads(row["state_json"])
+            return None
+
+    def get_run_info(self, thread_id: str) -> dict | None:
+        """Get run metadata without full state.
+
+        Args:
+            thread_id: Unique identifier for the run
+
+        Returns:
+            Dictionary with id, thread_id, created_at, updated_at, status
+        """
+        with self._get_connection() as conn:
+            row = conn.execute(
+                """SELECT id, thread_id, created_at, updated_at, status
+                   FROM pipeline_runs WHERE thread_id = ?""",
+                (thread_id,),
+            ).fetchone()
+
+            if row:
+                return dict(row)
+            return None
+
+    def list_runs(self, limit: int = 10) -> list[dict]:
+        """List recent pipeline runs.
+
+        Args:
+            limit: Maximum number of runs to return
+
+        Returns:
+            List of run metadata dictionaries
+        """
+        with self._get_connection() as conn:
+            rows = conn.execute(
+                """SELECT id, thread_id, created_at, updated_at, status
+                   FROM pipeline_runs
+                   ORDER BY updated_at DESC
+                   LIMIT ?""",
+                (limit,),
+            ).fetchall()
+
+            return [dict(row) for row in rows]
+
+    def delete_run(self, thread_id: str) -> bool:
+        """Delete a pipeline run.
+
+        Args:
+            thread_id: Unique identifier for the run
+
+        Returns:
+            True if deleted, False if not found
+        """
+        with self._get_connection() as conn:
+            cursor = conn.execute(
+                "DELETE FROM pipeline_runs WHERE thread_id = ?", (thread_id,)
+            )
+            conn.commit()
+            return cursor.rowcount > 0
+
+    def _serialize_state(self, state: dict) -> dict:
+        """Convert state to JSON-serializable format.
+
+        Handles Pydantic models and other complex types.
+
+        Args:
+            state: State dictionary
+
+        Returns:
+            JSON-serializable dictionary
+        """
+        result = {}
+        for key, value in state.items():
+            if isinstance(value, BaseModel):
+                result[key] = value.model_dump()
+            elif hasattr(value, "__dict__"):
+                result[key] = vars(value)
+            else:
+                result[key] = value
+        return result