Add exclude option to to_serializable() (#2479)

This commit is contained in:
Vini Brasil
2025-03-26 11:35:12 -03:00
committed by GitHub
parent ebb585e494
commit a25a27c3d3
2 changed files with 45 additions and 5 deletions

View File

@@ -1,4 +1,5 @@
import json
import uuid
from datetime import date, datetime
from typing import Any, Dict, List, Union
@@ -32,7 +33,7 @@ def export_state(flow: Flow) -> dict[str, 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:
"""Converts a Python object into a JSON-compatible representation.
@@ -42,6 +43,7 @@ def to_serializable(
Args:
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.
Returns:
@@ -50,21 +52,39 @@ def to_serializable(
if _current_depth >= max_depth:
return repr(obj)
if exclude is None:
exclude = set()
if isinstance(obj, (str, int, float, bool, type(None))):
return obj
elif isinstance(obj, uuid.UUID):
return str(obj)
elif isinstance(obj, (date, datetime)):
return obj.isoformat()
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):
return {
_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()
if key not in exclude
}
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:
return repr(obj)

View File

@@ -6,7 +6,7 @@ import pytest
from pydantic import BaseModel
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):
@@ -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"],
}