refactor: Refactor Agent logic for easier readability (#9726)

* Start refactor

* Update run_async to use the new code

* Slight updates

* Refactoring of tests

* Remove messages from execution context

* Cleanup

* More cleanup

* Formatting

* Fix some typing

* ignore typing issues

* Add reno

* Adding docstrings

* Small changes

* docstrings

* Updates

* Update haystack/components/agents/agent.py

Co-authored-by: Stefano Fiorucci <[email protected]>

* PR comments

* PR comments

---------

Co-authored-by: Stefano Fiorucci <[email protected]>

Author: sjrl

haystack/components/agents/agent.py

@@ -340,15 +340,15 @@ def _validate_tool_breakpoint_is_valid(agent_breakpoint: AgentBreakpoint, tools:
         raise ValueError(f"Tool '{tool_breakpoint.tool_name}' is not available in the agent's tools")
 
 
-def _check_chat_generator_breakpoint(
+def _trigger_chat_generator_breakpoint(
     *, agent_snapshot: AgentSnapshot, parent_snapshot: Optional[PipelineSnapshot]
 ) -> None:
     """
-    Check for breakpoint before calling the ChatGenerator.
+    Trigger a breakpoint before ChatGenerator execution in Agent.
 
     :param agent_snapshot: AgentSnapshot object containing the agent's state and breakpoints
     :param parent_snapshot: Optional parent snapshot containing the state of the pipeline that houses the agent.
-    :raises BreakpointException: If a breakpoint is triggered
+    :raises BreakpointException: Always raised when this function is called, indicating a breakpoint has been triggered.
     """
 
     break_point = agent_snapshot.break_point.break_point
@@ -381,16 +381,17 @@ def _check_chat_generator_breakpoint(
     )
 
 
-def _check_tool_invoker_breakpoint(
+def _handle_tool_invoker_breakpoint(
     *, llm_messages: list[ChatMessage], agent_snapshot: AgentSnapshot, parent_snapshot: Optional[PipelineSnapshot]
 ) -> None:
     """
-    Check for breakpoint before calling the ToolInvoker.
+    Check if a tool call breakpoint should be triggered before executing the tool invoker.
 
     :param llm_messages: List of ChatMessage objects containing potential tool calls.
     :param agent_snapshot: AgentSnapshot object containing the agent's state and breakpoints.
     :param parent_snapshot: Optional parent snapshot containing the state of the pipeline that houses the agent.
-    :raises BreakpointException: If a breakpoint is triggered
+    :raises BreakpointException: If the breakpoint is triggered, indicating a breakpoint has been reached for a tool
+        call.
     """
     if not isinstance(agent_snapshot.break_point.break_point, ToolBreakpoint):
         return

haystack/core/pipeline/breakpoint.py

@@ -0,0 +1,4 @@
+---
+enhancements:
+  - |
+    The internal Agent logic was refactored to help with readability and maintanability. This should help developers understand and extend the internal Agent logic moving forward.

releasenotes/notes/refactor-agent-b8d3b690f52c3bb5.yaml

@@ -12,13 +12,12 @@
 from openai import Stream
 from openai.types.chat import ChatCompletionChunk, chat_completion_chunk
 
-from haystack import Pipeline, tracing
+from haystack import Pipeline, component, tracing
 from haystack.components.agents import Agent
 from haystack.components.agents.state import merge_lists
 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
@@ -101,59 +100,56 @@ def openai_mock_chat_completion_chunk():
         yield mock_chat_completion_create
 
 
-class MockChatGeneratorWithoutTools(ChatGenerator):
+@component
+class MockChatGeneratorWithoutTools:
     """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()
 
+    @component.output_types(replies=list[ChatMessage])
     def run(self, messages: list[ChatMessage]) -> dict[str, Any]:
         return {"replies": [ChatMessage.from_assistant("Hello")]}
 
 
-class MockChatGeneratorWithoutRunAsync(ChatGenerator):
+@component
+class MockChatGeneratorWithoutRunAsync:
     """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()
 
+    @component.output_types(replies=list[ChatMessage])
     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={})
-
+@component
+class MockChatGenerator:
     def to_dict(self) -> dict[str, Any]:
         return {"type": "MockChatGeneratorWithoutRunAsync", "data": {}}
 
     @classmethod
     def from_dict(cls, data: dict[str, Any]) -> "MockChatGeneratorWithoutRunAsync":
         return cls()
 
+    @component.output_types(replies=list[ChatMessage])
     def run(
         self, messages: list[ChatMessage], tools: Optional[Union[list[Tool], Toolset]] = 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: Optional[Union[list[Tool], Toolset]] = None, **kwargs
     ) -> dict[str, Any]:
@@ -818,7 +814,7 @@ async def test_run_async_falls_back_to_run_when_chat_generator_has_no_run_async(
 
     @pytest.mark.asyncio
     async def test_run_async_uses_chat_generator_run_async_when_available(self, weather_tool):
-        chat_generator = MockChatGeneratorWithRunAsync()
+        chat_generator = MockChatGenerator()
         agent = Agent(chat_generator=chat_generator, tools=[weather_tool])
         agent.warm_up()
 
@@ -904,16 +900,16 @@ def test_agent_tracing_span_run(self, caplog, monkeypatch, weather_tool):
             "chat_generator",
             "MockChatGeneratorWithoutRunAsync",
             '{"messages": "list", "tools": "list"}',
-            "{}",
-            "{}",
+            '{"messages": {"type": "list", "senders": []}, "tools": {"type": "typing.Union[list[haystack.tools.tool.Tool], haystack.tools.toolset.Toolset, NoneType]", "senders": []}}',  # noqa: E501
+            '{"replies": {"type": "list", "receivers": []}}',
             '{"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}}]}',  # noqa: E501
             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}}]',  # noqa: E501
             '["text"]',
             '{"messages": {"type": "list[haystack.dataclasses.chat_message.ChatMessage]", "handler": "haystack.components.agents.state.state_utils.merge_lists"}}',  # noqa: E501
-            '{"messages": [{"role": "user", "meta": {}, "name": null, "content": [{"text": "What\'s the weather in Paris?"}]}], "streaming_callback": null}',  # noqa: E501
+            '{"messages": [{"role": "user", "meta": {}, "name": null, "content": [{"text": "What\'s the weather in Paris?"}]}], "streaming_callback": null, "break_point": null, "snapshot": null}',  # noqa: E501
             '{"messages": [{"role": "user", "meta": {}, "name": null, "content": [{"text": "What\'s the weather in Paris?"}]}, {"role": "assistant", "meta": {}, "name": null, "content": [{"text": "Hello"}]}]}',  # noqa: E501
             1,
         ]
