deepset-ai/haystack
1
0
Fork 0
mirror of https://github.com/deepset-ai/haystack.git synced 2026-10-12 02:12:50 +00:00
Code Issues Packages Projects Releases Wiki Activity

test: adopt MockChatGenerator in tests (#11994)

Browse Source 
Co-authored-by: David S. Batista <[email protected]>
This commit is contained in:
Stefano Fiorucci
2026-07-14 10:42:28 +02:00
committed by GitHub
co-authored by David S. Batista
parent bffe6ce8c2
commit af7ec4a48b
14 changed files with 215 additions and 443 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/components/agents/test_agent.py
+103 -175
View File
@@ -21,6 +21,7 @@ from haystack.components.agents.state import merge_lists, replace_values
from haystack.components.agents.tool_calling import _run_tool
from haystack.components.builders.chat_prompt_builder import ChatPromptBuilder
from haystack.components.builders.prompt_builder import PromptBuilder
from haystack.components.generators.chat import MockChatGenerator
from haystack.components.generators.chat.openai import OpenAIChatGenerator
from haystack.components.joiners.branch import BranchJoiner
from haystack.components.joiners.list_joiner import ListJoiner
@@ -101,7 +102,7 @@ def component_tool():
@pytest.fixture
def make_agent(weather_tool):
def _factory(**kwargs):
return Agent(chat_generator=MockChatGenerator(), tools=[weather_tool], **kwargs)
return Agent(chat_generator=MockChatGenerator("Hello"), tools=[weather_tool], **kwargs)
return _factory
@@ -177,46 +178,37 @@ class MockChatGeneratorWithoutRunAsync:
@component
class MockChatGenerator:
def to_dict(self) -> dict[str, Any]:
return {"type": "MockChatGeneratorWithoutRunAsync", "data": {}}
class ToolAssertingChatGenerator:
"""Asserts the Agent forwards the expected tools, then drives one tool call before a plain reply."""
@classmethod
def from_dict(cls, data: dict[str, Any]) -> "MockChatGenerator":
return cls()
def __init__(self, expected_tools):
self.expected_tools = expected_tools
self.tool_invoked = False
@component.output_types(replies=list[ChatMessage])
def run(self, messages: list[ChatMessage], tools: list[Tool] | Toolset | None = None, **kwargs) -> dict[str, Any]:
return {"replies": [ChatMessage.from_assistant("Hello")]}
@component.output_types(replies=list[ChatMessage])
async def run_async(
self, messages: list[ChatMessage], tools: list[Tool] | Toolset | None = None, **kwargs
) -> dict[str, Any]:
return {"replies": [ChatMessage.from_assistant("Hello from run_async")]}
@component
class ParallelToolCallingChatGenerator:
"""Requests two `weather_tool` calls on the first turn, then returns a plain reply so the agent loop exits."""
tool_invoked = False
@component.output_types(replies=list[ChatMessage])
def run(self, messages: list[ChatMessage], tools: list[Tool] | Toolset | None = None, **kwargs) -> dict[str, Any]:
if self.tool_invoked:
return {"replies": [ChatMessage.from_assistant("done")]}
assert tools == self.expected_tools
tool_message = ChatMessage.from_assistant(
tool_calls=[ToolCall(tool_name="weather_tool", arguments={"location": "Berlin"})]
)
message = tool_message if not self.tool_invoked else ChatMessage.from_assistant("Hello")
self.tool_invoked = True
return {
"replies": [
ChatMessage.from_assistant(
tool_calls=[
ToolCall(tool_name="weather_tool", arguments={"location": "Berlin"}),
ToolCall(tool_name="weather_tool", arguments={"location": "Paris"}),
]
)
]
}
return {"replies": [message]}
def _parallel_tool_calling_generator() -> MockChatGenerator:
"""Requests two `weather_tool` calls on the first turn, then returns a plain reply so the agent loop exits."""
return MockChatGenerator(
[
ChatMessage.from_assistant(
tool_calls=[
ToolCall(tool_name="weather_tool", arguments={"location": "Berlin"}),
ToolCall(tool_name="weather_tool", arguments={"location": "Paris"}),
]
),
"done",
]
)
class TestAgent:
@@ -791,21 +783,13 @@ class TestAgent:
assert "Agent reached maximum agent steps" in caplog.text
assert result["step_count"] == 0
def test_exit_condition_exits(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(
tool_calls=[ToolCall(tool_name="weather_tool", arguments={"location": "Berlin"})]
)
]
agent = Agent(chat_generator=generator, tools=[weather_tool], exit_conditions=["weather_tool"])
# Patch agent.chat_generator.run to return mock_messages
agent.chat_generator.run = MagicMock(return_value={"replies": mock_messages})
def test_exit_condition_exits(self, weather_tool):
tool_call_message = ChatMessage.from_assistant(
tool_calls=[ToolCall(tool_name="weather_tool", arguments={"location": "Berlin"})]
)
agent = Agent(
chat_generator=MockChatGenerator(tool_call_message), tools=[weather_tool], exit_conditions=["weather_tool"]
)
result = agent.run([ChatMessage.from_user("Hello")])
@@ -820,22 +804,19 @@ class TestAgent:
assert isinstance(result["last_message"], ChatMessage)
assert result["messages"][-1] == result["last_message"]
def test_exit_condition_on_tool_provided_at_runtime(self, monkeypatch, weather_tool):
def test_exit_condition_on_tool_provided_at_runtime(self, weather_tool):
"""An exit condition naming a tool absent at init still triggers once that tool is provided at runtime."""
monkeypatch.setenv("OPENAI_API_KEY", "fake-key")
# weather_tool is NOT among the init tools, but it is named as an exit condition.
# weather_tool is NOT among the init tools, but it is named as an exit condition. The model calls
# weather_tool, which is supplied only at runtime.
tool_call_message = ChatMessage.from_assistant(
tool_calls=[ToolCall(tool_name="weather_tool", arguments={"location": "Berlin"})]
)
agent = Agent(
chat_generator=OpenAIChatGenerator(), tools=[], exit_conditions=["weather_tool"], max_agent_steps=5
chat_generator=MockChatGenerator(tool_call_message),
tools=[],
exit_conditions=["weather_tool"],
max_agent_steps=5,
)
# The model calls weather_tool, which is supplied only at runtime.
mock_messages = [
ChatMessage.from_assistant(
tool_calls=[ToolCall(tool_name="weather_tool", arguments={"location": "Berlin"})]
)
]
agent.chat_generator.run = MagicMock(return_value={"replies": mock_messages})
result = agent.run([ChatMessage.from_user("What's the weather in Berlin?")], tools=[weather_tool])
@@ -911,17 +892,8 @@ class TestAgent:
assert agent._check_exit_conditions(llm_messages, tool_messages) is True
def test_agent_with_no_tools(self, monkeypatch):
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")]
agent = Agent(chat_generator=generator, tools=[], max_agent_steps=3)
# Patch agent.chat_generator.run to return mock_messages
agent.chat_generator.run = MagicMock(return_value={"replies": mock_messages})
def test_agent_with_no_tools(self):
agent = Agent(chat_generator=MockChatGenerator("Berlin"), tools=[], max_agent_steps=3)
response = agent.run([ChatMessage.from_user("What is the capital of Germany?")])
@@ -943,23 +915,7 @@ class TestAgent:
assert response["messages"][0].text == "This is a system prompt."
def test_run_with_tools_run_param(self, weather_tool: Tool, component_tool: Tool, monkeypatch):
@component
class MockChatGenerator:
tool_invoked = False
@component.output_types(replies=list[ChatMessage])
def run(
self, messages: list[ChatMessage], tools: list[Tool] | Toolset | None = None, **kwargs
) -> dict[str, Any]:
assert tools == [weather_tool]
tool_message = ChatMessage.from_assistant(
tool_calls=[ToolCall(tool_name="weather_tool", arguments={"location": "Berlin"})]
)
message = tool_message if not self.tool_invoked else ChatMessage.from_assistant("Hello")
self.tool_invoked = True
return {"replies": [message]}
chat_generator = MockChatGenerator()
chat_generator = ToolAssertingChatGenerator(expected_tools=[weather_tool])
agent = Agent(
chat_generator=chat_generator,
tools=[component_tool],
@@ -975,23 +931,7 @@ class TestAgent:
assert run_tool_mock.call_args.kwargs["enable_streaming_callback_passthrough"] is True
def test_run_with_tools_run_param_for_tool_selection(self, weather_tool: Tool, component_tool: Tool, monkeypatch):
@component
class MockChatGenerator:
tool_invoked = False
@component.output_types(replies=list[ChatMessage])
def run(
self, messages: list[ChatMessage], tools: list[Tool] | Toolset | None = None, **kwargs
) -> dict[str, Any]:
assert tools == [weather_tool]
tool_message = ChatMessage.from_assistant(
tool_calls=[ToolCall(tool_name="weather_tool", arguments={"location": "Berlin"})]
)
message = tool_message if not self.tool_invoked else ChatMessage.from_assistant("Hello")
self.tool_invoked = True
return {"replies": [message]}
chat_generator = MockChatGenerator()
chat_generator = ToolAssertingChatGenerator(expected_tools=[weather_tool])
agent = Agent(
chat_generator=chat_generator,
tools=[weather_tool, component_tool],
@@ -1055,7 +995,7 @@ class TestAgent:
@pytest.mark.asyncio
async def test_generation_kwargs(self):
chat_generator = MockChatGenerator()
chat_generator = MockChatGenerator("Hello")
agent = Agent(chat_generator=chat_generator)
@@ -1073,7 +1013,7 @@ class TestAgent:
@pytest.mark.asyncio
async def test_run_async_uses_chat_generator_run_async_when_available(self, weather_tool):
chat_generator = MockChatGenerator()
chat_generator = MockChatGenerator("Hello")
agent = Agent(chat_generator=chat_generator, tools=[weather_tool])
chat_generator.run_async = AsyncMock(
@@ -1163,7 +1103,7 @@ class TestAgent:
@pytest.mark.asyncio
async def test_run_async_with_async_streaming_callback(self, weather_tool):
chat_generator = MockChatGenerator()
chat_generator = MockChatGenerator("Hello")
agent = Agent(chat_generator=chat_generator, tools=[weather_tool], streaming_callback=async_streaming_callback)
# This should not raise any exception
@@ -1171,10 +1111,10 @@ class TestAgent:
assert "messages" in result
assert len(result["messages"]) == 2
assert result["messages"][1].text == "Hello from run_async"
assert result["messages"][1].text == "Hello"
def test_run_with_async_streaming_callback_fails(self, weather_tool):
chat_generator = MockChatGenerator()
chat_generator = MockChatGenerator("Hello")
agent = Agent(chat_generator=chat_generator, tools=[weather_tool], streaming_callback=async_streaming_callback)
with pytest.raises(ValueError, match="The init callback cannot be a coroutine"):
@@ -1182,7 +1122,7 @@ class TestAgent:
@pytest.mark.asyncio
async def test_run_async_with_sync_streaming_callback_warns(self, weather_tool, caplog):
chat_generator = MockChatGenerator()
chat_generator = MockChatGenerator("Hello")
agent = Agent(chat_generator=chat_generator, tools=[weather_tool], streaming_callback=sync_streaming_callback)
with caplog.at_level(logging.WARNING):
@@ -1200,10 +1140,8 @@ class TestAgent:
chat_generator=OpenAIChatGenerator(), tools=[weather_tool], state_schema={reserved: {"type": int}}
)
def test_run_populates_token_usage_and_tool_call_counts(self, monkeypatch, weather_tool, component_tool):
def test_run_populates_token_usage_and_tool_call_counts(self, weather_tool, component_tool):
"""A multi-step run aggregates step_count, token_usage (incl. nested details), and tool_call_counts."""
monkeypatch.setenv("OPENAI_API_KEY", "fake-key")
agent = Agent(chat_generator=OpenAIChatGenerator(), tools=[weather_tool, component_tool])
# Step 1: two parallel tool calls + usage with nested detail dicts.
# Step 2: one more weather_tool call + flat usage.
# Step 3: final text answer + usage.
@@ -1238,8 +1176,9 @@ class TestAgent:
},
)
]
agent.chat_generator.run = MagicMock(
side_effect=[{"replies": first_step}, {"replies": second_step}, {"replies": third_step}]
agent = Agent(
chat_generator=MockChatGenerator(first_step + second_step + third_step),
tools=[weather_tool, component_tool],
)
result = agent.run([ChatMessage.from_user("Hi")])
@@ -1253,9 +1192,7 @@ class TestAgent:
}
@pytest.mark.asyncio
async def test_run_async_populates_token_usage_and_tool_call_counts(self, monkeypatch, weather_tool):
monkeypatch.setenv("OPENAI_API_KEY", "fake-key")
agent = Agent(chat_generator=OpenAIChatGenerator(), tools=[weather_tool])
async def test_run_async_populates_token_usage_and_tool_call_counts(self, weather_tool):
first_step = [
_assistant_with_usage(
tool_calls=[ToolCall(tool_name="weather_tool", arguments={"location": "Berlin"})],
@@ -1265,7 +1202,7 @@ class TestAgent:
second_step = [
_assistant_with_usage("Done.", usage={"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4})
]
agent.chat_generator.run_async = AsyncMock(side_effect=[{"replies": first_step}, {"replies": second_step}])
agent = Agent(chat_generator=MockChatGenerator(first_step + second_step), tools=[weather_tool])
result = await agent.run_async([ChatMessage.from_user("Hi")])
assert result["step_count"] == 2
@@ -1274,9 +1211,10 @@ class TestAgent:
def test_metadata_outputs_show_defaults_when_no_data(self, weather_tool):
"""`token_usage` stays empty and `tool_call_counts` reports zero for every tool when nothing happens."""
agent = Agent(chat_generator=MockChatGenerator(), tools=[weather_tool])
# A text-only reply whose `usage` meta is empty leaves `token_usage` empty after aggregation.
chat_generator = MockChatGenerator(ChatMessage.from_assistant("Hello", meta={"usage": {}}))
agent = Agent(chat_generator=chat_generator, tools=[weather_tool])
result = agent.run([ChatMessage.from_user("Hi")])
# MockChatGenerator returns a text-only reply with no `usage` meta and no tool calls.
assert result["step_count"] == 1
assert result["token_usage"] == {}
assert result["tool_call_counts"] == {"weather_tool": 0}
@@ -1348,7 +1286,7 @@ class TestAgentTracing:
elif hasattr(record, "tag_name") and spans:
spans[-1][1][record.tag_name] = record.tag_value
# Keep only the agent's own spans. With the MockChatGenerator returning no tool calls, the inner
# Keep only the agent's own spans. With the chat generator returning no tool calls, the inner
# `haystack.agent.step.tool` span never fires - the loop exits after the LLM call.
agent_spans = [(op, tags) for op, tags in spans if op.startswith("haystack.agent")]
@@ -1415,26 +1353,15 @@ class TestAgentTracing:
tracing.disable_tracing()
def test_agent_tracing_span_run_with_tool_call(self, caplog, monkeypatch, weather_tool):
@component
class ToolCallingChatGenerator:
tool_invoked = False
@component.output_types(replies=list[ChatMessage])
def run(
self, messages: list[ChatMessage], tools: list[Tool] | Toolset | None = None, **kwargs
) -> dict[str, Any]:
if self.tool_invoked:
return {"replies": [ChatMessage.from_assistant("done")]}
self.tool_invoked = True
return {
"replies": [
ChatMessage.from_assistant(
tool_calls=[ToolCall(tool_name="weather_tool", arguments={"location": "Berlin"})]
)
]
}
agent = Agent(chat_generator=ToolCallingChatGenerator(), tools=[weather_tool])
chat_generator = MockChatGenerator(
[
ChatMessage.from_assistant(
tool_calls=[ToolCall(tool_name="weather_tool", arguments={"location": "Berlin"})]
),
"done",
]
)
agent = Agent(chat_generator=chat_generator, tools=[weather_tool])
tracing.tracer.is_content_tracing_enabled = True
tracing.enable_tracing(LoggingTracer())
@@ -1478,7 +1405,7 @@ class TestAgentTracing:
def test_agent_tracing_span_run_with_parallel_tool_calls(self, caplog, monkeypatch, weather_tool):
"""Each tool call in a step gets its own `haystack.agent.step.tool` span instead of one grouped span."""
agent = Agent(chat_generator=ParallelToolCallingChatGenerator(), tools=[weather_tool])
agent = Agent(chat_generator=_parallel_tool_calling_generator(), tools=[weather_tool])
tracing.tracer.is_content_tracing_enabled = True
tracing.enable_tracing(LoggingTracer())
@@ -1515,7 +1442,7 @@ class TestAgentTracing:
@pytest.mark.asyncio
async def test_agent_tracing_span_async_run(self, caplog, monkeypatch, weather_tool):
chat_generator = MockChatGenerator()
chat_generator = MockChatGenerator("Hello")
agent = Agent(chat_generator=chat_generator, tools=[weather_tool])
tracing.tracer.is_content_tracing_enabled = True
@@ -1544,7 +1471,7 @@ class TestAgentTracing:
)
assert (
llm_tags["haystack.agent.step.llm.output"]
== '{"replies": [{"role": "assistant", "meta": {}, "name": null, "content": [{"text": "Hello from run_async"}]}]}' # noqa: E501
== '{"replies": [{"role": "assistant", "meta": {"model": "mock-model", "index": 0, "finish_reason": "stop", "usage": {"prompt_tokens": 5, "completion_tokens": 1, "total_tokens": 6}}, "name": null, "content": [{"text": "Hello"}]}]}' # noqa: E501
)
_, step_tags = agent_spans[1]
@@ -1557,7 +1484,7 @@ class TestAgentTracing:
"haystack.agent.exit_conditions": '["text"]',
"haystack.agent.state_schema": '{"messages": {"type": "list[haystack.dataclasses.chat_message.ChatMessage]", "handler": "haystack.components.agents.state.state_utils.merge_lists"}, "step_count": {"type": "int", "handler": "haystack.components.agents.state.state_utils.replace_values"}, "token_usage": {"type": "dict[str, typing.Any]", "handler": "haystack.components.agents.state.state_utils.replace_values"}, "tool_call_counts": {"type": "dict[str, int]", "handler": "haystack.components.agents.state.state_utils.replace_values"}, "continue_run": {"type": "bool", "handler": "haystack.components.agents.state.state_utils.replace_values"}, "tools": {"type": "list", "handler": "haystack.components.agents.state.state_utils.replace_values"}, "hook_context": {"type": "dict[str, typing.Any]", "handler": "haystack.components.agents.state.state_utils.replace_values"}}', # noqa: E501
"haystack.agent.input": '{"messages": [{"role": "user", "meta": {}, "name": null, "content": [{"text": "What\'s the weather in Paris?"}]}], "streaming_callback": null}', # noqa: E501
"haystack.agent.output": '{"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"}]}], "step_count": 1, "token_usage": {}, "tool_call_counts": {"weather_tool": 0}, "last_message": {"role": "assistant", "meta": {}, "name": null, "content": [{"text": "Hello from run_async"}]}}', # noqa: E501
"haystack.agent.output": '{"messages": [{"role": "user", "meta": {}, "name": null, "content": [{"text": "What\'s the weather in Paris?"}]}, {"role": "assistant", "meta": {"model": "mock-model", "index": 0, "finish_reason": "stop", "usage": {"prompt_tokens": 5, "completion_tokens": 1, "total_tokens": 6}}, "name": null, "content": [{"text": "Hello"}]}], "step_count": 1, "token_usage": {"prompt_tokens": 5, "completion_tokens": 1, "total_tokens": 6}, "tool_call_counts": {"weather_tool": 0}, "last_message": {"role": "assistant", "meta": {"model": "mock-model", "index": 0, "finish_reason": "stop", "usage": {"prompt_tokens": 5, "completion_tokens": 1, "total_tokens": 6}}, "name": null, "content": [{"text": "Hello"}]}}', # noqa: E501
"haystack.agent.steps_taken": 1,
}
@@ -1568,7 +1495,7 @@ class TestAgentTracing:
@pytest.mark.asyncio
async def test_agent_tracing_span_async_run_with_parallel_tool_calls(self, caplog, monkeypatch, weather_tool):
"""The async path also emits one `haystack.agent.step.tool` span per tool call."""
agent = Agent(chat_generator=ParallelToolCallingChatGenerator(), tools=[weather_tool])
agent = Agent(chat_generator=_parallel_tool_calling_generator(), tools=[weather_tool])
tracing.tracer.is_content_tracing_enabled = True
tracing.enable_tracing(LoggingTracer())
@@ -1756,13 +1683,13 @@ class TestSelectToolsByName:
class TestAgentToolSelection:
def test_tool_selection_new_tool(self, weather_tool: Tool, component_tool: Tool):
chat_generator = MockChatGenerator()
chat_generator = MockChatGenerator("Hello")
agent = Agent(chat_generator=chat_generator, tools=[weather_tool], system_prompt="This is a system prompt.")
result = agent._select_tools([component_tool])
assert result == [component_tool]
def test_tool_selection_existing_tools(self, weather_tool: Tool, component_tool: Tool):
chat_generator = MockChatGenerator()
chat_generator = MockChatGenerator("Hello")
agent = Agent(
chat_generator=chat_generator,
tools=[weather_tool, component_tool],
@@ -1772,7 +1699,7 @@ class TestAgentToolSelection:
assert result == [weather_tool, component_tool]
def test_tool_selection_invalid_type(self, weather_tool: Tool, component_tool: Tool):
chat_generator = MockChatGenerator()
chat_generator = MockChatGenerator("Hello")
agent = Agent(
chat_generator=chat_generator,
tools=[weather_tool, component_tool],
@@ -1790,7 +1717,7 @@ class TestAgentToolSelection:
def test_tool_selection_with_list_of_toolsets(self, weather_tool: Tool, component_tool: Tool):
"""Test that list of Toolsets and Tools can be passed to agent."""
chat_generator = MockChatGenerator()
chat_generator = MockChatGenerator("Hello")
toolset1 = Toolset([weather_tool])
standalone_tool = Tool(
name="standalone",
@@ -2054,7 +1981,7 @@ class TestAgentUserPromptInPipeline:
assert "Documents:" in rendered
def test_rag_pipeline_messages_plus_user_prompt(self, document_store_with_docs, weather_tool):
chat_generator = MockChatGenerator()
chat_generator = MockChatGenerator("Hello")
agent = Agent(
chat_generator=chat_generator,
@@ -2142,7 +2069,7 @@ class TestAgentWaitsForBlockedPredecessor:
def run(self, processed_files: list[str]) -> dict:
return {"prompt": [ChatMessage.from_user(f"Files: {processed_files}")]}
chat_generator = MockChatGenerator()
chat_generator = MockChatGenerator("Hello")
agent = Agent(
chat_generator=chat_generator,
tools=[weather_tool],
@@ -2209,7 +2136,7 @@ class TestAgentWarmUp:
def test_warm_up_single_tool(self):
tool = self._make_tracking_tool()
agent = Agent(chat_generator=MockChatGenerator(), tools=[tool])
agent = Agent(chat_generator=MockChatGenerator("Hello"), tools=[tool])
assert not tool.was_warmed_up
agent.warm_up()
@@ -2218,7 +2145,7 @@ class TestAgentWarmUp:
def test_warm_up_multiple_tools(self):
tool1 = self._make_tracking_tool("tool1")
tool2 = self._make_tracking_tool("tool2")
agent = Agent(chat_generator=MockChatGenerator(), tools=[tool1, tool2])
agent = Agent(chat_generator=MockChatGenerator("Hello"), tools=[tool1, tool2])
assert not tool1.was_warmed_up
assert not tool2.was_warmed_up
@@ -2229,7 +2156,7 @@ class TestAgentWarmUp:
def test_warm_up_toolset(self):
inner_tool = self._make_tracking_tool()
toolset = self._make_tracking_toolset([inner_tool])
agent = Agent(chat_generator=MockChatGenerator(), tools=toolset)
agent = Agent(chat_generator=MockChatGenerator("Hello"), tools=toolset)
assert not toolset.was_warmed_up
agent.warm_up()
@@ -2241,7 +2168,7 @@ class TestAgentWarmUp:
tool2 = self._make_tracking_tool("tool2")
toolset2 = self._make_tracking_toolset([tool2])
agent = Agent(chat_generator=MockChatGenerator(), tools=toolset1 + toolset2)
agent = Agent(chat_generator=MockChatGenerator("Hello"), tools=toolset1 + toolset2)
assert not toolset1.was_warmed_up
assert not toolset2.was_warmed_up
@@ -2257,7 +2184,7 @@ class TestAgentWarmUp:
tool4 = self._make_tracking_tool("toolset_tool2")
toolset2 = self._make_tracking_toolset([tool4])
agent = Agent(chat_generator=MockChatGenerator(), tools=[tool1, toolset1, tool2, toolset2])
agent = Agent(chat_generator=MockChatGenerator("Hello"), tools=[tool1, toolset1, tool2, toolset2])
assert not tool1.was_warmed_up
assert not tool2.was_warmed_up
@@ -2285,7 +2212,7 @@ class TestAgentWarmUp:
tool.warm_up = counting_warm_up
agent = Agent(chat_generator=MockChatGenerator(), tools=[tool])
agent = Agent(chat_generator=MockChatGenerator("Hello"), tools=[tool])
agent.warm_up()
agent.warm_up()
agent.warm_up()
@@ -2318,7 +2245,7 @@ class TestAgentWarmUp:
self._connected = True
mcp_toolset = MockMCPToolset()
agent = Agent(chat_generator=MockChatGenerator(), tools=mcp_toolset)
agent = Agent(chat_generator=MockChatGenerator("Hello"), tools=mcp_toolset)
assert mcp_toolset.tools == [placeholder_tool]
agent.warm_up()
assert mcp_toolset.tools == [actual_tool]
@@ -2377,7 +2304,7 @@ class TestAgentWarmUp:
def test_run_warms_up_per_run_toolset(self):
"""Per-run tools passed to run() are not covered by Agent.warm_up() and must be warmed up at run time."""
init_tool = self._make_tracking_tool("init_tool")
agent = Agent(chat_generator=MockChatGenerator(), tools=Toolset([init_tool]))
agent = Agent(chat_generator=MockChatGenerator("Hello"), tools=Toolset([init_tool]))
agent.warm_up()
per_run_tool = self._make_tracking_tool("per_run_tool")
@@ -2393,7 +2320,7 @@ class TestAgentWarmUp:
def test_run_warms_up_per_run_list_of_tools_and_toolsets(self):
"""A per-run list of Tools and Toolsets must be warmed up at run time."""
init_tool = self._make_tracking_tool("init_tool")
agent = Agent(chat_generator=MockChatGenerator(), tools=[init_tool])
agent = Agent(chat_generator=MockChatGenerator("Hello"), tools=[init_tool])
agent.warm_up()
per_run_tool = self._make_tracking_tool("per_run_tool")
@@ -2410,7 +2337,7 @@ class TestAgentWarmUp:
async def test_run_async_warms_up_per_run_toolset(self):
"""The async run path must also warm up per-run tools."""
init_tool = self._make_tracking_tool("init_tool")
agent = Agent(chat_generator=MockChatGenerator(), tools=Toolset([init_tool]))
agent = Agent(chat_generator=MockChatGenerator("Hello"), tools=Toolset([init_tool]))
agent.warm_up()
per_run_tool = self._make_tracking_tool("per_run_tool")
@@ -2424,7 +2351,7 @@ class TestAgentWarmUp:
class TestComponentLifecycle:
def test_warm_up_delegates_to_chat_generator(self, weather_tool):
chat_generator = MockChatGenerator()
chat_generator = MockChatGenerator("Hello")
chat_generator.warm_up = MagicMock()
agent = Agent(chat_generator=chat_generator, tools=[weather_tool], system_prompt="This is a system prompt.")
@@ -2434,11 +2361,12 @@ class TestComponentLifecycle:
chat_generator.warm_up.reset_mock()
agent.run([ChatMessage.from_user("What is the weather in Berlin?")])
assert agent._tools_warmed_up is True
chat_generator.warm_up.assert_called_once()
# warm_up runs twice here: the Agent delegates to the generator, and the generator's own run() self-warms
assert chat_generator.warm_up.call_count == 2
@pytest.mark.asyncio
async def test_warm_up_async_delegates_to_chat_generator(self):
chat_generator = MockChatGenerator()
chat_generator = MockChatGenerator("Hello")
chat_generator.warm_up_async = AsyncMock()
chat_generator.warm_up = MagicMock()
agent = Agent(chat_generator=chat_generator, tools=[])
@@ -2455,7 +2383,7 @@ class TestComponentLifecycle:
chat_generator.warm_up.assert_called_once()
def test_close_delegates_to_chat_generator(self):
chat_generator = MockChatGenerator()
chat_generator = MockChatGenerator("Hello")
chat_generator.close = MagicMock()
agent = Agent(chat_generator=chat_generator, tools=[])
agent.close()
@@ -2463,7 +2391,7 @@ class TestComponentLifecycle:
@pytest.mark.asyncio
async def test_close_async_delegates_to_chat_generator(self):
chat_generator = MockChatGenerator()
chat_generator = MockChatGenerator("Hello")
chat_generator.close_async = AsyncMock()
agent = Agent(chat_generator=chat_generator, tools=[])
await agent.close_async()
@@ -2471,7 +2399,7 @@ class TestComponentLifecycle:
@pytest.mark.asyncio
async def test_close_async_falls_back_to_sync_close(self):
chat_generator = MockChatGenerator()
chat_generator = MockChatGenerator("Hello")
chat_generator.close = MagicMock()
agent = Agent(chat_generator=chat_generator, tools=[])
await agent.close_async()
@@ -2503,7 +2431,7 @@ class TestAgentNotTriggeredByInjectedInput:
def run(self) -> dict:
return {"messages": [ChatMessage.from_assistant("?")], "last_role": "assistant"}
chat_generator = MockChatGenerator()
chat_generator = MockChatGenerator("Hello")
agent = Agent(chat_generator=chat_generator, tools=[weather_tool])
chat_generator.run = MagicMock(return_value={"replies": [ChatMessage.from_assistant("x")]})
test/components/agents/test_agent_hitl.py
+2 -10
View File
@@ -8,9 +8,8 @@ from unittest.mock import MagicMock
import pytest
from haystack import component
from haystack.components.agents import Agent
from haystack.components.generators.chat import OpenAIChatGenerator
from haystack.components.generators.chat import MockChatGenerator, OpenAIChatGenerator
from haystack.dataclasses import ChatMessage, ToolCall
from haystack.hooks.human_in_the_loop import (
AlwaysAskPolicy,
@@ -21,7 +20,7 @@ from haystack.hooks.human_in_the_loop import (
SimpleConsoleUI,
)
from haystack.hooks.human_in_the_loop.types import ConfirmationStrategy, ConfirmationUI
from haystack.tools import Tool, Toolset, create_tool_from_function
from haystack.tools import Tool, create_tool_from_function
class MockUserInterface(ConfirmationUI):
@@ -234,13 +233,6 @@ class TestAgent:
assert result["token_usage"]["total_tokens"] > 0
@component
class MockChatGenerator:
@component.output_types(replies=list[ChatMessage])
def run(self, messages: list[ChatMessage], tools: list[Tool] | Toolset | None = None, **kwargs) -> dict[str, Any]:
return {"replies": [ChatMessage.from_assistant("Hello")]}
def _producer() -> dict[str, str]:
return {"value": "PRODUCED"}
test/components/agents/test_agent_hooks.py
+3 -16
View File
@@ -2,30 +2,17 @@
#
# SPDX-License-Identifier: Apache-2.0
from typing import Annotated, Any
from typing import Annotated
from unittest.mock import AsyncMock, MagicMock
import pytest
from haystack import component
from haystack.components.agents import Agent
from haystack.components.agents.state import State, replace_values
from haystack.components.generators.chat import MockChatGenerator
from haystack.dataclasses import ChatMessage, ToolCall
from haystack.hooks import hook
from haystack.tools import Tool, Toolset, tool
@component
class MockChatGenerator:
@component.output_types(replies=list[ChatMessage])
def run(self, messages: list[ChatMessage], tools: list[Tool] | Toolset | None = None, **kwargs) -> dict[str, Any]:
return {"replies": [ChatMessage.from_assistant("Hello")]}
@component.output_types(replies=list[ChatMessage])
async def run_async(
self, messages: list[ChatMessage], tools: list[Tool] | Toolset | None = None, **kwargs
) -> dict[str, Any]:
return {"replies": [ChatMessage.from_assistant("Hello")]}
from haystack.tools import tool
@tool
test/components/extractors/image/test_llm_document_content_extractor.py
+12 -28
View File
@@ -11,7 +11,7 @@ import pytest
from haystack import Document, Pipeline
from haystack.components.converters.image.document_to_image import DocumentToImageContent
from haystack.components.extractors.image import LLMDocumentContentExtractor
from haystack.components.generators.chat import OpenAIChatGenerator
from haystack.components.generators.chat import MockChatGenerator, OpenAIChatGenerator
from haystack.components.writers import DocumentWriter
from haystack.core.serialization import component_to_dict
from haystack.dataclasses.chat_message import ChatMessage, ImageContent
@@ -191,12 +191,10 @@ class TestLLMDocumentContentExtractor:
@patch.object(DocumentToImageContent, "run")
def test_run_plain_string_response_goes_to_content(self, mock_doc_to_image_run):
"""When LLM returns plain string (non-JSON), it is written to document content."""
mock_chat_generator = Mock(spec=OpenAIChatGenerator)
mock_chat_generator.run.return_value = {"replies": [ChatMessage.from_assistant(text="Plain text, not JSON")]}
mock_doc_to_image_run.return_value = {
"image_contents": [ImageContent.from_file_path("./test/test_files/images/apple.jpg")]
}
extractor = LLMDocumentContentExtractor(chat_generator=mock_chat_generator)
extractor = LLMDocumentContentExtractor(chat_generator=MockChatGenerator("Plain text, not JSON"))
docs = [Document(content="", meta={"file_path": "/path/to/image.pdf"})]
result = extractor.run(documents=docs)
assert len(result["documents"]) == 1
@@ -206,14 +204,10 @@ class TestLLMDocumentContentExtractor:
@patch.object(DocumentToImageContent, "run")
def test_run_valid_json_not_object_reports_error(self, mock_doc_to_image_run):
"""When LLM returns valid JSON that is not an object (e.g. array or primitive), report error."""
mock_chat_generator = Mock(spec=OpenAIChatGenerator)
mock_chat_generator.run.return_value = {
"replies": [ChatMessage.from_assistant(text='["array", "not", "object"]')]
}
mock_doc_to_image_run.return_value = {
"image_contents": [ImageContent.from_file_path("./test/test_files/images/apple.jpg")]
}
extractor = LLMDocumentContentExtractor(chat_generator=mock_chat_generator)
extractor = LLMDocumentContentExtractor(chat_generator=MockChatGenerator('["array", "not", "object"]'))
docs = [Document(content="", meta={"file_path": "/path/to/image.pdf"})]
result = extractor.run(documents=docs)
assert len(result["documents"]) == 0
@@ -224,16 +218,12 @@ class TestLLMDocumentContentExtractor:
@patch.object(DocumentToImageContent, "run")
def test_run_with_content_and_metadata_extraction(self, mock_doc_to_image_run):
"""When content mode gets JSON with document_content and other keys, other keys are merged into metadata."""
mock_chat_generator = Mock(spec=OpenAIChatGenerator)
mock_chat_generator.run.return_value = {
"replies": [
ChatMessage.from_assistant(text='{"document_content": "Main text", "title": "Doc Title", "page": "1"}')
]
}
mock_doc_to_image_run.return_value = {
"image_contents": [ImageContent.from_file_path("./test/test_files/images/apple.jpg")]
}
extractor = LLMDocumentContentExtractor(chat_generator=mock_chat_generator)
extractor = LLMDocumentContentExtractor(
chat_generator=MockChatGenerator('{"document_content": "Main text", "title": "Doc Title", "page": "1"}')
)
docs = [Document(content="", meta={"file_path": "/path/to/image.pdf"})]
result = extractor.run(documents=docs)
assert len(result["documents"]) == 1
@@ -285,16 +275,14 @@ class TestLLMDocumentContentExtractor:
@patch.object(DocumentToImageContent, "run")
def test_run_removes_extraction_error_from_previous_runs(self, mock_doc_to_image_run):
mock_chat_generator = Mock(spec=OpenAIChatGenerator)
mock_chat_generator.run.return_value = {
"replies": [ChatMessage.from_assistant(text='{"document_content": "Successfully extracted content"}')]
}
# Mock DocumentToImageContent to return valid image content
mock_doc_to_image_run.return_value = {
"image_contents": [ImageContent.from_file_path("./test/test_files/images/apple.jpg")]
}
extractor = LLMDocumentContentExtractor(chat_generator=mock_chat_generator)
extractor = LLMDocumentContentExtractor(
chat_generator=MockChatGenerator('{"document_content": "Successfully extracted content"}')
)
# Document with previous extraction error
docs = [
@@ -356,16 +344,12 @@ class TestLLMDocumentContentExtractor:
@patch.object(DocumentToImageContent, "run")
def test_run_json_multiple_keys_metadata_merged(self, mock_doc_to_image_run):
"""When LLM returns JSON with multiple keys and no document_content, all keys are merged into metadata."""
mock_chat_generator = Mock(spec=OpenAIChatGenerator)
mock_chat_generator.run.return_value = {
"replies": [
ChatMessage.from_assistant(text='{"title": "Sample Doc", "author": "Test", "document_type": "report"}')
]
}
mock_doc_to_image_run.return_value = {
"image_contents": [ImageContent.from_file_path("./test/test_files/images/apple.jpg")]
}
extractor = LLMDocumentContentExtractor(chat_generator=mock_chat_generator)
extractor = LLMDocumentContentExtractor(
chat_generator=MockChatGenerator('{"title": "Sample Doc", "author": "Test", "document_type": "report"}')
)
original_content = "Original content"
docs = [Document(content=original_content, meta={"file_path": "/path/to/image.pdf"})]
result = extractor.run(documents=docs)
test/components/query/test_query_expander.py
+14 -27
View File
@@ -8,6 +8,7 @@ from unittest.mock import AsyncMock, Mock
import pytest
from haystack.components.generators.chat import MockChatGenerator
from haystack.components.generators.chat.openai import OpenAIChatGenerator
from haystack.components.query.query_expander import DEFAULT_PROMPT_TEMPLATE, QueryExpander
from haystack.dataclasses.chat_message import ChatMessage
@@ -104,12 +105,9 @@ class TestQueryExpander:
]
mock_chat_generator.run.assert_called_once()
def test_run_without_including_original(self, mock_chat_generator):
mock_chat_generator.run.return_value = {
"replies": [ChatMessage.from_assistant('{"queries": ["alt1", "alt2"]}')]
}
expander = QueryExpander(chat_generator=mock_chat_generator, include_original_query=False)
def test_run_without_including_original(self):
chat_generator = MockChatGenerator('{"queries": ["alt1", "alt2"]}')
expander = QueryExpander(chat_generator=chat_generator, include_original_query=False)
result = expander.run("original")
assert result["queries"] == ["alt1", "alt2"]
@@ -149,9 +147,8 @@ class TestQueryExpander:
result = expander.run("test query")
assert result["queries"] == ["test query"]
def test_run_invalid_json_response(self, mock_chat_generator):
mock_chat_generator.run.return_value = {"replies": [ChatMessage.from_assistant("invalid json response")]}
expander = QueryExpander(chat_generator=mock_chat_generator)
def test_run_invalid_json_response(self):
expander = QueryExpander(chat_generator=MockChatGenerator("invalid json response"))
result = expander.run("test query")
assert result["queries"] == ["test query"]
@@ -196,20 +193,16 @@ class TestQueryExpander:
queries = expander._parse_expanded_queries('{"queries": ["valid query", 123, "", "another valid"]}')
assert queries == ["valid query", "another valid"]
def test_run_query_deduplication(self, mock_chat_generator):
mock_chat_generator.run.return_value = {
"replies": [ChatMessage.from_assistant('{"queries": ["original query", "alt1", "alt2"]}')]
}
expander = QueryExpander(chat_generator=mock_chat_generator, include_original_query=True)
def test_run_query_deduplication(self):
chat_generator = MockChatGenerator('{"queries": ["original query", "alt1", "alt2"]}')
expander = QueryExpander(chat_generator=chat_generator, include_original_query=True)
result = expander.run("original query")
assert result["queries"] == ["original query", "alt1", "alt2"]
assert len(result["queries"]) == 3
def test_run_truncates_excess_queries(self, mock_chat_generator, caplog):
mock_chat_generator.run.return_value = {
"replies": [ChatMessage.from_assistant('{"queries": ["q1", "q2", "q3", "q4", "q5"]}')]
}
expander = QueryExpander(chat_generator=mock_chat_generator, n_expansions=3, include_original_query=False)
def test_run_truncates_excess_queries(self, caplog):
chat_generator = MockChatGenerator('{"queries": ["q1", "q2", "q3", "q4", "q5"]}')
expander = QueryExpander(chat_generator=chat_generator, n_expansions=3, include_original_query=False)
with caplog.at_level(logging.WARNING):
result = expander.run("test query")
@@ -244,14 +237,8 @@ class TestQueryExpander:
assert "Create 2 alternative search queries for: test query" in call_args
assert "Return as JSON" in call_args
def test_component_output_types(self, mock_chat_generator, monkeypatch):
monkeypatch.setenv("OPENAI_API_KEY", "test-key-12345")
expander = QueryExpander()
mock_chat_generator.run.return_value = {
"replies": [ChatMessage.from_assistant('{"queries": ["test1", "test2"]}')]
}
expander.chat_generator = mock_chat_generator
def test_component_output_types(self):
expander = QueryExpander(chat_generator=MockChatGenerator('{"queries": ["test1", "test2"]}'))
result = expander.run("test")
assert "queries" in result
test/components/rankers/test_llm_ranker.py
+40 -53
View File
@@ -9,6 +9,7 @@ import pytest
from jinja2 import TemplateSyntaxError
from haystack import Document
from haystack.components.generators.chat import MockChatGenerator
from haystack.components.generators.chat.openai import OpenAIChatGenerator
from haystack.components.rankers.llm_ranker import DEFAULT_PROMPT_TEMPLATE, LLMRanker
from haystack.dataclasses import ChatMessage
@@ -109,84 +110,78 @@ def test_run_whitespace_query_returns_fallback(mock_chat_generator):
mock_chat_generator.run.assert_not_called()
def test_run_successful_ranking(mock_chat_generator):
def test_run_successful_ranking():
documents = [
Document(id="1", content="first"),
Document(id="2", content="second"),
Document(id="3", content="third"),
]
mock_chat_generator.run.return_value = {
"replies": [ChatMessage.from_assistant('{"documents": [{"index": 2}, {"index": 1}, {"index": 3}]}')]
}
ranker = LLMRanker(chat_generator=mock_chat_generator, top_k=2)
chat_generator = MockChatGenerator('{"documents": [{"index": 2}, {"index": 1}, {"index": 3}]}')
ranker = LLMRanker(chat_generator=chat_generator, top_k=2)
result = ranker.run(query="test query", documents=documents)
assert [document.id for document in result["documents"]] == ["2", "1"]
def test_run_returns_only_documents_listed_by_the_llm(mock_chat_generator):
def test_run_returns_only_documents_listed_by_the_llm():
documents = [Document(id="1", content="first"), Document(id="2", content="second")]
mock_chat_generator.run.return_value = {"replies": [ChatMessage.from_assistant('{"documents": [{"index": 2}]}')]}
ranker = LLMRanker(chat_generator=mock_chat_generator, top_k=2)
chat_generator = MockChatGenerator('{"documents": [{"index": 2}]}')
ranker = LLMRanker(chat_generator=chat_generator, top_k=2)
result = ranker.run(query="test query", documents=documents)
assert [document.id for document in result["documents"]] == ["2"]
def test_run_runtime_top_k_overrides_instance_top_k(mock_chat_generator):
def test_run_runtime_top_k_overrides_instance_top_k():
documents = [
Document(id="doc_1", content="first"),
Document(id="doc_2", content="second"),
Document(id="doc_3", content="third"),
]
mock_chat_generator.run.return_value = {
"replies": [ChatMessage.from_assistant('{"documents": [{"index": 3}, {"index": 2}, {"index": 1}]}')]
}
ranker = LLMRanker(chat_generator=mock_chat_generator, top_k=3)
chat_generator = MockChatGenerator('{"documents": [{"index": 3}, {"index": 2}, {"index": 1}]}')
ranker = LLMRanker(chat_generator=chat_generator, top_k=3)
result = ranker.run(query="test query", documents=documents, top_k=1)
assert [document.id for document in result["documents"]] == ["doc_3"]
def test_run_ignores_out_of_range_indices(mock_chat_generator):
def test_run_ignores_out_of_range_indices():
documents = [Document(id="1", content="first"), Document(id="2", content="second")]
mock_chat_generator.run.return_value = {
"replies": [ChatMessage.from_assistant('{"documents": [{"index": 99}, {"index": 2}, {"index": 1}]}')]
}
ranker = LLMRanker(chat_generator=mock_chat_generator)
chat_generator = MockChatGenerator('{"documents": [{"index": 99}, {"index": 2}, {"index": 1}]}')
ranker = LLMRanker(chat_generator=chat_generator)
result = ranker.run(query="test query", documents=documents)
assert [document.id for document in result["documents"]] == ["2", "1"]
def test_run_empty_ranking_result_returns_empty_documents(mock_chat_generator):
def test_run_empty_ranking_result_returns_empty_documents():
documents = [Document(id="1", content="first"), Document(id="2", content="second")]
mock_chat_generator.run.return_value = {"replies": [ChatMessage.from_assistant('{"documents": []}')]}
ranker = LLMRanker(chat_generator=mock_chat_generator)
chat_generator = MockChatGenerator('{"documents": []}')
ranker = LLMRanker(chat_generator=chat_generator)
result = ranker.run(query="test query", documents=documents)
assert result == {"documents": []}
def test_run_invalid_json_falls_back(mock_chat_generator):
def test_run_invalid_json_falls_back():
documents = [Document(id="1", content="first"), Document(id="2", content="second")]
mock_chat_generator.run.return_value = {"replies": [ChatMessage.from_assistant("not-json")]}
ranker = LLMRanker(chat_generator=mock_chat_generator, top_k=1, raise_on_failure=False)
chat_generator = MockChatGenerator("not-json")
ranker = LLMRanker(chat_generator=chat_generator, top_k=1, raise_on_failure=False)
result = ranker.run(query="test query", documents=documents)
assert result == {"documents": documents}
def test_run_invalid_json_raises(mock_chat_generator):
def test_run_invalid_json_raises():
documents = [Document(id="1", content="first")]
mock_chat_generator.run.return_value = {"replies": [ChatMessage.from_assistant("not-json")]}
ranker = LLMRanker(chat_generator=mock_chat_generator, raise_on_failure=True)
chat_generator = MockChatGenerator("not-json")
ranker = LLMRanker(chat_generator=chat_generator, raise_on_failure=True)
with pytest.raises(ValueError):
ranker.run(query="test query", documents=documents)
@@ -221,72 +216,64 @@ def test_run_no_replies_falls_back(mock_chat_generator):
assert result == {"documents": documents}
def test_run_reply_without_text_falls_back(mock_chat_generator):
def test_run_reply_without_text_falls_back():
documents = [Document(id="1", content="first"), Document(id="2", content="second")]
mock_chat_generator.run.return_value = {"replies": [ChatMessage.from_assistant(tool_calls=[])]}
ranker = LLMRanker(chat_generator=mock_chat_generator, top_k=1)
chat_generator = MockChatGenerator(ChatMessage.from_assistant(tool_calls=[]))
ranker = LLMRanker(chat_generator=chat_generator, top_k=1)
result = ranker.run(query="test query", documents=documents)
assert result == {"documents": documents}
def test_run_no_valid_document_indices_falls_back(mock_chat_generator):
def test_run_no_valid_document_indices_falls_back():
documents = [Document(id="1", content="first"), Document(id="2", content="second")]
mock_chat_generator.run.return_value = {
"replies": [ChatMessage.from_assistant('{"documents": [{"index": 0}, {"index": 3}]}')]
}
ranker = LLMRanker(chat_generator=mock_chat_generator, top_k=1)
chat_generator = MockChatGenerator('{"documents": [{"index": 0}, {"index": 3}]}')
ranker = LLMRanker(chat_generator=chat_generator, top_k=1)
result = ranker.run(query="test query", documents=documents)
assert result == {"documents": documents}
def test_run_deduplicates_documents_before_ranking(mock_chat_generator):
def test_run_deduplicates_documents_before_ranking():
documents = [
Document(id="duplicate", content="keep me", score=0.9),
Document(id="duplicate", content="drop me", score=0.1),
Document(id="unique", content="unique", score=0.2),
]
mock_chat_generator.run.return_value = {
"replies": [ChatMessage.from_assistant('{"documents": [{"index": 2}, {"index": 1}]}')]
}
ranker = LLMRanker(chat_generator=mock_chat_generator)
chat_generator = MockChatGenerator('{"documents": [{"index": 2}, {"index": 1}]}')
ranker = LLMRanker(chat_generator=chat_generator)
result = ranker.run(query="test query", documents=documents)
assert [document.content for document in result["documents"]] == ["unique", "keep me"]
def test_run_preserves_duplicate_indices(mock_chat_generator):
def test_run_preserves_duplicate_indices():
documents = [Document(id="1", content="first"), Document(id="2", content="second")]
mock_chat_generator.run.return_value = {
"replies": [ChatMessage.from_assistant('{"documents": [{"index": 2}, {"index": 2}, {"index": 1}]}')]
}
ranker = LLMRanker(chat_generator=mock_chat_generator)
chat_generator = MockChatGenerator('{"documents": [{"index": 2}, {"index": 2}, {"index": 1}]}')
ranker = LLMRanker(chat_generator=chat_generator)
result = ranker.run(query="test query", documents=documents)
assert [document.id for document in result["documents"]] == ["2", "2", "1"]
def test_run_numeric_string_index_is_accepted(mock_chat_generator):
def test_run_numeric_string_index_is_accepted():
documents = [Document(id="1", content="first"), Document(id="2", content="second")]
mock_chat_generator.run.return_value = {"replies": [ChatMessage.from_assistant('{"documents": [{"index": "2"}]}')]}
ranker = LLMRanker(chat_generator=mock_chat_generator)
chat_generator = MockChatGenerator('{"documents": [{"index": "2"}]}')
ranker = LLMRanker(chat_generator=chat_generator)
result = ranker.run(query="test query", documents=documents)
assert result == {"documents": [documents[1]]}
def test_run_invalid_index_type_falls_back(mock_chat_generator):
def test_run_invalid_index_type_falls_back():
documents = [Document(id="1", content="first"), Document(id="2", content="second")]
mock_chat_generator.run.return_value = {
"replies": [ChatMessage.from_assistant('{"documents": [{"index": "invalid"}]}')]
}
ranker = LLMRanker(chat_generator=mock_chat_generator)
chat_generator = MockChatGenerator('{"documents": [{"index": "invalid"}]}')
ranker = LLMRanker(chat_generator=chat_generator)
result = ranker.run(query="test query", documents=documents)
test/core/pipeline/breakpoints/test_pipeline_breakpoints_answer_joiner.py
+3 -14
View File
@@ -4,8 +4,8 @@
import pytest
from haystack import component
from haystack.components.builders.answer_builder import AnswerBuilder
from haystack.components.generators.chat import MockChatGenerator
from haystack.components.joiners import AnswerJoiner
from haystack.core.errors import BreakpointException
from haystack.core.pipeline.pipeline import Pipeline
@@ -13,24 +13,13 @@ from haystack.dataclasses import ChatMessage
from haystack.dataclasses.breakpoints import Breakpoint
@component
class FakeChatGenerator:
def __init__(self, content: str, model_name: str):
self.content = content
self.model_name = model_name
@component.output_types(replies=list[ChatMessage])
def run(self, messages: list[ChatMessage]) -> dict[str, list[ChatMessage]]:
return {"replies": [ChatMessage.from_assistant(self.content)]}
class TestPipelineBreakpoints:
@pytest.fixture
def answer_join_pipeline(self):
"""Creates a pipeline with fake components."""
pipeline = Pipeline()
pipeline.add_component("gpt-4o", FakeChatGenerator("GPT-4 response", "gpt-4o"))
pipeline.add_component("gpt-3", FakeChatGenerator("GPT-3 response", "gpt-3.5-turbo"))
pipeline.add_component("gpt-4o", MockChatGenerator("GPT-4 response", model="gpt-4o"))
pipeline.add_component("gpt-3", MockChatGenerator("GPT-3 response", model="gpt-3.5-turbo"))
pipeline.add_component("answer_builder_a", AnswerBuilder())
pipeline.add_component("answer_builder_b", AnswerBuilder())
pipeline.add_component("answer_joiner", AnswerJoiner())
test/core/pipeline/breakpoints/test_pipeline_breakpoints_branch_joiner.py
+2 -15
View File
@@ -3,12 +3,11 @@
# SPDX-License-Identifier: Apache-2.0
import json
from typing import Any
import pytest
from haystack import component
from haystack.components.converters import OutputAdapter
from haystack.components.generators.chat import MockChatGenerator
from haystack.components.joiners import BranchJoiner
from haystack.components.validators import JsonSchemaValidator
from haystack.core.errors import BreakpointException
@@ -17,18 +16,6 @@ from haystack.dataclasses import ChatMessage
from haystack.dataclasses.breakpoints import Breakpoint
@component
class FakeChatGenerator:
def __init__(self, content: str):
self.content = content
@component.output_types(replies=list[ChatMessage])
def run(
self, messages: list[ChatMessage], generation_kwargs: dict | None = None, **kwargs: Any
) -> dict[str, list[ChatMessage]]:
return {"replies": [ChatMessage.from_assistant(self.content)]}
class TestPipelineBreakpoints:
@pytest.fixture
def branch_joiner_pipeline(self):
@@ -46,7 +33,7 @@ class TestPipelineBreakpoints:
pipe = Pipeline()
pipe.add_component("joiner", BranchJoiner(list[ChatMessage]))
pipe.add_component("fc_llm", FakeChatGenerator(content))
pipe.add_component("fc_llm", MockChatGenerator(content))
pipe.add_component("validator", JsonSchemaValidator(json_schema=person_schema))
pipe.add_component("adapter", OutputAdapter("{{chat_message}}", list[ChatMessage], unsafe=True))
test/core/pipeline/breakpoints/test_pipeline_breakpoints_list_joiner.py
+4 -15
View File
@@ -2,28 +2,17 @@
#
# SPDX-License-Identifier: Apache-2.0
from typing import Any
import pytest
from haystack import Pipeline, component
from haystack import Pipeline
from haystack.components.builders import ChatPromptBuilder
from haystack.components.generators.chat import MockChatGenerator
from haystack.components.joiners import ListJoiner
from haystack.core.errors import BreakpointException
from haystack.dataclasses import ChatMessage
from haystack.dataclasses.breakpoints import Breakpoint
@component
class FakeChatGenerator:
def __init__(self, response: str):
self.response = response
@component.output_types(replies=list[ChatMessage])
def run(self, messages: list[ChatMessage], **kwargs: Any) -> dict[str, list[ChatMessage]]:
return {"replies": [ChatMessage.from_assistant(self.response)]}
class TestPipelineBreakpoints:
@pytest.fixture
def list_joiner_pipeline(self):
@@ -39,11 +28,11 @@ class TestPipelineBreakpoints:
pipe = Pipeline()
pipe.add_component("prompt_builder", ChatPromptBuilder(template=user_message, required_variables=None))
pipe.add_component("llm", FakeChatGenerator("Nuclear physics is the study of atomic nuclei."))
pipe.add_component("llm", MockChatGenerator("Nuclear physics is the study of atomic nuclei."))
pipe.add_component(
"feedback_prompt_builder", ChatPromptBuilder(template=feedback_message, required_variables=None)
)
pipe.add_component("feedback_llm", FakeChatGenerator("Score: 8/10. Concise and accurate."))
pipe.add_component("feedback_llm", MockChatGenerator("Score: 8/10. Concise and accurate."))
pipe.add_component("list_joiner", ListJoiner(list[ChatMessage]))
pipe.connect("prompt_builder.prompt", "llm.messages")
test/core/pipeline/breakpoints/test_pipeline_breakpoints_loops.py
+2 -11
View File
@@ -10,6 +10,7 @@ from pydantic import BaseModel, ValidationError
from haystack import component
from haystack.components.builders import ChatPromptBuilder
from haystack.components.generators.chat import MockChatGenerator
from haystack.core.errors import BreakpointException
from haystack.core.pipeline.pipeline import Pipeline
from haystack.dataclasses import ChatMessage
@@ -34,16 +35,6 @@ class OutputValidator:
return {"invalid_replies": replies, "error_message": str(e)}
@component
class FakeChatGenerator:
def __init__(self, response: str):
self.response = response
@component.output_types(replies=list[ChatMessage])
def run(self, messages: list[ChatMessage]) -> dict[str, list[ChatMessage]]:
return {"replies": [ChatMessage.from_assistant(self.response)]}
class City(BaseModel):
name: str
country: str
@@ -95,7 +86,7 @@ class TestPipelineBreakpointsLoops:
instance=ChatPromptBuilder(template=prompt_template, required_variables=["passage", "schema"]),
name="prompt_builder",
)
pipeline.add_component(instance=FakeChatGenerator(response=response_json), name="llm")
pipeline.add_component(instance=MockChatGenerator(response_json), name="llm")
pipeline.add_component(instance=OutputValidator(pydantic_model=CitiesData), name="output_validator")
pipeline.connect("prompt_builder.prompt", "llm.messages")
test/core/pipeline/test_pipeline.py
+8 -14
View File
@@ -8,18 +8,12 @@ from concurrent.futures import ThreadPoolExecutor
import pytest
from haystack.components.agents import Agent
from haystack.components.generators.chat import MockChatGenerator
from haystack.components.joiners import BranchJoiner
from haystack.core.component import component
from haystack.core.errors import PipelineConnectError, PipelineRuntimeError
from haystack.core.pipeline import Pipeline
from haystack.dataclasses import ChatMessage, Document
@component
class MockChatGenerator:
@component.output_types(replies=list[ChatMessage])
def run(self, messages: list[ChatMessage]) -> dict[str, list[ChatMessage]]:
return {"replies": [ChatMessage.from_assistant("Hello, world!")]}
from haystack.dataclasses import ChatMessage, ChatRole, Document
@component
@@ -400,16 +394,16 @@ class TestPipeline:
p = Pipeline()
p.add_component("message_producer", MessageProducer())
p.add_component("message_producer2", MessageProducer())
p.add_component("agent", Agent(chat_generator=MockChatGenerator()))
p.add_component("agent", Agent(chat_generator=MockChatGenerator("Hello, world!")))
p.connect("message_producer", "agent.messages")
p.connect("message_producer2", "agent.messages")
result = p.run({})
assert result["agent"]["messages"] == [
ChatMessage.from_user("Hello, world!"),
ChatMessage.from_user("Hello, world!"),
ChatMessage.from_assistant("Hello, world!"),
]
messages = result["agent"]["messages"]
assert messages[:2] == [ChatMessage.from_user("Hello, world!"), ChatMessage.from_user("Hello, world!")]
assert messages[2].role == ChatRole.ASSISTANT
assert messages[2].text == "Hello, world!"
def test_run_auto_variadic_str_to_list_str(self):
"""Two str producers connected to a list[str] input are auto-joined and flattened at runtime."""
test/hooks/tool_result_offloading/test_hooks.py
+5 -18
View File
@@ -3,14 +3,14 @@
# SPDX-License-Identifier: Apache-2.0
from pathlib import Path
from typing import Annotated, Any
from typing import Annotated
from unittest.mock import AsyncMock, MagicMock
import pytest
from haystack import component
from haystack.components.agents import Agent
from haystack.components.agents.state.state import State
from haystack.components.generators.chat import MockChatGenerator
from haystack.dataclasses import ChatMessage, ImageContent, TextContent, ToolCall
from haystack.hooks.tool_result_offloading import (
RESULT_STORE_CONTEXT_KEY,
@@ -20,20 +20,7 @@ from haystack.hooks.tool_result_offloading import (
OffloadOverChars,
ToolResultOffloadHook,
)
from haystack.tools import Tool, Toolset, tool
@component
class MockChatGenerator:
@component.output_types(replies=list[ChatMessage])
def run(self, messages: list[ChatMessage], tools: list[Tool] | Toolset | None = None, **kwargs) -> dict[str, Any]:
return {"replies": [ChatMessage.from_assistant("done")]}
@component.output_types(replies=list[ChatMessage])
async def run_async(
self, messages: list[ChatMessage], tools: list[Tool] | Toolset | None = None, **kwargs
) -> dict[str, Any]:
return {"replies": [ChatMessage.from_assistant("done")]}
from haystack.tools import tool
@tool
@@ -237,7 +224,7 @@ class TestToolResultOffloadHookInAgent:
hook = ToolResultOffloadHook(
store=FileSystemToolResultStore(root=tmp_path), offload_strategies={"*": AlwaysOffload()}
)
agent = Agent(chat_generator=MockChatGenerator(), tools=[big_tool], hooks={"after_tool": [hook]})
agent = Agent(chat_generator=MockChatGenerator("done"), tools=[big_tool], hooks={"after_tool": [hook]})
agent.warm_up()
agent.chat_generator.run = MagicMock(
side_effect=[
@@ -258,7 +245,7 @@ class TestToolResultOffloadHookInAgentAsync:
hook = ToolResultOffloadHook(
store=FileSystemToolResultStore(root=tmp_path), offload_strategies={"*": AlwaysOffload()}
)
agent = Agent(chat_generator=MockChatGenerator(), tools=[big_tool], hooks={"after_tool": [hook]})
agent = Agent(chat_generator=MockChatGenerator("done"), tools=[big_tool], hooks={"after_tool": [hook]})
agent.warm_up()
agent.chat_generator.run_async = AsyncMock(
side_effect=[
test/tools/test_component_tool.py
+4 -20
View File
@@ -15,10 +15,10 @@ from openai.types.chat.chat_completion import Choice
from haystack import Pipeline, SuperComponent, component
from haystack.components.agents import Agent, State
from haystack.components.builders import PromptBuilder
from haystack.components.generators.chat import OpenAIChatGenerator
from haystack.components.generators.chat import MockChatGenerator, OpenAIChatGenerator
from haystack.core.pipeline.utils import _deepcopy_with_exceptions
from haystack.dataclasses import ChatMessage, ChatRole, Document
from haystack.tools import ComponentTool, ToolsType
from haystack.tools import ComponentTool
from test.tools.test_parameters_schema_utils import BYTE_STREAM_SCHEMA, DOCUMENT_SCHEMA, SPARSE_EMBEDDING_SCHEMA
# Component and Model Definitions
@@ -148,22 +148,6 @@ class DocumentProcessor:
return {"concatenated": "\n".join(doc.content for doc in documents[:top_k] if doc.content)}
@component
class FakeChatGenerator:
def __init__(self, messages: list[ChatMessage]):
self.messages = messages
@component.output_types(replies=list[ChatMessage])
def run(
self,
messages: list[ChatMessage],
generation_kwargs: dict[str, Any] | None = None,
*,
tools: ToolsType | None = None,
) -> dict[str, list[ChatMessage]]:
return {"replies": self.messages}
def output_handler(old, new):
"""
Output handler to test serialization.
@@ -555,7 +539,7 @@ class TestComponentTool:
def test_component_invoker_with_agent(self):
"""Tests that Agent as a ComponentTool can be invoked when calling it with a list of dicts"""
agent = Agent(chat_generator=FakeChatGenerator(messages=[ChatMessage.from_assistant("Answer")]))
agent = Agent(chat_generator=MockChatGenerator("Answer"))
tool = ComponentTool(
component=agent,
name="agent_tool",
@@ -563,7 +547,7 @@ class TestComponentTool:
outputs_to_string={"source": "last_message"},
)
result = tool.invoke(messages=[{"role": "user", "content": [{"text": "A 4-day trip in the south of France"}]}])
assert result["last_message"] == ChatMessage.from_assistant("Answer")
assert result["last_message"].text == "Answer"
def test_convert_param_union_with_list_arm(self):
@component
test/tools/test_searchable_toolset.py
+13 -27
View File
@@ -11,7 +11,7 @@ import pytest
from haystack import component
from haystack.components.agents import Agent
from haystack.components.generators.chat import OpenAIChatGenerator
from haystack.components.generators.chat import MockChatGenerator, OpenAIChatGenerator
from haystack.dataclasses import ChatMessage, ToolCall
from haystack.tools import SearchableToolset, Tool, Toolset, flatten_tools_or_toolsets
from haystack.tools.from_function import create_tool_from_function
@@ -923,32 +923,18 @@ class TestSearchableToolsetAgentToolSelection:
"""tool_call_counts seeds only the search tool up front; a discovered+called tool is added lazily."""
toolset = SearchableToolset(catalog=large_catalog, search_threshold=3, top_k=5)
@component
class SearchThenWeatherGenerator:
step = 0
@component.output_types(replies=list[ChatMessage])
def run(self, messages, tools=None, **kwargs):
self.step += 1
if self.step == 1:
return {
"replies": [
ChatMessage.from_assistant(
tool_calls=[ToolCall(tool_name="search_tools", arguments={"tool_keywords": "weather"})]
)
]
}
if self.step == 2:
return {
"replies": [
ChatMessage.from_assistant(
tool_calls=[ToolCall(tool_name="get_weather", arguments={"city": "Berlin"})]
)
]
}
return {"replies": [ChatMessage.from_assistant("done")]}
agent = Agent(chat_generator=SearchThenWeatherGenerator(), tools=toolset, max_agent_steps=6)
chat_generator = MockChatGenerator(
[
ChatMessage.from_assistant(
tool_calls=[ToolCall(tool_name="search_tools", arguments={"tool_keywords": "weather"})]
),
ChatMessage.from_assistant(
tool_calls=[ToolCall(tool_name="get_weather", arguments={"city": "Berlin"})]
),
"done",
]
)
agent = Agent(chat_generator=chat_generator, tools=toolset, max_agent_steps=6)
result = agent.run(messages=[ChatMessage.from_user("What's the weather in Berlin?")])
counts = result["tool_call_counts"]