Skip to content

Commit c0aac99

Browse files
sjrlsjrl
authored
chore: replace ToolInvoker with Agent in tests, examples and docstrings (#4108)
1 parent 61f2470 commit c0aac99

9 files changed

Lines changed: 65 additions & 197 deletions

File tree

‎integrations/aimlapi/examples/aimlapi_with_tools_example.py‎

Lines changed: 11 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66

77
"""Run AIMLAPI chat generation with tool calling and execution."""
88

9-
from haystack.components.tools import ToolInvoker
9+
from haystack.components.agents import Agent
1010
from haystack.dataclasses import ChatMessage
1111
from haystack.tools import Tool
1212

@@ -20,7 +20,7 @@ def weather(city: str) -> str:
2020

2121

2222
def main() -> None:
23-
"""Demonstrate calling a tool and feeding the result back to AIMLAPI."""
23+
"""Demonstrate an Agent calling a tool and feeding the result back to AIMLAPI."""
2424

2525
tool_parameters = {
2626
"type": "object",
@@ -35,42 +35,19 @@ def main() -> None:
3535
function=weather,
3636
)
3737

38-
tool_invoker = ToolInvoker(tools=[weather_tool])
39-
40-
client = AIMLAPIChatGenerator(model="openai/gpt-5-mini-2025-08-07")
41-
42-
messages = [
43-
ChatMessage.from_system("You help users by calling the provided tools when they are relevant."),
44-
ChatMessage.from_user("What's the weather in Tokyo today?"),
45-
]
46-
47-
print("Requesting a tool call from the model...")
48-
tool_request = client.run(
49-
messages=messages,
38+
agent = Agent(
39+
chat_generator=AIMLAPIChatGenerator(model="openai/gpt-5-mini-2025-08-07"),
40+
system_prompt="You help users by calling the provided tools when they are relevant.",
5041
tools=[weather_tool],
51-
generation_kwargs={"tool_choice": {"type": "function", "function": {"name": "weather"}}},
52-
)["replies"][0]
53-
54-
print(f"assistant tool request: {tool_request}")
55-
56-
if not tool_request.tool_calls:
57-
print("No tool call was produced by the model.")
58-
return
59-
60-
tool_messages = tool_invoker.run(messages=[tool_request])["tool_messages"]
61-
for tool_message in tool_messages:
62-
for tool_result in tool_message.tool_call_results:
63-
print(f"tool output: {tool_result.result}")
42+
)
6443

65-
follow_up_messages = [*messages, tool_request, *tool_messages]
44+
result = agent.run(messages=[ChatMessage.from_user("What's the weather in Tokyo today?")])
6645

67-
final_reply = client.run(
68-
messages=follow_up_messages,
69-
tools=[weather_tool],
70-
generation_kwargs={"tool_choice": "none"},
71-
)["replies"][0]
46+
for message in result["messages"]:
47+
if message.tool_call_result:
48+
print(f"tool output: {message.tool_call_result.result}")
7249

73-
print(f"assistant final answer: {final_reply.text}")
50+
print(f"assistant final answer: {result['last_message'].text}")
7451

7552

7653
if __name__ == "__main__":

‎integrations/aimlapi/tests/test_aimlapi_chat_generator.py‎

Lines changed: 7 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -5,12 +5,8 @@
55
import pytest
66
import pytz
77
from haystack import Pipeline
8+
from haystack.components.agents import Agent
89
from haystack.components.generators.utils import print_streaming_chunk
9-
10-
try:
11-
from haystack.components.tools import ToolInvoker
12-
except ImportError: # ToolInvoker was removed in Haystack 3.0
13-
ToolInvoker = None
1410
from haystack.dataclasses import ChatMessage, ChatRole, StreamingChunk, ToolCall
1511
from haystack.tools import Tool
1612
from haystack.utils.auth import Secret
@@ -409,30 +405,15 @@ def test_live_run_with_tools_streaming(self, tools):
409405
reason="Export an env var called AIMLAPI_API_KEY containing the AIMLAPI API key to run this test.",
410406
)
411407
@pytest.mark.integration
412-
@pytest.mark.skipif(ToolInvoker is None, reason="ToolInvoker is not available in the installed haystack-ai version")
413-
def test_pipeline_with_aimlapi_chat_generator(self, tools):
408+
def test_agent_with_aimlapi_chat_generator(self, tools):
414409
"""
415-
Test that the AIMLAPIChatGenerator component can be used in a pipeline
410+
Test that the AIMLAPIChatGenerator component can drive tool calls in an Agent
416411
"""
417-
pipeline = Pipeline()
418-
pipeline.add_component("generator", AIMLAPIChatGenerator(tools=tools, model="openai/gpt-5-mini-2025-08-07"))
419-
pipeline.add_component("tool_invoker", ToolInvoker(tools=tools))
420-
421-
pipeline.connect("generator", "tool_invoker")
422-
423-
results = pipeline.run(
424-
data={
425-
"generator": {
426-
"messages": [ChatMessage.from_user("What's the weather like in Paris?")],
427-
"generation_kwargs": {"tool_choice": "auto"},
428-
}
429-
}
430-
)
412+
agent = Agent(chat_generator=AIMLAPIChatGenerator(model="openai/gpt-5-mini-2025-08-07"), tools=tools)
413+
result = agent.run(messages=[ChatMessage.from_user("What's the weather like in Paris?")])
431414

