andre/crewAI
1
0
Fork 0
mirror of https://github.com/crewAIInc/crewAI.git synced 2026-01-09 16:18:30 +00:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

Add @persist decorator with FlowPersistence interface (#1892)

Browse Source 
* 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]
2025-01-16 10:23:46 -03:00
committed by GitHub
parent 3dc442801f
commit 294f2cc3a9
10 changed files with 1061 additions and 138 deletions
Download Patch File Download Diff File Expand all files Collapse all files

384
src/crewai/flow/flow.py
View File

@@ -1,5 +1,6 @@
import asyncio
import inspect
import uuid
from typing import (
Any,
Callable,
@@ -12,6 +13,7 @@ from typing import (
TypeVar,
Union,
cast,
overload,
)
from uuid import uuid4
@@ -25,6 +27,8 @@ from crewai.flow.flow_events import (
MethodExecutionStartedEvent,
)
from crewai.flow.flow_visualizer import plot_flow
from crewai.flow.persistence import FlowPersistence
from crewai.flow.persistence.base import FlowPersistence
from crewai.flow.utils import get_possible_return_constants
from crewai.telemetry import Telemetry
@@ -33,7 +37,46 @@ class FlowState(BaseModel):
"""Base model for all flow states, ensuring each state has a unique ID."""
id: str = Field(default_factory=lambda: str(uuid4()), description="Unique identifier for the flow state")
T = TypeVar("T", bound=Union[FlowState, Dict[str, Any]])
# Type variables with explicit bounds
T = TypeVar("T", bound=Union[Dict[str, Any], BaseModel]) # Generic flow state type parameter
StateT = TypeVar("StateT", bound=Union[Dict[str, Any], BaseModel]) # State validation type parameter
def ensure_state_type(state: Any, expected_type: Type[StateT]) -> StateT:
"""Ensure state matches expected type with proper validation.
Args:
state: State instance to validate
expected_type: Expected type for the state
Returns:
Validated state instance
Raises:
TypeError: If state doesn't match expected type
ValueError: If state validation fails
"""
"""Ensure state matches expected type with proper validation.
Args:
state: State instance to validate
expected_type: Expected type for the state
Returns:
Validated state instance
Raises:
TypeError: If state doesn't match expected type
ValueError: If state validation fails
"""
if expected_type == dict:
if not isinstance(state, dict):
raise TypeError(f"Expected dict, got {type(state).__name__}")
return cast(StateT, state)
if isinstance(expected_type, type) and issubclass(expected_type, BaseModel):
if not isinstance(state, expected_type):
raise TypeError(f"Expected {expected_type.__name__}, got {type(state).__name__}")
return cast(StateT, state)
raise TypeError(f"Invalid expected_type: {expected_type}")
def start(condition: Optional[Union[str, dict, Callable]] = None) -> Callable:
@@ -326,21 +369,27 @@ class FlowMeta(type):
routers = set()
for attr_name, attr_value in dct.items():
if hasattr(attr_value, "__is_start_method__"):
start_methods.append(attr_name)
# Check for any flow-related attributes
if (hasattr(attr_value, "__is_flow_method__") or
hasattr(attr_value, "__is_start_method__") or
hasattr(attr_value, "__trigger_methods__") or
hasattr(attr_value, "__is_router__")):
# Register start methods
if hasattr(attr_value, "__is_start_method__"):
start_methods.append(attr_name)
# Register listeners and routers
if hasattr(attr_value, "__trigger_methods__"):
methods = attr_value.__trigger_methods__
condition_type = getattr(attr_value, "__condition_type__", "OR")
listeners[attr_name] = (condition_type, methods)
elif hasattr(attr_value, "__trigger_methods__"):
methods = attr_value.__trigger_methods__
condition_type = getattr(attr_value, "__condition_type__", "OR")
listeners[attr_name] = (condition_type, methods)
if hasattr(attr_value, "__is_router__") and attr_value.__is_router__:
routers.add(attr_name)
possible_returns = get_possible_return_constants(attr_value)
if possible_returns:
router_paths[attr_name] = possible_returns
if hasattr(attr_value, "__is_router__") and attr_value.__is_router__:
routers.add(attr_name)
possible_returns = get_possible_return_constants(attr_value)
if possible_returns:
router_paths[attr_name] = possible_returns
setattr(cls, "_start_methods", start_methods)
setattr(cls, "_listeners", listeners)
@@ -351,6 +400,9 @@ class FlowMeta(type):
class Flow(Generic[T], metaclass=FlowMeta):
"""Base class for all flows.
Type parameter T must be either Dict[str, Any] or a subclass of BaseModel."""
_telemetry = Telemetry()
_start_methods: List[str] = []
@@ -367,53 +419,220 @@ class Flow(Generic[T], metaclass=FlowMeta):
_FlowGeneric.__name__ = f"{cls.__name__}[{item.__name__}]"
return _FlowGeneric
def __init__(self) -> None:
def __init__(
self,
persistence: Optional[FlowPersistence] = None,
restore_uuid: Optional[str] = None,
**kwargs: Any,
) -> None:
"""Initialize a new Flow instance.
Args:
persistence: Optional persistence backend for storing flow states
restore_uuid: Optional UUID to restore state from persistence
**kwargs: Additional state values to initialize or override
"""
# Initialize basic instance attributes
self._methods: Dict[str, Callable] = {}
self._state: T = self._create_initial_state()
self._method_execution_counts: Dict[str, int] = {}
self._pending_and_listeners: Dict[str, Set[str]] = {}
self._method_outputs: List[Any] = [] # List to store all method outputs
self._persistence: Optional[FlowPersistence] = persistence
# Validate state model before initialization
if isinstance(self.initial_state, type):
if issubclass(self.initial_state, BaseModel) and not issubclass(self.initial_state, FlowState):
# Check if model has id field
model_fields = getattr(self.initial_state, "model_fields", None)
if not model_fields or "id" not in model_fields:
raise ValueError("Flow state model must have an 'id' field")
# Handle persistence and potential ID conflicts
stored_state = None
if self._persistence is not None:
if restore_uuid and kwargs and "id" in kwargs and restore_uuid != kwargs["id"]:
raise ValueError(
f"Conflicting IDs provided: restore_uuid='{restore_uuid}' "
f"vs kwargs['id']='{kwargs['id']}'. Use only one ID for restoration."
)
# Attempt to load state, prioritizing restore_uuid
if restore_uuid:
stored_state = self._persistence.load_state(restore_uuid)
if not stored_state:
raise ValueError(f"No state found for restore_uuid='{restore_uuid}'")
elif kwargs and "id" in kwargs:
stored_state = self._persistence.load_state(kwargs["id"])
if not stored_state:
# For kwargs["id"], we allow creating new state if not found
self._state = self._create_initial_state()
if kwargs:
self._initialize_state(kwargs)
return
# Initialize state based on persistence and kwargs
if stored_state:
# Create initial state and restore from persistence
self._state = self._create_initial_state()
self._restore_state(stored_state)
# Apply any additional kwargs to override specific fields
if kwargs:
filtered_kwargs = {k: v for k, v in kwargs.items() if k != "id"}
if filtered_kwargs:
self._initialize_state(filtered_kwargs)
else:
# No stored state, create new state with initial values
self._state = self._create_initial_state()
# Apply any additional kwargs
if kwargs:
self._initialize_state(kwargs)
self._telemetry.flow_creation_span(self.__class__.__name__)
# Register all flow-related methods
for method_name in dir(self):
if callable(getattr(self, method_name)) and not method_name.startswith(
"__"
):
self._methods[method_name] = getattr(self, method_name)
if not method_name.startswith("_"):
method = getattr(self, method_name)
# Check for any flow-related attributes
if (hasattr(method, "__is_flow_method__") or
hasattr(method, "__is_start_method__") or
hasattr(method, "__trigger_methods__") or
hasattr(method, "__is_router__")):
# Ensure method is bound to this instance
if not hasattr(method, "__self__"):
method = method.__get__(self, self.__class__)
self._methods[method_name] = method
def _create_initial_state(self) -> T:
"""Create and initialize flow state with UUID and default values.
Returns:
New state instance with UUID and default values initialized
Raises:
ValueError: If structured state model lacks 'id' field
TypeError: If state is neither BaseModel nor dictionary
"""
# Handle case where initial_state is None but we have a type parameter
if self.initial_state is None and hasattr(self, "_initial_state_T"):
state_type = getattr(self, "_initial_state_T")
if isinstance(state_type, type):
if issubclass(state_type, FlowState):
return state_type() # type: ignore
# Create instance without id, then set it
instance = state_type()
if not hasattr(instance, 'id'):
setattr(instance, 'id', str(uuid4()))
return cast(T, instance)
elif issubclass(state_type, BaseModel):
# Create a new type that includes the ID field
class StateWithId(state_type, FlowState): # type: ignore
pass
return StateWithId() # type: ignore
instance = StateWithId()
if not hasattr(instance, 'id'):
setattr(instance, 'id', str(uuid4()))
return cast(T, instance)
elif state_type == dict:
return cast(T, {"id": str(uuid4())}) # Minimal dict state
# Handle case where no initial state is provided
if self.initial_state is None:
return cast(T, {"id": str(uuid4())})
# Handle case where initial_state is a type (class)
if isinstance(self.initial_state, type):
if issubclass(self.initial_state, FlowState):
return cast(T, self.initial_state()) # Uses model defaults
elif issubclass(self.initial_state, BaseModel):
# Validate that the model has an id field
model_fields = getattr(self.initial_state, "model_fields", None)
if not model_fields or "id" not in model_fields:
raise ValueError("Flow state model must have an 'id' field")
return cast(T, self.initial_state()) # Uses model defaults
elif self.initial_state == dict:
return cast(T, {"id": str(uuid4())})
# Handle dictionary instance case
if isinstance(self.initial_state, dict):
new_state = dict(self.initial_state) # Copy to avoid mutations
if "id" not in new_state:
new_state["id"] = str(uuid4())
return cast(T, new_state)
# Handle BaseModel instance case
if isinstance(self.initial_state, BaseModel):
model = cast(BaseModel, self.initial_state)
if not hasattr(model, "id"):
raise ValueError("Flow state model must have an 'id' field")
# Create new instance with same values to avoid mutations
if hasattr(model, "model_dump"):
# Pydantic v2
state_dict = model.model_dump()
elif hasattr(model, "dict"):
# Pydantic v1
state_dict = model.dict()
else:
# Fallback for other BaseModel implementations
state_dict = {
k: v for k, v in model.__dict__.items()
if not k.startswith("_")
}
# Create new instance of the same class
model_class = type(model)
return cast(T, model_class(**state_dict))
raise TypeError(
f"Initial state must be dict or BaseModel, got {type(self.initial_state)}"
)
# Handle case where initial_state is None but we have a type parameter
if self.initial_state is None and hasattr(self, "_initial_state_T"):
state_type = getattr(self, "_initial_state_T")
if isinstance(state_type, type):
if issubclass(state_type, FlowState):
return cast(T, state_type())
elif issubclass(state_type, BaseModel):
# Create a new type that includes the ID field
class StateWithId(state_type, FlowState): # type: ignore
pass
return cast(T, StateWithId())
elif state_type == dict:
return cast(T, {"id": str(uuid4())})
# Handle case where no initial state is provided
if self.initial_state is None:
return {"id": str(uuid4())} # type: ignore
return cast(T, {"id": str(uuid4())})
# Handle case where initial_state is a type (class)
if isinstance(self.initial_state, type):
if issubclass(self.initial_state, FlowState):
return self.initial_state() # type: ignore
return cast(T, self.initial_state())
elif issubclass(self.initial_state, BaseModel):
# Create a new type that includes the ID field
class StateWithId(self.initial_state, FlowState): # type: ignore
pass
return StateWithId() # type: ignore
# Validate that the model has an id field
model_fields = getattr(self.initial_state, "model_fields", None)
if not model_fields or "id" not in model_fields:
raise ValueError("Flow state model must have an 'id' field")
return cast(T, self.initial_state())
elif self.initial_state == dict:
return cast(T, {"id": str(uuid4())})
# Handle dictionary case
if isinstance(self.initial_state, dict) and "id" not in self.initial_state:
self.initial_state["id"] = str(uuid4())
# Handle dictionary instance case
if isinstance(self.initial_state, dict):
if "id" not in self.initial_state:
self.initial_state["id"] = str(uuid4())
return cast(T, dict(self.initial_state)) # Create new dict to avoid mutations
return self.initial_state # type: ignore
# Handle BaseModel instance case
if isinstance(self.initial_state, BaseModel):
if not hasattr(self.initial_state, "id"):
raise ValueError("Flow state model must have an 'id' field")
return cast(T, self.initial_state)
raise TypeError(
f"Initial state must be dict or BaseModel, got {type(self.initial_state)}"
)
@property
def state(self) -> T:
@@ -425,50 +644,95 @@ class Flow(Generic[T], metaclass=FlowMeta):
return self._method_outputs
def _initialize_state(self, inputs: Dict[str, Any]) -> None:
"""Initialize or update flow state with new inputs.
Args:
inputs: Dictionary of state values to set/update
Raises:
ValueError: If validation fails for structured state
TypeError: If state is neither BaseModel nor dictionary
"""
if isinstance(self._state, dict):
# Preserve the ID when updating unstructured state
# For dict states, preserve existing fields unless overridden
current_id = self._state.get("id")
self._state.update(inputs)
# Only update specified fields
for k, v in inputs.items():
self._state[k] = v
# Ensure ID is preserved or generated
if current_id:
self._state["id"] = current_id
elif "id" not in self._state:
self._state["id"] = str(uuid4())
elif isinstance(self._state, BaseModel):
# Structured state
# For BaseModel states, preserve existing fields unless overridden
try:
def create_model_with_extra_forbid(
base_model: Type[BaseModel],
) -> Type[BaseModel]:
class ModelWithExtraForbid(base_model): # type: ignore
model_config = base_model.model_config.copy()
model_config["extra"] = "forbid"
return ModelWithExtraForbid
# Get current state as dict, preserving the ID if it exists
state_model = cast(BaseModel, self._state)
current_state = (
state_model.model_dump()
if hasattr(state_model, "model_dump")
else state_model.dict()
if hasattr(state_model, "dict")
else {
k: v
for k, v in state_model.__dict__.items()
model = cast(BaseModel, self._state)
# Get current state as dict
if hasattr(model, "model_dump"):
current_state = model.model_dump()
elif hasattr(model, "dict"):
current_state = model.dict()
else:
current_state = {
k: v for k, v in model.__dict__.items()
if not k.startswith("_")
}
)
ModelWithExtraForbid = create_model_with_extra_forbid(
self._state.__class__
)
self._state = cast(
T, ModelWithExtraForbid(**{**current_state, **inputs})
)
# Create new state with preserved fields and updates
new_state = {**current_state, **inputs}
# Create new instance with merged state
model_class = type(model)
if hasattr(model_class, "model_validate"):
# Pydantic v2
self._state = cast(T, model_class.model_validate(new_state))
elif hasattr(model_class, "parse_obj"):
# Pydantic v1
self._state = cast(T, model_class.parse_obj(new_state))
else:
# Fallback for other BaseModel implementations
self._state = cast(T, model_class(**new_state))
except ValidationError as e:
raise ValueError(f"Invalid inputs for structured state: {e}") from e
else:
raise TypeError("State must be a BaseModel instance or a dictionary.")
def _restore_state(self, stored_state: Dict[str, Any]) -> None:
"""Restore flow state from persistence.
Args:
stored_state: Previously stored state to restore
Raises:
ValueError: If validation fails for structured state
TypeError: If state is neither BaseModel nor dictionary
"""
# When restoring from persistence, use the stored ID
stored_id = stored_state.get("id")
if not stored_id:
raise ValueError("Stored state must have an 'id' field")
if isinstance(self._state, dict):
# For dict states, update all fields from stored state
self._state.clear()
self._state.update(stored_state)
elif isinstance(self._state, BaseModel):
# For BaseModel states, create new instance with stored values
model = cast(BaseModel, self._state)
if hasattr(model, "model_validate"):
# Pydantic v2
self._state = cast(T, type(model).model_validate(stored_state))
elif hasattr(model, "parse_obj"):
# Pydantic v1
self._state = cast(T, type(model).parse_obj(stored_state))
else:
# Fallback for other BaseModel implementations
self._state = cast(T, type(model)(**stored_state))
else:
raise TypeError(
f"State must be dict or BaseModel, got {type(self._state)}"
)
def kickoff(self, inputs: Optional[Dict[str, Any]] = None) -> Any:
self.event_emitter.send(