mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-10 00:28:31 +00:00
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>
This commit is contained in:
devin-ai-integration[bot]
committed by
GitHub
GitHub
parent
3dc442801f
commit
294f2cc3a9
177
src/crewai/flow/persistence/decorators.py
Normal file
177
src/crewai/flow/persistence/decorators.py
Normal file
@@ -0,0 +1,177 @@
|
||||
"""
|
||||
Decorators for flow state persistence.
|
||||
|
||||
Example:
|
||||
```python
|
||||
from crewai.flow.flow import Flow, start
|
||||
from crewai.flow.persistence import persist, SQLiteFlowPersistence
|
||||
|
||||
class MyFlow(Flow):
|
||||
@start()
|
||||
@persist(SQLiteFlowPersistence())
|
||||
def sync_method(self):
|
||||
# Synchronous method implementation
|
||||
pass
|
||||
|
||||
@start()
|
||||
@persist(SQLiteFlowPersistence())
|
||||
async def async_method(self):
|
||||
# Asynchronous method implementation
|
||||
await some_async_operation()
|
||||
```
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import functools
|
||||
import inspect
|
||||
import logging
|
||||
from typing import (
|
||||
Any,
|
||||
Callable,
|
||||
Dict,
|
||||
Optional,
|
||||
Type,
|
||||
TypeVar,
|
||||
Union,
|
||||
cast,
|
||||
get_type_hints,
|
||||
)
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
from crewai.flow.persistence.base import FlowPersistence
|
||||
from crewai.flow.persistence.sqlite import SQLiteFlowPersistence
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
def persist(persistence: Optional[FlowPersistence] = None):
|
||||
"""Decorator to persist flow state.
|
||||
|
||||
This decorator can be applied at either the class level or method level.
|
||||
When applied at the class level, it automatically persists all flow method
|
||||
states. When applied at the method level, it persists only that method's
|
||||
state.
|
||||
|
||||
Args:
|
||||
persistence: Optional FlowPersistence implementation to use.
|
||||
If not provided, uses SQLiteFlowPersistence.
|
||||
|
||||
Returns:
|
||||
A decorator that can be applied to either a class or method
|
||||
|
||||
Raises:
|
||||
ValueError: If the flow state doesn't have an 'id' field
|
||||
RuntimeError: If state persistence fails
|
||||
|
||||
Example:
|
||||
@persist # Class-level persistence with default SQLite
|
||||
class MyFlow(Flow[MyState]):
|
||||
@start()
|
||||
def begin(self):
|
||||
pass
|
||||
"""
|
||||
def _persist_state(flow_instance: Any, method_name: str, persistence_instance: FlowPersistence) -> None:
|
||||
"""Helper to persist state with error handling."""
|
||||
try:
|
||||
# Get flow UUID from state
|
||||
state = getattr(flow_instance, 'state', None)
|
||||
if state is None:
|
||||
raise ValueError("Flow instance has no state")
|
||||
|
||||
flow_uuid: Optional[str] = None
|
||||
if isinstance(state, dict):
|
||||
flow_uuid = state.get('id')
|
||||
elif isinstance(state, BaseModel):
|
||||
flow_uuid = getattr(state, 'id', None)
|
||||
|
||||
if not flow_uuid:
|
||||
raise ValueError(
|
||||
"Flow state must have an 'id' field for persistence"
|
||||
)
|
||||
|
||||
# Persist the state
|
||||
persistence_instance.save_state(
|
||||
flow_uuid=flow_uuid,
|
||||
method_name=method_name,
|
||||
state_data=state,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Failed to persist state for method {method_name}: {str(e)}"
|
||||
)
|
||||
raise RuntimeError(f"State persistence failed: {str(e)}") from e
|
||||
|
||||
def decorator(target: Union[Type, Callable[..., T]]) -> Union[Type, Callable[..., T]]:
|
||||
"""Decorator that handles both class and method decoration."""
|
||||
actual_persistence = persistence or SQLiteFlowPersistence()
|
||||
|
||||
if isinstance(target, type):
|
||||
# Class decoration
|
||||
class_methods = {}
|
||||
for name, method in target.__dict__.items():
|
||||
if callable(method) and hasattr(method, "__is_flow_method__"):
|
||||
# Wrap each flow method with persistence
|
||||
if asyncio.iscoroutinefunction(method):
|
||||
@functools.wraps(method)
|
||||
async def class_async_wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
|
||||
method_coro = method(self, *args, **kwargs)
|
||||
if asyncio.iscoroutine(method_coro):
|
||||
result = await method_coro
|
||||
else:
|
||||
result = method_coro
|
||||
_persist_state(self, method.__name__, actual_persistence)
|
||||
return result
|
||||
class_methods[name] = class_async_wrapper
|
||||
else:
|
||||
@functools.wraps(method)
|
||||
def class_sync_wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
|
||||
result = method(self, *args, **kwargs)
|
||||
_persist_state(self, method.__name__, actual_persistence)
|
||||
return result
|
||||
class_methods[name] = class_sync_wrapper
|
||||
|
||||
# Preserve flow-specific attributes
|
||||
for attr in ["__is_start_method__", "__trigger_methods__", "__condition_type__", "__is_router__"]:
|
||||
if hasattr(method, attr):
|
||||
setattr(class_methods[name], attr, getattr(method, attr))
|
||||
setattr(class_methods[name], "__is_flow_method__", True)
|
||||
|
||||
# Update class with wrapped methods
|
||||
for name, method in class_methods.items():
|
||||
setattr(target, name, method)
|
||||
return target
|
||||
else:
|
||||
# Method decoration
|
||||
method = target
|
||||
setattr(method, "__is_flow_method__", True)
|
||||
|
||||
if asyncio.iscoroutinefunction(method):
|
||||
@functools.wraps(method)
|
||||
async def method_async_wrapper(flow_instance: Any, *args: Any, **kwargs: Any) -> T:
|
||||
method_coro = method(flow_instance, *args, **kwargs)
|
||||
if asyncio.iscoroutine(method_coro):
|
||||
result = await method_coro
|
||||
else:
|
||||
result = method_coro
|
||||
_persist_state(flow_instance, method.__name__, actual_persistence)
|
||||
return result
|
||||
for attr in ["__is_start_method__", "__trigger_methods__", "__condition_type__", "__is_router__"]:
|
||||
if hasattr(method, attr):
|
||||
setattr(method_async_wrapper, attr, getattr(method, attr))
|
||||
setattr(method_async_wrapper, "__is_flow_method__", True)
|
||||
return cast(Callable[..., T], method_async_wrapper)
|
||||
else:
|
||||
@functools.wraps(method)
|
||||
def method_sync_wrapper(flow_instance: Any, *args: Any, **kwargs: Any) -> T:
|
||||
result = method(flow_instance, *args, **kwargs)
|
||||
_persist_state(flow_instance, method.__name__, actual_persistence)
|
||||
return result
|
||||
for attr in ["__is_start_method__", "__trigger_methods__", "__condition_type__", "__is_router__"]:
|
||||
if hasattr(method, attr):
|
||||
setattr(method_sync_wrapper, attr, getattr(method, attr))
|
||||
setattr(method_sync_wrapper, "__is_flow_method__", True)
|
||||
return cast(Callable[..., T], method_sync_wrapper)
|
||||
|
||||
return decorator
|
||||
Reference in New Issue
Block a user