432-
assert (
433-
"The weather in Paris is sunny and 32°C"
434-
== results["tool_invoker"]["tool_messages"][0].tool_call_result.result
435-
)
415+
tool_results = [msg.tool_call_result.result for msg in result["messages"] if msg.tool_call_result]
416+
assert "The weather in Paris is sunny and 32°C" in tool_results
436417

437418
def test_serde_in_pipeline(self, monkeypatch):
438419
"""

‎integrations/amazon_bedrock/tests/test_chat_generator.py‎

Lines changed: 6 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -8,11 +8,6 @@
88
from haystack import Pipeline
99
from haystack.components.agents import Agent
1010
from haystack.components.generators.utils import print_streaming_chunk
11-
12-
try:
13-
from haystack.components.tools import ToolInvoker
14-
except ImportError: # ToolInvoker was removed in Haystack 3.0
15-
ToolInvoker = None
1611
from haystack.dataclasses import (
1712
ChatMessage,
1813
ChatRole,
@@ -1648,26 +1643,15 @@ def test_prompt_caching_live_run_with_user_message(self, model_name, streaming_c
16481643
assert usage["cache_details"] == expected, f"Unexpected usage stats: {usage}"
16491644

16501645
@pytest.mark.parametrize("model_name", MODELS_TO_TEST_WITH_TOOLS[:1]) # just one model is enough
1651-
@pytest.mark.skipif(ToolInvoker is None, reason="ToolInvoker is not available in the installed haystack-ai version")
1652-
def test_pipeline_with_amazon_bedrock_chat_generator(self, model_name, tools):
1646+
def test_agent_with_amazon_bedrock_chat_generator(self, model_name, tools):
16531647
"""
1654-
Test that the AmazonBedrockChatGenerator component can be used in a pipeline
1648+
Test that the AmazonBedrockChatGenerator component can drive tool calls in an Agent
16551649
"""
1650+
agent = Agent(chat_generator=AmazonBedrockChatGenerator(model=model_name), tools=tools)
1651+
result = agent.run(messages=[ChatMessage.from_user("What's the weather like in Paris?")])
16561652

1657-
pipeline = Pipeline()
1658-
pipeline.add_component("generator", AmazonBedrockChatGenerator(model=model_name, tools=tools))
1659-
pipeline.add_component("tool_invoker", ToolInvoker(tools=tools))
1660-
1661-
pipeline.connect("generator", "tool_invoker")
1662-
1663-
results = pipeline.run(
1664-
data={"generator": {"messages": [ChatMessage.from_user("What's the weather like in Paris?")]}}
1665-
)
1666-
1667-
assert (
1668-
"The weather in Paris is sunny and 32°C"
1669-
== results["tool_invoker"]["tool_messages"][0].tool_call_result.result
1670-
)
1653+
tool_results = [msg.tool_call_result.result for msg in result["messages"] if msg.tool_call_result]
1654+
assert "The weather in Paris is sunny and 32°C" in tool_results
16711655

16721656

16731657
# In the CI, those tests are skipped if AWS Authentication fails

‎integrations/cohere/src/haystack_integrations/components/generators/cohere/chat/chat_generator.py‎

Lines changed: 7 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -460,13 +460,11 @@ class CohereChatGenerator:
460460
461461
### Advanced example
462462
463-
CohereChatGenerator can be integrated into pipelines and supports Haystack's tooling
464-
architecture, enabling tools to be invoked seamlessly across various generators.
463+
CohereChatGenerator supports Haystack's tooling architecture, so it can drive tool calls in an `Agent`.
465464
466465
```python
467-
from haystack import Pipeline
466+
from haystack.components.agents import Agent
468467
from haystack.dataclasses import ChatMessage
469-
from haystack.components.tools import ToolInvoker
470468
from haystack.tools import Tool
471469
from haystack_integrations.components.generators.cohere import CohereChatGenerator
472470
@@ -490,20 +488,12 @@ def weather(city: str) -> str:
490488
function=weather,
491489
)
492490
493-
# Create and set up the pipeline
494-
pipeline = Pipeline()
495-
pipeline.add_component("generator", CohereChatGenerator(tools=[weather_tool]))
496-
pipeline.add_component("tool_invoker", ToolInvoker(tools=[weather_tool]))
497-
pipeline.connect("generator", "tool_invoker")
491+
# The Agent calls the tool and feeds the result back to the model
492+
agent = Agent(chat_generator=CohereChatGenerator(), tools=[weather_tool])
493+
result = agent.run(messages=[ChatMessage.from_user("What's the weather like in Paris?")])
498494
499-
# Run the pipeline with a weather query
500-
results = pipeline.run(
501-
data={"generator": {"messages": [ChatMessage.from_user("What's the weather like in Paris?")]}}
502-
)
503-
504-
# The tool result will be available in the pipeline output
505-
print(results["tool_invoker"]["tool_messages"][0].tool_call_result.result)
506-
# Output: "The weather in Paris is sunny and 32°C"
495+
print(result["last_message"].text)
496+
# Output: "The weather in Paris is sunny and 32°C."
507497
```
508498
"""
509499

