Coverage for haystack/components/agents/state/state.py: 97%
77 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-21 13:53 +0000
1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
2#
3# SPDX-License-Identifier: Apache-2.0
5from collections.abc import Callable
6from copy import deepcopy
7from typing import Any, get_args
9from haystack.dataclasses import ChatMessage
10from haystack.utils import _deserialize_value_with_schema
11from haystack.utils.base_serialization import _serialize_with_field_fallback
12from haystack.utils.callable_serialization import deserialize_callable, serialize_callable
13from haystack.utils.type_serialization import deserialize_type, serialize_type
15from .state_utils import _is_list_type, _is_valid_type, merge_lists, replace_values
18def _schema_to_dict(schema: dict[str, Any]) -> dict[str, Any]:
19 """
20 Convert a schema dictionary to a serializable format.
22 Converts each parameter's type and optional handler function into a serializable
23 format using type and callable serialization utilities.
25 :param schema: Dictionary mapping parameter names to their type and handler configs
26 :returns: Dictionary with serialized type and handler information
27 """
28 serialized_schema = {}
29 for param, config in schema.items():
30 serialized_schema[param] = {"type": serialize_type(config["type"])}
31 if config.get("handler"):
32 serialized_schema[param]["handler"] = serialize_callable(config["handler"])
34 return serialized_schema
37def _schema_from_dict(schema: dict[str, Any]) -> dict[str, Any]:
38 """
39 Convert a serialized schema dictionary back to its original format.
41 Deserializes the type and optional handler function for each parameter from their
42 serialized format back into Python types and callables.
44 :param schema: Dictionary containing serialized schema information
45 :returns: Dictionary with deserialized type and handler configurations
46 """
47 deserialized_schema = {}
48 for param, config in schema.items():
49 deserialized_schema[param] = {"type": deserialize_type(config["type"])}
51 if config.get("handler"):
52 deserialized_schema[param]["handler"] = deserialize_callable(config["handler"])
54 return deserialized_schema
57def _validate_schema(schema: dict[str, Any]) -> None:
58 """
59 Validate that a schema dictionary meets all required constraints.
61 Checks that each parameter definition has a valid type field and that any handler
62 specified is a callable function.
64 :param schema: Dictionary mapping parameter names to their type and handler configs
65 :raises ValueError: If schema validation fails due to missing or invalid fields
66 """
67 for param, definition in schema.items():
68 if "type" not in definition:
69 raise ValueError(f"StateSchema: Key '{param}' is missing a 'type' entry.")
70 if not _is_valid_type(definition["type"]):
71 raise ValueError(f"StateSchema: 'type' for key '{param}' must be a Python type, got {definition['type']}")
72 if definition.get("handler") is not None and not callable(definition["handler"]):
73 raise ValueError(f"StateSchema: 'handler' for key '{param}' must be callable or None")
74 if param == "messages": # definition["type"] != list[ChatMessage] but split to cover also List[ChatMessage]
75 if not _is_list_type(definition["type"]):
76 raise ValueError(f"StateSchema: 'messages' must be of type list[ChatMessage], got {definition['type']}")
77 # Check if the list contains ChatMessage elements
78 args = get_args(definition["type"])
79 if not args or not issubclass(args[0], ChatMessage):
80 raise ValueError(f"StateSchema: 'messages' must be of type list[ChatMessage], got {definition['type']}")
83class State:
84 """
85 State is a container for storing shared information during the execution of an Agent and its tools.
87 For instance, State can be used to store documents, context, and intermediate results.
89 Internally it wraps a `_data` dictionary defined by a `schema`. Each schema entry has:
90 ```json
91 "parameter_name": {
92 "type": SomeType, # expected type
93 "handler": Optional[Callable[[Any, Any], Any]] # merge/update function
94 }
95 ```
97 Handlers control how values are merged when using the `set()` method:
98 - For list types: defaults to `merge_lists` (concatenates lists)
99 - For other types: defaults to `replace_values` (overwrites existing value)
101 A `messages` field with type `list[ChatMessage]` is automatically added to the schema.
103 This makes it possible for the Agent to read from and write to the same context.
105 ### Usage example
106 ```python
107 from haystack.components.agents.state import State
109 my_state = State(
110 schema={"gh_repo_name": {"type": str}, "user_name": {"type": str}},
111 data={"gh_repo_name": "my_repo", "user_name": "my_user_name"}
112 )
113 ```
114 """
116 def __init__(self, schema: dict[str, Any], data: dict[str, Any] | None = None) -> None:
117 """
118 Initialize a State object with a schema and optional data.
120 :param schema: Dictionary mapping parameter names to their type and handler configs.
121 Type must be a valid Python type, and handler must be a callable function or None.
122 If handler is None, the default handler for the type will be used. The default handlers are:
123 - For list types: `haystack.agents.state.state_utils.merge_lists`
124 - For all other types: `haystack.agents.state.state_utils.replace_values`
125 :param data: Optional dictionary of initial data to populate the state
126 """
127 _validate_schema(schema)
128 self.schema = deepcopy(schema)
129 if self.schema.get("messages") is None:
130 self.schema["messages"] = {"type": list[ChatMessage], "handler": merge_lists}
131 self._data = data or {}
133 # Set default handlers if not provided in schema
134 for definition in self.schema.values():
135 # Skip if handler is already defined and not None
136 if definition.get("handler") is not None:
137 continue
138 # Set default handler based on type
139 if _is_list_type(definition["type"]):
140 definition["handler"] = merge_lists
141 else:
142 definition["handler"] = replace_values
144 def get(self, key: str, default: Any = None) -> Any:
145 """
146 Retrieve a value from the state by key.
148 :param key: Key to look up in the state
149 :param default: Value to return if key is not found
150 :returns: Value associated with key or default if not found
151 """
152 return deepcopy(self._data.get(key, default))
154 def set(self, key: str, value: Any, handler_override: Callable[[Any, Any], Any] | None = None) -> None:
155 """
156 Set or merge a value in the state according to schema rules.
158 Value is merged or overwritten according to these rules:
159 - if handler_override is given, use that
160 - else use the handler defined in the schema for 'key'
162 :param key: Key to store the value under
163 :param value: Value to store or merge
164 :param handler_override: Optional function to override the default merge behavior
165 """
166 # If key not in schema, we throw an error
167 definition = self.schema.get(key, None)
168 if definition is None:
169 raise ValueError(f"State: Key '{key}' not found in schema. Schema: {self.schema}")
171 # Get current value from state and apply handler
172 current_value = self._data.get(key, None)
173 handler = handler_override or definition["handler"]
174 self._data[key] = handler(current_value, value)
176 @property
177 def data(self) -> dict[str, Any]:
178 """
179 All current data of the state.
180 """
181 return self._data
183 def has(self, key: str) -> bool:
184 """
185 Check if a key exists in the state.
187 :param key: Key to check for existence
188 :returns: True if key exists in state, False otherwise
189 """
190 return key in self._data
192 def to_dict(self, skip_keys: list[str] | None = None) -> dict[str, Any]:
193 """
194 Convert the State object to a dictionary.
196 :param skip_keys: List of keys to skip during serialization
197 :returns: Dictionary representation of the State object
198 """
199 skip_keys = skip_keys or []
200 schema = {key: definition for key, definition in self.schema.items() if key not in skip_keys}
201 data = {key: value for key, value in self._data.items() if key not in skip_keys}
203 serialized = {}
204 serialized["schema"] = _schema_to_dict(schema=schema)
205 # Field-level fallback so a single non-serializable value omits only that field instead of
206 # failing the whole State serialization.
207 serialized["data"] = _serialize_with_field_fallback(payload=data, description="the agent's State data")
209 return serialized
211 @classmethod
212 def from_dict(cls, data: dict[str, Any]) -> "State":
213 """
214 Convert a dictionary back to a State object.
215 """
216 schema = _schema_from_dict(data.get("schema", {}))
217 deserialized_data = _deserialize_value_with_schema(data.get("data", {}))
218 return State(schema, deserialized_data)