refactor: Some refactoring and finish TODO for component tool (#9751)

* Some refactoring and finish TODO for component tool

* Add docstrings

Author: sjrl

haystack/tools/component_tool.py

@@ -18,6 +18,7 @@
 from haystack.tools.errors import SchemaGenerationError
 from haystack.tools.from_function import _remove_title_from_schema
 from haystack.tools.parameters_schema_utils import _get_component_param_descriptions, _resolve_type
+from haystack.tools.tool import _deserialize_outputs_to_state, _serialize_outputs_to_state
 from haystack.utils.callable_serialization import deserialize_callable, serialize_callable
 
 logger = logging.getLogger(__name__)
@@ -46,6 +47,8 @@ class ComponentTool(Tool):
     You can create a ComponentTool from the component by passing the component to the ComponentTool constructor.
     Below is an example of creating a ComponentTool from an existing SerperDevWebSearch component.
 
+    ## Usage Example:
+
     ```python
     from haystack import component, Pipeline
     from haystack.tools import ComponentTool
@@ -93,7 +96,7 @@ def __init__(
         outputs_to_string: Optional[dict[str, Union[str, Callable[[Any], str]]]] = None,
         inputs_from_state: Optional[dict[str, str]] = None,
         outputs_to_state: Optional[dict[str, dict[str, Union[str, Callable]]]] = None,
-    ):
+    ) -> None:
         """
         Create a Tool instance from a Haystack component.
 
@@ -213,10 +216,8 @@ def to_dict(self) -> dict[str, Any]:
         """
         Serializes the ComponentTool to a dictionary.
         """
-        serialized_component = component_to_dict(obj=self._component, name=self.name)
-
         serialized: dict[str, Any] = {
-            "component": serialized_component,
+            "component": component_to_dict(obj=self._component, name=self.name),
             "name": self.name,
             "description": self.description,
             "parameters": self._unresolved_parameters,
@@ -226,13 +227,7 @@ def to_dict(self) -> dict[str, Any]:
         }
 
         if self.outputs_to_state is not None:
-            serialized_outputs = {}
-            for key, config in self.outputs_to_state.items():
-                serialized_config = config.copy()
-                if "handler" in config:
-                    serialized_config["handler"] = serialize_callable(config["handler"])
-                serialized_outputs[key] = serialized_config
-            serialized["outputs_to_state"] = serialized_outputs
+            serialized["outputs_to_state"] = _serialize_outputs_to_state(self.outputs_to_state)
 
         if self.outputs_to_string is not None and self.outputs_to_string.get("handler") is not None:
             # This is soft-copied as to not modify the attributes in place
@@ -253,13 +248,7 @@ def from_dict(cls, data: dict[str, Any]) -> "Tool":
         component = component_from_dict(cls=component_class, data=inner_data["component"], name=inner_data["name"])
 
         if "outputs_to_state" in inner_data and inner_data["outputs_to_state"]:
-            deserialized_outputs = {}
-            for key, config in inner_data["outputs_to_state"].items():
-                deserialized_config = config.copy()
-                if "handler" in config:
-                    deserialized_config["handler"] = deserialize_callable(config["handler"])
-                deserialized_outputs[key] = deserialized_config
-            inner_data["outputs_to_state"] = deserialized_outputs
+            inner_data["outputs_to_state"] = _deserialize_outputs_to_state(inner_data["outputs_to_state"])
 
         if (
             inner_data.get("outputs_to_string") is not None

haystack/tools/tool.py

@@ -10,6 +10,7 @@
 
 from haystack.core.serialization import generate_qualified_class_name
 from haystack.tools.errors import ToolInvocationError
+from haystack.utils.callable_serialization import deserialize_callable, serialize_callable
 
 
 @dataclass
@@ -116,24 +117,18 @@ def to_dict(self) -> dict[str, Any]:
         :returns:
             Dictionary with serialized data.
         """
-        # Import here to avoid circular dependency with utils.callable_serialization
-        from haystack.utils import serialize_callable
-
         data = asdict(self)
         data["function"] = serialize_callable(self.function)
 
-        # Serialize output handlers if they exist
-        if self.outputs_to_state:
-            serialized_outputs = {}
-            for key, config in self.outputs_to_state.items():
-                serialized_config = config.copy()
-                if "handler" in config:
-                    serialized_config["handler"] = serialize_callable(config["handler"])
-                serialized_outputs[key] = serialized_config
-            data["outputs_to_state"] = serialized_outputs
+        if self.outputs_to_state is not None:
+            data["outputs_to_state"] = _serialize_outputs_to_state(self.outputs_to_state)
 
         if self.outputs_to_string is not None and self.outputs_to_string.get("handler") is not None:
+            # This is soft-copied as to not modify the attributes in place
+            data["outputs_to_string"] = self.outputs_to_string.copy()
             data["outputs_to_string"]["handler"] = serialize_callable(self.outputs_to_string["handler"])
+        else:
+            data["outputs_to_string"] = None
 
         return {"type": generate_qualified_class_name(type(self)), "data": data}
 
@@ -147,21 +142,10 @@ def from_dict(cls, data: dict[str, Any]) -> "Tool":
         :returns:
             Deserialized Tool.
         """
-        # Import here to avoid circular dependency with utils.callable_serialization
-        from haystack.utils import deserialize_callable
-
         init_parameters = data["data"]
         init_parameters["function"] = deserialize_callable(init_parameters["function"])
-
-        # Deserialize output handlers if they exist
         if "outputs_to_state" in init_parameters and init_parameters["outputs_to_state"]:
-            deserialized_outputs = {}
-            for key, config in init_parameters["outputs_to_state"].items():
-                deserialized_config = config.copy()
-                if "handler" in config:
-                    deserialized_config["handler"] = deserialize_callable(config["handler"])
-                deserialized_outputs[key] = deserialized_config
-            init_parameters["outputs_to_state"] = deserialized_outputs
+            init_parameters["outputs_to_state"] = _deserialize_outputs_to_state(init_parameters["outputs_to_state"])
 
         if (
             init_parameters.get("outputs_to_string") is not None
@@ -187,3 +171,35 @@ def _check_duplicate_tool_names(tools: Optional[list[Tool]]) -> None:
     duplicate_tool_names = {name for name in tool_names if tool_names.count(name) > 1}
     if duplicate_tool_names:
         raise ValueError(f"Duplicate tool names found: {duplicate_tool_names}")
+
+
+def _serialize_outputs_to_state(outputs_to_state: dict[str, dict[str, Any]]) -> dict[str, dict[str, Any]]:
+    """
+    Serializes the outputs_to_state dictionary, converting any callable handlers to their string representation.
+
+    :param outputs_to_state: The outputs_to_state dictionary to serialize.
+    :returns: The serialized outputs_to_state dictionary.
+    """
+    serialized_outputs = {}
+    for key, config in outputs_to_state.items():
+        serialized_config = config.copy()
+        if "handler" in config:
+            serialized_config["handler"] = serialize_callable(config["handler"])
+        serialized_outputs[key] = serialized_config
+    return serialized_outputs
+
+
+def _deserialize_outputs_to_state(outputs_to_state: dict[str, dict[str, Any]]) -> dict[str, dict[str, Any]]:
+    """
+    Deserializes the outputs_to_state dictionary, converting any string handlers back to callables.
+
+    :param outputs_to_state: The outputs_to_state dictionary to deserialize.
+    :returns: The deserialized outputs_to_state dictionary.
+    """
+    deserialized_outputs = {}
+    for key, config in outputs_to_state.items():
+        deserialized_config = config.copy()
+        if "handler" in config:
+            deserialized_config["handler"] = deserialize_callable(config["handler"])
+        deserialized_outputs[key] = deserialized_config
+    return deserialized_outputs

haystack/utils/callable_serialization.py

@@ -6,7 +6,6 @@
 from typing import Any, Callable
 
 from haystack.core.errors import DeserializationError, SerializationError
-from haystack.tools.tool import Tool
 from haystack.utils.type_serialization import thread_safe_import
 
 
@@ -51,6 +50,9 @@ def deserialize_callable(callable_handle: str) -> Callable:
     :return: The callable
     :raises DeserializationError: If the callable cannot be found
     """
+    # Import here to avoid circular imports
+    from haystack.tools.tool import Tool
+
     parts = callable_handle.split(".")
 
     for i in range(len(parts), 0, -1):