‎integrations/cohere/tests/test_chat_generator.py‎

Lines changed: 7 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -4,12 +4,8 @@
44
import pytest
55
from cohere.core import ApiError
66
from haystack import Pipeline
7+
from haystack.components.agents import Agent
78
from haystack.components.generators.utils import print_streaming_chunk
8-
9-
try:
10-
from haystack.components.tools import ToolInvoker
11-
except ImportError: # ToolInvoker was removed in Haystack 3.0
12-
ToolInvoker = None
139
from haystack.dataclasses import ChatMessage, ChatRole, ImageContent, ToolCall
1410
from haystack.dataclasses.streaming_chunk import StreamingChunk
1511
from haystack.tools import Tool, Toolset
@@ -810,10 +806,9 @@ def test_live_run_with_tools_streaming(self):
810806
assert len(final_message.text) > 0
811807
assert "paris" in final_message.text.lower()
812808

813-
@pytest.mark.skipif(ToolInvoker is None, reason="ToolInvoker is not available in the installed haystack-ai version")
814-
def test_pipeline_with_cohere_chat_generator(self):
809+
def test_agent_with_cohere_chat_generator(self):
815810
"""
816-
Test that the CohereChatGenerator component can be used in a pipeline
811+
Test that the CohereChatGenerator component can drive tool calls in an Agent
817812
"""
818813
weather_tool = Tool(
819814
name="weather",
@@ -831,20 +826,11 @@ def test_pipeline_with_cohere_chat_generator(self):
831826
function=weather,
832827
)
833828

