wehub-resource-sync
c56bef871b
CodeQL / Analyze (python) (push) Has been cancelled
Update Platform Components Table / update (push) Has been cancelled
Docker image release / Build base image (push) Has been cancelled
Sync docs with Docusaurus / sync (push) Has been cancelled
Tests / Check if changed (push) Has been cancelled
Tests / format (push) Has been cancelled
Tests / check-imports (push) Has been cancelled
Tests / Unit / macos-latest (push) Has been cancelled
Tests / Unit / ubuntu-latest (push) Has been cancelled
Tests / Unit / windows-latest (push) Has been cancelled
Tests / mypy (push) Has been cancelled
Tests / Integration / ubuntu-latest (push) Has been cancelled
Tests / Integration / macos-latest (push) Has been cancelled
Tests / Integration / windows-latest (push) Has been cancelled
Tests / notify-slack-on-failure (push) Has been cancelled
Tests / Mark tests as completed (push) Has been cancelled
505 lines
20 KiB
Python
505 lines
20 KiB
Python
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
|
#
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
from typing import Annotated, Any
|
|
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.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")]}
|
|
|
|
|
|
@tool
|
|
def save(content: Annotated[str, "what to save"]) -> str:
|
|
"""Save content."""
|
|
return "saved"
|
|
|
|
|
|
@tool
|
|
def final_answer(answer: Annotated[str, "the answer"]) -> str:
|
|
"""Provide the final answer."""
|
|
return answer
|
|
|
|
|
|
# --- Module-level hooks (importable so they can be serialized) ---
|
|
|
|
|
|
@hook
|
|
def record_a(state: State) -> None:
|
|
state.set("trace", ["a"])
|
|
|
|
|
|
@hook
|
|
def record_b(state: State) -> None:
|
|
state.set("trace", ["b"])
|
|
|
|
|
|
@hook
|
|
def build_context(state: State) -> None:
|
|
if state.get("step_count") == 0:
|
|
state.set("messages", [ChatMessage.from_system("INJECTED")])
|
|
|
|
|
|
@hook
|
|
def record_tool_calls(state: State) -> None:
|
|
# Realistic before_tool use case: audit the tool calls the model is about to run.
|
|
pending = state.data["messages"][-1].tool_calls
|
|
state.set("trace", [tc.tool_name for tc in pending])
|
|
|
|
|
|
@hook
|
|
def replace_pending_save_call(state: State) -> None:
|
|
messages = state.data["messages"]
|
|
replacement = ChatMessage.from_assistant(tool_calls=[ToolCall("save", {"content": "changed"})])
|
|
state.set("messages", [*messages[:-1], replacement], handler_override=replace_values)
|
|
|
|
|
|
@hook
|
|
def replace_pending_call_with_non_tool_message(state: State) -> None:
|
|
messages = state.data["messages"]
|
|
replacement = ChatMessage.from_assistant("Skipping tool execution.")
|
|
state.set("messages", [*messages[:-1], replacement], handler_override=replace_values)
|
|
|
|
|
|
@hook
|
|
def record_after_tool(state: State) -> None:
|
|
state.set("trace", ["after_tool"])
|
|
|
|
|
|
@hook
|
|
def offload_tool_results(state: State) -> None:
|
|
# Realistic after_tool use case: replace freshly produced tool results with a compact pointer, preserving origin.
|
|
rewritten = [
|
|
ChatMessage.from_tool(tool_result="OFFLOADED", origin=m.tool_call_result.origin, error=m.tool_call_result.error)
|
|
if m.tool_call_result is not None
|
|
else m
|
|
for m in state.data["messages"]
|
|
]
|
|
state.set("messages", rewritten, handler_override=replace_values)
|
|
|
|
|
|
@hook
|
|
def require_save(state: State) -> None:
|
|
if state.get("tool_call_counts", {}).get("save", 0) == 0:
|
|
state.set("messages", [ChatMessage.from_system("You must call save before finishing.")])
|
|
state.set("continue_run", True)
|
|
|
|
|
|
@hook
|
|
def always_continue(state: State) -> None:
|
|
state.set("messages", [ChatMessage.from_system("keep going")])
|
|
state.set("continue_run", True)
|
|
|
|
|
|
@hook
|
|
def critique(state: State) -> None:
|
|
# Push back on the first final answer to force one more loop.
|
|
if state.get("tool_call_counts", {}).get("final_answer", 0) < 2:
|
|
state.set("messages", [ChatMessage.from_user("Please expand.")])
|
|
state.set("continue_run", True)
|
|
|
|
|
|
class LifecycleHook:
|
|
"""A class-based hook that records its lifecycle calls (e.g. opening/closing a client)."""
|
|
|
|
def __init__(self) -> None:
|
|
self.warmed = 0
|
|
self.warmed_async = 0
|
|
self.closed = 0
|
|
self.closed_async = 0
|
|
|
|
def run(self, state: State) -> None:
|
|
pass
|
|
|
|
def warm_up(self) -> None:
|
|
self.warmed += 1
|
|
|
|
async def warm_up_async(self) -> None:
|
|
self.warmed_async += 1
|
|
|
|
def close(self) -> None:
|
|
self.closed += 1
|
|
|
|
async def close_async(self) -> None:
|
|
self.closed_async += 1
|
|
|
|
|
|
def _agent(generator, **kwargs):
|
|
agent = Agent(chat_generator=generator, **kwargs)
|
|
agent.warm_up()
|
|
return agent
|
|
|
|
|
|
class TestAgentHooksValidation:
|
|
def test_invalid_hook_point_raises(self):
|
|
with pytest.raises(ValueError):
|
|
Agent(chat_generator=MockChatGenerator(), hooks={"bogus": [record_a]})
|
|
|
|
def test_non_callable_object_raises(self):
|
|
with pytest.raises(TypeError, match="must have a callable 'run"):
|
|
Agent(chat_generator=MockChatGenerator(), hooks={"before_llm": [object()]})
|
|
|
|
def test_unwrapped_function_hints_at_hook_decorator(self):
|
|
def my_hook(state: State) -> None:
|
|
pass
|
|
|
|
with pytest.raises(TypeError, match="@hook decorator"):
|
|
Agent(chat_generator=MockChatGenerator(), hooks={"before_llm": [my_hook]})
|
|
|
|
def test_continue_run_is_a_reserved_state_schema_key(self):
|
|
with pytest.raises(ValueError):
|
|
Agent(chat_generator=MockChatGenerator(), state_schema={"continue_run": {"type": bool}})
|
|
|
|
def test_continue_run_is_not_exposed_as_output(self):
|
|
agent = _agent(MockChatGenerator())
|
|
agent.chat_generator.run = MagicMock(return_value={"replies": [ChatMessage.from_assistant("done")]})
|
|
result = agent.run(messages=[ChatMessage.from_user("hi")])
|
|
assert "continue_run" not in result
|
|
|
|
|
|
class TestBeforeLlmHook:
|
|
def test_step_0_only_context_injected_once_across_loops(self):
|
|
agent = _agent(MockChatGenerator(), tools=[save], hooks={"before_llm": [build_context]})
|
|
agent.chat_generator.run = MagicMock(
|
|
side_effect=[
|
|
{"replies": [ChatMessage.from_assistant(tool_calls=[ToolCall("save", {"content": "x"})])]},
|
|
{"replies": [ChatMessage.from_assistant("done")]},
|
|
]
|
|
)
|
|
result = agent.run(messages=[ChatMessage.from_user("hi")])
|
|
# build_context only injects on step_count == 0, so it appears once despite the two-loop run
|
|
injected = [m for m in result["messages"] if m.is_from("system") and m.text == "INJECTED"]
|
|
assert len(injected) == 1
|
|
assert agent.chat_generator.run.call_count == 2
|
|
|
|
def test_hooks_run_in_list_order(self):
|
|
agent = _agent(
|
|
MockChatGenerator(), state_schema={"trace": {"type": list}}, hooks={"before_llm": [record_a, record_b]}
|
|
)
|
|
agent.chat_generator.run = MagicMock(return_value={"replies": [ChatMessage.from_assistant("done")]})
|
|
result = agent.run(messages=[ChatMessage.from_user("hi")])
|
|
assert result["trace"] == ["a", "b"]
|
|
|
|
|
|
class TestBeforeToolHook:
|
|
def test_runs_before_tools_when_model_calls_a_tool(self):
|
|
agent = _agent(
|
|
MockChatGenerator(),
|
|
tools=[save],
|
|
state_schema={"trace": {"type": list}},
|
|
hooks={"before_tool": [record_tool_calls]},
|
|
)
|
|
agent.chat_generator.run = MagicMock(
|
|
side_effect=[
|
|
{"replies": [ChatMessage.from_assistant(tool_calls=[ToolCall("save", {"content": "x"})])]},
|
|
{"replies": [ChatMessage.from_assistant("done")]},
|
|
]
|
|
)
|
|
result = agent.run(messages=[ChatMessage.from_user("hi")])
|
|
# before_tool sees the pending tool call and does not fire on the final text-only step
|
|
assert result["trace"] == ["save"]
|
|
assert result["tool_call_counts"]["save"] == 1
|
|
|
|
def test_rereads_last_state_message_after_hook(self):
|
|
agent = _agent(MockChatGenerator(), tools=[save], hooks={"before_tool": [replace_pending_save_call]})
|
|
agent.chat_generator.run = MagicMock(
|
|
side_effect=[
|
|
{"replies": [ChatMessage.from_assistant(tool_calls=[ToolCall("save", {"content": "original"})])]},
|
|
{"replies": [ChatMessage.from_assistant("done")]},
|
|
]
|
|
)
|
|
result = agent.run(messages=[ChatMessage.from_user("hi")])
|
|
tool_messages = [m for m in result["messages"] if m.tool_call_result is not None]
|
|
assert len(tool_messages) == 1
|
|
assert tool_messages[0].tool_call_result.origin.arguments == {"content": "changed"}
|
|
|
|
def test_skips_tool_execution_when_hook_leaves_last_message_without_tool_calls(self):
|
|
agent = _agent(
|
|
MockChatGenerator(), tools=[save], hooks={"before_tool": [replace_pending_call_with_non_tool_message]}
|
|
)
|
|
agent.chat_generator.run = MagicMock(
|
|
side_effect=[
|
|
{"replies": [ChatMessage.from_assistant(tool_calls=[ToolCall("save", {"content": "x"})])]},
|
|
{"replies": [ChatMessage.from_assistant("done")]},
|
|
]
|
|
)
|
|
result = agent.run(messages=[ChatMessage.from_user("hi")])
|
|
assert agent.chat_generator.run.call_count == 2
|
|
assert result["tool_call_counts"]["save"] == 0
|
|
assert not [m for m in result["messages"] if m.tool_call_result is not None]
|
|
assert [m.text for m in result["messages"]][-2:] == ["Skipping tool execution.", "done"]
|
|
|
|
|
|
class TestAfterToolHook:
|
|
def test_runs_after_tools_and_not_on_text_only_exit(self):
|
|
agent = _agent(
|
|
MockChatGenerator(),
|
|
tools=[save],
|
|
state_schema={"trace": {"type": list}},
|
|
hooks={"after_tool": [record_after_tool]},
|
|
)
|
|
agent.chat_generator.run = MagicMock(
|
|
side_effect=[
|
|
{"replies": [ChatMessage.from_assistant(tool_calls=[ToolCall("save", {"content": "x"})])]},
|
|
{"replies": [ChatMessage.from_assistant("done")]},
|
|
]
|
|
)
|
|
result = agent.run(messages=[ChatMessage.from_user("hi")])
|
|
# Fires once for the single tool step; does not fire on the final text-only step (no tools ran there).
|
|
assert result["trace"] == ["after_tool"]
|
|
|
|
def test_does_not_run_when_no_tools_are_called(self):
|
|
agent = _agent(
|
|
MockChatGenerator(),
|
|
tools=[save],
|
|
state_schema={"trace": {"type": list}},
|
|
hooks={"after_tool": [record_after_tool]},
|
|
)
|
|
agent.chat_generator.run = MagicMock(return_value={"replies": [ChatMessage.from_assistant("done")]})
|
|
result = agent.run(messages=[ChatMessage.from_user("hi")])
|
|
assert result.get("trace", []) == []
|
|
|
|
def test_rewrites_tool_results_before_next_llm_call(self):
|
|
agent = _agent(MockChatGenerator(), tools=[save], hooks={"after_tool": [offload_tool_results]})
|
|
agent.chat_generator.run = MagicMock(
|
|
side_effect=[
|
|
{"replies": [ChatMessage.from_assistant(tool_calls=[ToolCall("save", {"content": "x"})])]},
|
|
{"replies": [ChatMessage.from_assistant("done")]},
|
|
]
|
|
)
|
|
result = agent.run(messages=[ChatMessage.from_user("hi")])
|
|
# The rewritten result is what the next LLM call sees...
|
|
second_call_messages = agent.chat_generator.run.call_args_list[1].kwargs["messages"]
|
|
assert [m.tool_call_result.result for m in second_call_messages if m.tool_call_result is not None] == [
|
|
"OFFLOADED"
|
|
]
|
|
# ...and what ends up in the final conversation.
|
|
assert [m.tool_call_result.result for m in result["messages"] if m.tool_call_result is not None] == [
|
|
"OFFLOADED"
|
|
]
|
|
|
|
def test_rewriting_results_does_not_change_exit_matching(self):
|
|
# Exit conditions match on tool name + error, not result content, so offloading the result still exits.
|
|
agent = _agent(
|
|
MockChatGenerator(),
|
|
tools=[final_answer],
|
|
exit_conditions=["final_answer"],
|
|
hooks={"after_tool": [offload_tool_results]},
|
|
)
|
|
agent.chat_generator.run = MagicMock(
|
|
return_value={
|
|
"replies": [ChatMessage.from_assistant(tool_calls=[ToolCall("final_answer", {"answer": "42"})])]
|
|
}
|
|
)
|
|
result = agent.run(messages=[ChatMessage.from_user("hi")])
|
|
assert agent.chat_generator.run.call_count == 1
|
|
assert [m.tool_call_result.result for m in result["messages"] if m.tool_call_result is not None] == [
|
|
"OFFLOADED"
|
|
]
|
|
|
|
|
|
class TestOnExitHook:
|
|
def test_mandatory_tool_forces_extra_step(self):
|
|
agent = _agent(MockChatGenerator(), tools=[save], hooks={"on_exit": [require_save]})
|
|
agent.chat_generator.run = MagicMock(
|
|
side_effect=[
|
|
{"replies": [ChatMessage.from_assistant("Done")]}, # tries to exit without saving
|
|
{"replies": [ChatMessage.from_assistant(tool_calls=[ToolCall("save", {"content": "x"})])]},
|
|
{"replies": [ChatMessage.from_assistant("All finished")]},
|
|
]
|
|
)
|
|
result = agent.run(messages=[ChatMessage.from_user("hi")])
|
|
assert agent.chat_generator.run.call_count == 3
|
|
assert result["tool_call_counts"]["save"] == 1
|
|
assert result["last_message"].text == "All finished"
|
|
|
|
def test_readonly_hook_does_not_change_exit(self):
|
|
fired = []
|
|
|
|
def record(state: State) -> None:
|
|
fired.append(1)
|
|
|
|
agent = _agent(MockChatGenerator(), tools=[save], hooks={"on_exit": [hook(record)]})
|
|
agent.chat_generator.run = MagicMock(return_value={"replies": [ChatMessage.from_assistant("Final")]})
|
|
agent.run(messages=[ChatMessage.from_user("hi")])
|
|
assert agent.chat_generator.run.call_count == 1
|
|
assert fired == [1]
|
|
|
|
def test_critique_on_tool_based_exit(self):
|
|
agent = _agent(
|
|
MockChatGenerator(), tools=[final_answer], exit_conditions=["final_answer"], hooks={"on_exit": [critique]}
|
|
)
|
|
agent.chat_generator.run = MagicMock(
|
|
side_effect=[
|
|
{"replies": [ChatMessage.from_assistant(tool_calls=[ToolCall("final_answer", {"answer": "short"})])]},
|
|
{"replies": [ChatMessage.from_assistant(tool_calls=[ToolCall("final_answer", {"answer": "longer"})])]},
|
|
]
|
|
)
|
|
result = agent.run(messages=[ChatMessage.from_user("q")])
|
|
# critique forces one extra loop before the final answer is accepted
|
|
assert agent.chat_generator.run.call_count == 2
|
|
assert result["tool_call_counts"]["final_answer"] == 2
|
|
|
|
def test_max_agent_steps_bounds_always_cancel_hook(self):
|
|
agent = _agent(MockChatGenerator(), max_agent_steps=3, hooks={"on_exit": [always_continue]})
|
|
agent.chat_generator.run = MagicMock(return_value={"replies": [ChatMessage.from_assistant("text")]})
|
|
agent.run(messages=[ChatMessage.from_user("hi")])
|
|
assert agent.chat_generator.run.call_count == 3
|
|
|
|
def test_continue_run_from_before_llm_hook_does_not_force_continuation(self):
|
|
def leak_continue(state: State) -> None:
|
|
state.set("continue_run", True)
|
|
|
|
def audit(state: State) -> None:
|
|
pass
|
|
|
|
agent = _agent(
|
|
MockChatGenerator(),
|
|
max_agent_steps=5,
|
|
hooks={"before_llm": [hook(leak_continue)], "on_exit": [hook(audit)]},
|
|
)
|
|
agent.chat_generator.run = MagicMock(return_value={"replies": [ChatMessage.from_assistant("Final")]})
|
|
agent.run(messages=[ChatMessage.from_user("hi")])
|
|
assert agent.chat_generator.run.call_count == 1
|
|
|
|
|
|
class TestHookReuseAcrossHookPoints:
|
|
def test_same_hook_under_two_hook_points(self):
|
|
counter = []
|
|
|
|
def count_call(state: State) -> None:
|
|
counter.append(1)
|
|
|
|
counting = hook(count_call)
|
|
agent = _agent(MockChatGenerator(), hooks={"before_llm": [counting], "on_exit": [counting]})
|
|
agent.chat_generator.run = MagicMock(return_value={"replies": [ChatMessage.from_assistant("done")]})
|
|
agent.run(messages=[ChatMessage.from_user("hi")])
|
|
# once for before_llm, once for on_exit (which appends nothing, so the run still exits)
|
|
assert len(counter) == 2
|
|
|
|
|
|
class TestAgentHooksSerde:
|
|
def test_to_dict_from_dict_roundtrip(self, monkeypatch):
|
|
monkeypatch.setenv("OPENAI_API_KEY", "fake-key")
|
|
from haystack.components.generators.chat import OpenAIChatGenerator
|
|
|
|
agent = Agent(
|
|
chat_generator=OpenAIChatGenerator(),
|
|
tools=[save],
|
|
hooks={"before_llm": [build_context], "on_exit": [require_save]},
|
|
)
|
|
restored = Agent.from_dict(agent.to_dict())
|
|
assert set(restored.hooks) == {"before_llm", "on_exit"}
|
|
assert restored.hooks["before_llm"][0].function is build_context.function
|
|
assert restored.hooks["on_exit"][0].function is require_save.function
|
|
|
|
|
|
class TestAgentHooksAsync:
|
|
@pytest.mark.asyncio
|
|
async def test_sync_hook_runs_in_async_run(self):
|
|
fired = []
|
|
|
|
def record(state: State) -> None:
|
|
fired.append(1)
|
|
|
|
agent = _agent(MockChatGenerator(), hooks={"before_llm": [hook(record)]})
|
|
agent.chat_generator.run_async = AsyncMock(return_value={"replies": [ChatMessage.from_assistant("done")]})
|
|
await agent.run_async(messages=[ChatMessage.from_user("hi")])
|
|
assert fired == [1]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_after_tool_hook_rewrites_results(self):
|
|
agent = _agent(MockChatGenerator(), tools=[save], hooks={"after_tool": [offload_tool_results]})
|
|
agent.chat_generator.run_async = AsyncMock(
|
|
side_effect=[
|
|
{"replies": [ChatMessage.from_assistant(tool_calls=[ToolCall("save", {"content": "x"})])]},
|
|
{"replies": [ChatMessage.from_assistant("done")]},
|
|
]
|
|
)
|
|
result = await agent.run_async(messages=[ChatMessage.from_user("hi")])
|
|
assert [m.tool_call_result.result for m in result["messages"] if m.tool_call_result is not None] == [
|
|
"OFFLOADED"
|
|
]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_on_exit_hook_forces_extra_step(self):
|
|
async def require_save_async(state: State) -> None:
|
|
if state.get("tool_call_counts", {}).get("save", 0) == 0:
|
|
state.set("messages", [ChatMessage.from_system("call save first")])
|
|
state.set("continue_run", True)
|
|
|
|
agent = _agent(MockChatGenerator(), tools=[save], hooks={"on_exit": [hook(require_save_async)]})
|
|
agent.chat_generator.run_async = AsyncMock(
|
|
side_effect=[
|
|
{"replies": [ChatMessage.from_assistant("Done")]},
|
|
{"replies": [ChatMessage.from_assistant(tool_calls=[ToolCall("save", {"content": "x"})])]},
|
|
{"replies": [ChatMessage.from_assistant("All finished")]},
|
|
]
|
|
)
|
|
result = await agent.run_async(messages=[ChatMessage.from_user("hi")])
|
|
assert agent.chat_generator.run_async.call_count == 3
|
|
assert result["tool_call_counts"]["save"] == 1
|
|
|
|
|
|
class TestAgentHookLifecycle:
|
|
def test_init_does_not_warm_up(self):
|
|
h = LifecycleHook()
|
|
Agent(chat_generator=MockChatGenerator(), hooks={"before_llm": [h]})
|
|
assert h.warmed == 0
|
|
|
|
def test_warm_up_warms_hooks_once(self):
|
|
h = LifecycleHook()
|
|
agent = Agent(chat_generator=MockChatGenerator(), hooks={"before_llm": [h]})
|
|
agent.warm_up()
|
|
agent.warm_up()
|
|
assert h.warmed == 1
|
|
|
|
def test_reused_hook_warmed_once(self):
|
|
h = LifecycleHook()
|
|
agent = Agent(chat_generator=MockChatGenerator(), hooks={"before_llm": [h], "on_exit": [h]})
|
|
agent.warm_up()
|
|
assert h.warmed == 1
|
|
|
|
def test_close_closes_hooks(self):
|
|
h = LifecycleHook()
|
|
agent = Agent(chat_generator=MockChatGenerator(), hooks={"before_llm": [h]})
|
|
agent.close()
|
|
assert h.closed == 1
|
|
|
|
|
|
class TestAgentHookLifecycleAsync:
|
|
@pytest.mark.asyncio
|
|
async def test_warm_up_async_prefers_async(self):
|
|
h = LifecycleHook()
|
|
agent = Agent(chat_generator=MockChatGenerator(), hooks={"before_llm": [h]})
|
|
await agent.warm_up_async()
|
|
assert h.warmed_async == 1
|
|
assert h.warmed == 0
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_close_async_prefers_async(self):
|
|
h = LifecycleHook()
|
|
agent = Agent(chat_generator=MockChatGenerator(), hooks={"before_llm": [h]})
|
|
await agent.close_async()
|
|
assert h.closed_async == 1
|
|
assert h.closed == 0
|