diff --git a/haystack/components/agents/state/state.py b/haystack/components/agents/state/state.py index 7db24cdcd47..2764508dd54 100644 --- a/haystack/components/agents/state/state.py +++ b/haystack/components/agents/state/state.py @@ -189,15 +189,23 @@ def has(self, key: str) -> bool: """ return key in self._data - def to_dict(self) -> dict[str, Any]: + def to_dict(self, skip_keys: list[str] | None = None) -> dict[str, Any]: """ Convert the State object to a dictionary. + + :param skip_keys: List of keys to skip during serialization + :returns: Dictionary representation of the State object """ + skip_keys = skip_keys or [] + schema = {key: definition for key, definition in self.schema.items() if key not in skip_keys} + data = {key: value for key, value in self._data.items() if key not in skip_keys} + serialized = {} - serialized["schema"] = _schema_to_dict(self.schema) + serialized["schema"] = _schema_to_dict(schema=schema) # Field-level fallback so a single non-serializable value omits only that field instead of # failing the whole State serialization. - serialized["data"] = _serialize_with_field_fallback(self._data, description="the agent's State data") + serialized["data"] = _serialize_with_field_fallback(payload=data, description="the agent's State data") + return serialized @classmethod diff --git a/releasenotes/notes/state-todict-skipkeys-ec8c5dc61354c015.yaml b/releasenotes/notes/state-todict-skipkeys-ec8c5dc61354c015.yaml new file mode 100644 index 00000000000..86adc8ca4a1 --- /dev/null +++ b/releasenotes/notes/state-todict-skipkeys-ec8c5dc61354c015.yaml @@ -0,0 +1,4 @@ +--- +enhancements: + - | + ``State.to_dict`` now accepts a ``skip_keys`` parameter to exclude specific keys from the output. diff --git a/test/components/agents/test_state_class.py b/test/components/agents/test_state_class.py index cd9cb686e9e..79f0bbf7c26 100644 --- a/test/components/agents/test_state_class.py +++ b/test/components/agents/test_state_class.py @@ -424,6 +424,30 @@ def test_state_to_dict(self): }, } + def test_state_to_dict_skip_keys(self): + state_schema = {"numbers": {"type": int}, "messages": {"type": list[ChatMessage]}} + + data = {"numbers": 1, "messages": [ChatMessage.from_user(text="Hello, world!")]} + state = State(state_schema, data) + state_dict = state.to_dict(skip_keys=["numbers"]) + assert state_dict["schema"] == { + "messages": { + "type": "list[haystack.dataclasses.chat_message.ChatMessage]", + "handler": "haystack.components.agents.state.state_utils.merge_lists", + } + } + assert state_dict["data"] == { + "serialization_schema": { + "type": "object", + "properties": { + "messages": {"type": "array", "items": {"type": "haystack.dataclasses.chat_message.ChatMessage"}} + }, + }, + "serialized_data": { + "messages": [{"role": "user", "meta": {}, "name": None, "content": [{"text": "Hello, world!"}]}] + }, + } + def test_state_from_dict(self): state_dict = { "schema": {