Add @persist decorator with FlowPersistence interface (#1892)

* Add @persist decorator with SQLite persistence - Add FlowPersistence abstract base class - Implement SQLiteFlowPersistence backend - Add @persist decorator for flow state persistence - Add tests for flow persistence functionality Co-Authored-By: Joe Moura <joao@crewai.com> * Fix remaining merge conflicts in uv.lock - Remove stray merge conflict markers - Keep main's comprehensive platform-specific resolution markers - Preserve all required dependencies for persistence functionality Co-Authored-By: Joe Moura <joao@crewai.com> * Fix final CUDA dependency conflicts in uv.lock - Resolve NVIDIA CUDA solver dependency conflicts - Use main's comprehensive platform checks - Ensure all merge conflict markers are removed - Preserve persistence-related dependencies Co-Authored-By: Joe Moura <joao@crewai.com> * Fix nvidia-cusparse-cu12 dependency conflicts in uv.lock - Resolve NVIDIA CUSPARSE dependency conflicts - Use main's comprehensive platform checks - Complete systematic check of entire uv.lock file - Ensure all merge conflict markers are removed Co-Authored-By: Joe Moura <joao@crewai.com> * Fix triton filelock dependency conflicts in uv.lock - Resolve triton package filelock dependency conflict - Use main's comprehensive platform checks - Complete final systematic check of entire uv.lock file - Ensure TOML file structure is valid Co-Authored-By: Joe Moura <joao@crewai.com> * Fix merge conflict in crew_test.py - Remove duplicate assertion in test_multimodal_agent_live_image_analysis - Clean up conflict markers - Preserve test functionality Co-Authored-By: Joe Moura <joao@crewai.com> * Clean up trailing merge conflict marker in crew_test.py - Remove remaining conflict marker at end of file - Preserve test functionality - Complete conflict resolution Co-Authored-By: Joe Moura <joao@crewai.com> * Improve type safety in persistence implementation and resolve merge conflicts Co-Authored-By: Joe Moura <joao@crewai.com> * fix: Add explicit type casting in _create_initial_state method Co-Authored-By: Joe Moura <joao@crewai.com> * fix: Improve type safety in flow state handling with proper validation Co-Authored-By: Joe Moura <joao@crewai.com> * fix: Improve type system with proper TypeVar scoping and validation Co-Authored-By: Joe Moura <joao@crewai.com> * fix: Improve state restoration logic and add comprehensive tests Co-Authored-By: Joe Moura <joao@crewai.com> * fix: Initialize FlowState instances without passing id to constructor Co-Authored-By: Joe Moura <joao@crewai.com> * feat: Add class-level flow persistence decorator with SQLite default - Add class-level @persist decorator support - Set SQLiteFlowPersistence as default backend - Use db_storage_path for consistent database location - Improve async method handling and type safety - Add comprehensive docstrings and examples Co-Authored-By: Joe Moura <joao@crewai.com> * fix: Sort imports in decorators.py to fix lint error Co-Authored-By: Joe Moura <joao@crewai.com> * style: Organize imports according to PEP 8 standard Co-Authored-By: Joe Moura <joao@crewai.com> * style: Format typing imports with line breaks for better readability Co-Authored-By: Joe Moura <joao@crewai.com> * style: Simplify import organization to fix lint error Co-Authored-By: Joe Moura <joao@crewai.com> * style: Fix import sorting using Ruff auto-fix Co-Authored-By: Joe Moura <joao@crewai.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: Joe Moura <joao@crewai.com> Co-authored-by: João Moura <joaomdmoura@gmail.com>
2026-01-21 13:58:15 +00:00 · 2025-01-16 10:23:46 -03:00
parent 209f852300
commit eda8e673d7
10 changed files with 1061 additions and 138 deletions
--- a/src/crewai/flow/persistence/sqlite.py
+++ b/src/crewai/flow/persistence/sqlite.py
@@ -0,0 +1,124 @@
+"""
+SQLite-based implementation of flow state persistence.
+"""
+
+import json
+import os
+import sqlite3
+import tempfile
+from datetime import datetime
+from typing import Any, Dict, Optional, Union
+
+from pydantic import BaseModel
+
+from crewai.flow.persistence.base import FlowPersistence
+
+
+class SQLiteFlowPersistence(FlowPersistence):
+    """SQLite-based implementation of flow state persistence.
+    
+    This class provides a simple, file-based persistence implementation using SQLite.
+    It's suitable for development and testing, or for production use cases with
+    moderate performance requirements.
+    """
+    
+    db_path: str  # Type annotation for instance variable
+    
+    def __init__(self, db_path: Optional[str] = None):
+        """Initialize SQLite persistence.
+        
+        Args:
+            db_path: Path to the SQLite database file. If not provided, uses
+                    db_storage_path() from utilities.paths.
+                    
+        Raises:
+            ValueError: If db_path is invalid
+        """
+        from crewai.utilities.paths import db_storage_path
+        # Get path from argument or default location
+        path = db_path or db_storage_path()
+        
+        if not path:
+            raise ValueError("Database path must be provided")
+            
+        self.db_path = path  # Now mypy knows this is str
+        self.init_db()
+    
+    def init_db(self) -> None:
+        """Create the necessary tables if they don't exist."""
+        with sqlite3.connect(self.db_path) as conn:
+            conn.execute("""
+            CREATE TABLE IF NOT EXISTS flow_states (
+                id INTEGER PRIMARY KEY AUTOINCREMENT,
+                flow_uuid TEXT NOT NULL,
+                method_name TEXT NOT NULL,
+                timestamp DATETIME NOT NULL,
+                state_json TEXT NOT NULL
+            )
+            """)
+            # Add index for faster UUID lookups
+            conn.execute("""
+            CREATE INDEX IF NOT EXISTS idx_flow_states_uuid 
+            ON flow_states(flow_uuid)
+            """)
+    
+    def save_state(
+        self,
+        flow_uuid: str,
+        method_name: str,
+        state_data: Union[Dict[str, Any], BaseModel],
+    ) -> None:
+        """Save the current flow state to SQLite.
+        
+        Args:
+            flow_uuid: Unique identifier for the flow instance
+            method_name: Name of the method that just completed
+            state_data: Current state data (either dict or Pydantic model)
+        """
+        # Convert state_data to dict, handling both Pydantic and dict cases
+        if isinstance(state_data, BaseModel):
+            state_dict = dict(state_data)  # Use dict() for better type compatibility
+        elif isinstance(state_data, dict):
+            state_dict = state_data
+        else:
+            raise ValueError(
+                f"state_data must be either a Pydantic BaseModel or dict, got {type(state_data)}"
+            )
+            
+        with sqlite3.connect(self.db_path) as conn:
+            conn.execute("""
+            INSERT INTO flow_states (
+                flow_uuid,
+                method_name,
+                timestamp,
+                state_json
+            ) VALUES (?, ?, ?, ?)
+            """, (
+                flow_uuid,
+                method_name,
+                datetime.utcnow().isoformat(),
+                json.dumps(state_dict),
+            ))
+    
+    def load_state(self, flow_uuid: str) -> Optional[Dict[str, Any]]:
+        """Load the most recent state for a given flow UUID.
+        
+        Args:
+            flow_uuid: Unique identifier for the flow instance
+            
+        Returns:
+            The most recent state as a dictionary, or None if no state exists
+        """
+        with sqlite3.connect(self.db_path) as conn:
+            cursor = conn.execute("""
+            SELECT state_json
+            FROM flow_states
+            WHERE flow_uuid = ?
+            ORDER BY id DESC
+            LIMIT 1
+            """, (flow_uuid,))
+            row = cursor.fetchone()
+            
+        if row:
+            return json.loads(row[0])
+        return None