From 9f1d7d1aa91bba036a9d69677243f3362541945e Mon Sep 17 00:00:00 2001 From: Lucas Gomide Date: Wed, 6 Aug 2025 10:04:59 -0300 Subject: [PATCH] fix: allow persist Flow state with BaseModel entries (#3276) Co-authored-by: Lorenze Jay <63378463+lorenzejay@users.noreply.github.com> --- src/crewai/flow/persistence/sqlite.py | 2 +- tests/test_flow_persistence.py | 44 +++++++++++++++++++++++++-- 2 files changed, 43 insertions(+), 3 deletions(-) diff --git a/src/crewai/flow/persistence/sqlite.py b/src/crewai/flow/persistence/sqlite.py index 8b2a0f3f2..ee38c614d 100644 --- a/src/crewai/flow/persistence/sqlite.py +++ b/src/crewai/flow/persistence/sqlite.py @@ -81,7 +81,7 @@ class SQLiteFlowPersistence(FlowPersistence): """ # 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 + state_dict = state_data.model_dump() elif isinstance(state_data, dict): state_dict = state_data else: diff --git a/tests/test_flow_persistence.py b/tests/test_flow_persistence.py index cf3bb22f0..667a1f058 100644 --- a/tests/test_flow_persistence.py +++ b/tests/test_flow_persistence.py @@ -1,9 +1,8 @@ """Test flow state persistence functionality.""" import os -from typing import Dict +from typing import Dict, List -import pytest from pydantic import BaseModel from crewai.flow.flow import Flow, FlowState, listen, start @@ -208,3 +207,44 @@ def test_persist_decorator_verbose_logging(tmp_path, caplog): flow = VerboseFlow(persistence=persistence) flow.kickoff() assert "Saving flow state" in caplog.text + + + +def test_persistence_with_base_model(tmp_path): + db_path = os.path.join(tmp_path, "test_flows.db") + persistence = SQLiteFlowPersistence(db_path) + + class Message(BaseModel): + role: str + type: str + content: str + + class State(FlowState): + latest_message: Message | None = None + history: List[Message] = [] + + @persist(persistence) + class BaseModelFlow(Flow[State]): + initial_state = State(latest_message=None, history=[]) + + @start() + def init_step(self): + self.state.latest_message = Message(role="user", type="text", content="Hello, World!") + self.state.history.append(self.state.latest_message) + + flow = BaseModelFlow(persistence=persistence) + flow.kickoff() + + latest_message = flow.state.latest_message + message, = flow.state.history + + assert latest_message is not None + assert latest_message.role == "user" + assert latest_message.type == "text" + assert latest_message.content == "Hello, World!" + + assert len(flow.state.history) == 1 + assert message.role == "user" + assert message.type == "text" + assert message.content == "Hello, World!" + assert isinstance(flow.state, State)