mirror of
https://github.com/deepset-ai/haystack.git
synced 2026-01-01 01:27:28 +00:00
691 lines
24 KiB
Python
691 lines
24 KiB
Python
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
|
#
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
import pytest
|
|
|
|
from haystack import Pipeline
|
|
from haystack.dataclasses import ChatMessage
|
|
from haystack.components.tools import ToolInvoker
|
|
from haystack.components.generators.chat import OpenAIChatGenerator
|
|
from haystack.components.converters import OutputAdapter
|
|
from haystack.dataclasses.chat_message import ToolCall
|
|
from haystack.tools import Tool, Toolset
|
|
|
|
|
|
def test_toolset_with_tool_invoker():
|
|
"""Test that a Toolset works with ToolInvoker."""
|
|
|
|
# Define a simple function for a tool
|
|
def add_numbers(a: int, b: int) -> int:
|
|
return a + b
|
|
|
|
# Create a tool
|
|
add_tool = Tool(
|
|
name="add",
|
|
description="Add two numbers",
|
|
parameters={
|
|
"type": "object",
|
|
"properties": {"a": {"type": "integer"}, "b": {"type": "integer"}},
|
|
"required": ["a", "b"],
|
|
},
|
|
function=add_numbers,
|
|
)
|
|
|
|
# Create a toolset with the tool
|
|
toolset = Toolset([add_tool])
|
|
|
|
# Create a ToolInvoker with the toolset
|
|
invoker = ToolInvoker(tools=toolset)
|
|
|
|
# Create a message with a tool call
|
|
tool_call = ToolCall(tool_name="add", arguments={"a": 2, "b": 3})
|
|
message = ChatMessage.from_assistant(tool_calls=[tool_call])
|
|
|
|
# Invoke the tool
|
|
result = invoker.run(messages=[message])
|
|
|
|
# Check that the tool was invoked correctly
|
|
assert len(result["tool_messages"]) == 1
|
|
assert result["tool_messages"][0].tool_call_result.result == "5"
|
|
|
|
|
|
def test_toolset_with_multiple_tools():
|
|
"""Test that a Toolset with multiple tools works properly."""
|
|
|
|
# Define simple functions for tools
|
|
def add_numbers(a: int, b: int) -> int:
|
|
return a + b
|
|
|
|
def multiply_numbers(a: int, b: int) -> int:
|
|
return a * b
|
|
|
|
# Create tools
|
|
add_tool = Tool(
|
|
name="add",
|
|
description="Add two numbers",
|
|
parameters={
|
|
"type": "object",
|
|
"properties": {"a": {"type": "integer"}, "b": {"type": "integer"}},
|
|
"required": ["a", "b"],
|
|
},
|
|
function=add_numbers,
|
|
)
|
|
|
|
multiply_tool = Tool(
|
|
name="multiply",
|
|
description="Multiply two numbers",
|
|
parameters={
|
|
"type": "object",
|
|
"properties": {"a": {"type": "integer"}, "b": {"type": "integer"}},
|
|
"required": ["a", "b"],
|
|
},
|
|
function=multiply_numbers,
|
|
)
|
|
|
|
# Create a toolset with the tools
|
|
toolset = Toolset([add_tool, multiply_tool])
|
|
|
|
# Check that both tools are in the toolset
|
|
assert len(toolset) == 2
|
|
assert toolset[0].name == "add"
|
|
assert toolset[1].name == "multiply"
|
|
|
|
# Create a ToolInvoker with the toolset
|
|
invoker = ToolInvoker(tools=list(toolset))
|
|
|
|
# Create messages with tool calls
|
|
add_call = ToolCall(tool_name="add", arguments={"a": 2, "b": 3})
|
|
add_message = ChatMessage.from_assistant(tool_calls=[add_call])
|
|
|
|
multiply_call = ToolCall(tool_name="multiply", arguments={"a": 4, "b": 5})
|
|
multiply_message = ChatMessage.from_assistant(tool_calls=[multiply_call])
|
|
|
|
# Invoke the tools
|
|
result = invoker.run(messages=[add_message, multiply_message])
|
|
|
|
# Check that the tools were invoked correctly
|
|
assert len(result["tool_messages"]) == 2
|
|
tool_results = [message.tool_call_result.result for message in result["tool_messages"]]
|
|
assert "5" in tool_results
|
|
assert "20" in tool_results
|
|
|
|
|
|
def test_toolset_registration():
|
|
"""Test that tools can be registered with a Toolset."""
|
|
|
|
# Create an empty toolset
|
|
toolset = Toolset()
|
|
assert len(toolset) == 0
|
|
|
|
# Define a simple function for a tool
|
|
def add_numbers(a: int, b: int) -> int:
|
|
return a + b
|
|
|
|
# Create a tool
|
|
add_tool = Tool(
|
|
name="add",
|
|
description="Add two numbers",
|
|
parameters={
|
|
"type": "object",
|
|
"properties": {"a": {"type": "integer"}, "b": {"type": "integer"}},
|
|
"required": ["a", "b"],
|
|
},
|
|
function=add_numbers,
|
|
)
|
|
|
|
# Register the tool
|
|
toolset.register_tool(add_tool)
|
|
assert len(toolset) == 1
|
|
assert toolset[0].name == "add"
|
|
|
|
# Test with ToolInvoker
|
|
invoker = ToolInvoker(tools=list(toolset))
|
|
tool_call = ToolCall(tool_name="add", arguments={"a": 2, "b": 3})
|
|
message = ChatMessage.from_assistant(tool_calls=[tool_call])
|
|
result = invoker.run(messages=[message])
|
|
|
|
assert len(result["tool_messages"]) == 1
|
|
assert result["tool_messages"][0].tool_call_result.result == "5"
|
|
|
|
|
|
def test_toolset_addition():
|
|
"""Test that toolsets can be combined."""
|
|
|
|
# Define simple functions for tools
|
|
def add_numbers(a: int, b: int) -> int:
|
|
return a + b
|
|
|
|
def multiply_numbers(a: int, b: int) -> int:
|
|
return a * b
|
|
|
|
def subtract_numbers(a: int, b: int) -> int:
|
|
return a - b
|
|
|
|
# Create tools
|
|
add_tool = Tool(
|
|
name="add",
|
|
description="Add two numbers",
|
|
parameters={
|
|
"type": "object",
|
|
"properties": {"a": {"type": "integer"}, "b": {"type": "integer"}},
|
|
"required": ["a", "b"],
|
|
},
|
|
function=add_numbers,
|
|
)
|
|
|
|
multiply_tool = Tool(
|
|
name="multiply",
|
|
description="Multiply two numbers",
|
|
parameters={
|
|
"type": "object",
|
|
"properties": {"a": {"type": "integer"}, "b": {"type": "integer"}},
|
|
"required": ["a", "b"],
|
|
},
|
|
function=multiply_numbers,
|
|
)
|
|
|
|
subtract_tool = Tool(
|
|
name="subtract",
|
|
description="Subtract two numbers",
|
|
parameters={
|
|
"type": "object",
|
|
"properties": {"a": {"type": "integer"}, "b": {"type": "integer"}},
|
|
"required": ["a", "b"],
|
|
},
|
|
function=subtract_numbers,
|
|
)
|
|
|
|
# Create toolsets
|
|
toolset1 = Toolset([add_tool])
|
|
toolset2 = Toolset([multiply_tool])
|
|
|
|
# Combine toolsets
|
|
combined_toolset = toolset1 + toolset2
|
|
assert len(combined_toolset) == 2
|
|
|
|
# Add a single tool
|
|
combined_toolset = combined_toolset + subtract_tool
|
|
assert len(combined_toolset) == 3
|
|
|
|
# Check that all tools are accessible
|
|
tool_names = [tool.name for tool in combined_toolset]
|
|
assert "add" in tool_names
|
|
assert "multiply" in tool_names
|
|
assert "subtract" in tool_names
|
|
|
|
# Test with ToolInvoker
|
|
invoker = ToolInvoker(tools=list(combined_toolset))
|
|
|
|
# Create messages with tool calls
|
|
add_call = ToolCall(tool_name="add", arguments={"a": 10, "b": 5})
|
|
multiply_call = ToolCall(tool_name="multiply", arguments={"a": 10, "b": 5})
|
|
subtract_call = ToolCall(tool_name="subtract", arguments={"a": 10, "b": 5})
|
|
|
|
message = ChatMessage.from_assistant(tool_calls=[add_call, multiply_call, subtract_call])
|
|
|
|
# Invoke the tools
|
|
result = invoker.run(messages=[message])
|
|
|
|
# Check that all tools were invoked correctly
|
|
assert len(result["tool_messages"]) == 3
|
|
tool_results = [message.tool_call_result.result for message in result["tool_messages"]]
|
|
assert "15" in tool_results
|
|
assert "50" in tool_results
|
|
assert "5" in tool_results
|
|
|
|
|
|
def add_numbers_for_serialization(a: int, b: int) -> int:
|
|
return a + b
|
|
|
|
|
|
def test_toolset_serialization():
|
|
"""Test that a Toolset can be serialized and deserialized."""
|
|
|
|
# Create a tool
|
|
add_tool = Tool(
|
|
name="add",
|
|
description="Add two numbers",
|
|
parameters={
|
|
"type": "object",
|
|
"properties": {"a": {"type": "integer"}, "b": {"type": "integer"}},
|
|
"required": ["a", "b"],
|
|
},
|
|
function=add_numbers_for_serialization,
|
|
)
|
|
|
|
# Create a toolset with the tool
|
|
toolset = Toolset([add_tool])
|
|
|
|
# Serialize the toolset
|
|
serialized = toolset.to_dict()
|
|
|
|
# Deserialize the toolset
|
|
deserialized = Toolset.from_dict(serialized)
|
|
|
|
# Check that the deserialized toolset has the same tool
|
|
assert len(deserialized) == 1
|
|
assert deserialized[0].name == "add"
|
|
assert deserialized[0].description == "Add two numbers"
|
|
|
|
# Test that the deserialized toolset works with ToolInvoker
|
|
invoker = ToolInvoker(tools=list(deserialized))
|
|
tool_call = ToolCall(tool_name="add", arguments={"a": 2, "b": 3})
|
|
message = ChatMessage.from_assistant(tool_calls=[tool_call])
|
|
result = invoker.run(messages=[message])
|
|
|
|
assert len(result["tool_messages"]) == 1
|
|
assert result["tool_messages"][0].tool_call_result.result == "5"
|
|
|
|
|
|
def test_toolset_duplicate_tool_names():
|
|
"""Test that a Toolset raises an error for duplicate tool names."""
|
|
|
|
# Define a simple function for tools
|
|
def add_numbers(a: int, b: int) -> int:
|
|
return a + b
|
|
|
|
# Create tools with the same name
|
|
add_tool1 = Tool(
|
|
name="add",
|
|
description="Add two numbers (first)",
|
|
parameters={
|
|
"type": "object",
|
|
"properties": {"a": {"type": "integer"}, "b": {"type": "integer"}},
|
|
"required": ["a", "b"],
|
|
},
|
|
function=add_numbers,
|
|
)
|
|
|
|
add_tool2 = Tool(
|
|
name="add",
|
|
description="Add two numbers (second)",
|
|
parameters={
|
|
"type": "object",
|
|
"properties": {"a": {"type": "integer"}, "b": {"type": "integer"}},
|
|
"required": ["a", "b"],
|
|
},
|
|
function=add_numbers,
|
|
)
|
|
|
|
# Check that creating a toolset with duplicate tool names raises an error
|
|
with pytest.raises(ValueError, match="Duplicate tool names found"):
|
|
Toolset([add_tool1, add_tool2])
|
|
|
|
# Create a toolset with one tool
|
|
toolset = Toolset([add_tool1])
|
|
|
|
# Check that registering a tool with a duplicate name raises an error
|
|
with pytest.raises(ValueError, match="Duplicate tool names found"):
|
|
toolset.register_tool(add_tool2)
|
|
|
|
# Check that combining toolsets with duplicate tool names raises an error
|
|
toolset2 = Toolset([add_tool2])
|
|
with pytest.raises(ValueError, match="Duplicate tool names found"):
|
|
combined = toolset + toolset2
|
|
|
|
|
|
def add_numbers_for_pipeline(a: int, b: int) -> int:
|
|
return a + b
|
|
|
|
|
|
def multiply_numbers_for_pipeline(a: int, b: int) -> int:
|
|
return a * b
|
|
|
|
|
|
def add_numbers_for_calculator(a: int, b: int) -> int:
|
|
return a + b
|
|
|
|
|
|
# Integration tests using full pipelines
|
|
@pytest.mark.integration
|
|
class TestToolsetIntegration:
|
|
"""Integration tests for Toolset in complete pipelines."""
|
|
|
|
def test_basic_math_pipeline(self):
|
|
"""Test basic math operations through a complete pipeline flow."""
|
|
|
|
def add_numbers(a: int, b: int) -> int:
|
|
return a + b
|
|
|
|
def subtract_numbers(a: int, b: int) -> int:
|
|
return a - b
|
|
|
|
# Create the tools
|
|
add_tool = Tool(
|
|
name="add",
|
|
description="Add two numbers",
|
|
parameters={
|
|
"type": "object",
|
|
"properties": {"a": {"type": "integer"}, "b": {"type": "integer"}},
|
|
"required": ["a", "b"],
|
|
},
|
|
function=add_numbers,
|
|
)
|
|
|
|
subtract_tool = Tool(
|
|
name="subtract",
|
|
description="Subtract two numbers",
|
|
parameters={
|
|
"type": "object",
|
|
"properties": {"a": {"type": "integer"}, "b": {"type": "integer"}},
|
|
"required": ["a", "b"],
|
|
},
|
|
function=subtract_numbers,
|
|
)
|
|
|
|
# Create a toolset
|
|
math_toolset = Toolset([add_tool, subtract_tool])
|
|
|
|
# Create a complete pipeline
|
|
pipeline = Pipeline()
|
|
pipeline.add_component("llm", OpenAIChatGenerator(model="gpt-4o-mini", tools=list(math_toolset)))
|
|
pipeline.add_component("tool_invoker", ToolInvoker(tools=list(math_toolset)))
|
|
pipeline.add_component(
|
|
"adapter",
|
|
OutputAdapter(
|
|
template="{{ initial_msg + initial_tool_messages + tool_messages }}",
|
|
output_type=list[ChatMessage],
|
|
unsafe=True,
|
|
),
|
|
)
|
|
pipeline.add_component("response_llm", OpenAIChatGenerator(model="gpt-4o-mini"))
|
|
|
|
# Connect the components
|
|
pipeline.connect("llm.replies", "tool_invoker.messages")
|
|
pipeline.connect("llm.replies", "adapter.initial_tool_messages")
|
|
pipeline.connect("tool_invoker.tool_messages", "adapter.tool_messages")
|
|
pipeline.connect("adapter.output", "response_llm.messages")
|
|
|
|
# Test addition through the pipeline
|
|
user_input = "What is 5 plus 3?"
|
|
user_input_msg = ChatMessage.from_user(text=user_input)
|
|
|
|
result = pipeline.run({"llm": {"messages": [user_input_msg]}, "adapter": {"initial_msg": [user_input_msg]}})
|
|
|
|
# Verify the complete flow worked
|
|
pipe_result = result["response_llm"]["replies"][0].text
|
|
assert "8" in pipe_result or "eight" in pipe_result
|
|
|
|
def test_combined_toolsets_pipeline(self):
|
|
"""Test combining multiple toolsets in a complete pipeline."""
|
|
|
|
# Create basic and advanced math tools
|
|
def add_numbers(a: int, b: int) -> int:
|
|
return a + b
|
|
|
|
def multiply_numbers(a: int, b: int) -> int:
|
|
return a * b
|
|
|
|
# Create the tools
|
|
add_tool = Tool(
|
|
name="add",
|
|
description="Add two numbers",
|
|
parameters={
|
|
"type": "object",
|
|
"properties": {"a": {"type": "integer"}, "b": {"type": "integer"}},
|
|
"required": ["a", "b"],
|
|
},
|
|
function=add_numbers,
|
|
)
|
|
|
|
multiply_tool = Tool(
|
|
name="multiply",
|
|
description="Multiply two numbers",
|
|
parameters={
|
|
"type": "object",
|
|
"properties": {"a": {"type": "integer"}, "b": {"type": "integer"}},
|
|
"required": ["a", "b"],
|
|
},
|
|
function=multiply_numbers,
|
|
)
|
|
|
|
# Create separate toolsets
|
|
basic_math_toolset = Toolset([add_tool])
|
|
advanced_math_toolset = Toolset([multiply_tool])
|
|
|
|
# Combine toolsets
|
|
combined_toolset = basic_math_toolset + advanced_math_toolset
|
|
|
|
# Create a complete pipeline
|
|
pipeline = Pipeline()
|
|
pipeline.add_component("llm", OpenAIChatGenerator(model="gpt-4o-mini", tools=combined_toolset))
|
|
pipeline.add_component("tool_invoker", ToolInvoker(tools=list(combined_toolset)))
|
|
pipeline.add_component(
|
|
"adapter",
|
|
OutputAdapter(
|
|
template="{{ initial_msg + initial_tool_messages + tool_messages }}",
|
|
output_type=list[ChatMessage],
|
|
unsafe=True,
|
|
),
|
|
)
|
|
pipeline.add_component("response_llm", OpenAIChatGenerator(model="gpt-4o-mini"))
|
|
|
|
# Connect the components
|
|
pipeline.connect("llm.replies", "tool_invoker.messages")
|
|
pipeline.connect("llm.replies", "adapter.initial_tool_messages")
|
|
pipeline.connect("tool_invoker.tool_messages", "adapter.tool_messages")
|
|
pipeline.connect("adapter.output", "response_llm.messages")
|
|
|
|
# Test multiplication through the pipeline
|
|
user_input = "What is 6 times 7?"
|
|
user_input_msg = ChatMessage.from_user(text=user_input)
|
|
|
|
result = pipeline.run({"llm": {"messages": [user_input_msg]}, "adapter": {"initial_msg": [user_input_msg]}})
|
|
|
|
# Verify the complete flow worked
|
|
pipe_result = result["response_llm"]["replies"][0].text
|
|
assert "42" in pipe_result or "forty two" in pipe_result
|
|
|
|
def test_custom_calculator_pipeline(self):
|
|
"""Test a custom calculator toolset in a complete pipeline."""
|
|
|
|
class CalculatorToolset(Toolset):
|
|
"""A toolset for calculator operations."""
|
|
|
|
def __init__(self):
|
|
"""Create a toolset with calculator operations."""
|
|
super().__init__([])
|
|
self._create_tools()
|
|
|
|
def _create_tools(self):
|
|
"""Create calculator operation tools."""
|
|
|
|
def add_numbers(a: int, b: int) -> int:
|
|
return a + b
|
|
|
|
add_tool = Tool(
|
|
name="add",
|
|
description="Add two numbers",
|
|
parameters={
|
|
"type": "object",
|
|
"properties": {"a": {"type": "integer"}, "b": {"type": "integer"}},
|
|
"required": ["a", "b"],
|
|
},
|
|
function=add_numbers,
|
|
)
|
|
|
|
def multiply_numbers(a: int, b: int) -> int:
|
|
return a * b
|
|
|
|
multiply_tool = Tool(
|
|
name="multiply",
|
|
description="Multiply two numbers",
|
|
parameters={
|
|
"type": "object",
|
|
"properties": {"a": {"type": "integer"}, "b": {"type": "integer"}},
|
|
"required": ["a", "b"],
|
|
},
|
|
function=multiply_numbers,
|
|
)
|
|
|
|
# Register the tools
|
|
self.register_tool(add_tool)
|
|
self.register_tool(multiply_tool)
|
|
|
|
# Create a CalculatorToolset instance
|
|
calculator_toolset = CalculatorToolset()
|
|
|
|
# Create a complete pipeline
|
|
pipeline = Pipeline()
|
|
pipeline.add_component("llm", OpenAIChatGenerator(model="gpt-4o-mini", tools=list(calculator_toolset)))
|
|
pipeline.add_component("tool_invoker", ToolInvoker(tools=list(calculator_toolset)))
|
|
pipeline.add_component(
|
|
"adapter",
|
|
OutputAdapter(
|
|
template="{{ initial_msg + initial_tool_messages + tool_messages }}",
|
|
output_type=list[ChatMessage],
|
|
unsafe=True,
|
|
),
|
|
)
|
|
pipeline.add_component("response_llm", OpenAIChatGenerator(model="gpt-4o-mini"))
|
|
|
|
# Connect the components
|
|
pipeline.connect("llm.replies", "tool_invoker.messages")
|
|
pipeline.connect("llm.replies", "adapter.initial_tool_messages")
|
|
pipeline.connect("tool_invoker.tool_messages", "adapter.tool_messages")
|
|
pipeline.connect("adapter.output", "response_llm.messages")
|
|
|
|
# Test the calculator through the pipeline
|
|
user_input = "What is 15 plus 7?"
|
|
user_input_msg = ChatMessage.from_user(text=user_input)
|
|
|
|
result = pipeline.run({"llm": {"messages": [user_input_msg]}, "adapter": {"initial_msg": [user_input_msg]}})
|
|
|
|
# Verify the complete flow worked
|
|
pipe_result = result["response_llm"]["replies"][0].text
|
|
assert "22" in pipe_result or "twenty two" in pipe_result
|
|
|
|
def test_serde_in_pipeline(self, monkeypatch):
|
|
"""Test serialization and deserialization of a pipeline with toolsets."""
|
|
monkeypatch.setenv("OPENAI_API_KEY", "test-key")
|
|
|
|
# Create tools
|
|
add_tool = Tool(
|
|
name="add",
|
|
description="Add two numbers",
|
|
parameters={
|
|
"type": "object",
|
|
"properties": {"a": {"type": "integer"}, "b": {"type": "integer"}},
|
|
"required": ["a", "b"],
|
|
},
|
|
function=add_numbers_for_pipeline,
|
|
)
|
|
|
|
multiply_tool = Tool(
|
|
name="multiply",
|
|
description="Multiply two numbers",
|
|
parameters={
|
|
"type": "object",
|
|
"properties": {"a": {"type": "integer"}, "b": {"type": "integer"}},
|
|
"required": ["a", "b"],
|
|
},
|
|
function=multiply_numbers_for_pipeline,
|
|
)
|
|
|
|
# Create toolsets and combine them
|
|
basic_toolset = Toolset([add_tool])
|
|
advanced_toolset = Toolset([multiply_tool])
|
|
combined_toolset = basic_toolset + advanced_toolset
|
|
|
|
# Create and configure the pipeline
|
|
pipeline = Pipeline()
|
|
pipeline.add_component("tool_invoker", ToolInvoker(tools=list(combined_toolset)))
|
|
pipeline.add_component("llm", OpenAIChatGenerator(model="gpt-3.5-turbo", tools=list(combined_toolset)))
|
|
pipeline.add_component(
|
|
"adapter",
|
|
OutputAdapter(
|
|
template="{{ initial_msg + initial_tool_messages + tool_messages }}",
|
|
output_type=list[ChatMessage],
|
|
unsafe=True,
|
|
),
|
|
)
|
|
pipeline.add_component("response_llm", OpenAIChatGenerator(model="gpt-3.5-turbo"))
|
|
|
|
# Connect components
|
|
pipeline.connect("llm.replies", "tool_invoker.messages")
|
|
pipeline.connect("llm.replies", "adapter.initial_tool_messages")
|
|
pipeline.connect("tool_invoker.tool_messages", "adapter.tool_messages")
|
|
pipeline.connect("adapter.output", "response_llm.messages")
|
|
|
|
# Serialize to dict and verify structure
|
|
pipeline_dict = pipeline.to_dict()
|
|
assert (
|
|
pipeline_dict["components"]["tool_invoker"]["type"] == "haystack.components.tools.tool_invoker.ToolInvoker"
|
|
)
|
|
assert len(pipeline_dict["components"]["tool_invoker"]["init_parameters"]["tools"]) == 2
|
|
|
|
# Verify tool serialization
|
|
tools_dict = pipeline_dict["components"]["tool_invoker"]["init_parameters"]["tools"]
|
|
tool_names = [tool["data"]["name"] for tool in tools_dict]
|
|
assert "add" in tool_names
|
|
assert "multiply" in tool_names
|
|
|
|
# Test round-trip serialization
|
|
pipeline_yaml = pipeline.dumps()
|
|
new_pipeline = Pipeline.loads(pipeline_yaml)
|
|
|
|
# Compare components
|
|
original_components = set(pipeline_dict["components"].keys())
|
|
new_components = set(new_pipeline.to_dict()["components"].keys())
|
|
assert original_components == new_components
|
|
|
|
for name in original_components:
|
|
assert type(pipeline.get_component(name)) == type(new_pipeline.get_component(name))
|
|
|
|
# Compare connections
|
|
original_connections = {(edge[0], edge[1]) for edge in pipeline.graph.edges()}
|
|
new_connections = {(edge[0], edge[1]) for edge in new_pipeline.graph.edges()}
|
|
assert original_connections == new_connections
|
|
|
|
def test_toolset_serde_in_pipeline(self):
|
|
"""Test serialization and deserialization of toolsets within a pipeline."""
|
|
|
|
class CalculatorToolset(Toolset):
|
|
"""A toolset for calculator operations."""
|
|
|
|
def __init__(self):
|
|
"""Create a toolset with calculator operations."""
|
|
super().__init__([])
|
|
self._create_tools()
|
|
|
|
def _create_tools(self):
|
|
"""Create calculator operation tools."""
|
|
add_tool = Tool(
|
|
name="add",
|
|
description="Add two numbers",
|
|
parameters={
|
|
"type": "object",
|
|
"properties": {"a": {"type": "integer"}, "b": {"type": "integer"}},
|
|
"required": ["a", "b"],
|
|
},
|
|
function=add_numbers_for_calculator,
|
|
)
|
|
|
|
self.register_tool(add_tool)
|
|
|
|
# Create a CalculatorToolset instance
|
|
calculator_toolset = CalculatorToolset()
|
|
|
|
# Create a pipeline with the toolset
|
|
pipeline = Pipeline()
|
|
pipeline.add_component("tool_invoker", ToolInvoker(tools=list(calculator_toolset)))
|
|
|
|
# Serialize the pipeline
|
|
pipeline_dict = pipeline.to_dict()
|
|
|
|
# Verify toolset serialization
|
|
tool_invoker_dict = pipeline_dict["components"]["tool_invoker"]
|
|
assert tool_invoker_dict["type"] == "haystack.components.tools.tool_invoker.ToolInvoker"
|
|
assert len(tool_invoker_dict["init_parameters"]["tools"]) == 1
|
|
|
|
# Verify tool details
|
|
tool_dict = tool_invoker_dict["init_parameters"]["tools"][0]
|
|
assert tool_dict["data"]["name"] == "add"
|
|
assert tool_dict["data"]["description"] == "Add two numbers"
|
|
|
|
# Test round-trip serialization
|
|
pipeline_yaml = pipeline.dumps()
|
|
new_pipeline = Pipeline.loads(pipeline_yaml)
|
|
assert new_pipeline == pipeline
|