15131674881

Committed 20 May 2025 07:35AM UTC coverage: 90.156% (-0.3%) from 90.471%

Build # 15131674881

Build Type

Pull #9407

github

Committed by

web-flow

Commit Message

Merge b382eca10 into 6ad23f822

Pull Request Pull Request #9407: feat: stream `ToolResult` from run_async in Agent

Run Details

10972 of 12170 relevant lines covered (90.16%)

0.9 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

89.66

haystack/tools/tool.py

# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
#
# SPDX-License-Identifier: Apache-2.0

import asyncio
from dataclasses import asdict, dataclass
from functools import partial
from typing import Any, Callable, Dict, List, Optional

from jsonschema import Draft202012Validator
from jsonschema.exceptions import SchemaError

from haystack.core.serialization import generate_qualified_class_name
from haystack.tools.errors import ToolInvocationError
from haystack.utils import deserialize_callable, serialize_callable


@dataclass
class Tool:
    """
    Data class representing a Tool that Language Models can prepare a call for.

    Accurate definitions of the textual attributes such as `name` and `description`
    are important for the Language Model to correctly prepare the call.

    :param name:
        Name of the Tool.
    :param description:
        Description of the Tool.
    :param parameters:
        A JSON schema defining the parameters expected by the Tool.
    :param function:
        The function that will be invoked when the Tool is called.
    :param outputs_to_string:
        Optional dictionary defining how a tool outputs should be converted into a string.
        If the source is provided only the specified output key is sent to the handler.
        If the source is omitted the whole tool result is sent to the handler.
        Example: {
            "source": "docs", "handler": format_documents
        }
    :param inputs_from_state:
        Optional dictionary mapping state keys to tool parameter names.
        Example: {"repository": "repo"} maps state's "repository" to tool's "repo" parameter.
    :param outputs_to_state:
        Optional dictionary defining how tool outputs map to keys within state as well as optional handlers.
        If the source is provided only the specified output key is sent to the handler.
        Example: {
            "documents": {"source": "docs", "handler": custom_handler}
        }
        If the source is omitted the whole tool result is sent to the handler.
        Example: {
            "documents": {"handler": custom_handler}
        }
    """

    name: str
    description: str
    parameters: Dict[str, Any]
    function: Callable
    outputs_to_string: Optional[Dict[str, Any]] = None
    inputs_from_state: Optional[Dict[str, str]] = None
    outputs_to_state: Optional[Dict[str, Dict[str, Any]]] = None

    def __post_init__(self):
        # Check that the parameters define a valid JSON schema
        try:
            Draft202012Validator.check_schema(self.parameters)
        except SchemaError as e:
            raise ValueError("The provided parameters do not define a valid JSON schema") from e

        # Validate outputs structure if provided
        if self.outputs_to_state is not None:
            for key, config in self.outputs_to_state.items():
                if not isinstance(config, dict):
                    raise ValueError(f"outputs_to_state configuration for key '{key}' must be a dictionary")
                if "source" in config and not isinstance(config["source"], str):
                    raise ValueError(f"outputs_to_state source for key '{key}' must be a string.")
                if "handler" in config and not callable(config["handler"]):
                    raise ValueError(f"outputs_to_state handler for key '{key}' must be callable")

        if self.outputs_to_string is not None:
            if "source" in self.outputs_to_string and not isinstance(self.outputs_to_string["source"], str):
                raise ValueError("outputs_to_string source must be a string.")
            if "handler" in self.outputs_to_string and not callable(self.outputs_to_string["handler"]):
                raise ValueError("outputs_to_string handler must be callable")

    @property
    def tool_spec(self) -> Dict[str, Any]:
        """
        Return the Tool specification to be used by the Language Model.
        """
        return {"name": self.name, "description": self.description, "parameters": self.parameters}

    def invoke(self, **kwargs) -> Any:
        """
        Invoke the Tool with the provided keyword arguments.
        """
        try:
            result = self.function(**kwargs)
        except Exception as e:
            raise ToolInvocationError(
                f"Failed to invoke Tool `{self.name}` with parameters {kwargs}. Error: {e}"
            ) from e
        return result

    async def invoke_async(self, **kwargs) -> Any:
        """
        Asynchronously invoke the Tool with the provided keyword arguments.
        """
        loop = asyncio.get_running_loop()
        try:
            return await loop.run_in_executor(None, partial(self.function, **kwargs))
        except Exception as e:
            raise ToolInvocationError(
                f"Failed to invoke Tool `{self.name}` with parameters {kwargs}. Error: {e}"
            ) from e

    def to_dict(self) -> Dict[str, Any]:
        """
        Serializes the Tool to a dictionary.

        :returns:
            Dictionary with serialized data.
        """
        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_string is not None and self.outputs_to_string.get("handler") is not None:
            data["outputs_to_string"] = serialize_callable(self.outputs_to_string["handler"])

        return {"type": generate_qualified_class_name(type(self)), "data": data}

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "Tool":
        """
        Deserializes the Tool from a dictionary.

        :param data:
            Dictionary to deserialize from.
        :returns:
            Deserialized Tool.
        """
        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

        if (
            init_parameters.get("outputs_to_string") is not None
            and init_parameters["outputs_to_string"].get("handler") is not None
        ):
            init_parameters["outputs_to_string"]["handler"] = deserialize_callable(
                init_parameters["outputs_to_string"]["handler"]
            )

        return cls(**init_parameters)


