From 6651dada6e27eb5718afe86708f862525128a897 Mon Sep 17 00:00:00 2001 From: Cameron G <156701171+camgrimsec@users.noreply.github.com> Date: Thu, 25 Jun 2026 14:48:17 -0400 Subject: [PATCH 1/2] chore: apply agent runtime budget patch --- .../workflows/apply-agent-runtime-budget.yml | 132 ++++++++++++++++++ 1 file changed, 132 insertions(+) create mode 100644 .github/workflows/apply-agent-runtime-budget.yml diff --git a/.github/workflows/apply-agent-runtime-budget.yml b/.github/workflows/apply-agent-runtime-budget.yml new file mode 100644 index 00000000000..66d3aa671f6 --- /dev/null +++ b/.github/workflows/apply-agent-runtime-budget.yml @@ -0,0 +1,132 @@ +name: Apply agent runtime budget patch + +on: + push: + branches: + - feat/agent-runtime-budget + +permissions: + contents: write + +jobs: + apply-patch: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + with: + ref: feat/agent-runtime-budget + + - name: Apply focused runtime-budget patch + run: | + python - <<'PY' + from pathlib import Path + + agent_path = Path("haystack/components/agents/agent.py") + text = agent_path.read_text() + + def replace_once(old: str, new: str) -> None: + global text + count = text.count(old) + if count != 1: + raise RuntimeError(f"Expected one occurrence, found {count}: {old[:80]!r}") + text = text.replace(old, new, 1) + + def replace_all(old: str, new: str, expected: int) -> None: + global text + count = text.count(old) + if count != expected: + raise RuntimeError(f"Expected {expected} occurrences, found {count}: {old[:80]!r}") + text = text.replace(old, new) + + replace_once("import inspect\nimport re\n", "import inspect\nimport re\nimport time\n") + replace_once( + " max_agent_steps: int = 100,\n streaming_callback: StreamingCallbackT | None = None,", + " max_agent_steps: int = 100,\n max_agent_time_seconds: float | None = None,\n streaming_callback: StreamingCallbackT | None = None,", + ) + replace_once( + " :param max_agent_steps: Maximum number of steps the agent will run before stopping. Defaults to 100.\n" + " If the agent exceeds this number of steps, it will stop and return the current state.\n" + " :param streaming_callback:", + " :param max_agent_steps: Maximum number of steps the agent will run before stopping. Defaults to 100.\n" + " If the agent exceeds this number of steps, it will stop and return the current state.\n" + " :param max_agent_time_seconds: Optional maximum wall-clock runtime for one agent invocation, in seconds.\n" + " The limit is checked between component invocations. A currently running synchronous tool must\n" + " enforce its own timeout. Defaults to None.\n" + " :param streaming_callback:", + ) + replace_once( + " # Validate state schema if provided\n", + " if max_agent_time_seconds is not None and max_agent_time_seconds <= 0:\n" + " raise ValueError(\"max_agent_time_seconds must be greater than zero when provided.\")\n\n" + " # Validate state schema if provided\n", + ) + replace_once( + " self.max_agent_steps = max_agent_steps\n self.raise_on_tool_invocation_failure", + " self.max_agent_steps = max_agent_steps\n self.max_agent_time_seconds = max_agent_time_seconds\n self.raise_on_tool_invocation_failure", + ) + replace_once( + " max_agent_steps=self.max_agent_steps,\n streaming_callback=", + " max_agent_steps=self.max_agent_steps,\n max_agent_time_seconds=self.max_agent_time_seconds,\n streaming_callback=", + ) + replace_once( + ' "haystack.agent.max_steps": self.max_agent_steps,\n "haystack.agent.tools": self.tools,', + ' "haystack.agent.max_steps": self.max_agent_steps,\n "haystack.agent.max_time_seconds": self.max_agent_time_seconds,\n "haystack.agent.tools": self.tools,', + ) + replace_once( + " parent_span=parent_span,\n )\n\n def _initialize_fresh_execution(", + " parent_span=parent_span,\n )\n\n" + " def _is_time_limit_exceeded(self, deadline: float | None) -> bool:\n" + " if deadline is None or time.monotonic() < deadline:\n" + " return False\n\n" + " logger.warning(\n" + " \"Agent reached maximum execution time of {max_agent_time_seconds} seconds, stopping.\",\n" + " max_agent_time_seconds=self.max_agent_time_seconds,\n" + " )\n" + " return True\n\n" + " def _initialize_fresh_execution(", + ) + replace_all( + " span.set_content_tag(\"haystack.agent.input\", agent_inputs)\n\n while exe_context.counter < self.max_agent_steps:\n", + " span.set_content_tag(\"haystack.agent.input\", agent_inputs)\n\n" + " deadline = (\n" + " time.monotonic() + self.max_agent_time_seconds\n" + " if self.max_agent_time_seconds is not None\n" + " else None\n" + " )\n\n" + " while exe_context.counter < self.max_agent_steps:\n" + " if self._is_time_limit_exceeded(deadline):\n" + " break\n", + 2, + ) + replace_all( + " # Exit for `exit_conditions=[\"text\"]` behavior: the agent stops when there is no tool invoker, or when\n", + " if self._is_time_limit_exceeded(deadline):\n" + " break\n\n" + " # Exit for `exit_conditions=[\"text\"]` behavior: the agent stops when there is no tool invoker, or when\n", + 2, + ) + agent_path.write_text(text) + + test_path = Path("test/components/agents/test_agent.py") + test_text = test_path.read_text() + expected = ' "max_agent_steps": 100,\n' + replacements = test_text.count(expected) + if replacements < 2: + raise RuntimeError(f"Expected serialized Agent fixtures, found {replacements}") + test_path.write_text( + test_text.replace(expected, expected + ' "max_agent_time_seconds": None,\n') + ) + + Path("test/components/agents/test_agent_time_limit.py").write_text('''# SPDX-FileCopyrightText: 2022-present deepset GmbH \n#\n# SPDX-License-Identifier: Apache-2.0\n\nimport logging\nfrom typing import Any\n\nimport pytest\n\nfrom haystack import component\nfrom haystack.components.agents import Agent\nfrom haystack.dataclasses import ChatMessage, ToolCall\nfrom haystack.tools import Tool, Toolset\n\n\n@component\nclass LoopingChatGenerator:\n def __init__(self) -> None:\n self.calls = 0\n\n def to_dict(self) -> dict[str, Any]:\n return {"type": "LoopingChatGenerator", "data": {}}\n\n @classmethod\n def from_dict(cls, data: dict[str, Any]) -> "LoopingChatGenerator":\n return cls()\n\n @component.output_types(replies=list[ChatMessage])\n def run(self, messages: list[ChatMessage], tools: list[Tool] | Toolset | None = None, **kwargs: Any) -> dict[str, Any]:\n self.calls += 1\n return {\n "replies": [\n ChatMessage.from_assistant(\n "Calling the loop tool.",\n tool_calls=[ToolCall(tool_name="loop", arguments={})],\n )\n ]\n }\n\n @component.output_types(replies=list[ChatMessage])\n async def run_async(\n self, messages: list[ChatMessage], tools: list[Tool] | Toolset | None = None, **kwargs: Any\n ) -> dict[str, Any]:\n return self.run(messages=messages, tools=tools, **kwargs)\n\n\ndef loop_tool() -> str:\n return "looped"\n\n\ndef make_agent() -> tuple[Agent, LoopingChatGenerator]:\n chat_generator = LoopingChatGenerator()\n tool = Tool(\n name="loop",\n description="Return a loop marker.",\n parameters={"type": "object", "properties": {}},\n function=loop_tool,\n )\n return (\n Agent(chat_generator=chat_generator, tools=[tool], max_agent_steps=10, max_agent_time_seconds=1.0),\n chat_generator,\n )\n\n\nclass TestAgentTimeLimit:\n def test_requires_positive_time_limit(self) -> None:\n with pytest.raises(ValueError, match="max_agent_time_seconds must be greater than zero"):\n Agent(chat_generator=LoopingChatGenerator(), max_agent_time_seconds=0)\n\n def test_serializes_time_limit(self) -> None:\n agent, _ = make_agent()\n\n assert agent.to_dict()["init_parameters"]["max_agent_time_seconds"] == 1.0\n\n def test_stops_before_tool_invocation_when_time_limit_expires(\n self, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture\n ) -> None:\n agent, chat_generator = make_agent()\n values = iter([10.0, 10.0, 11.0])\n monkeypatch.setattr("haystack.components.agents.agent.time.monotonic", lambda: next(values))\n\n with caplog.at_level(logging.WARNING):\n result = agent.run(messages=[ChatMessage.from_user("Start")])\n\n assert chat_generator.calls == 1\n assert result["last_message"].tool_calls[0].tool_name == "loop"\n assert "Agent reached maximum execution time of 1.0 seconds, stopping." in caplog.text\n\n @pytest.mark.asyncio\n async def test_stops_before_tool_invocation_when_time_limit_expires_async(\n self, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture\n ) -> None:\n agent, chat_generator = make_agent()\n values = iter([10.0, 10.0, 11.0])\n monkeypatch.setattr("haystack.components.agents.agent.time.monotonic", lambda: next(values))\n\n with caplog.at_level(logging.WARNING):\n result = await agent.run_async(messages=[ChatMessage.from_user("Start")])\n\n assert chat_generator.calls == 1\n assert result["last_message"].tool_calls[0].tool_name == "loop"\n assert "Agent reached maximum execution time of 1.0 seconds, stopping." in caplog.text\n''') + + Path("releasenotes/notes/add-agent-runtime-budget-f3b6ce29.yaml").write_text('''---\nfeatures:\n - |\n Added an optional ``max_agent_time_seconds`` runtime budget to ``Agent``.\n When configured, the agent stops before starting another component invocation\n after the budget expires, complementing ``max_agent_steps`` for deployments\n with slow or remote tools.\n''') + PY + + python -m compileall -q haystack/components/agents/agent.py test/components/agents/test_agent_time_limit.py + git diff --check + rm .github/workflows/apply-agent-runtime-budget.yml + git add haystack/components/agents/agent.py test/components/agents/test_agent.py test/components/agents/test_agent_time_limit.py releasenotes/notes/add-agent-runtime-budget-f3b6ce29.yaml .github/workflows/apply-agent-runtime-budget.yml + git config user.name "camgrimsec" + git config user.email "156701171+camgrimsec@users.noreply.github.com" + git commit -m "feat: add agent runtime budget" + git push origin HEAD:feat/agent-runtime-budget From d5a3e26cda78b8612b15cfbff5527ac2d0943a45 Mon Sep 17 00:00:00 2001 From: camgrimsec <156701171+camgrimsec@users.noreply.github.com> Date: Thu, 25 Jun 2026 18:48:31 +0000 Subject: [PATCH 2/2] feat: add agent runtime budget --- .../workflows/apply-agent-runtime-budget.yml | 132 ------------------ haystack/components/agents/agent.py | 43 ++++++ .../add-agent-runtime-budget-f3b6ce29.yaml | 7 + test/components/agents/test_agent.py | 5 + .../agents/test_agent_time_limit.py | 102 ++++++++++++++ 5 files changed, 157 insertions(+), 132 deletions(-) delete mode 100644 .github/workflows/apply-agent-runtime-budget.yml create mode 100644 releasenotes/notes/add-agent-runtime-budget-f3b6ce29.yaml create mode 100644 test/components/agents/test_agent_time_limit.py diff --git a/.github/workflows/apply-agent-runtime-budget.yml b/.github/workflows/apply-agent-runtime-budget.yml deleted file mode 100644 index 66d3aa671f6..00000000000 --- a/.github/workflows/apply-agent-runtime-budget.yml +++ /dev/null @@ -1,132 +0,0 @@ -name: Apply agent runtime budget patch - -on: - push: - branches: - - feat/agent-runtime-budget - -permissions: - contents: write - -jobs: - apply-patch: - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v4 - with: - ref: feat/agent-runtime-budget - - - name: Apply focused runtime-budget patch - run: | - python - <<'PY' - from pathlib import Path - - agent_path = Path("haystack/components/agents/agent.py") - text = agent_path.read_text() - - def replace_once(old: str, new: str) -> None: - global text - count = text.count(old) - if count != 1: - raise RuntimeError(f"Expected one occurrence, found {count}: {old[:80]!r}") - text = text.replace(old, new, 1) - - def replace_all(old: str, new: str, expected: int) -> None: - global text - count = text.count(old) - if count != expected: - raise RuntimeError(f"Expected {expected} occurrences, found {count}: {old[:80]!r}") - text = text.replace(old, new) - - replace_once("import inspect\nimport re\n", "import inspect\nimport re\nimport time\n") - replace_once( - " max_agent_steps: int = 100,\n streaming_callback: StreamingCallbackT | None = None,", - " max_agent_steps: int = 100,\n max_agent_time_seconds: float | None = None,\n streaming_callback: StreamingCallbackT | None = None,", - ) - replace_once( - " :param max_agent_steps: Maximum number of steps the agent will run before stopping. Defaults to 100.\n" - " If the agent exceeds this number of steps, it will stop and return the current state.\n" - " :param streaming_callback:", - " :param max_agent_steps: Maximum number of steps the agent will run before stopping. Defaults to 100.\n" - " If the agent exceeds this number of steps, it will stop and return the current state.\n" - " :param max_agent_time_seconds: Optional maximum wall-clock runtime for one agent invocation, in seconds.\n" - " The limit is checked between component invocations. A currently running synchronous tool must\n" - " enforce its own timeout. Defaults to None.\n" - " :param streaming_callback:", - ) - replace_once( - " # Validate state schema if provided\n", - " if max_agent_time_seconds is not None and max_agent_time_seconds <= 0:\n" - " raise ValueError(\"max_agent_time_seconds must be greater than zero when provided.\")\n\n" - " # Validate state schema if provided\n", - ) - replace_once( - " self.max_agent_steps = max_agent_steps\n self.raise_on_tool_invocation_failure", - " self.max_agent_steps = max_agent_steps\n self.max_agent_time_seconds = max_agent_time_seconds\n self.raise_on_tool_invocation_failure", - ) - replace_once( - " max_agent_steps=self.max_agent_steps,\n streaming_callback=", - " max_agent_steps=self.max_agent_steps,\n max_agent_time_seconds=self.max_agent_time_seconds,\n streaming_callback=", - ) - replace_once( - ' "haystack.agent.max_steps": self.max_agent_steps,\n "haystack.agent.tools": self.tools,', - ' "haystack.agent.max_steps": self.max_agent_steps,\n "haystack.agent.max_time_seconds": self.max_agent_time_seconds,\n "haystack.agent.tools": self.tools,', - ) - replace_once( - " parent_span=parent_span,\n )\n\n def _initialize_fresh_execution(", - " parent_span=parent_span,\n )\n\n" - " def _is_time_limit_exceeded(self, deadline: float | None) -> bool:\n" - " if deadline is None or time.monotonic() < deadline:\n" - " return False\n\n" - " logger.warning(\n" - " \"Agent reached maximum execution time of {max_agent_time_seconds} seconds, stopping.\",\n" - " max_agent_time_seconds=self.max_agent_time_seconds,\n" - " )\n" - " return True\n\n" - " def _initialize_fresh_execution(", - ) - replace_all( - " span.set_content_tag(\"haystack.agent.input\", agent_inputs)\n\n while exe_context.counter < self.max_agent_steps:\n", - " span.set_content_tag(\"haystack.agent.input\", agent_inputs)\n\n" - " deadline = (\n" - " time.monotonic() + self.max_agent_time_seconds\n" - " if self.max_agent_time_seconds is not None\n" - " else None\n" - " )\n\n" - " while exe_context.counter < self.max_agent_steps:\n" - " if self._is_time_limit_exceeded(deadline):\n" - " break\n", - 2, - ) - replace_all( - " # Exit for `exit_conditions=[\"text\"]` behavior: the agent stops when there is no tool invoker, or when\n", - " if self._is_time_limit_exceeded(deadline):\n" - " break\n\n" - " # Exit for `exit_conditions=[\"text\"]` behavior: the agent stops when there is no tool invoker, or when\n", - 2, - ) - agent_path.write_text(text) - - test_path = Path("test/components/agents/test_agent.py") - test_text = test_path.read_text() - expected = ' "max_agent_steps": 100,\n' - replacements = test_text.count(expected) - if replacements < 2: - raise RuntimeError(f"Expected serialized Agent fixtures, found {replacements}") - test_path.write_text( - test_text.replace(expected, expected + ' "max_agent_time_seconds": None,\n') - ) - - Path("test/components/agents/test_agent_time_limit.py").write_text('''# SPDX-FileCopyrightText: 2022-present deepset GmbH \n#\n# SPDX-License-Identifier: Apache-2.0\n\nimport logging\nfrom typing import Any\n\nimport pytest\n\nfrom haystack import component\nfrom haystack.components.agents import Agent\nfrom haystack.dataclasses import ChatMessage, ToolCall\nfrom haystack.tools import Tool, Toolset\n\n\n@component\nclass LoopingChatGenerator:\n def __init__(self) -> None:\n self.calls = 0\n\n def to_dict(self) -> dict[str, Any]:\n return {"type": "LoopingChatGenerator", "data": {}}\n\n @classmethod\n def from_dict(cls, data: dict[str, Any]) -> "LoopingChatGenerator":\n return cls()\n\n @component.output_types(replies=list[ChatMessage])\n def run(self, messages: list[ChatMessage], tools: list[Tool] | Toolset | None = None, **kwargs: Any) -> dict[str, Any]:\n self.calls += 1\n return {\n "replies": [\n ChatMessage.from_assistant(\n "Calling the loop tool.",\n tool_calls=[ToolCall(tool_name="loop", arguments={})],\n )\n ]\n }\n\n @component.output_types(replies=list[ChatMessage])\n async def run_async(\n self, messages: list[ChatMessage], tools: list[Tool] | Toolset | None = None, **kwargs: Any\n ) -> dict[str, Any]:\n return self.run(messages=messages, tools=tools, **kwargs)\n\n\ndef loop_tool() -> str:\n return "looped"\n\n\ndef make_agent() -> tuple[Agent, LoopingChatGenerator]:\n chat_generator = LoopingChatGenerator()\n tool = Tool(\n name="loop",\n description="Return a loop marker.",\n parameters={"type": "object", "properties": {}},\n function=loop_tool,\n )\n return (\n Agent(chat_generator=chat_generator, tools=[tool], max_agent_steps=10, max_agent_time_seconds=1.0),\n chat_generator,\n )\n\n\nclass TestAgentTimeLimit:\n def test_requires_positive_time_limit(self) -> None:\n with pytest.raises(ValueError, match="max_agent_time_seconds must be greater than zero"):\n Agent(chat_generator=LoopingChatGenerator(), max_agent_time_seconds=0)\n\n def test_serializes_time_limit(self) -> None:\n agent, _ = make_agent()\n\n assert agent.to_dict()["init_parameters"]["max_agent_time_seconds"] == 1.0\n\n def test_stops_before_tool_invocation_when_time_limit_expires(\n self, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture\n ) -> None:\n agent, chat_generator = make_agent()\n values = iter([10.0, 10.0, 11.0])\n monkeypatch.setattr("haystack.components.agents.agent.time.monotonic", lambda: next(values))\n\n with caplog.at_level(logging.WARNING):\n result = agent.run(messages=[ChatMessage.from_user("Start")])\n\n assert chat_generator.calls == 1\n assert result["last_message"].tool_calls[0].tool_name == "loop"\n assert "Agent reached maximum execution time of 1.0 seconds, stopping." in caplog.text\n\n @pytest.mark.asyncio\n async def test_stops_before_tool_invocation_when_time_limit_expires_async(\n self, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture\n ) -> None:\n agent, chat_generator = make_agent()\n values = iter([10.0, 10.0, 11.0])\n monkeypatch.setattr("haystack.components.agents.agent.time.monotonic", lambda: next(values))\n\n with caplog.at_level(logging.WARNING):\n result = await agent.run_async(messages=[ChatMessage.from_user("Start")])\n\n assert chat_generator.calls == 1\n assert result["last_message"].tool_calls[0].tool_name == "loop"\n assert "Agent reached maximum execution time of 1.0 seconds, stopping." in caplog.text\n''') - - Path("releasenotes/notes/add-agent-runtime-budget-f3b6ce29.yaml").write_text('''---\nfeatures:\n - |\n Added an optional ``max_agent_time_seconds`` runtime budget to ``Agent``.\n When configured, the agent stops before starting another component invocation\n after the budget expires, complementing ``max_agent_steps`` for deployments\n with slow or remote tools.\n''') - PY - - python -m compileall -q haystack/components/agents/agent.py test/components/agents/test_agent_time_limit.py - git diff --check - rm .github/workflows/apply-agent-runtime-budget.yml - git add haystack/components/agents/agent.py test/components/agents/test_agent.py test/components/agents/test_agent_time_limit.py releasenotes/notes/add-agent-runtime-budget-f3b6ce29.yaml .github/workflows/apply-agent-runtime-budget.yml - git config user.name "camgrimsec" - git config user.email "156701171+camgrimsec@users.noreply.github.com" - git commit -m "feat: add agent runtime budget" - git push origin HEAD:feat/agent-runtime-budget 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