mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-09-25 04:10:04 +00:00
fix(flows): align MongoDB state serialization
This commit is contained in:
@@ -3,6 +3,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from datetime import date, datetime
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
@@ -15,6 +16,20 @@ if TYPE_CHECKING:
|
||||
_persistence_registry: dict[str, type[FlowPersistence]] = {}
|
||||
|
||||
|
||||
def _json_default(obj: Any) -> Any:
|
||||
"""Serialize flow-state values that are not JSON-native."""
|
||||
if isinstance(obj, BaseModel):
|
||||
try:
|
||||
return obj.model_dump(mode="json")
|
||||
except Exception:
|
||||
return obj.model_dump(mode="python")
|
||||
if isinstance(obj, (set, tuple)):
|
||||
return list(obj)
|
||||
if isinstance(obj, (date, datetime)):
|
||||
return obj.isoformat()
|
||||
return str(obj)
|
||||
|
||||
|
||||
class FlowPersistence(BaseModel, ABC):
|
||||
"""Abstract base class for flow state persistence.
|
||||
|
||||
|
||||
@@ -31,7 +31,7 @@ from typing import TYPE_CHECKING, Any
|
||||
|
||||
from pydantic import BaseModel, Field, PrivateAttr
|
||||
|
||||
from crewai.flow.persistence.base import FlowPersistence
|
||||
from crewai.flow.persistence.base import FlowPersistence, _json_default
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -190,7 +190,7 @@ class MongoDbFlowPersistence(FlowPersistence):
|
||||
"flow_uuid": flow_uuid,
|
||||
"method_name": method_name,
|
||||
"timestamp": datetime.now(timezone.utc).isoformat(),
|
||||
"state_json": json.dumps(state_dict),
|
||||
"state_json": json.dumps(state_dict, default=_json_default),
|
||||
"seq": self._next_sequence(self.states_collection),
|
||||
}
|
||||
)
|
||||
@@ -221,8 +221,8 @@ class MongoDbFlowPersistence(FlowPersistence):
|
||||
{"flow_uuid": flow_uuid},
|
||||
{
|
||||
"flow_uuid": flow_uuid,
|
||||
"context_json": json.dumps(context.to_dict()),
|
||||
"state_json": json.dumps(state_dict),
|
||||
"context_json": json.dumps(context.to_dict(), default=_json_default),
|
||||
"state_json": json.dumps(state_dict, default=_json_default),
|
||||
"created_at": datetime.now(timezone.utc).isoformat(),
|
||||
},
|
||||
upsert=True,
|
||||
|
||||
@@ -3,7 +3,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from contextlib import closing
|
||||
from datetime import date, datetime, timezone
|
||||
from datetime import datetime, timezone
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
@@ -15,27 +15,13 @@ from crewai_core.paths import db_storage_path
|
||||
from pydantic import BaseModel, Field, PrivateAttr, model_validator
|
||||
from typing_extensions import Self
|
||||
|
||||
from crewai.flow.persistence.base import FlowPersistence
|
||||
from crewai.flow.persistence.base import FlowPersistence, _json_default
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from crewai.flow.async_feedback.types import PendingFeedbackContext
|
||||
|
||||
|
||||
def _json_default(obj: Any) -> Any:
|
||||
"""Fallback serializer for non-primitive types in JSON dumps."""
|
||||
if isinstance(obj, BaseModel):
|
||||
try:
|
||||
return obj.model_dump(mode="json")
|
||||
except Exception:
|
||||
return obj.model_dump(mode="python")
|
||||
if isinstance(obj, (set, tuple)):
|
||||
return list(obj)
|
||||
if isinstance(obj, (date, datetime)):
|
||||
return obj.isoformat()
|
||||
return str(obj)
|
||||
|
||||
|
||||
class SQLiteFlowPersistence(FlowPersistence):
|
||||
"""SQLite-based implementation of flow state persistence.
|
||||
|
||||
|
||||
@@ -176,6 +176,36 @@ def test_basemodel_state_serialized_as_json(
|
||||
assert loaded["when"].startswith("2026-01-02T03:04:05")
|
||||
|
||||
|
||||
def test_dict_state_serializes_non_json_values(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
_patch_client(monkeypatch)
|
||||
|
||||
class NestedState(BaseModel):
|
||||
when: datetime
|
||||
|
||||
when = datetime(2026, 1, 2, 3, 4, 5, tzinfo=timezone.utc)
|
||||
persistence = MongoDbFlowPersistence(CONN)
|
||||
|
||||
persistence.save_state(
|
||||
"flow-1",
|
||||
"s",
|
||||
{
|
||||
"when": when,
|
||||
"tags": {"a", "b"},
|
||||
"items": (1, 2),
|
||||
"nested": NestedState(when=when),
|
||||
},
|
||||
)
|
||||
|
||||
assert persistence.load_state("flow-1") == {
|
||||
"when": "2026-01-02T03:04:05+00:00",
|
||||
"tags": ["a", "b"],
|
||||
"items": [1, 2],
|
||||
"nested": {"when": "2026-01-02T03:04:05Z"},
|
||||
}
|
||||
|
||||
|
||||
def test_missing_connection_string_raises(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
|
||||
Reference in New Issue
Block a user