diff --git a/haystack/components/agents/agent.py b/haystack/components/agents/agent.py index dbc2f2f9344..d899bf2d89c 100644 --- a/haystack/components/agents/agent.py +++ b/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,6 +489,7 @@ 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), @@ -486,6 +497,16 @@ def _create_agent_span(self) -> Any: 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 diff --git a/releasenotes/notes/add-agent-runtime-budget-f3b6ce29.yaml b/releasenotes/notes/add-agent-runtime-budget-f3b6ce29.yaml new file mode 100644 index 00000000000..93c7dbce504 --- /dev/null +++ b/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. diff --git a/test/components/agents/test_agent.py b/test/components/agents/test_agent.py index 155f70e82be..8b8e00acb47 100644 --- a/test/components/agents/test_agent.py +++ b/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}, diff --git a/test/components/agents/test_agent_time_limit.py b/test/components/agents/test_agent_time_limit.py new file mode 100644 index 00000000000..29a4d09a236 --- /dev/null +++ b/test/components/agents/test_agent_time_limit.py @@ -0,0 +1,102 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# 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