834-
pipeline = Pipeline()
835-
pipeline.add_component("generator", CohereChatGenerator(model="command-r7b-12-2024", tools=[weather_tool]))
836-
pipeline.add_component("tool_invoker", ToolInvoker(tools=[weather_tool]))
837-
838-
pipeline.connect("generator", "tool_invoker")
829+
agent = Agent(chat_generator=CohereChatGenerator(model="command-r7b-12-2024"), tools=[weather_tool])
830+
result = agent.run(messages=[ChatMessage.from_user("What's the weather like in Paris?")])
839831

840-
results = pipeline.run(
841-
data={"generator": {"messages": [ChatMessage.from_user("What's the weather like in Paris?")]}}
842-
)
843-
844-
assert (
845-
"The weather in Paris is sunny and 32°C"
846-
== results["tool_invoker"]["tool_messages"][0].tool_call_result.result
847-
)
832+
tool_results = [msg.tool_call_result.result for msg in result["messages"] if msg.tool_call_result]
833+
assert "The weather in Paris is sunny and 32°C" in tool_results
848834

849835
def test_live_run_multimodal(self):
850836
generator = CohereChatGenerator(

‎integrations/cometapi/tests/test_cometapi_chat_generator.py‎

Lines changed: 7 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -5,12 +5,8 @@
55
import pytest
66
import pytz
77
from haystack import Pipeline
8+
from haystack.components.agents import Agent
89
from haystack.components.generators.utils import print_streaming_chunk
9-
10-
try:
11-
from haystack.components.tools import ToolInvoker
12-
except ImportError: # ToolInvoker was removed in Haystack 3.0
13-
ToolInvoker = None
1410
from haystack.dataclasses import ChatMessage, ChatRole, StreamingChunk, ToolCall
1511
from haystack.tools import Tool
1612
from haystack.utils.auth import Secret
@@ -385,30 +381,15 @@ def test_live_run_with_tools_streaming(self, tools):
385381
reason="Export an env var called COMET_API_KEY containing the OpenAI API key to run this test.",
386382
)
387383
@pytest.mark.integration
388-
@pytest.mark.skipif(ToolInvoker is None, reason="ToolInvoker is not available in the installed haystack-ai version")
389-
def test_pipeline_with_cometapi_chat_generator(self, tools):
384+
def test_agent_with_cometapi_chat_generator(self, tools):
390385
"""
391-
Test that the CometAPIChatGenerator component can be used in a pipeline
386+
Test that the CometAPIChatGenerator component can drive tool calls in an Agent
392387
"""
393-
pipeline = Pipeline()
394-
pipeline.add_component("generator", CometAPIChatGenerator(tools=tools))
395-
pipeline.add_component("tool_invoker", ToolInvoker(tools=tools))
396-
397-
pipeline.connect("generator", "tool_invoker")
388+
agent = Agent(chat_generator=CometAPIChatGenerator(), tools=tools)
389+
result = agent.run(messages=[ChatMessage.from_user("What's the weather like in Paris?")])
398390

399-
results = pipeline.run(
400-
data={
401-
"generator": {
402-
"messages": [ChatMessage.from_user("What's the weather like in Paris?")],
403-
"generation_kwargs": {"tool_choice": "auto"},
404-
}
405-
}
406-
)
407-
408-
assert (
409-
"The weather in Paris is sunny and 32°C"
410-
== results["tool_invoker"]["tool_messages"][0].tool_call_result.result
411-
)
391+
tool_results = [msg.tool_call_result.result for msg in result["messages"] if msg.tool_call_result]
392+
assert "The weather in Paris is sunny and 32°C" in tool_results
412393