@@ -927,7 +923,7 @@ def test_agent_tracing_span_run(self, caplog, monkeypatch, weather_tool):
 
     @pytest.mark.asyncio
     async def test_agent_tracing_span_async_run(self, caplog, monkeypatch, weather_tool):
-        chat_generator = MockChatGeneratorWithRunAsync()
+        chat_generator = MockChatGenerator()
         agent = Agent(chat_generator=chat_generator, tools=[weather_tool])
 
         tracing.tracer.is_content_tracing_enabled = True
@@ -962,18 +958,18 @@ async def test_agent_tracing_span_async_run(self, caplog, monkeypatch, weather_t
 
         expected_tag_values = [
             "chat_generator",
-            "MockChatGeneratorWithRunAsync",
+            "MockChatGenerator",
             '{"messages": "list", "tools": "list"}',
-            "{}",
-            "{}",
+            '{"messages": {"type": "list", "senders": []}, "tools": {"type": "typing.Union[list[haystack.tools.tool.Tool], haystack.tools.toolset.Toolset, NoneType]", "senders": []}}',  # noqa: E501
+            '{"replies": {"type": "list", "receivers": []}}',
             '{"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}}]}',  # noqa: E501
             1,
             '{"replies": [{"role": "assistant", "meta": {}, "name": null, "content": [{"text": "Hello from run_async"}]}]}',  # noqa: E501
             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}}]',  # noqa: E501
             '["text"]',
             '{"messages": {"type": "list[haystack.dataclasses.chat_message.ChatMessage]", "handler": "haystack.components.agents.state.state_utils.merge_lists"}}',  # noqa: E501
-            '{"messages": [{"role": "user", "meta": {}, "name": null, "content": [{"text": "What\'s the weather in Paris?"}]}], "streaming_callback": null}',  # noqa: E501
+            '{"messages": [{"role": "user", "meta": {}, "name": null, "content": [{"text": "What\'s the weather in Paris?"}]}], "streaming_callback": null, "break_point": null, "snapshot": null}',  # noqa: E501
             '{"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"}]}]}',  # noqa: E501
             1,
         ]