mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-05-03 08:12:39 +00:00
Add exclude option to to_serializable() (#2479)
This commit is contained in:
@@ -1,4 +1,5 @@
|
|||||||
import json
|
import json
|
||||||
|
import uuid
|
||||||
from datetime import date, datetime
|
from datetime import date, datetime
|
||||||
from typing import Any, Dict, List, Union
|
from typing import Any, Dict, List, Union
|
||||||
|
|
||||||
@@ -32,7 +33,7 @@ def export_state(flow: Flow) -> dict[str, Serializable]:
|
|||||||
|
|
||||||
|
|
||||||
def to_serializable(
|
def to_serializable(
|
||||||
obj: Any, max_depth: int = 5, _current_depth: int = 0
|
obj: Any, exclude: set[str] | None = None, max_depth: int = 5, _current_depth: int = 0
|
||||||
) -> Serializable:
|
) -> Serializable:
|
||||||
"""Converts a Python object into a JSON-compatible representation.
|
"""Converts a Python object into a JSON-compatible representation.
|
||||||
|
|
||||||
@@ -42,6 +43,7 @@ def to_serializable(
|
|||||||
|
|
||||||
Args:
|
Args:
|
||||||
obj (Any): Object to transform.
|
obj (Any): Object to transform.
|
||||||
|
exclude (set[str], optional): Set of keys to exclude from the result.
|
||||||
max_depth (int, optional): Maximum recursion depth. Defaults to 5.
|
max_depth (int, optional): Maximum recursion depth. Defaults to 5.
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
@@ -50,21 +52,39 @@ def to_serializable(
|
|||||||
if _current_depth >= max_depth:
|
if _current_depth >= max_depth:
|
||||||
return repr(obj)
|
return repr(obj)
|
||||||
|
|
||||||
|
if exclude is None:
|
||||||
|
exclude = set()
|
||||||
|
|
||||||
if isinstance(obj, (str, int, float, bool, type(None))):
|
if isinstance(obj, (str, int, float, bool, type(None))):
|
||||||
return obj
|
return obj
|
||||||
|
elif isinstance(obj, uuid.UUID):
|
||||||
|
return str(obj)
|
||||||
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=max_depth, _current_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
|
obj=value,
|
||||||
|
exclude=exclude,
|
||||||
|
max_depth=max_depth,
|
||||||
|
_current_depth=_current_depth + 1,
|
||||||
)
|
)
|
||||||
for key, value in obj.items()
|
for key, value in obj.items()
|
||||||
|
if key not in exclude
|
||||||
}
|
}
|
||||||
elif isinstance(obj, BaseModel):
|
elif isinstance(obj, BaseModel):
|
||||||
return to_serializable(obj.model_dump(), max_depth, _current_depth + 1)
|
return to_serializable(
|
||||||
|
obj=obj.model_dump(exclude=exclude),
|
||||||
|
max_depth=max_depth,
|
||||||
|
_current_depth=_current_depth + 1,
|
||||||
|
)
|
||||||
else:
|
else:
|
||||||
return repr(obj)
|
return repr(obj)
|
||||||
|
|
||||||
|
|||||||
@@ -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, to_string
|
from crewai.flow.state_utils import export_state, to_serializable, to_string
|
||||||
|
|
||||||
|
|
||||||
class Address(BaseModel):
|
class Address(BaseModel):
|
||||||
@@ -148,3 +148,23 @@ def test_depth_limit(mock_flow):
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def test_exclude_keys():
|
||||||
|
result = to_serializable({"key1": "value1", "key2": "value2"}, exclude={"key1"})
|
||||||
|
assert result == {"key2": "value2"}
|
||||||
|
|
||||||
|
model = Person(
|
||||||
|
name="John Doe",
|
||||||
|
age=30,
|
||||||
|
address=Address(street="123 Main St", city="Tech City", country="Pythonia"),
|
||||||
|
birthday=date(1994, 1, 1),
|
||||||
|
skills=["Python", "Testing"],
|
||||||
|
)
|
||||||
|
result = to_serializable(model, exclude={"address"})
|
||||||
|
assert result == {
|
||||||
|
"name": "John Doe",
|
||||||
|
"age": 30,
|
||||||
|
"birthday": "1994-01-01",
|
||||||
|
"skills": ["Python", "Testing"],
|
||||||
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user