mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-10 08:38:30 +00:00
Implement flow.state_utils.to_string method and improve types (#2161)
This commit is contained in:
@@ -1,12 +1,18 @@
|
|||||||
|
import json
|
||||||
from datetime import date, datetime
|
from datetime import date, datetime
|
||||||
from typing import Any
|
from typing import Any, Dict, List, Union
|
||||||
|
|
||||||
from pydantic import BaseModel
|
from pydantic import BaseModel
|
||||||
|
|
||||||
from crewai.flow import Flow
|
from crewai.flow import Flow
|
||||||
|
|
||||||
|
SerializablePrimitive = Union[str, int, float, bool, None]
|
||||||
|
Serializable = Union[
|
||||||
|
SerializablePrimitive, List["Serializable"], Dict[str, "Serializable"]
|
||||||
|
]
|
||||||
|
|
||||||
def export_state(flow: Flow) -> dict[str, Any]:
|
|
||||||
|
def export_state(flow: Flow) -> dict[str, Serializable]:
|
||||||
"""Exports the Flow's internal state as JSON-compatible data structures.
|
"""Exports the Flow's internal state as JSON-compatible data structures.
|
||||||
|
|
||||||
Performs a one-way transformation of a Flow's state into basic Python types
|
Performs a one-way transformation of a Flow's state into basic Python types
|
||||||
@@ -20,10 +26,27 @@ def export_state(flow: Flow) -> dict[str, Any]:
|
|||||||
dict[str, Any]: The transformed state using JSON-compatible Python
|
dict[str, Any]: The transformed state using JSON-compatible Python
|
||||||
types.
|
types.
|
||||||
"""
|
"""
|
||||||
return _to_serializable(flow._state)
|
result = to_serializable(flow._state)
|
||||||
|
assert isinstance(result, dict)
|
||||||
|
return result
|
||||||
|
|
||||||
|
|
||||||
def _to_serializable(obj: Any, max_depth: int = 5, _current_depth: int = 0) -> Any:
|
def to_serializable(
|
||||||
|
obj: Any, max_depth: int = 5, _current_depth: int = 0
|
||||||
|
) -> Serializable:
|
||||||
|
"""Converts a Python object into a JSON-compatible representation.
|
||||||
|
|
||||||
|
Supports primitives, datetime objects, collections, dictionaries, and
|
||||||
|
Pydantic models. Recursion depth is limited to prevent infinite nesting.
|
||||||
|
Non-convertible objects default to their string representations.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
obj (Any): Object to transform.
|
||||||
|
max_depth (int, optional): Maximum recursion depth. Defaults to 5.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Serializable: A JSON-compatible structure.
|
||||||
|
"""
|
||||||
if _current_depth >= max_depth:
|
if _current_depth >= max_depth:
|
||||||
return repr(obj)
|
return repr(obj)
|
||||||
|
|
||||||
@@ -32,16 +55,16 @@ def _to_serializable(obj: Any, max_depth: int = 5, _current_depth: int = 0) -> A
|
|||||||
elif isinstance(obj, (date, datetime)):
|
elif isinstance(obj, (date, datetime)):
|
||||||
return obj.isoformat()
|
return obj.isoformat()
|
||||||
elif isinstance(obj, (list, tuple, set)):
|
elif isinstance(obj, (list, tuple, set)):
|
||||||
return [_to_serializable(item, max_depth, _current_depth + 1) for item in obj]
|
return [to_serializable(item, max_depth, _current_depth + 1) for item in obj]
|
||||||
elif isinstance(obj, dict):
|
elif isinstance(obj, dict):
|
||||||
return {
|
return {
|
||||||
_to_serializable_key(key): _to_serializable(
|
_to_serializable_key(key): to_serializable(
|
||||||
value, max_depth, _current_depth + 1
|
value, max_depth, _current_depth + 1
|
||||||
)
|
)
|
||||||
for key, value in obj.items()
|
for key, value in obj.items()
|
||||||
}
|
}
|
||||||
elif isinstance(obj, BaseModel):
|
elif isinstance(obj, BaseModel):
|
||||||
return _to_serializable(obj.model_dump(), max_depth, _current_depth + 1)
|
return to_serializable(obj.model_dump(), max_depth, _current_depth + 1)
|
||||||
else:
|
else:
|
||||||
return repr(obj)
|
return repr(obj)
|
||||||
|
|
||||||
@@ -50,3 +73,19 @@ def _to_serializable_key(key: Any) -> str:
|
|||||||
if isinstance(key, (str, int)):
|
if isinstance(key, (str, int)):
|
||||||
return str(key)
|
return str(key)
|
||||||
return f"key_{id(key)}_{repr(key)}"
|
return f"key_{id(key)}_{repr(key)}"
|
||||||
|
|
||||||
|
|
||||||
|
def to_string(obj: Any) -> str | None:
|
||||||
|
"""Serializes an object into a JSON string.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
obj (Any): Object to serialize.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
str | None: A JSON-formatted string or `None` if empty.
|
||||||
|
"""
|
||||||
|
serializable = to_serializable(obj)
|
||||||
|
if serializable is None:
|
||||||
|
return None
|
||||||
|
else:
|
||||||
|
return json.dumps(serializable)
|
||||||
|
|||||||
@@ -6,7 +6,7 @@ import pytest
|
|||||||
from pydantic import BaseModel
|
from pydantic import BaseModel
|
||||||
|
|
||||||
from crewai.flow import Flow
|
from crewai.flow import Flow
|
||||||
from crewai.flow.state_utils import export_state
|
from crewai.flow.state_utils import export_state, to_string
|
||||||
|
|
||||||
|
|
||||||
class Address(BaseModel):
|
class Address(BaseModel):
|
||||||
@@ -119,16 +119,10 @@ def test_pydantic_model_serialization(mock_flow):
|
|||||||
)
|
)
|
||||||
|
|
||||||
result = export_state(flow)
|
result = export_state(flow)
|
||||||
|
assert (
|
||||||
assert result["single_model"]["street"] == "123 Main St"
|
to_string(result)
|
||||||
|
== '{"single_model": {"street": "123 Main St", "city": "Tech City", "country": "Pythonia"}, "nested_model": {"name": "John Doe", "age": 30, "address": {"street": "123 Main St", "city": "Tech City", "country": "Pythonia"}, "birthday": "1994-01-01", "skills": ["Python", "Testing"]}, "model_list": [{"street": "123 Main St", "city": "Tech City", "country": "Pythonia"}, {"street": "123 Main St", "city": "Tech City", "country": "Pythonia"}], "model_dict": {"home": {"street": "123 Main St", "city": "Tech City", "country": "Pythonia"}}}'
|
||||||
assert result["nested_model"]["name"] == "John Doe"
|
)
|
||||||
assert result["nested_model"]["address"]["city"] == "Tech City"
|
|
||||||
assert result["nested_model"]["birthday"] == "1994-01-01"
|
|
||||||
|
|
||||||
assert len(result["model_list"]) == 2
|
|
||||||
assert all(m["street"] == "123 Main St" for m in result["model_list"])
|
|
||||||
assert result["model_dict"]["home"]["city"] == "Tech City"
|
|
||||||
|
|
||||||
|
|
||||||
def test_depth_limit(mock_flow):
|
def test_depth_limit(mock_flow):
|
||||||
|
|||||||
Reference in New Issue
Block a user