# SPDX-FileCopyrightText: 2022-present deepset GmbH # # SPDX-License-Identifier: Apache-2.0 import logging import os from datetime import datetime from typing import Iterator, Dict, Any, List, Optional, Union from unittest.mock import MagicMock, patch, AsyncMock import pytest from openai import Stream from openai.types.chat import ChatCompletionChunk, chat_completion_chunk from haystack.tracing.logging_tracer import LoggingTracer from haystack import Pipeline, tracing from haystack.components.agents import Agent from haystack.components.builders.chat_prompt_builder import ChatPromptBuilder from haystack.components.builders.prompt_builder import PromptBuilder from haystack.components.generators.chat.openai import OpenAIChatGenerator from haystack.components.generators.chat.types import ChatGenerator from haystack.core.component.types import OutputSocket from haystack.dataclasses import ChatMessage, ToolCall from haystack.dataclasses.chat_message import ChatRole, TextContent from haystack.dataclasses.streaming_chunk import StreamingChunk from haystack.tools import Tool, ComponentTool from haystack.tools.toolset import Toolset from haystack.utils import serialize_callable, Secret from haystack.components.agents.state import merge_lists def streaming_callback_for_serde(chunk: StreamingChunk): pass def weather_function(location): weather_info = { "Berlin": {"weather": "mostly sunny", "temperature": 7, "unit": "celsius"}, "Paris": {"weather": "mostly cloudy", "temperature": 8, "unit": "celsius"}, "Rome": {"weather": "sunny", "temperature": 14, "unit": "celsius"}, } return weather_info.get(location, {"weather": "unknown", "temperature": 0, "unit": "celsius"}) @pytest.fixture def weather_tool(): return Tool( name="weather_tool", description="Provides weather information for a given location.", parameters={"type": "object", "properties": {"location": {"type": "string"}}, "required": ["location"]}, function=weather_function, ) @pytest.fixture def component_tool(): return ComponentTool(name="parrot", description="This is a parrot.", component=PromptBuilder(template="{{parrot}}")) class OpenAIMockStream(Stream[ChatCompletionChunk]): def __init__(self, mock_chunk: ChatCompletionChunk, client=None, *args, **kwargs): client = client or MagicMock() super().__init__(client=client, *args, **kwargs) self.mock_chunk = mock_chunk def __stream__(self) -> Iterator[ChatCompletionChunk]: yield self.mock_chunk @pytest.fixture def openai_mock_chat_completion_chunk(): """ Mock the OpenAI API completion chunk response and reuse it for tests """ with patch("openai.resources.chat.completions.Completions.create") as mock_chat_completion_create: completion = ChatCompletionChunk( id="foo", model="gpt-4", object="chat.completion.chunk", choices=[ chat_completion_chunk.Choice( finish_reason="stop", logprobs=None, index=0, delta=chat_completion_chunk.ChoiceDelta(content="Hello", role="assistant"), ) ], created=int(datetime.now().timestamp()), usage=None, ) mock_chat_completion_create.return_value = OpenAIMockStream( completion, cast_to=None, response=None, client=None ) yield mock_chat_completion_create class MockChatGeneratorWithoutTools(ChatGenerator): """A mock chat generator that implements ChatGenerator protocol but doesn't support tools.""" __haystack_input__ = MagicMock(_sockets_dict={}) __haystack_output__ = MagicMock(_sockets_dict={}) def to_dict(self) -> Dict[str, Any]: return {"type": "MockChatGeneratorWithoutTools", "data": {}} @classmethod def from_dict(cls, data: Dict[str, Any]) -> "MockChatGeneratorWithoutTools": return cls() def run(self, messages: List[ChatMessage]) -> Dict[str, Any]: return {"replies": [ChatMessage.from_assistant("Hello")]} class MockChatGeneratorWithoutRunAsync(ChatGenerator): """A mock chat generator that implements ChatGenerator protocol but doesn't have run_async method.""" __haystack_input__ = MagicMock(_sockets_dict={}) __haystack_output__ = MagicMock(_sockets_dict={}) def to_dict(self) -> Dict[str, Any]: return {"type": "MockChatGeneratorWithoutRunAsync", "data": {}} @classmethod def from_dict(cls, data: Dict[str, Any]) -> "MockChatGeneratorWithoutRunAsync": return cls() def run( self, messages: List[ChatMessage], tools: Optional[Union[List[Tool], Toolset]] = None, **kwargs ) -> Dict[str, Any]: return {"replies": [ChatMessage.from_assistant("Hello")]} class MockChatGeneratorWithRunAsync(ChatGenerator): __haystack_supports_async__ = True __haystack_input__ = MagicMock(_sockets_dict={}) __haystack_output__ = MagicMock(_sockets_dict={}) def to_dict(self) -> Dict[str, Any]: return {"type": "MockChatGeneratorWithoutRunAsync", "data": {}} @classmethod def from_dict(cls, data: Dict[str, Any]) -> "MockChatGeneratorWithoutRunAsync": return cls() def run( self, messages: List[ChatMessage], tools: Optional[Union[List[Tool], Toolset]] = None, **kwargs ) -> Dict[str, Any]: return {"replies": [ChatMessage.from_assistant("Hello")]} async def run_async( self, messages: List[ChatMessage], tools: Optional[Union[List[Tool], Toolset]] = None, **kwargs ) -> Dict[str, Any]: return {"replies": [ChatMessage.from_assistant("Hello from run_async")]} class TestAgent: def test_output_types(self, weather_tool, component_tool, monkeypatch): monkeypatch.setenv("OPENAI_API_KEY", "fake-key") chat_generator = OpenAIChatGenerator() agent = Agent(chat_generator=chat_generator, tools=[weather_tool, component_tool]) assert agent.__haystack_output__._sockets_dict == { "messages": OutputSocket(name="messages", type=List[ChatMessage], receivers=[]), "last_message": OutputSocket(name="last_message", type=ChatMessage, receivers=[]), } def test_to_dict(self, weather_tool, component_tool, monkeypatch): monkeypatch.setenv("OPENAI_API_KEY", "fake-key") generator = OpenAIChatGenerator() agent = Agent( chat_generator=generator, tools=[weather_tool, component_tool], exit_conditions=["text", "weather_tool"], state_schema={"foo": {"type": str}}, ) serialized_agent = agent.to_dict() assert serialized_agent == { "type": "haystack.components.agents.agent.Agent", "init_parameters": { "chat_generator": { "type": "haystack.components.generators.chat.openai.OpenAIChatGenerator", "init_parameters": { "model": "gpt-4o-mini", "streaming_callback": None, "api_base_url": None, "organization": None, "generation_kwargs": {}, "api_key": {"type": "env_var", "env_vars": ["OPENAI_API_KEY"], "strict": True}, "timeout": None, "max_retries": None, "tools": None, "tools_strict": False, "http_client_kwargs": None, }, }, "tools": [ { "type": "haystack.tools.tool.Tool", "data": { "name": "weather_tool", "description": "Provides weather information for a given location.", "parameters": { "type": "object", "properties": {"location": {"type": "string"}}, "required": ["location"], }, "function": "test_agent.weather_function", "outputs_to_string": None, "inputs_from_state": None, "outputs_to_state": None, }, }, { "type": "haystack.tools.component_tool.ComponentTool", "data": { "component": { "type": "haystack.components.builders.prompt_builder.PromptBuilder", "init_parameters": { "template": "{{parrot}}", "variables": None, "required_variables": None, }, }, "name": "parrot", "description": "This is a parrot.", "parameters": None, "outputs_to_string": None, "inputs_from_state": None, "outputs_to_state": None, }, }, ], "system_prompt": None, "exit_conditions": ["text", "weather_tool"], "state_schema": {"foo": {"type": "str"}}, "max_agent_steps": 100, "raise_on_tool_invocation_failure": False, "streaming_callback": None, }, } def test_to_dict_with_toolset(self, monkeypatch, weather_tool): monkeypatch.setenv("OPENAI_API_KEY", "fake-key") toolset = Toolset(tools=[weather_tool]) agent = Agent(chat_generator=OpenAIChatGenerator(), tools=toolset) serialized_agent = agent.to_dict() assert serialized_agent == { "type": "haystack.components.agents.agent.Agent", "init_parameters": { "chat_generator": { "type": "haystack.components.generators.chat.openai.OpenAIChatGenerator", "init_parameters": { "model": "gpt-4o-mini", "streaming_callback": None, "api_base_url": None, "organization": None, "generation_kwargs": {}, "api_key": {"type": "env_var", "env_vars": ["OPENAI_API_KEY"], "strict": True}, "timeout": None, "max_retries": None, "tools": None, "tools_strict": False, "http_client_kwargs": None, }, }, "tools": { "type": "haystack.tools.toolset.Toolset", "data": { "tools": [ { "type": "haystack.tools.tool.Tool", "data": { "name": "weather_tool", "description": "Provides weather information for a given location.", "parameters": { "type": "object", "properties": {"location": {"type": "string"}}, "required": ["location"], }, "function": "test_agent.weather_function", "outputs_to_string": None, "inputs_from_state": None, "outputs_to_state": None, }, } ] }, }, "system_prompt": None, "exit_conditions": ["text"], "state_schema": {}, "max_agent_steps": 100, "raise_on_tool_invocation_failure": False, "streaming_callback": None, }, } def test_from_dict(self, monkeypatch): monkeypatch.setenv("OPENAI_API_KEY", "fake-key") data = { "type": "haystack.components.agents.agent.Agent", "init_parameters": { "chat_generator": { "type": "haystack.components.generators.chat.openai.OpenAIChatGenerator", "init_parameters": { "model": "gpt-4o-mini", "streaming_callback": None, "api_base_url": None, "organization": None, "generation_kwargs": {}, "api_key": {"type": "env_var", "env_vars": ["OPENAI_API_KEY"], "strict": True}, "timeout": None, "max_retries": None, "tools": None, "tools_strict": False, "http_client_kwargs": None, }, }, "tools": [ { "type": "haystack.tools.tool.Tool", "data": { "name": "weather_tool", "description": "Provides weather information for a given location.", "parameters": { "type": "object", "properties": {"location": {"type": "string"}}, "required": ["location"], }, "function": "test_agent.weather_function", "outputs_to_string": None, "inputs_from_state": None, "outputs_to_state": None, }, }, { "type": "haystack.tools.component_tool.ComponentTool", "data": { "component": { "type": "haystack.components.builders.prompt_builder.PromptBuilder", "init_parameters": { "template": "{{parrot}}", "variables": None, "required_variables": None, }, }, "name": "parrot", "description": "This is a parrot.", "parameters": None, "outputs_to_string": None, "inputs_from_state": None, "outputs_to_state": None, }, }, ], "system_prompt": None, "exit_conditions": ["text", "weather_tool"], "state_schema": {"foo": {"type": "str"}}, "max_agent_steps": 100, "raise_on_tool_invocation_failure": False, "streaming_callback": None, }, } agent = Agent.from_dict(data) assert isinstance(agent, Agent) assert isinstance(agent.chat_generator, OpenAIChatGenerator) assert agent.chat_generator.model == "gpt-4o-mini" assert agent.chat_generator.api_key == Secret.from_env_var("OPENAI_API_KEY") assert agent.tools[0].function is weather_function assert isinstance(agent.tools[1]._component, PromptBuilder) assert agent.exit_conditions == ["text", "weather_tool"] assert agent.state_schema == { "foo": {"type": str}, "messages": {"handler": merge_lists, "type": List[ChatMessage]}, } def test_from_dict_with_toolset(self, monkeypatch): monkeypatch.setenv("OPENAI_API_KEY", "fake-key") data = { "type": "haystack.components.agents.agent.Agent", "init_parameters": { "chat_generator": { "type": "haystack.components.generators.chat.openai.OpenAIChatGenerator", "init_parameters": { "model": "gpt-4o-mini", "streaming_callback": None, "api_base_url": None, "organization": None, "generation_kwargs": {}, "api_key": {"type": "env_var", "env_vars": ["OPENAI_API_KEY"], "strict": True}, "timeout": None, "max_retries": None, "tools": None, "tools_strict": False, "http_client_kwargs": None, }, }, "tools": { "type": "haystack.tools.toolset.Toolset", "data": { "tools": [ { "type": "haystack.tools.tool.Tool", "data": { "name": "weather_tool", "description": "Provides weather information for a given location.", "parameters": { "type": "object", "properties": {"location": {"type": "string"}}, "required": ["location"], }, "function": "test_agent.weather_function", "outputs_to_string": None, "inputs_from_state": None, "outputs_to_state": None, }, } ] }, }, "system_prompt": None, "exit_conditions": ["text"], "state_schema": {}, "max_agent_steps": 100, "raise_on_tool_invocation_failure": False, "streaming_callback": None, }, } agent = Agent.from_dict(data) assert isinstance(agent, Agent) assert isinstance(agent.chat_generator, OpenAIChatGenerator) assert agent.chat_generator.model == "gpt-4o-mini" assert agent.chat_generator.api_key == Secret.from_env_var("OPENAI_API_KEY") assert isinstance(agent.tools, Toolset) assert agent.tools[0].function is weather_function assert agent.exit_conditions == ["text"] def test_serde(self, weather_tool, component_tool, monkeypatch): monkeypatch.setenv("FAKE_OPENAI_KEY", "fake-key") generator = OpenAIChatGenerator(api_key=Secret.from_env_var("FAKE_OPENAI_KEY")) agent = Agent( chat_generator=generator, tools=[weather_tool, component_tool], exit_conditions=["text", "weather_tool"], state_schema={"foo": {"type": str}}, ) serialized_agent = agent.to_dict() init_parameters = serialized_agent["init_parameters"] assert serialized_agent["type"] == "haystack.components.agents.agent.Agent" assert ( init_parameters["chat_generator"]["type"] == "haystack.components.generators.chat.openai.OpenAIChatGenerator" ) assert init_parameters["streaming_callback"] is None assert init_parameters["tools"][0]["data"]["function"] == serialize_callable(weather_function) assert ( init_parameters["tools"][1]["data"]["component"]["type"] == "haystack.components.builders.prompt_builder.PromptBuilder" ) assert init_parameters["exit_conditions"] == ["text", "weather_tool"] deserialized_agent = Agent.from_dict(serialized_agent) assert isinstance(deserialized_agent, Agent) assert isinstance(deserialized_agent.chat_generator, OpenAIChatGenerator) assert deserialized_agent.tools[0].function is weather_function assert isinstance(deserialized_agent.tools[1]._component, PromptBuilder) assert deserialized_agent.exit_conditions == ["text", "weather_tool"] assert deserialized_agent.state_schema == { "foo": {"type": str}, "messages": {"handler": merge_lists, "type": List[ChatMessage]}, } def test_serde_with_streaming_callback(self, weather_tool, component_tool, monkeypatch): monkeypatch.setenv("FAKE_OPENAI_KEY", "fake-key") generator = OpenAIChatGenerator(api_key=Secret.from_env_var("FAKE_OPENAI_KEY")) agent = Agent( chat_generator=generator, tools=[weather_tool, component_tool], streaming_callback=streaming_callback_for_serde, ) serialized_agent = agent.to_dict() init_parameters = serialized_agent["init_parameters"] assert init_parameters["streaming_callback"] == "test_agent.streaming_callback_for_serde" deserialized_agent = Agent.from_dict(serialized_agent) assert deserialized_agent.streaming_callback is streaming_callback_for_serde def test_exit_conditions_validation(self, weather_tool, component_tool, monkeypatch): monkeypatch.setenv("FAKE_OPENAI_KEY", "fake-key") generator = OpenAIChatGenerator(api_key=Secret.from_env_var("FAKE_OPENAI_KEY")) # Test invalid exit condition with pytest.raises(ValueError, match="Invalid exit conditions provided:"): Agent(chat_generator=generator, tools=[weather_tool, component_tool], exit_conditions=["invalid_tool"]) # Test default exit condition agent = Agent(chat_generator=generator, tools=[weather_tool, component_tool]) assert agent.exit_conditions == ["text"] # Test multiple valid exit conditions agent = Agent( chat_generator=generator, tools=[weather_tool, component_tool], exit_conditions=["text", "weather_tool"] ) assert agent.exit_conditions == ["text", "weather_tool"] def test_run_with_params_streaming(self, openai_mock_chat_completion_chunk, weather_tool): chat_generator = OpenAIChatGenerator(api_key=Secret.from_token("test-api-key")) streaming_callback_called = False def streaming_callback(chunk: StreamingChunk) -> None: nonlocal streaming_callback_called streaming_callback_called = True agent = Agent(chat_generator=chat_generator, streaming_callback=streaming_callback, tools=[weather_tool]) agent.warm_up() response = agent.run([ChatMessage.from_user("Hello")]) # check we called the streaming callback assert streaming_callback_called is True # check that the component still returns the correct response assert isinstance(response, dict) assert "messages" in response assert isinstance(response["messages"], list) assert len(response["messages"]) == 2 assert [isinstance(reply, ChatMessage) for reply in response["messages"]] assert "Hello" in response["messages"][1].text # see openai_mock_chat_completion_chunk assert "last_message" in response assert isinstance(response["last_message"], ChatMessage) def test_run_with_run_streaming(self, openai_mock_chat_completion_chunk, weather_tool): chat_generator = OpenAIChatGenerator(api_key=Secret.from_token("test-api-key")) streaming_callback_called = False def streaming_callback(chunk: StreamingChunk) -> None: nonlocal streaming_callback_called streaming_callback_called = True agent = Agent(chat_generator=chat_generator, tools=[weather_tool]) agent.warm_up() response = agent.run([ChatMessage.from_user("Hello")], streaming_callback=streaming_callback) # check we called the streaming callback assert streaming_callback_called is True # check that the component still returns the correct response assert isinstance(response, dict) assert "messages" in response assert isinstance(response["messages"], list) assert len(response["messages"]) == 2 assert [isinstance(reply, ChatMessage) for reply in response["messages"]] assert "Hello" in response["messages"][1].text # see openai_mock_chat_completion_chunk assert "last_message" in response assert isinstance(response["last_message"], ChatMessage) def test_keep_generator_streaming(self, openai_mock_chat_completion_chunk, weather_tool): streaming_callback_called = False def streaming_callback(chunk: StreamingChunk) -> None: nonlocal streaming_callback_called streaming_callback_called = True chat_generator = OpenAIChatGenerator( api_key=Secret.from_token("test-api-key"), streaming_callback=streaming_callback ) agent = Agent(chat_generator=chat_generator, tools=[weather_tool]) agent.warm_up() response = agent.run([ChatMessage.from_user("Hello")]) # check we called the streaming callback assert streaming_callback_called is True # check that the component still returns the correct response assert isinstance(response, dict) assert "messages" in response assert isinstance(response["messages"], list) assert len(response["messages"]) == 2 assert [isinstance(reply, ChatMessage) for reply in response["messages"]] assert "Hello" in response["messages"][1].text # see openai_mock_chat_completion_chunk assert "last_message" in response assert isinstance(response["last_message"], ChatMessage) def test_chat_generator_must_support_tools(self, weather_tool): chat_generator = MockChatGeneratorWithoutTools() with pytest.raises(TypeError, match="MockChatGeneratorWithoutTools does not accept tools"): Agent(chat_generator=chat_generator, tools=[weather_tool]) def test_multiple_llm_responses_with_tool_call(self, monkeypatch, weather_tool): monkeypatch.setenv("OPENAI_API_KEY", "fake-key") generator = OpenAIChatGenerator() mock_messages = [ ChatMessage.from_assistant("First response"), ChatMessage.from_assistant( tool_calls=[ToolCall(tool_name="weather_tool", arguments={"location": "Berlin"})] ), ] agent = Agent(chat_generator=generator, tools=[weather_tool], max_agent_steps=1) agent.warm_up() # Patch agent.chat_generator.run to return mock_messages agent.chat_generator.run = MagicMock(return_value={"replies": mock_messages}) result = agent.run([ChatMessage.from_user("Hello")]) assert "messages" in result assert len(result["messages"]) == 4 assert ( result["messages"][-1].tool_call_result.result == "{'weather': 'mostly sunny', 'temperature': 7, 'unit': 'celsius'}" ) assert "last_message" in result assert isinstance(result["last_message"], ChatMessage) assert result["messages"][-1] == result["last_message"] def test_exceed_max_steps(self, monkeypatch, weather_tool, caplog): monkeypatch.setenv("OPENAI_API_KEY", "fake-key") generator = OpenAIChatGenerator() mock_messages = [ ChatMessage.from_assistant("First response"), ChatMessage.from_assistant( tool_calls=[ToolCall(tool_name="weather_tool", arguments={"location": "Berlin"})] ), ] agent = Agent(chat_generator=generator, tools=[weather_tool], max_agent_steps=0) agent.warm_up() # Patch agent.chat_generator.run to return mock_messages agent.chat_generator.run = MagicMock(return_value={"replies": mock_messages}) with caplog.at_level(logging.WARNING): agent.run([ChatMessage.from_user("Hello")]) assert "Agent reached maximum agent steps" in caplog.text def test_exit_conditions_checked_across_all_llm_messages(self, monkeypatch, weather_tool): monkeypatch.setenv("OPENAI_API_KEY", "fake-key") generator = OpenAIChatGenerator() # Mock messages where the exit condition appears in the second message mock_messages = [ ChatMessage.from_assistant("First response"), ChatMessage.from_assistant( tool_calls=[ToolCall(tool_name="weather_tool", arguments={"location": "Berlin"})] ), ] agent = Agent(chat_generator=generator, tools=[weather_tool], exit_conditions=["weather_tool"]) agent.warm_up() # Patch agent.chat_generator.run to return mock_messages agent.chat_generator.run = MagicMock(return_value={"replies": mock_messages}) result = agent.run([ChatMessage.from_user("Hello")]) assert "messages" in result assert len(result["messages"]) == 4 assert result["messages"][-2].tool_call.tool_name == "weather_tool" assert ( result["messages"][-1].tool_call_result.result == "{'weather': 'mostly sunny', 'temperature': 7, 'unit': 'celsius'}" ) assert "last_message" in result assert isinstance(result["last_message"], ChatMessage) assert result["messages"][-1] == result["last_message"] def test_agent_with_no_tools(self, monkeypatch, caplog): monkeypatch.setenv("OPENAI_API_KEY", "fake-key") generator = OpenAIChatGenerator() # Mock messages where the exit condition appears in the second message mock_messages = [ChatMessage.from_assistant("Berlin")] with caplog.at_level("WARNING"): agent = Agent(chat_generator=generator, tools=[], max_agent_steps=3) agent.warm_up() assert "No tools provided to the Agent." in caplog.text # Patch agent.chat_generator.run to return mock_messages agent.chat_generator.run = MagicMock(return_value={"replies": mock_messages}) response = agent.run([ChatMessage.from_user("What is the capital of Germany?")]) assert isinstance(response, dict) assert "messages" in response assert isinstance(response["messages"], list) assert len(response["messages"]) == 2 assert [isinstance(reply, ChatMessage) for reply in response["messages"]] assert response["messages"][0].text == "What is the capital of Germany?" assert response["messages"][1].text == "Berlin" assert "last_message" in response assert isinstance(response["last_message"], ChatMessage) assert response["messages"][-1] == response["last_message"] def test_run_with_system_prompt(self, weather_tool): chat_generator = MockChatGeneratorWithoutRunAsync() agent = Agent(chat_generator=chat_generator, tools=[weather_tool], system_prompt="This is a system prompt.") agent.warm_up() response = agent.run([ChatMessage.from_user("What is the weather in Berlin?")]) assert response["messages"][0].text == "This is a system prompt." def test_run_not_warmed_up(self, weather_tool): chat_generator = MockChatGeneratorWithoutRunAsync() chat_generator.warm_up = MagicMock() agent = Agent(chat_generator=chat_generator, tools=[weather_tool], system_prompt="This is a system prompt.") with pytest.raises(RuntimeError, match="The component Agent wasn't warmed up."): agent.run([ChatMessage.from_user("What is the weather in Berlin?")]) @pytest.mark.skipif(not os.environ.get("OPENAI_API_KEY"), reason="OPENAI_API_KEY not set") @pytest.mark.integration def test_run(self, weather_tool): chat_generator = OpenAIChatGenerator(model="gpt-4o-mini") agent = Agent(chat_generator=chat_generator, tools=[weather_tool], max_agent_steps=3) agent.warm_up() response = agent.run([ChatMessage.from_user("What is the weather in Berlin?")]) assert isinstance(response, dict) assert "messages" in response assert isinstance(response["messages"], list) assert len(response["messages"]) == 4 assert [isinstance(reply, ChatMessage) for reply in response["messages"]] # Loose check of message texts assert response["messages"][0].text == "What is the weather in Berlin?" assert response["messages"][1].text is None assert response["messages"][2].text is None assert response["messages"][3].text is not None # Loose check of message metadata assert response["messages"][0].meta == {} assert response["messages"][1].meta.get("model") is not None assert response["messages"][2].meta == {} assert response["messages"][3].meta.get("model") is not None # Loose check of tool calls and results assert response["messages"][1].tool_calls[0].tool_name == "weather_tool" assert response["messages"][1].tool_calls[0].arguments is not None assert response["messages"][2].tool_call_results[0].result is not None assert response["messages"][2].tool_call_results[0].origin is not None assert "last_message" in response assert isinstance(response["last_message"], ChatMessage) assert response["messages"][-1] == response["last_message"] @pytest.mark.asyncio async def test_run_async_falls_back_to_run_when_chat_generator_has_no_run_async(self, weather_tool): chat_generator = MockChatGeneratorWithoutRunAsync() agent = Agent(chat_generator=chat_generator, tools=[weather_tool]) agent.warm_up() chat_generator.run = MagicMock(return_value={"replies": [ChatMessage.from_assistant("Hello")]}) result = await agent.run_async([ChatMessage.from_user("Hello")]) expected_messages = [ ChatMessage(_role=ChatRole.USER, _content=[TextContent(text="Hello")], _name=None, _meta={}) ] chat_generator.run.assert_called_once_with(messages=expected_messages, tools=[weather_tool]) assert isinstance(result, dict) assert "messages" in result assert isinstance(result["messages"], list) assert len(result["messages"]) == 2 assert [isinstance(reply, ChatMessage) for reply in result["messages"]] assert "Hello" in result["messages"][1].text assert "last_message" in result assert isinstance(result["last_message"], ChatMessage) assert result["messages"][-1] == result["last_message"] @pytest.mark.asyncio async def test_run_async_uses_chat_generator_run_async_when_available(self, weather_tool): chat_generator = MockChatGeneratorWithRunAsync() agent = Agent(chat_generator=chat_generator, tools=[weather_tool]) agent.warm_up() chat_generator.run_async = AsyncMock( return_value={"replies": [ChatMessage.from_assistant("Hello from run_async")]} ) result = await agent.run_async([ChatMessage.from_user("Hello")]) expected_messages = [ ChatMessage(_role=ChatRole.USER, _content=[TextContent(text="Hello")], _name=None, _meta={}) ] chat_generator.run_async.assert_called_once_with(messages=expected_messages, tools=[weather_tool]) assert isinstance(result, dict) assert "messages" in result assert isinstance(result["messages"], list) assert len(result["messages"]) == 2 assert [isinstance(reply, ChatMessage) for reply in result["messages"]] assert "Hello from run_async" in result["messages"][1].text assert "last_message" in result assert isinstance(result["last_message"], ChatMessage) assert result["messages"][-1] == result["last_message"] @pytest.mark.integration @pytest.mark.skipif(not os.environ.get("OPENAI_API_KEY"), reason="OPENAI_API_KEY not set") def test_agent_streaming_with_tool_call(self, weather_tool): chat_generator = OpenAIChatGenerator() agent = Agent(chat_generator=chat_generator, tools=[weather_tool]) agent.warm_up() streaming_callback_called = False def streaming_callback(chunk: StreamingChunk) -> None: nonlocal streaming_callback_called streaming_callback_called = True result = agent.run( [ChatMessage.from_user("What's the weather in Paris?")], streaming_callback=streaming_callback ) assert result is not None assert result["messages"] is not None assert result["last_message"] is not None assert streaming_callback_called class TestAgentTracing: def test_agent_tracing_span_run(self, caplog, monkeypatch, weather_tool): chat_generator = MockChatGeneratorWithoutRunAsync() agent = Agent(chat_generator=chat_generator, tools=[weather_tool]) tracing.tracer.is_content_tracing_enabled = True tracing.enable_tracing(LoggingTracer()) caplog.set_level(logging.DEBUG) _ = agent.run([ChatMessage.from_user("What's the weather in Paris?")]) # Ensure tracing span was emitted assert any("Operation: haystack.component.run" in record.message for record in caplog.records) # Check specific tags tags_records = [r for r in caplog.records if hasattr(r, "tag_name")] expected_tag_names = [ "haystack.component.name", "haystack.component.type", "haystack.component.input_types", "haystack.component.input_spec", "haystack.component.output_spec", "haystack.component.input", "haystack.component.visits", "haystack.component.output", "haystack.agent.max_steps", "haystack.agent.tools", "haystack.agent.exit_conditions", "haystack.agent.state_schema", "haystack.agent.input", "haystack.agent.output", "haystack.agent.steps_taken", ] expected_tag_values = [ "chat_generator", "MockChatGeneratorWithoutRunAsync", '{"messages": "list", "tools": "list"}', "{}", "{}", '{"messages": [{"role": "user", "meta": {}, "name": null, "content": [{"text": "What\'s the weather in Paris?"}]}], "tools": [{"type": "haystack.tools.tool.Tool", "data": {"name": "weather_tool", "description": "Provides weather information for a given location.", "parameters": {"type": "object", "properties": {"location": {"type": "string"}}, "required": ["location"]}, "function": "test_agent.weather_function", "outputs_to_string": null, "inputs_from_state": null, "outputs_to_state": null}}]}', 1, '{"replies": [{"role": "assistant", "meta": {}, "name": null, "content": [{"text": "Hello"}]}]}', 100, '[{"type": "haystack.tools.tool.Tool", "data": {"name": "weather_tool", "description": "Provides weather information for a given location.", "parameters": {"type": "object", "properties": {"location": {"type": "string"}}, "required": ["location"]}, "function": "test_agent.weather_function", "outputs_to_string": null, "inputs_from_state": null, "outputs_to_state": null}}]', '["text"]', '{"messages": {"type": "typing.List[haystack.dataclasses.chat_message.ChatMessage]", "handler": "haystack.components.agents.state.state_utils.merge_lists"}}', '{"messages": [{"role": "user", "meta": {}, "name": null, "content": [{"text": "What\'s the weather in Paris?"}]}], "streaming_callback": null}', '{"messages": [{"role": "user", "meta": {}, "name": null, "content": [{"text": "What\'s the weather in Paris?"}]}, {"role": "assistant", "meta": {}, "name": null, "content": [{"text": "Hello"}]}]}', 1, ] for idx, record in enumerate(tags_records): assert record.tag_name == expected_tag_names[idx] assert record.tag_value == expected_tag_values[idx] # Clean up tracing.tracer.is_content_tracing_enabled = False tracing.disable_tracing() @pytest.mark.asyncio async def test_agent_tracing_span_async_run(self, caplog, monkeypatch, weather_tool): chat_generator = MockChatGeneratorWithRunAsync() agent = Agent(chat_generator=chat_generator, tools=[weather_tool]) tracing.tracer.is_content_tracing_enabled = True tracing.enable_tracing(LoggingTracer()) caplog.set_level(logging.DEBUG) _ = await agent.run_async([ChatMessage.from_user("What's the weather in Paris?")]) # Ensure tracing span was emitted assert any("Operation: haystack.component.run" in record.message for record in caplog.records) # Check specific tags tags_records = [r for r in caplog.records if hasattr(r, "tag_name")] expected_tag_names = [ "haystack.component.name", "haystack.component.type", "haystack.component.input_types", "haystack.component.input_spec", "haystack.component.output_spec", "haystack.component.input", "haystack.component.visits", "haystack.component.output", "haystack.agent.max_steps", "haystack.agent.tools", "haystack.agent.exit_conditions", "haystack.agent.state_schema", "haystack.agent.input", "haystack.agent.output", "haystack.agent.steps_taken", ] expected_tag_values = [ "chat_generator", "MockChatGeneratorWithRunAsync", '{"messages": "list", "tools": "list"}', "{}", "{}", '{"messages": [{"role": "user", "meta": {}, "name": null, "content": [{"text": "What\'s the weather in Paris?"}]}], "tools": [{"type": "haystack.tools.tool.Tool", "data": {"name": "weather_tool", "description": "Provides weather information for a given location.", "parameters": {"type": "object", "properties": {"location": {"type": "string"}}, "required": ["location"]}, "function": "test_agent.weather_function", "outputs_to_string": null, "inputs_from_state": null, "outputs_to_state": null}}]}', 1, '{"replies": [{"role": "assistant", "meta": {}, "name": null, "content": [{"text": "Hello from run_async"}]}]}', 100, '[{"type": "haystack.tools.tool.Tool", "data": {"name": "weather_tool", "description": "Provides weather information for a given location.", "parameters": {"type": "object", "properties": {"location": {"type": "string"}}, "required": ["location"]}, "function": "test_agent.weather_function", "outputs_to_string": null, "inputs_from_state": null, "outputs_to_state": null}}]', '["text"]', '{"messages": {"type": "typing.List[haystack.dataclasses.chat_message.ChatMessage]", "handler": "haystack.components.agents.state.state_utils.merge_lists"}}', '{"messages": [{"role": "user", "meta": {}, "name": null, "content": [{"text": "What\'s the weather in Paris?"}]}], "streaming_callback": null}', '{"messages": [{"role": "user", "meta": {}, "name": null, "content": [{"text": "What\'s the weather in Paris?"}]}, {"role": "assistant", "meta": {}, "name": null, "content": [{"text": "Hello from run_async"}]}]}', 1, ] for idx, record in enumerate(tags_records): assert record.tag_name == expected_tag_names[idx] assert record.tag_value == expected_tag_values[idx] # Clean up tracing.tracer.is_content_tracing_enabled = False tracing.disable_tracing() def test_agent_tracing_in_pipeline(self, caplog, monkeypatch, weather_tool): chat_generator = MockChatGeneratorWithoutRunAsync() agent = Agent(chat_generator=chat_generator, tools=[weather_tool]) agent.warm_up() tracing.tracer.is_content_tracing_enabled = True tracing.enable_tracing(LoggingTracer()) caplog.set_level(logging.DEBUG) pipeline = Pipeline() pipeline.add_component( "prompt_builder", ChatPromptBuilder(template=[ChatMessage.from_user("Hello {{location}}")]) ) pipeline.add_component("agent", agent) pipeline.connect("prompt_builder.prompt", "agent.messages") pipeline.run(data={"prompt_builder": {"location": "Berlin"}}) assert any("Operation: haystack.pipeline.run" in record.message for record in caplog.records) tags_records = [r for r in caplog.records if hasattr(r, "tag_name")] expected_tag_names = [ "haystack.component.name", "haystack.component.type", "haystack.component.input_types", "haystack.component.input_spec", "haystack.component.output_spec", "haystack.component.input", "haystack.component.visits", "haystack.component.output", "haystack.component.name", "haystack.component.type", "haystack.component.input_types", "haystack.component.input_spec", "haystack.component.output_spec", "haystack.component.input", "haystack.component.visits", "haystack.component.output", "haystack.agent.max_steps", "haystack.agent.tools", "haystack.agent.exit_conditions", "haystack.agent.state_schema", "haystack.agent.input", "haystack.agent.output", "haystack.agent.steps_taken", "haystack.component.name", "haystack.component.type", "haystack.component.input_types", "haystack.component.input_spec", "haystack.component.output_spec", "haystack.component.input", "haystack.component.visits", "haystack.component.output", "haystack.pipeline.input_data", "haystack.pipeline.output_data", "haystack.pipeline.metadata", "haystack.pipeline.max_runs_per_component", ] for idx, record in enumerate(tags_records): assert record.tag_name == expected_tag_names[idx] # Clean up tracing.tracer.is_content_tracing_enabled = False tracing.disable_tracing()