feat: add serialization to State / move State to utils
#9345
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Amnah199
requested review from
dfokina and
julian-risch
and removed request for
a team
Pull Request Test Coverage Report for Build 14859791142Warning: This coverage report may be inaccurate.This pull request's base commit is no longer the HEAD commit of its target branch. This means it includes changes from outside the original pull request, including, potentially, unrelated coverage changes.
Details
💛 - Coveralls |
|
I’ve kept the current |
Amnah199
commented
May 6, 2025
Comment on lines
+361
to
+424
| def test_state_to_dict(): | ||
| # we test dict, a python type and a haystack dataclass | ||
| state_schema = {"numbers": {"type": int}, "messages": {"type": List[ChatMessage]}, "dict_of_lists": {"type": dict}} | ||
|
|
||
| data = { | ||
| "numbers": 1, | ||
| "messages": [ChatMessage.from_user(text="Hello, world!")], | ||
| "dict_of_lists": {"numbers": [1, 2, 3]}, | ||
| } | ||
| state = State(state_schema, data) | ||
| state_dict = state.to_dict() | ||
| assert state_dict["schema"] == { | ||
| "numbers": {"type": "int", "handler": "haystack.utils.state_utils.replace_values"}, | ||
| "messages": { | ||
| "type": "typing.List[haystack.dataclasses.chat_message.ChatMessage]", | ||
| "handler": "haystack.utils.state_utils.merge_lists", | ||
| }, | ||
| "dict_of_lists": {"type": "dict", "handler": "haystack.utils.state_utils.replace_values"}, | ||
| } | ||
| assert state_dict["data"] == { | ||
| "numbers": 1, | ||
| "messages": [ | ||
| {"role": "user", "meta": {}, "name": None, "content": [{"text": "Hello, world!"}], "_type": "ChatMessage"} | ||
| ], | ||
| "dict_of_lists": {"numbers": [1, 2, 3]}, | ||
| } | ||
|
|
||
|
|
||
| def test_state_from_dict(): | ||
| state_dict = { | ||
| "schema": { | ||
| "numbers": {"type": "int", "handler": "haystack.utils.state_utils.replace_values"}, | ||
| "messages": { | ||
| "type": "typing.List[haystack.dataclasses.chat_message.ChatMessage]", | ||
| "handler": "haystack.utils.state_utils.merge_lists", | ||
| }, | ||
| "dict_of_lists": {"type": "dict", "handler": "haystack.utils.state_utils.replace_values"}, | ||
| }, | ||
| "data": { | ||
| "numbers": 1, | ||
| "messages": [ | ||
| { | ||
| "role": "user", | ||
| "meta": {}, | ||
| "name": None, | ||
| "content": [{"text": "Hello, world!"}], | ||
| "_type": "ChatMessage", | ||
| } | ||
| ], | ||
| "dict_of_lists": {"numbers": [1, 2, 3]}, | ||
| }, | ||
| } | ||
| state = State.from_dict(state_dict) | ||
| # Check types are correctly converted | ||
| assert state.schema["numbers"]["type"] == int | ||
| assert state.schema["dict_of_lists"]["type"] == dict | ||
| # Check handlers are functions, not comparing exact functions as they might be different references | ||
| assert callable(state.schema["numbers"]["handler"]) | ||
| assert callable(state.schema["messages"]["handler"]) | ||
| assert callable(state.schema["dict_of_lists"]["handler"]) | ||
| # Check data is correct | ||
| assert state.data["numbers"] == 1 | ||
| assert state.data["messages"] == [ChatMessage.from_user(text="Hello, world!")] | ||
| assert state.data["dict_of_lists"] == {"numbers": [1, 2, 3]} |
There was a problem hiding this comment.
For ease in review, these are the two new tests for serialization. Rest are just moved from old test file.
sjrl
reviewed
May 9, 2025
Comment on lines
+126
to
+133
| _type_deserializers = { | ||
| "Answer": Answer.from_dict, | ||
| "ChatMessage": ChatMessage.from_dict, | ||
| "Document": Document.from_dict, | ||
| "ExtractedAnswer": ExtractedAnswer.from_dict, | ||
| "GeneratedAnswer": GeneratedAnswer.from_dict, | ||
| "SparseEmbedding": SparseEmbedding.from_dict, | ||
| } |
There was a problem hiding this comment.
Instead of using a predefined set of deserializers couldn't we follow a similar methodology like we do for component_from_dict where we import the _type field, check if the imported class has a from_dict attribute and if so use that? That way we wouldn't need to create a list of hard-coded deserializers.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Related Issues
State#9286Proposed Changes:
serialize_valueanddeserialize_valueutility methods in theutilsmodule. These methods encapsulate logic that is also used in the breakpoints de/serialization logic andtracing.utils.coerce_tag_value(except the json load). Once the breakpoints feature is merged in haystack, it can reuse these centralized utility functions.Stateclass to theutilsmodule, as it is not actually a data class. A deprecation warning is added to the existing State class in the dataclasses module to guide users toward the updated implementation.How did you test it?
Moved the existing tests to
test_utils_state.pyand added two new tests for serialization and deserialization.