413394
def test_serde_in_pipeline(self, monkeypatch):
414395
"""

‎integrations/llama_stack/examples/llama-stack-with-tools.py‎

Lines changed: 7 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -8,11 +8,11 @@
88
# set up Llama Stack Server through an inference provider and have a model available.
99
# For a quick start on how to setup server with Ollama, see [Llama Stack docs](https://llama-stack.readthedocs.io/en/latest/getting_started/index.html).
1010

11-
from haystack.components.tools import ToolInvoker
11+
from haystack.components.agents import Agent
1212
from haystack.dataclasses import ChatMessage
1313
from haystack.tools import Tool
1414

15-
from .components.generators.llama_stack.chat.chat_generator import LlamaStackChatGenerator
15+
from haystack_integrations.components.generators.llama_stack import LlamaStackChatGenerator
1616

1717

1818
# Define a tool that models can call
@@ -30,20 +30,14 @@ def weather(city: str):
3030
function=weather,
3131
)
3232

33-
# Create a tool invoker with the weather tool
34-
tool_invoker = ToolInvoker(tools=[weather_tool])
35-
3633
client = LlamaStackChatGenerator(
3734
model="ollama/llama3.2:3b", # model depends on the inference provider used for the Llama Stack Server.
3835
api_base_url="http://localhost:8321/v1",
3936
)
40-
messages = [ChatMessage.from_user("What's the weather in Tokyo?")]
41-
42-
response = client.run(messages=messages, tools=[weather_tool])["replies"]
4337

44-
print(f"assistant messages: {response[0]}\n")
38+
# The Agent calls the tool and feeds the result back to the model
39+
agent = Agent(chat_generator=client, tools=[weather_tool])
40+
result = agent.run(messages=[ChatMessage.from_user("What's the weather in Tokyo?")])
4541

46-
# If the assistant message contains a tool call, run the tool invoker
47-
if response[0].tool_calls:
48-
tool_messages = tool_invoker.run(messages=response)["tool_messages"]
49-
print(f"tool messages: {tool_messages}")
42+
print(f"messages: {result['messages']}\n")
43+
print(f"assistant final answer: {result['last_message'].text}")

‎integrations/openrouter/examples/openrouter_with_tools_example.py‎

Lines changed: 6 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
# To run this example, you will need to
88
# set `OPENROUTER_API_KEY` environment variable
99

10-
from haystack.components.tools import ToolInvoker
10+
from haystack.components.agents import Agent
1111
from haystack.dataclasses import ChatMessage
1212
from haystack.tools import Tool
1313

@@ -29,9 +29,6 @@ def weather(city: str):
2929
function=weather,
3030
)
3131

32-
# Create a tool invoker with the weather tool
33-
tool_invoker = ToolInvoker(tools=[weather_tool])
34-
3532
# We can setup model routing by setting the `model` parameter to `openrouter/auto`
3633
# and providing a list of models to route to.
3734
client = OpenRouterChatGenerator(
@@ -40,13 +37,10 @@ def weather(city: str):
4037
"models": ["openai/gpt-4o-mini", "anthropic/claude-3-haiku"],
4138
},
4239
)
43-
messages = [ChatMessage.from_user("What's the weather in Tokyo?")]
44-
45-
response = client.run(messages=messages, tools=[weather_tool])["replies"]
4640

47-
print(f"assistant messages: {response[0]}\n")
41+
# The Agent calls the tool and feeds the result back to the model
42+
agent = Agent(chat_generator=client, tools=[weather_tool])
43+
result = agent.run(messages=[ChatMessage.from_user("What's the weather in Tokyo?")])
4844

49-
# If the assistant message contains a tool call, run the tool invoker
50-
if response[0].tool_calls:
51-
tool_messages = tool_invoker.run(messages=response)["tool_messages"]
52-
print(f"tool messages: {tool_messages}")
45+
print(f"messages: {result['messages']}\n")
46+
print(f"assistant final answer: {result['last_message'].text}")

0 commit comments

Comments
 (0)