def _check_duplicate_tool_names(tools: Optional[List[Tool]]) -> None:
    """
    Checks for duplicate tool names and raises a ValueError if they are found.

    :param tools: The list of tools to check.
    :raises ValueError: If duplicate tool names are found.
    """
    if tools is None:
        return
    tool_names = [tool.name for tool in tools]
    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}")

1	# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
2	#
3	# SPDX-License-Identifier: Apache-2.0
4
5	import asyncio	1✔
6	from dataclasses import asdict, dataclass	1✔
7	from functools import partial	1✔
8	from typing import Any, Callable, Dict, List, Optional	1✔
9
10	from jsonschema import Draft202012Validator	1✔
11	from jsonschema.exceptions import SchemaError	1✔
12
13	from haystack.core.serialization import generate_qualified_class_name	1✔
14	from haystack.tools.errors import ToolInvocationError	1✔
15	from haystack.utils import deserialize_callable, serialize_callable	1✔
16
17
18	@dataclass	1✔
19	class Tool:	1✔
20	"""
21	Data class representing a Tool that Language Models can prepare a call for.
22
23	Accurate definitions of the textual attributes such as `name` and `description`
24	are important for the Language Model to correctly prepare the call.
25
26	:param name:
27	Name of the Tool.
28	:param description:
29	Description of the Tool.
30	:param parameters:
31	A JSON schema defining the parameters expected by the Tool.
32	:param function:
33	The function that will be invoked when the Tool is called.
34	:param outputs_to_string:
35	Optional dictionary defining how a tool outputs should be converted into a string.
36	If the source is provided only the specified output key is sent to the handler.
37	If the source is omitted the whole tool result is sent to the handler.
38	Example: {
39	"source": "docs", "handler": format_documents
40	}
41	:param inputs_from_state:
42	Optional dictionary mapping state keys to tool parameter names.
43	Example: {"repository": "repo"} maps state's "repository" to tool's "repo" parameter.
44	:param outputs_to_state:
45	Optional dictionary defining how tool outputs map to keys within state as well as optional handlers.
46	If the source is provided only the specified output key is sent to the handler.
47	Example: {
48	"documents": {"source": "docs", "handler": custom_handler}
49	}
50	If the source is omitted the whole tool result is sent to the handler.
51	Example: {
52	"documents": {"handler": custom_handler}
53	}
54	"""
55
56	name: str	1✔
57	description: str	1✔
58	parameters: Dict[str, Any]	1✔
59	function: Callable	1✔
60	outputs_to_string: Optional[Dict[str, Any]] = None	1✔
61	inputs_from_state: Optional[Dict[str, str]] = None	1✔
62	outputs_to_state: Optional[Dict[str, Dict[str, Any]]] = None	1✔
63
64	def __post_init__(self):	1✔
65	# Check that the parameters define a valid JSON schema
66	try:	1✔
67	Draft202012Validator.check_schema(self.parameters)	1✔
68	except SchemaError as e:	1✔
69	raise ValueError("The provided parameters do not define a valid JSON schema") from e	1✔
70
71	# Validate outputs structure if provided
72	if self.outputs_to_state is not None:	1✔
73	for key, config in self.outputs_to_state.items():	1✔
74	if not isinstance(config, dict):	1✔
75	raise ValueError(f"outputs_to_state configuration for key '{key}' must be a dictionary")	1✔
76	if "source" in config and not isinstance(config["source"], str):	1✔
77	raise ValueError(f"outputs_to_state source for key '{key}' must be a string.")	1✔
78	if "handler" in config and not callable(config["handler"]):	1✔
79	raise ValueError(f"outputs_to_state handler for key '{key}' must be callable")	1✔
80
81	if self.outputs_to_string is not None:	1✔
82	if "source" in self.outputs_to_string and not isinstance(self.outputs_to_string["source"], str):	1✔
83	raise ValueError("outputs_to_string source must be a string.")	×
84	if "handler" in self.outputs_to_string and not callable(self.outputs_to_string["handler"]):	1✔
85	raise ValueError("outputs_to_string handler must be callable")	×
86
87	@property	1✔
88	def tool_spec(self) -> Dict[str, Any]:	1✔
89	"""
90	Return the Tool specification to be used by the Language Model.
91	"""
92	return {"name": self.name, "description": self.description, "parameters": self.parameters}	1✔
93
94	def invoke(self, **kwargs) -> Any:	1✔
95	"""
96	Invoke the Tool with the provided keyword arguments.
97	"""
98	try:	1✔
99	result = self.function(**kwargs)	1✔
100	except Exception as e:	1✔
101	raise ToolInvocationError(	1✔
102	f"Failed to invoke Tool `{self.name}` with parameters {kwargs}. Error: {e}"
103	) from e
104	return result	1✔
105
106	async def invoke_async(self, **kwargs) -> Any:	1✔
107	"""
108	Asynchronously invoke the Tool with the provided keyword arguments.
109	"""
110	loop = asyncio.get_running_loop()	×
111	try:	×
112	return await loop.run_in_executor(None, partial(self.function, **kwargs))	×
113	except Exception as e:	×
114	raise ToolInvocationError(	×
115	f"Failed to invoke Tool `{self.name}` with parameters {kwargs}. Error: {e}"
116	) from e
117
118	def to_dict(self) -> Dict[str, Any]:	1✔
119	"""
120	Serializes the Tool to a dictionary.
121
122	:returns:
123	Dictionary with serialized data.
124	"""
125	data = asdict(self)	1✔
126	data["function"] = serialize_callable(self.function)	1✔
127
128	# Serialize output handlers if they exist
129	if self.outputs_to_state:	1✔
130	serialized_outputs = {}	1✔
131	for key, config in self.outputs_to_state.items():	1✔
132	serialized_config = config.copy()	1✔
133	if "handler" in config:	1✔
134	serialized_config["handler"] = serialize_callable(config["handler"])	1✔
135	serialized_outputs[key] = serialized_config	1✔
136	data["outputs_to_state"] = serialized_outputs	1✔
137
138	if self.outputs_to_string is not None and self.outputs_to_string.get("handler") is not None:	1✔
139	data["outputs_to_string"] = serialize_callable(self.outputs_to_string["handler"])	×
140
141	return {"type": generate_qualified_class_name(type(self)), "data": data}	1✔
142
143	@classmethod	1✔
144	def from_dict(cls, data: Dict[str, Any]) -> "Tool":	1✔
145	"""
146	Deserializes the Tool from a dictionary.
147
148	:param data:
149	Dictionary to deserialize from.
150	:returns:
151	Deserialized Tool.
152	"""
153	init_parameters = data["data"]	1✔
154	init_parameters["function"] = deserialize_callable(init_parameters["function"])	1✔
155
156	# Deserialize output handlers if they exist
157	if "outputs_to_state" in init_parameters and init_parameters["outputs_to_state"]:	1✔
158	deserialized_outputs = {}	1✔
159	for key, config in init_parameters["outputs_to_state"].items():	1✔
160	deserialized_config = config.copy()	1✔
161	if "handler" in config:	1✔
162	deserialized_config["handler"] = deserialize_callable(config["handler"])	1✔
163	deserialized_outputs[key] = deserialized_config	1✔
164	init_parameters["outputs_to_state"] = deserialized_outputs	1✔
165
166	if (	1✔
167	init_parameters.get("outputs_to_string") is not None
168	and init_parameters["outputs_to_string"].get("handler") is not None
169	):
170	init_parameters["outputs_to_string"]["handler"] = deserialize_callable(	×
171	init_parameters["outputs_to_string"]["handler"]
172	)
173
174	return cls(**init_parameters)	1✔
175
176
177	def _check_duplicate_tool_names(tools: Optional[List[Tool]]) -> None:	1✔
178	"""
179	Checks for duplicate tool names and raises a ValueError if they are found.
180
181	:param tools: The list of tools to check.
182	:raises ValueError: If duplicate tool names are found.
183	"""
184	if tools is None:	1✔
185	return	1✔
186	tool_names = [tool.name for tool in tools]	1✔
187	duplicate_tool_names = {name for name in tool_names if tool_names.count(name) > 1}	1✔
188	if duplicate_tool_names:	1✔
189	raise ValueError(f"Duplicate tool names found: {duplicate_tool_names}")	1✔

deepset-ai / haystack / 15131674881

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous