Feat/agent runtime budget

haystack/components/agents/agent.py

@@ -4,6 +4,7 @@
 
 import inspect
 import re
+import time
 from dataclasses import dataclass
 from typing import Any, Literal, cast
 
@@ -219,6 +220,7 @@ def __init__(
         exit_conditions: list[str] | None = None,
         state_schema: dict[str, Any] | None = None,
         max_agent_steps: int = 100,
+        max_agent_time_seconds: float | None = None,
         streaming_callback: StreamingCallbackT | None = None,
         raise_on_tool_invocation_failure: bool = False,
         tool_invoker_kwargs: dict[str, Any] | None = None,
@@ -248,6 +250,9 @@ def __init__(
             Tools can read from and write to state keys using `inputs_from_state` and `outputs_to_state`.
         :param max_agent_steps: Maximum number of steps the agent will run before stopping. Defaults to 100.
             If the agent exceeds this number of steps, it will stop and return the current state.
+        :param max_agent_time_seconds: Optional maximum wall-clock runtime for one agent invocation, in seconds.
+            The limit is checked between component invocations. A currently running synchronous tool must
+            enforce its own timeout. Defaults to None.
         :param streaming_callback: A callback that will be invoked when a response is streamed from the LLM.
             The same callback can be configured to emit tool results when a tool is called.
         :param raise_on_tool_invocation_failure: Should the agent raise an exception when a tool invocation fails?
@@ -277,6 +282,9 @@ def __init__(
                 "Ensure that each exit condition corresponds to either 'text' or a valid tool name."
             )
 
+        if max_agent_time_seconds is not None and max_agent_time_seconds <= 0:
+            raise ValueError("max_agent_time_seconds must be greater than zero when provided.")
+
         # Validate state schema if provided
         if state_schema is not None:
             _validate_schema(state_schema)
@@ -296,6 +304,7 @@ def __init__(
         self.required_variables = required_variables
         self.exit_conditions = exit_conditions
         self.max_agent_steps = max_agent_steps
+        self.max_agent_time_seconds = max_agent_time_seconds
         self.raise_on_tool_invocation_failure = raise_on_tool_invocation_failure
         self.streaming_callback = streaming_callback
 
@@ -420,6 +429,7 @@ def to_dict(self) -> dict[str, Any]:
             # We serialize the original state schema, not the resolved one to reflect the original user input
             state_schema=_schema_to_dict(self._state_schema),
             max_agent_steps=self.max_agent_steps,
+            max_agent_time_seconds=self.max_agent_time_seconds,
             streaming_callback=serialize_callable(self.streaming_callback) if self.streaming_callback else None,
             raise_on_tool_invocation_failure=self.raise_on_tool_invocation_failure,
             tool_invoker_kwargs=self.tool_invoker_kwargs,
@@ -479,13 +489,24 @@ def _create_agent_span(self) -> Any:
             "haystack.agent.run",
             tags={
                 "haystack.agent.max_steps": self.max_agent_steps,
+                "haystack.agent.max_time_seconds": self.max_agent_time_seconds,
                 "haystack.agent.tools": self.tools,
                 "haystack.agent.exit_conditions": self.exit_conditions,
                 "haystack.agent.state_schema": _schema_to_dict(self.state_schema),
             },
             parent_span=parent_span,
         )
 
+    def _is_time_limit_exceeded(self, deadline: float | None) -> bool:
+        if deadline is None or time.monotonic() < deadline:
+            return False
+
+        logger.warning(
+            "Agent reached maximum execution time of {max_agent_time_seconds} seconds, stopping.",
+            max_agent_time_seconds=self.max_agent_time_seconds,
+        )
+        return True
+
     def _initialize_fresh_execution(
         self,
         messages: list[ChatMessage],
@@ -806,7 +827,15 @@ def run(  # noqa: PLR0915
             # agent_inputs is local and not used after this point, so we avoid deepcopying it
             span.set_content_tag("haystack.agent.input", agent_inputs)
 
+            deadline = (
+                time.monotonic() + self.max_agent_time_seconds
+                if self.max_agent_time_seconds is not None
+                else None
+            )
+
             while exe_context.counter < self.max_agent_steps:
+                if self._is_time_limit_exceeded(deadline):
+                    break
                 # We skip the chat generator when restarting from a snapshot from a ToolBreakpoint
                 if exe_context.skip_chat_generator:
                     llm_messages = exe_context.state.get("messages", [])[-1:]
@@ -857,6 +886,9 @@ def run(  # noqa: PLR0915
                     llm_messages = result["replies"]
                     exe_context.state.set("messages", llm_messages)
 
+                if self._is_time_limit_exceeded(deadline):
+                    break
+
                 # Exit for `exit_conditions=["text"]` behavior: the agent stops when there is no tool invoker, or when
                 # the model returns a plain text response (no tool calls). We require the last message to be a non-empty
                 # assistant text message so that an invalid response (e.g. a message with no tool calls or text) won't
@@ -1049,7 +1081,15 @@ async def run_async(  # noqa: PLR0915
             # agent_inputs is local and not used after this point, so we avoid deepcopying it
             span.set_content_tag("haystack.agent.input", agent_inputs)
 
+            deadline = (
+                time.monotonic() + self.max_agent_time_seconds
+                if self.max_agent_time_seconds is not None
+                else None
+            )
+
             while exe_context.counter < self.max_agent_steps:
+                if self._is_time_limit_exceeded(deadline):
+                    break
                 # We skip the chat generator when restarting from a snapshot from a ToolBreakpoint
                 if exe_context.skip_chat_generator:
                     llm_messages = exe_context.state.get("messages", [])[-1:]
@@ -1101,6 +1141,9 @@ async def run_async(  # noqa: PLR0915
                     llm_messages = result["replies"]
                     exe_context.state.set("messages", llm_messages)
 
+                if self._is_time_limit_exceeded(deadline):
+                    break
+
                 # Exit for `exit_conditions=["text"]` behavior: the agent stops when there is no tool invoker, or when
                 # the model returns a plain text response (no tool calls). We require the last message to be a non-empty
                 # assistant text message so that an invalid response (e.g. a message with no tool calls or text) won't

releasenotes/notes/add-agent-runtime-budget-f3b6ce29.yaml

@@ -0,0 +1,7 @@
+---
+features:
+  - |
+    Added an optional ``max_agent_time_seconds`` runtime budget to ``Agent``.
+    When configured, the agent stops before starting another component invocation
+    after the budget expires, complementing ``max_agent_steps`` for deployments
+    with slow or remote tools.

test/components/agents/test_agent.py

@@ -272,6 +272,7 @@ def test_to_dict(self, weather_tool, component_tool, monkeypatch):
                 "exit_conditions": ["text", "weather_tool"],
                 "state_schema": {"foo": {"type": "str"}},
                 "max_agent_steps": 100,
+                "max_agent_time_seconds": None,
                 "streaming_callback": None,
                 "raise_on_tool_invocation_failure": False,
                 "tool_invoker_kwargs": {"max_workers": 5, "enable_streaming_callback_passthrough": True},
@@ -337,6 +338,7 @@ def test_to_dict_with_toolset(self, monkeypatch, weather_tool):
                 "exit_conditions": ["text"],
                 "state_schema": {},
                 "max_agent_steps": 100,
+                "max_agent_time_seconds": None,
                 "raise_on_tool_invocation_failure": False,
                 "streaming_callback": None,
                 "tool_invoker_kwargs": None,
@@ -421,6 +423,7 @@ def test_from_dict(self, monkeypatch):
                 "exit_conditions": ["text", "weather_tool"],
                 "state_schema": {"foo": {"type": "str"}},
                 "max_agent_steps": 100,
+                "max_agent_time_seconds": None,
                 "raise_on_tool_invocation_failure": False,
                 "streaming_callback": None,
                 "tool_invoker_kwargs": {"max_workers": 5, "enable_streaming_callback_passthrough": True},
@@ -491,6 +494,7 @@ def test_from_dict_with_toolset(self, monkeypatch):
                 "exit_conditions": ["text"],
                 "state_schema": {},
                 "max_agent_steps": 100,
+                "max_agent_time_seconds": None,
                 "raise_on_tool_invocation_failure": False,
                 "streaming_callback": None,
                 "tool_invoker_kwargs": None,
@@ -568,6 +572,7 @@ def test_from_dict_state_schema_none(self, monkeypatch):
                 "exit_conditions": ["text", "weather_tool"],
                 "state_schema": None,
                 "max_agent_steps": 100,
+                "max_agent_time_seconds": None,
                 "raise_on_tool_invocation_failure": False,
                 "streaming_callback": None,
                 "tool_invoker_kwargs": {"max_workers": 5, "enable_streaming_callback_passthrough": True},

test/components/agents/test_agent_time_limit.py

@@ -0,0 +1,102 @@
+# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+from typing import Any
+
+import pytest
+
+from haystack import component
+from haystack.components.agents import Agent
+from haystack.dataclasses import ChatMessage, ToolCall
+from haystack.tools import Tool, Toolset
+
+
+@component
+class LoopingChatGenerator:
+    def __init__(self) -> None:
+        self.calls = 0
+
+    def to_dict(self) -> dict[str, Any]:
+        return {"type": "LoopingChatGenerator", "data": {}}
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "LoopingChatGenerator":
+        return cls()
+
+    @component.output_types(replies=list[ChatMessage])
+    def run(self, messages: list[ChatMessage], tools: list[Tool] | Toolset | None = None, **kwargs: Any) -> dict[str, Any]:
+        self.calls += 1
+        return {
+            "replies": [
+                ChatMessage.from_assistant(
+                    "Calling the loop tool.",
+                    tool_calls=[ToolCall(tool_name="loop", arguments={})],
+                )
+            ]
+        }
+
+    @component.output_types(replies=list[ChatMessage])
+    async def run_async(
+        self, messages: list[ChatMessage], tools: list[Tool] | Toolset | None = None, **kwargs: Any
+    ) -> dict[str, Any]:
+        return self.run(messages=messages, tools=tools, **kwargs)
+
+
+def loop_tool() -> str:
+    return "looped"
+
+
+def make_agent() -> tuple[Agent, LoopingChatGenerator]:
+    chat_generator = LoopingChatGenerator()
+    tool = Tool(
+        name="loop",
+        description="Return a loop marker.",
+        parameters={"type": "object", "properties": {}},
+        function=loop_tool,
+    )
+    return (
+        Agent(chat_generator=chat_generator, tools=[tool], max_agent_steps=10, max_agent_time_seconds=1.0),
+        chat_generator,
+    )
+
+
+class TestAgentTimeLimit:
+    def test_requires_positive_time_limit(self) -> None:
+        with pytest.raises(ValueError, match="max_agent_time_seconds must be greater than zero"):
+            Agent(chat_generator=LoopingChatGenerator(), max_agent_time_seconds=0)
+
+    def test_serializes_time_limit(self) -> None:
+        agent, _ = make_agent()
+
+        assert agent.to_dict()["init_parameters"]["max_agent_time_seconds"] == 1.0
+
+    def test_stops_before_tool_invocation_when_time_limit_expires(
+        self, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
+    ) -> None:
+        agent, chat_generator = make_agent()
+        values = iter([10.0, 10.0, 11.0])
+        monkeypatch.setattr("haystack.components.agents.agent.time.monotonic", lambda: next(values))
+
+        with caplog.at_level(logging.WARNING):
+            result = agent.run(messages=[ChatMessage.from_user("Start")])
+
+        assert chat_generator.calls == 1
+        assert result["last_message"].tool_calls[0].tool_name == "loop"
+        assert "Agent reached maximum execution time of 1.0 seconds, stopping." in caplog.text
+
+    @pytest.mark.asyncio
+    async def test_stops_before_tool_invocation_when_time_limit_expires_async(
+        self, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
+    ) -> None:
+        agent, chat_generator = make_agent()
+        values = iter([10.0, 10.0, 11.0])
+        monkeypatch.setattr("haystack.components.agents.agent.time.monotonic", lambda: next(values))
+
+        with caplog.at_level(logging.WARNING):
+            result = await agent.run_async(messages=[ChatMessage.from_user("Start")])
+
+        assert chat_generator.calls == 1
+        assert result["last_message"].tool_calls[0].tool_name == "loop"
+        assert "Agent reached maximum execution time of 1.0 seconds, stopping." in caplog.text