From 798a4b9e5a39848b9b7da0ab368d7ff22e081b50 Mon Sep 17 00:00:00 2001 From: Eric Evans <194135482+ericevans-nv@users.noreply.github.com> Date: Fri, 6 Mar 2026 06:42:35 -0600 Subject: [PATCH 1/2] Add LLM error recovery, retry logic, and diagnostic logging to ToolCallAgentGraph Signed-off-by: Eric Evans <194135482+ericevans-nv@users.noreply.github.com> --- .../src/nat/llm/utils/constants.py | 9 + .../agent/tool_calling_agent/agent.py | 230 ++++++++++++++- .../agent/tool_calling_agent/register.py | 38 ++- .../tests/agent/test_tool_calling.py | 272 ++++++++++++++---- 4 files changed, 478 insertions(+), 71 deletions(-) diff --git a/packages/nvidia_nat_core/src/nat/llm/utils/constants.py b/packages/nvidia_nat_core/src/nat/llm/utils/constants.py index 9ad6d062db..52cdf94e65 100644 --- a/packages/nvidia_nat_core/src/nat/llm/utils/constants.py +++ b/packages/nvidia_nat_core/src/nat/llm/utils/constants.py @@ -21,3 +21,12 @@ class LLMHeaderPrefix(StrEnum): PAYLOAD = "X-Payload" # Custom metadata propagation DYNAMO = "x-prefix" # Dynamo KV cache routing + + +class LLMFinishReason(StrEnum): + """LLM finish_reason values used for response validation.""" + + STOP = "stop" + LENGTH = "length" + CONTENT_FILTER = "content_filter" + TOOL_CALLS = "tool_calls" diff --git a/packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/tool_calling_agent/agent.py b/packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/tool_calling_agent/agent.py index 446613f0c6..09a69e3a84 100644 --- a/packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/tool_calling_agent/agent.py +++ b/packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/tool_calling_agent/agent.py @@ -22,6 +22,7 @@ from langchain_core.messages import AIMessageChunk from langchain_core.messages import SystemMessage from langchain_core.messages import ToolMessage +from langchain_core.messages.ai import UsageMetadata from langchain_core.messages.base import BaseMessage from langchain_core.messages.utils import convert_to_openai_messages from langchain_core.runnables import RunnableLambda @@ -32,6 +33,7 @@ from pydantic import BaseModel from pydantic import Field +from nat.llm.utils.constants import LLMFinishReason from nat.plugins.langchain.agent.base import AGENT_CALL_LOG_MESSAGE from nat.plugins.langchain.agent.base import AGENT_LOG_PREFIX from nat.plugins.langchain.agent.base import AgentDecision @@ -86,6 +88,9 @@ def __init__( log_response_max_chars: int = 1000, handle_tool_errors: bool = True, return_direct: list[BaseTool] | None = None, + max_truncation_retries: int = 0, + truncation_token_increment: int = 1024, + max_empty_response_retries: int = 0, ): super().__init__(llm=llm, tools=tools, @@ -117,24 +122,47 @@ def __init__( self.agent = prompt_runnable | self.bound_llm self.tool_caller = ToolNode(tools, handle_tool_errors=handle_tool_errors) self.return_direct = [tool.name for tool in return_direct] if return_direct else [] + + self._max_truncation_retries: int = max_truncation_retries + self._truncation_token_increment: int = truncation_token_increment + self._max_empty_response_retries: int = max_empty_response_retries + logger.debug("%s Initialized Tool Calling Agent Graph", AGENT_LOG_PREFIX) + async def _invoke_llm(self, state: ToolCallAgentGraphState): + """Stream the LLM and return the accumulated AIMessage response. + + Args: + state: Current agent graph state containing the conversation messages. + + Returns: + The fully accumulated AIMessage from the LLM. + + Raises: + RuntimeError: If the LLM returns no response. + """ + # Use astream so LangGraph's stream_mode="messages" can observe individual LLM tokens. + # Config is inherited from LangGraph's context, preserving streaming callbacks. + response = None + async for chunk in self.agent.astream({"messages": state.messages}): + response = chunk if response is None else response + chunk + if response is None: + raise RuntimeError('No response received from agent') + return response + async def agent_node(self, state: ToolCallAgentGraphState): try: logger.debug("%s Starting the Tool Calling Agent Node", AGENT_LOG_PREFIX) if len(state.messages) == 0: raise RuntimeError('No input received in state: "messages"') - # Use astream so LangGraph's stream_mode="messages" can observe individual LLM tokens. - # Config is inherited from LangGraph's context, preserving streaming callbacks. - response = None - async for chunk in self.agent.astream({"messages": state.messages}): - response = chunk if response is None else response + chunk - if response is None: - raise RuntimeError('No response received from agent') + + response = await self._invoke_llm(state) if isinstance(response, AIMessageChunk): response = _chunk_to_message(response) + response = await self._validate_llm_response(response, state) + if self.detailed_logs: agent_input = "\n".join(str(message.content) for message in state.messages) logger.info(AGENT_CALL_LOG_MESSAGE, agent_input, response) @@ -145,6 +173,193 @@ async def agent_node(self, state: ToolCallAgentGraphState): logger.error("%s Failed to call agent_node: %s", AGENT_LOG_PREFIX, ex) raise + def _get_token_usage(self, response: AIMessage) -> UsageMetadata: + """Extract token usage from an AIMessage, checking both LangChain and OpenAI formats. + + Returns: + UsageMetadata with ``input_tokens``, ``output_tokens``, ``total_tokens`` + (values default to 0 if unavailable). + """ + usage_meta: UsageMetadata | None = getattr(response, "usage_metadata", None) + if usage_meta: + return UsageMetadata( + input_tokens=usage_meta.get("input_tokens", 0), + output_tokens=usage_meta.get("output_tokens", 0), + total_tokens=usage_meta.get("total_tokens", 0), + ) + resp_meta: dict = getattr(response, "response_metadata", {}) + openai_usage: dict = resp_meta.get("usage", {}) + if openai_usage: + return UsageMetadata( + input_tokens=openai_usage.get("prompt_tokens", 0), + output_tokens=openai_usage.get("completion_tokens", 0), + total_tokens=openai_usage.get("total_tokens", 0), + ) + return UsageMetadata(input_tokens=0, output_tokens=0, total_tokens=0) + + async def _validate_llm_response(self, response: AIMessage, state: ToolCallAgentGraphState) -> AIMessage: + """Validate the LLM response and attempt recovery if configured. + + Args: + response: The accumulated AIMessage from the LLM. + state: Current agent graph state. + + Returns: + A validated (possibly retried) response. + """ + metadata: dict = getattr(response, "response_metadata", {}) + finish_reason: str | None = metadata.get("finish_reason") + + # Output truncated — can cause GraphRecursionError loop if unhandled + if finish_reason == LLMFinishReason.LENGTH: + if self._max_truncation_retries > 0: + response = await self._retry_on_truncation(response, state) + else: + usage: UsageMetadata = self._get_token_usage(response) + truncated_content: str = str(response.content)[:500] if response.content else "" + model_name: str = metadata.get("model_name", "unknown") + msg: str = (f"LLM output truncated (finish_reason='length'). " + f"model={model_name}, " + f"output_tokens={usage['output_tokens']}, " + f"input_tokens={usage['input_tokens']}, " + f"total_tokens={usage['total_tokens']}. " + f"Truncated output: {truncated_content!r}") + raise RuntimeError(msg) + + # Empty response — no content, no tool calls at all + if (not response.tool_calls and not getattr(response, "invalid_tool_calls", None) + and not (response.content and str(response.content).strip())): + if self._max_empty_response_retries > 0: + response = await self._retry_on_empty_response(state, metadata) + else: + logger.error( + "%s LLM returned an empty response (no content, no tool calls). " + "finish_reason=%s, response_metadata=%s", + AGENT_LOG_PREFIX, + finish_reason, + metadata, + ) + + # Invalid tool calls — LLM produced tool call JSON that could not be parsed + if not response.tool_calls and getattr(response, "invalid_tool_calls", None): + raw_calls: list = response.invalid_tool_calls + logger.warning( + "%s LLM produced %d invalid tool call(s) that could not be parsed. " + "raw_calls=%s, response_metadata=%s", + AGENT_LOG_PREFIX, + len(raw_calls), + raw_calls, + metadata, + ) + + # Content filter — LLM provider blocked the response + if finish_reason == LLMFinishReason.CONTENT_FILTER: + logger.error( + "%s LLM response blocked by content filter (finish_reason='content_filter'). " + "response_metadata=%s", + AGENT_LOG_PREFIX, + metadata, + ) + + return response + + async def _retry_on_truncation(self, first_response: AIMessage, state: ToolCallAgentGraphState) -> AIMessage: + """Retry the LLM call with a higher max_tokens when the output is truncated. + + Without recovery, a truncated response typically lacks valid tool calls + or a complete answer, causing the agent to loop until it hits a + GraphRecursionError. Each retry adds a fixed token increment to + max_tokens, giving the LLM room to finish its output. + + Args: + first_response: The truncated AIMessage from the LLM. + state: Current agent graph state. + + Returns: + The response from the successful retry, or the last truncated + response if all retries are exhausted. + """ + configured_max: int | None = getattr(self.llm, "max_tokens", None) + if configured_max is not None: + current_tokens: int = configured_max + else: + usage: UsageMetadata = self._get_token_usage(first_response) + current_tokens = usage.get("output_tokens", 0) or 4096 + + response: AIMessage | None = None + + for attempt in range(1, self._max_truncation_retries + 1): + new_limit: int = current_tokens + self._truncation_token_increment + logger.warning( + "%s Output truncated (finish_reason=length, current_max_tokens=%s). " + "Retry %d/%d with max_tokens=%d (+%d)", + AGENT_LOG_PREFIX, + current_tokens, + attempt, + self._max_truncation_retries, + new_limit, + self._truncation_token_increment, + ) + + self.bound_llm = self.llm.bind_tools(self.tools).bind(max_tokens=new_limit) + self.agent = self.agent.first | self.bound_llm # type: ignore[union-attr] + + response = await self._invoke_llm(state) + metadata: dict = getattr(response, "response_metadata", {}) + finish_reason: str | None = metadata.get("finish_reason") + + if finish_reason != LLMFinishReason.LENGTH: + logger.info("%s Truncation retry succeeded on attempt %d with max_tokens=%d", + AGENT_LOG_PREFIX, + attempt, + new_limit) + return response + + current_tokens = new_limit + + logger.warning("%s All %d truncation retries exhausted", AGENT_LOG_PREFIX, self._max_truncation_retries) + return response + + async def _retry_on_empty_response(self, state: ToolCallAgentGraphState, first_metadata: dict) -> AIMessage: + """Retry the LLM call when it returns an empty response. + + Args: + state: Current agent graph state. + first_metadata: The response_metadata from the empty response. + + Returns: + The response from the successful retry, or the last empty + response if all retries are exhausted. + """ + response: AIMessage | None = None + for attempt in range(1, self._max_empty_response_retries + 1): + logger.warning( + "%s LLM returned an empty response. Retry %d/%d. response_metadata=%s", + AGENT_LOG_PREFIX, + attempt, + self._max_empty_response_retries, + first_metadata, + ) + response = await self._invoke_llm(state) + + has_content: bool = bool(response.content and str(response.content).strip()) + if response.tool_calls or has_content: + logger.info( + "%s Empty response retry succeeded on attempt %d", + AGENT_LOG_PREFIX, + attempt, + ) + return response + + first_metadata = getattr(response, "response_metadata", {}) + + logger.warning( + "%s All %d empty response retries exhausted", + AGENT_LOG_PREFIX, + self._max_empty_response_retries, + ) + return response + async def conditional_edge(self, state: ToolCallAgentGraphState): try: logger.debug("%s Starting the Tool Calling Conditional Edge", AGENT_LOG_PREFIX) @@ -167,6 +382,7 @@ async def tool_node(self, state: ToolCallAgentGraphState): tool_calls = state.messages[-1].tool_calls tools = [tool.get("name") for tool in tool_calls] tool_input = state.messages[-1] + tool_response = await self.tool_caller.ainvoke( input={"messages": [tool_input]}, config=self._runnable_config, diff --git a/packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/tool_calling_agent/register.py b/packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/tool_calling_agent/register.py index f5d9ad2154..bcc1f1a34a 100644 --- a/packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/tool_calling_agent/register.py +++ b/packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/tool_calling_agent/register.py @@ -19,6 +19,7 @@ import uuid from collections.abc import AsyncGenerator +from pydantic import BaseModel from pydantic import Field from nat.builder.builder import Builder @@ -40,6 +41,15 @@ logger = logging.getLogger(__name__) +class TruncationRetryConfig(BaseModel): + """Configuration for retrying LLM calls that are truncated (finish_reason='length').""" + + max_retries: int = Field(default=0, + description="Number of retries when LLM output is truncated. " + "0 disables recovery (raises RuntimeError).") + token_increment: int = Field(default=1024, description="Fixed number of tokens added to max_tokens on each retry.") + + class ToolCallAgentWorkflowConfig(AgentBaseConfig, name="tool_calling_agent"): """ A Tool Calling Agent requires an LLM which supports tool calling. A tool Calling Agent utilizes the tool @@ -52,6 +62,12 @@ class ToolCallAgentWorkflowConfig(AgentBaseConfig, name="tool_calling_agent"): max_iterations: int = Field(default=15, description="Number of tool calls before stoping the tool calling agent.") max_history: int = Field(default=15, description="Maximum number of messages to keep in the conversation history.") + truncation_retry: TruncationRetryConfig = Field(default_factory=TruncationRetryConfig, + description="Configuration for retrying truncated LLM responses.") + max_empty_response_retries: int = Field( + default=0, + description="Number of retries when LLM returns an empty response (no content, no tool calls). " + "0 disables recovery (raises RuntimeError).") system_prompt: str | None = Field(default=None, description="Provides the system prompt to use with the agent.") additional_instructions: str | None = Field(default=None, description="Additional instructions appended to the system prompt.") @@ -87,13 +103,18 @@ async def tool_calling_agent_workflow(config: ToolCallAgentWorkflowConfig, build tool_names=config.return_direct, wrapper_type=LLMFrameworkEnum.LANGCHAIN) if config.return_direct else None # construct the Tool Calling Agent Graph from the configured llm, and tools - graph: CompiledStateGraph = await ToolCallAgentGraph(llm=llm, - tools=tools, - prompt=prompt, - detailed_logs=config.verbose, - log_response_max_chars=config.log_response_max_chars, - handle_tool_errors=config.handle_tool_errors, - return_direct=return_direct_tools).build_graph() + graph: CompiledStateGraph = await ToolCallAgentGraph( + llm=llm, + tools=tools, + prompt=prompt, + detailed_logs=config.verbose, + log_response_max_chars=config.log_response_max_chars, + handle_tool_errors=config.handle_tool_errors, + return_direct=return_direct_tools, + max_truncation_retries=config.truncation_retry.max_retries, + truncation_token_increment=config.truncation_retry.token_increment, + max_empty_response_retries=config.max_empty_response_retries, + ).build_graph() async def _response_fn(chat_request_or_message: ChatRequestOrMessage) -> str: """ @@ -129,14 +150,17 @@ async def _response_fn(chat_request_or_message: ChatRequestOrMessage) -> str: state = ToolCallAgentGraphState(**state) output_message = state.messages[-1] return str(output_message.content) + except GraphRecursionError: logger.warning( "%s Tool Calling Agent reached its maximum iteration limit (%d) without producing a final answer. " "This typically means the LLM kept calling tools instead of returning a response.", AGENT_LOG_PREFIX, config.max_iterations) + return (f"The tool calling agent could not produce a final answer within {config.max_iterations} " "iterations. The agent repeatedly called tools without converging on a response.") + except Exception as ex: logger.error("%s Tool Calling Agent failed with exception: %s", AGENT_LOG_PREFIX, ex) raise diff --git a/packages/nvidia_nat_langchain/tests/agent/test_tool_calling.py b/packages/nvidia_nat_langchain/tests/agent/test_tool_calling.py index f050f20615..4a02134a60 100644 --- a/packages/nvidia_nat_langchain/tests/agent/test_tool_calling.py +++ b/packages/nvidia_nat_langchain/tests/agent/test_tool_calling.py @@ -15,10 +15,12 @@ import datetime import json +import typing +from unittest.mock import AsyncMock +from unittest.mock import patch import pytest from langchain_core.messages import AIMessage -from langchain_core.messages import AIMessageChunk from langchain_core.messages import HumanMessage from langchain_core.messages import ToolMessage from langgraph.graph.state import CompiledStateGraph @@ -32,7 +34,6 @@ from nat.plugins.langchain.agent.base import AgentDecision from nat.plugins.langchain.agent.tool_calling_agent.agent import ToolCallAgentGraph from nat.plugins.langchain.agent.tool_calling_agent.agent import ToolCallAgentGraphState -from nat.plugins.langchain.agent.tool_calling_agent.agent import _chunk_to_message from nat.plugins.langchain.agent.tool_calling_agent.agent import create_tool_calling_agent_prompt from nat.plugins.langchain.agent.tool_calling_agent.register import ToolCallAgentWorkflowConfig @@ -380,66 +381,223 @@ def test_tool_call_chunk_serialization(): assert tc["function"]["arguments"] == "" -def test_tool_call_chunk_arguments_streaming(): - """Test that incremental tool call argument chunks serialize correctly (no id/type on follow-up chunks).""" - chunk = ChatResponseChunk( - id="test-chunk-id", - choices=[ - ChatResponseChunkChoice( - index=0, - delta=ChoiceDelta(tool_calls=[ - ChoiceDeltaToolCall(index=0, - id=None, - type=None, - function=ChoiceDeltaToolCallFunction( - name=None, - arguments='{"author":', - )) - ]), - finish_reason=None, - ) - ], - created=datetime.datetime(2026, 1, 1, tzinfo=datetime.UTC), +@pytest.fixture(name="error_mock_llm") +def fixture_error_mock_llm(): + """Mock LLM with a settable ``max_tokens`` field for truncation retry tests.""" + from langchain_core.callbacks import AsyncCallbackManagerForLLMRun + from langchain_core.callbacks import CallbackManagerForLLMRun + from langchain_core.language_models import BaseChatModel + from langchain_core.messages import BaseMessage + from langchain_core.outputs import ChatGeneration + from langchain_core.outputs import ChatResult + + class _ErrorMockLLM(BaseChatModel): + max_tokens: int | None = None + + async def _agenerate( + self, + messages: list[BaseMessage], + stop: list[str] | None = None, + run_manager: AsyncCallbackManagerForLLMRun | None = None, + **kwargs: typing.Any, + ) -> ChatResult: + msg = AIMessage(content=messages[-1].content, response_metadata={"mock_llm_response": True}) + return ChatResult(generations=[ChatGeneration(message=msg)]) + + def _generate( + self, + messages: list[BaseMessage], + stop: list[str] | None = None, + run_manager: CallbackManagerForLLMRun | None = None, + **kwargs: typing.Any, + ) -> ChatResult: + msg = AIMessage(content=messages[-1].content, response_metadata={"mock_llm_response": True}) + return ChatResult(generations=[ChatGeneration(message=msg)]) + + def bind_tools(self, tools, **kwargs): + return self + + @property + def _llm_type(self) -> str: + return "error-mock-llm" + + return _ErrorMockLLM() + + +def _make_agent(llm, mock_tool, **kwargs) -> ToolCallAgentGraph: + """Helper to build a ``ToolCallAgentGraph`` with a single tool and custom params.""" + return ToolCallAgentGraph(llm=llm, tools=[mock_tool("T")], **kwargs) + + +async def test_validate_truncation_raises_when_disabled(error_mock_llm, mock_tool): + """finish_reason=length with max_truncation_retries=0 must raise RuntimeError with token info.""" + agent = _make_agent(error_mock_llm, mock_tool, max_truncation_retries=0) + response = AIMessage( + content="partial output here", + response_metadata={ + "finish_reason": "length", "model_name": "test-model" + }, + usage_metadata={ + "output_tokens": 100, "input_tokens": 50, "total_tokens": 150 + }, ) + state = ToolCallAgentGraphState(messages=[HumanMessage(content="test")]) - payload = json.loads(chunk.model_dump_json()) - tc = payload["choices"][0]["delta"]["tool_calls"][0] - assert tc["index"] == 0 - assert tc["id"] is None - assert tc["type"] is None - assert tc["function"]["arguments"] == '{"author":' + with pytest.raises(RuntimeError, match="truncated") as exc_info: + await agent._validate_llm_response(response, state) + msg: str = str(exc_info.value) + assert "model=test-model" in msg + assert "output_tokens=100" in msg + assert "input_tokens=50" in msg + assert "total_tokens=150" in msg + assert "partial output here" in msg -def test_content_chunk_has_no_tool_calls(): - """Test that a content-only ChatResponseChunk does not include tool_calls in the delta.""" - chunk = ChatResponseChunk.create_streaming_chunk("Hello world", id_="test-id") - payload = json.loads(chunk.model_dump_json()) - delta = payload["choices"][0]["delta"] - assert delta["content"] == "Hello world" - assert delta.get("tool_calls") is None - - -def test_chunk_to_message_rehydrates_empty_tool_calls(): - """Test that _chunk_to_message rehydrates when additional_kwargs["tool_calls"] is an empty list.""" - chunk = AIMessageChunk( - content="", - tool_call_chunks=[{ - "name": "test_tool", - "args": '{"query": "hello"}', - "id": "call_abc123", - "index": 0, - "type": "tool_call_chunk", - }], - additional_kwargs={"tool_calls": []}, + +async def test_validate_truncation_delegates_to_retry(error_mock_llm, mock_tool): + """finish_reason=length with max_truncation_retries>0 delegates to _retry_on_truncation.""" + agent = _make_agent(error_mock_llm, mock_tool, max_truncation_retries=2, truncation_token_increment=512) + truncated = AIMessage(content="partial", response_metadata={"finish_reason": "length"}) + good = AIMessage(content="complete", response_metadata={"finish_reason": "stop"}) + state = ToolCallAgentGraphState(messages=[HumanMessage(content="test")]) + + with patch.object(agent, "_retry_on_truncation", new_callable=AsyncMock, return_value=good) as mock_retry: + result = await agent._validate_llm_response(truncated, state) + + mock_retry.assert_awaited_once() + assert result.content == "complete" + + +async def test_validate_empty_response_delegates_to_retry(error_mock_llm, mock_tool): + """Empty response with retries>0 delegates to _retry_on_empty_response.""" + agent = _make_agent(error_mock_llm, mock_tool, max_empty_response_retries=2) + empty = AIMessage(content="", response_metadata={"finish_reason": "stop"}) + good = AIMessage(content="actual answer", response_metadata={"finish_reason": "stop"}) + state = ToolCallAgentGraphState(messages=[HumanMessage(content="test")]) + + with patch.object(agent, "_retry_on_empty_response", new_callable=AsyncMock, return_value=good) as mock_retry: + result = await agent._validate_llm_response(empty, state) + + mock_retry.assert_awaited_once() + assert result.content == "actual answer" + + +def test_get_token_usage_from_openai_response_metadata(error_mock_llm, mock_tool): + """Falls back to response_metadata['usage'] (OpenAI format) when usage_metadata is absent.""" + agent = _make_agent(error_mock_llm, mock_tool) + response = AIMessage( + content="hi", + response_metadata={"usage": { + "prompt_tokens": 5, "completion_tokens": 15, "total_tokens": 20 + }}, ) + usage = agent._get_token_usage(response) + assert usage["input_tokens"] == 5 + assert usage["output_tokens"] == 15 + assert usage["total_tokens"] == 20 + + +async def test_retry_on_truncation_succeeds(error_mock_llm, mock_tool): + """Truncation retry succeeds when the retried LLM call finishes normally.""" + agent = _make_agent(error_mock_llm, mock_tool, max_truncation_retries=3, truncation_token_increment=512) + first_response = AIMessage( + content="partial", + response_metadata={"finish_reason": "length"}, + usage_metadata={ + "output_tokens": 100, "input_tokens": 50, "total_tokens": 150 + }, + ) + good = AIMessage(content="complete", response_metadata={"finish_reason": "stop"}) + state = ToolCallAgentGraphState(messages=[HumanMessage(content="test")]) - assert chunk.tool_calls, "Chunk should have parsed tool_calls" - assert not chunk.additional_kwargs.get("tool_calls"), "Wire format tool_calls should be empty" + with patch.object(agent, "_invoke_llm", new_callable=AsyncMock, return_value=good): + result = await agent._retry_on_truncation(first_response, state) - msg = _chunk_to_message(chunk) + assert result.content == "complete" + assert result.response_metadata["finish_reason"] == "stop" - assert isinstance(msg, AIMessage) - assert len(msg.additional_kwargs["tool_calls"]) == 1 - tc = msg.additional_kwargs["tool_calls"][0] - assert tc["function"]["name"] == "test_tool" - assert json.loads(tc["function"]["arguments"]) == {"query": "hello"} + +async def test_retry_on_truncation_exhausted(error_mock_llm, mock_tool): + """All truncation retries exhausted returns the last truncated response.""" + agent = _make_agent(error_mock_llm, mock_tool, max_truncation_retries=2, truncation_token_increment=512) + first_response = AIMessage( + content="partial", + response_metadata={"finish_reason": "length"}, + usage_metadata={ + "output_tokens": 100, "input_tokens": 50, "total_tokens": 150 + }, + ) + still_truncated = AIMessage(content="partial", response_metadata={"finish_reason": "length"}) + state = ToolCallAgentGraphState(messages=[HumanMessage(content="test")]) + + with patch.object(agent, "_invoke_llm", new_callable=AsyncMock, return_value=still_truncated): + result = await agent._retry_on_truncation(first_response, state) + + assert result.response_metadata["finish_reason"] == "length" + + +async def test_retry_on_truncation_increments_from_usage(error_mock_llm, mock_tool): + """When max_tokens is not configured, the base is taken from usage_metadata output_tokens.""" + increment: int = 256 + agent = _make_agent(error_mock_llm, mock_tool, max_truncation_retries=3, truncation_token_increment=increment) + assert agent.llm.max_tokens is None + first_response = AIMessage( + content="partial", + response_metadata={"finish_reason": "length"}, + usage_metadata={ + "output_tokens": 100, "input_tokens": 50, "total_tokens": 150 + }, + ) + still_truncated = AIMessage(content="partial", response_metadata={"finish_reason": "length"}) + state = ToolCallAgentGraphState(messages=[HumanMessage(content="test")]) + + observed_max_tokens: list[int | None] = [] + + async def _capture_and_invoke(s): + observed_max_tokens.append(getattr(agent.bound_llm, "kwargs", {}).get("max_tokens")) + return still_truncated + + with patch.object(agent, "_invoke_llm", side_effect=_capture_and_invoke): + await agent._retry_on_truncation(first_response, state) + + assert observed_max_tokens == [100 + increment, 100 + 2 * increment, 100 + 3 * increment] + + +async def test_retry_on_truncation_increments_from_configured_max_tokens(error_mock_llm, mock_tool): + """When max_tokens is already configured on the LLM, it is used as the base instead of usage data.""" + increment: int = 512 + agent = _make_agent(error_mock_llm, mock_tool, max_truncation_retries=2, truncation_token_increment=increment) + agent.llm.max_tokens = 50 + first_response = AIMessage( + content="partial", + response_metadata={"finish_reason": "length"}, + usage_metadata={ + "output_tokens": 48, "input_tokens": 30, "total_tokens": 78 + }, + ) + still_truncated = AIMessage(content="partial", response_metadata={"finish_reason": "length"}) + state = ToolCallAgentGraphState(messages=[HumanMessage(content="test")]) + + observed_max_tokens: list[int | None] = [] + + async def _capture_and_invoke(s): + observed_max_tokens.append(getattr(agent.bound_llm, "kwargs", {}).get("max_tokens")) + return still_truncated + + with patch.object(agent, "_invoke_llm", side_effect=_capture_and_invoke): + await agent._retry_on_truncation(first_response, state) + + assert observed_max_tokens == [50 + increment, 50 + 2 * increment] + + +async def test_retry_on_empty_succeeds(error_mock_llm, mock_tool): + """Empty-response retry succeeds when the retried LLM call returns content.""" + agent = _make_agent(error_mock_llm, mock_tool, max_empty_response_retries=3) + first_meta: dict = {"finish_reason": "stop"} + good = AIMessage(content="actual answer", response_metadata={"finish_reason": "stop"}) + state = ToolCallAgentGraphState(messages=[HumanMessage(content="test")]) + + with patch.object(agent, "_invoke_llm", new_callable=AsyncMock, return_value=good): + result = await agent._retry_on_empty_response(state, first_meta) + + assert result.content == "actual answer" From bf7438575e90d1319f8cdc3ee98cc682439ee1cb Mon Sep 17 00:00:00 2001 From: Eric Evans <194135482+ericevans-nv@users.noreply.github.com> Date: Fri, 6 Mar 2026 15:19:07 -0600 Subject: [PATCH 2/2] Add callable scaling strategy and persistent retry state for truncation recovery Signed-off-by: Eric Evans <194135482+ericevans-nv@users.noreply.github.com> --- .../agent/tool_calling_agent/agent.py | 92 +++++------ .../agent/tool_calling_agent/register.py | 28 +++- .../tests/agent/test_tool_calling.py | 143 ++++++++++++++++-- 3 files changed, 202 insertions(+), 61 deletions(-) diff --git a/packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/tool_calling_agent/agent.py b/packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/tool_calling_agent/agent.py index 09a69e3a84..b6eb73e062 100644 --- a/packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/tool_calling_agent/agent.py +++ b/packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/tool_calling_agent/agent.py @@ -89,7 +89,7 @@ def __init__( handle_tool_errors: bool = True, return_direct: list[BaseTool] | None = None, max_truncation_retries: int = 0, - truncation_token_increment: int = 1024, + truncation_scaling_fn: typing.Callable[[int], int] | None = None, max_empty_response_retries: int = 0, ): super().__init__(llm=llm, @@ -124,7 +124,9 @@ def __init__( self.return_direct = [tool.name for tool in return_direct] if return_direct else [] self._max_truncation_retries: int = max_truncation_retries - self._truncation_token_increment: int = truncation_token_increment + self._truncation_retries_remaining: int = max_truncation_retries + self._truncation_scaling_fn: typing.Callable[[int], int] = truncation_scaling_fn or (lambda c: c + 1024) + self._current_max_tokens: int | None = getattr(llm, "max_tokens", None) self._max_empty_response_retries: int = max_empty_response_retries logger.debug("%s Initialized Tool Calling Agent Graph", AGENT_LOG_PREFIX) @@ -143,11 +145,14 @@ async def _invoke_llm(self, state: ToolCallAgentGraphState): """ # Use astream so LangGraph's stream_mode="messages" can observe individual LLM tokens. # Config is inherited from LangGraph's context, preserving streaming callbacks. - response = None + chunks: list[AIMessageChunk] = [] async for chunk in self.agent.astream({"messages": state.messages}): - response = chunk if response is None else response + chunk - if response is None: + chunks.append(chunk) + if not chunks: raise RuntimeError('No response received from agent') + response: AIMessageChunk = chunks[0] + for c in chunks[1:]: + response = response + c return response async def agent_node(self, state: ToolCallAgentGraphState): @@ -232,13 +237,8 @@ async def _validate_llm_response(self, response: AIMessage, state: ToolCallAgent if self._max_empty_response_retries > 0: response = await self._retry_on_empty_response(state, metadata) else: - logger.error( - "%s LLM returned an empty response (no content, no tool calls). " - "finish_reason=%s, response_metadata=%s", - AGENT_LOG_PREFIX, - finish_reason, - metadata, - ) + raise RuntimeError(f"LLM returned an empty response (no content, no tool calls). " + f"finish_reason={finish_reason}, response_metadata={metadata}") # Invalid tool calls — LLM produced tool call JSON that could not be parsed if not response.tool_calls and getattr(response, "invalid_tool_calls", None): @@ -268,39 +268,41 @@ async def _retry_on_truncation(self, first_response: AIMessage, state: ToolCallA Without recovery, a truncated response typically lacks valid tool calls or a complete answer, causing the agent to loop until it hits a - GraphRecursionError. Each retry adds a fixed token increment to - max_tokens, giving the LLM room to finish its output. + GraphRecursionError. Each retry increases max_tokens via the configured + scaling callable, giving the LLM room to finish its output. Args: first_response: The truncated AIMessage from the LLM. state: Current agent graph state. Returns: - The response from the successful retry, or the last truncated - response if all retries are exhausted. + The response from the successful retry. + + Raises: + RuntimeError: If all retries are exhausted without a non-truncated response. """ - configured_max: int | None = getattr(self.llm, "max_tokens", None) - if configured_max is not None: - current_tokens: int = configured_max - else: + if self._current_max_tokens is None: usage: UsageMetadata = self._get_token_usage(first_response) - current_tokens = usage.get("output_tokens", 0) or 4096 + self._current_max_tokens = usage.get("output_tokens", 0) or 4096 response: AIMessage | None = None - for attempt in range(1, self._max_truncation_retries + 1): - new_limit: int = current_tokens + self._truncation_token_increment + while self._truncation_retries_remaining > 0: + self._truncation_retries_remaining -= 1 + new_limit: int = self._truncation_scaling_fn(self._current_max_tokens) + retries_used: int = self._max_truncation_retries - self._truncation_retries_remaining + logger.warning( - "%s Output truncated (finish_reason=length, current_max_tokens=%s). " - "Retry %d/%d with max_tokens=%d (+%d)", + "%s Output truncated (finish_reason=length, current_max_tokens=%d). " + "Retry %d/%d with max_tokens=%d", AGENT_LOG_PREFIX, - current_tokens, - attempt, + self._current_max_tokens, + retries_used, self._max_truncation_retries, new_limit, - self._truncation_token_increment, ) + self._current_max_tokens = new_limit self.bound_llm = self.llm.bind_tools(self.tools).bind(max_tokens=new_limit) self.agent = self.agent.first | self.bound_llm # type: ignore[union-attr] @@ -309,16 +311,20 @@ async def _retry_on_truncation(self, first_response: AIMessage, state: ToolCallA finish_reason: str | None = metadata.get("finish_reason") if finish_reason != LLMFinishReason.LENGTH: - logger.info("%s Truncation retry succeeded on attempt %d with max_tokens=%d", - AGENT_LOG_PREFIX, - attempt, - new_limit) + logger.info( + "%s Truncation retry succeeded on attempt %d with max_tokens=%d", + AGENT_LOG_PREFIX, + retries_used, + new_limit, + ) return response - current_tokens = new_limit - - logger.warning("%s All %d truncation retries exhausted", AGENT_LOG_PREFIX, self._max_truncation_retries) - return response + usage = self._get_token_usage(response or first_response) + raise RuntimeError(f"LLM output still truncated after {self._max_truncation_retries} retries " + f"(last max_tokens={self._current_max_tokens}). " + f"output_tokens={usage.get('output_tokens', 'N/A')}, " + f"input_tokens={usage.get('input_tokens', 'N/A')}, " + f"total_tokens={usage.get('total_tokens', 'N/A')}") async def _retry_on_empty_response(self, state: ToolCallAgentGraphState, first_metadata: dict) -> AIMessage: """Retry the LLM call when it returns an empty response. @@ -328,8 +334,10 @@ async def _retry_on_empty_response(self, state: ToolCallAgentGraphState, first_m first_metadata: The response_metadata from the empty response. Returns: - The response from the successful retry, or the last empty - response if all retries are exhausted. + The response from the successful retry. + + Raises: + RuntimeError: If all retries are exhausted without a non-empty response. """ response: AIMessage | None = None for attempt in range(1, self._max_empty_response_retries + 1): @@ -353,12 +361,8 @@ async def _retry_on_empty_response(self, state: ToolCallAgentGraphState, first_m first_metadata = getattr(response, "response_metadata", {}) - logger.warning( - "%s All %d empty response retries exhausted", - AGENT_LOG_PREFIX, - self._max_empty_response_retries, - ) - return response + raise RuntimeError(f"LLM still returning empty responses after {self._max_empty_response_retries} retries. " + f"response_metadata={first_metadata}") async def conditional_edge(self, state: ToolCallAgentGraphState): try: diff --git a/packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/tool_calling_agent/register.py b/packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/tool_calling_agent/register.py index bcc1f1a34a..3e1ae689da 100644 --- a/packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/tool_calling_agent/register.py +++ b/packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/tool_calling_agent/register.py @@ -18,9 +18,11 @@ import logging import uuid from collections.abc import AsyncGenerator +from collections.abc import Callable from pydantic import BaseModel from pydantic import Field +from pydantic import model_validator from nat.builder.builder import Builder from nat.builder.framework_enum import LLMFrameworkEnum @@ -47,7 +49,29 @@ class TruncationRetryConfig(BaseModel): max_retries: int = Field(default=0, description="Number of retries when LLM output is truncated. " "0 disables recovery (raises RuntimeError).") - token_increment: int = Field(default=1024, description="Fixed number of tokens added to max_tokens on each retry.") + token_increment: int | None = Field(default=None, + description="Fixed number of tokens added to max_tokens on each retry. " + "Mutually exclusive with token_scaling. Defaults to 1024 if neither is set.") + token_scaling: float | None = Field(default=None, + description="Multiplicative factor applied to max_tokens on each retry " + "(e.g. 1.5 = 50%% increase per retry). " + "Mutually exclusive with token_increment.") + + @model_validator(mode="after") + def _check_scaling_strategy(self) -> "TruncationRetryConfig": + if self.token_increment is not None and self.token_scaling is not None: + raise ValueError("Set token_increment or token_scaling, not both.") + if self.max_retries > 0 and self.token_increment is None and self.token_scaling is None: + self.token_increment = 1024 + return self + + def build_scaling_fn(self) -> Callable[[int], int]: + """Build a callable that computes the next max_tokens from the current value.""" + if self.token_scaling is not None: + factor: float = self.token_scaling + return lambda current: int(current * factor) + increment: int = self.token_increment or 1024 + return lambda current: current + increment class ToolCallAgentWorkflowConfig(AgentBaseConfig, name="tool_calling_agent"): @@ -112,7 +136,7 @@ async def tool_calling_agent_workflow(config: ToolCallAgentWorkflowConfig, build handle_tool_errors=config.handle_tool_errors, return_direct=return_direct_tools, max_truncation_retries=config.truncation_retry.max_retries, - truncation_token_increment=config.truncation_retry.token_increment, + truncation_scaling_fn=config.truncation_retry.build_scaling_fn(), max_empty_response_retries=config.max_empty_response_retries, ).build_graph() diff --git a/packages/nvidia_nat_langchain/tests/agent/test_tool_calling.py b/packages/nvidia_nat_langchain/tests/agent/test_tool_calling.py index 4a02134a60..7b890eb831 100644 --- a/packages/nvidia_nat_langchain/tests/agent/test_tool_calling.py +++ b/packages/nvidia_nat_langchain/tests/agent/test_tool_calling.py @@ -36,6 +36,34 @@ from nat.plugins.langchain.agent.tool_calling_agent.agent import ToolCallAgentGraphState from nat.plugins.langchain.agent.tool_calling_agent.agent import create_tool_calling_agent_prompt from nat.plugins.langchain.agent.tool_calling_agent.register import ToolCallAgentWorkflowConfig +from nat.plugins.langchain.agent.tool_calling_agent.register import TruncationRetryConfig + + +def test_truncation_retry_config_rejects_both_strategies(): + """Setting both token_increment and token_scaling must raise ValidationError.""" + with pytest.raises(Exception, match="token_increment or token_scaling"): + TruncationRetryConfig(max_retries=2, token_increment=1024, token_scaling=1.5) + + +def test_truncation_retry_config_defaults_to_increment_when_neither_set(): + """max_retries > 0 with neither set defaults token_increment to 1024.""" + config = TruncationRetryConfig(max_retries=2, token_increment=None, token_scaling=None) + assert config.token_increment == 1024 + assert config.token_scaling is None + + +def test_truncation_retry_config_accepts_scaling_only(): + """Setting only token_scaling (without token_increment) must succeed.""" + config = TruncationRetryConfig(max_retries=3, token_scaling=1.5) + assert config.token_scaling == 1.5 + assert config.token_increment is None + + +def test_truncation_retry_config_accepts_increment_only(): + """Setting only token_increment (without token_scaling) must succeed.""" + config = TruncationRetryConfig(max_retries=3, token_increment=512) + assert config.token_increment == 512 + assert config.token_scaling is None async def test_state_schema(): @@ -456,7 +484,7 @@ async def test_validate_truncation_raises_when_disabled(error_mock_llm, mock_too async def test_validate_truncation_delegates_to_retry(error_mock_llm, mock_tool): """finish_reason=length with max_truncation_retries>0 delegates to _retry_on_truncation.""" - agent = _make_agent(error_mock_llm, mock_tool, max_truncation_retries=2, truncation_token_increment=512) + agent = _make_agent(error_mock_llm, mock_tool, max_truncation_retries=2, truncation_scaling_fn=lambda c: c + 512) truncated = AIMessage(content="partial", response_metadata={"finish_reason": "length"}) good = AIMessage(content="complete", response_metadata={"finish_reason": "stop"}) state = ToolCallAgentGraphState(messages=[HumanMessage(content="test")]) @@ -468,6 +496,16 @@ async def test_validate_truncation_delegates_to_retry(error_mock_llm, mock_tool) assert result.content == "complete" +async def test_validate_empty_response_raises_when_disabled(error_mock_llm, mock_tool): + """Empty response with max_empty_response_retries=0 must raise RuntimeError.""" + agent = _make_agent(error_mock_llm, mock_tool, max_empty_response_retries=0) + empty = AIMessage(content="", response_metadata={"finish_reason": "stop"}) + state = ToolCallAgentGraphState(messages=[HumanMessage(content="test")]) + + with pytest.raises(RuntimeError, match="empty response"): + await agent._validate_llm_response(empty, state) + + async def test_validate_empty_response_delegates_to_retry(error_mock_llm, mock_tool): """Empty response with retries>0 delegates to _retry_on_empty_response.""" agent = _make_agent(error_mock_llm, mock_tool, max_empty_response_retries=2) @@ -499,7 +537,7 @@ def test_get_token_usage_from_openai_response_metadata(error_mock_llm, mock_tool async def test_retry_on_truncation_succeeds(error_mock_llm, mock_tool): """Truncation retry succeeds when the retried LLM call finishes normally.""" - agent = _make_agent(error_mock_llm, mock_tool, max_truncation_retries=3, truncation_token_increment=512) + agent = _make_agent(error_mock_llm, mock_tool, max_truncation_retries=3, truncation_scaling_fn=lambda c: c + 512) first_response = AIMessage( content="partial", response_metadata={"finish_reason": "length"}, @@ -515,11 +553,13 @@ async def test_retry_on_truncation_succeeds(error_mock_llm, mock_tool): assert result.content == "complete" assert result.response_metadata["finish_reason"] == "stop" + assert agent._current_max_tokens == 612 + assert agent._truncation_retries_remaining == 2 async def test_retry_on_truncation_exhausted(error_mock_llm, mock_tool): - """All truncation retries exhausted returns the last truncated response.""" - agent = _make_agent(error_mock_llm, mock_tool, max_truncation_retries=2, truncation_token_increment=512) + """All truncation retries exhausted raises RuntimeError.""" + agent = _make_agent(error_mock_llm, mock_tool, max_truncation_retries=2, truncation_scaling_fn=lambda c: c + 512) first_response = AIMessage( content="partial", response_metadata={"finish_reason": "length"}, @@ -527,19 +567,29 @@ async def test_retry_on_truncation_exhausted(error_mock_llm, mock_tool): "output_tokens": 100, "input_tokens": 50, "total_tokens": 150 }, ) - still_truncated = AIMessage(content="partial", response_metadata={"finish_reason": "length"}) + still_truncated = AIMessage( + content="partial", + response_metadata={"finish_reason": "length"}, + usage_metadata={ + "output_tokens": 200, "input_tokens": 50, "total_tokens": 250 + }, + ) state = ToolCallAgentGraphState(messages=[HumanMessage(content="test")]) with patch.object(agent, "_invoke_llm", new_callable=AsyncMock, return_value=still_truncated): - result = await agent._retry_on_truncation(first_response, state) - - assert result.response_metadata["finish_reason"] == "length" + with pytest.raises(RuntimeError, match="LLM output still truncated after 2 retries"): + await agent._retry_on_truncation(first_response, state) async def test_retry_on_truncation_increments_from_usage(error_mock_llm, mock_tool): """When max_tokens is not configured, the base is taken from usage_metadata output_tokens.""" increment: int = 256 - agent = _make_agent(error_mock_llm, mock_tool, max_truncation_retries=3, truncation_token_increment=increment) + agent = _make_agent( + error_mock_llm, + mock_tool, + max_truncation_retries=3, + truncation_scaling_fn=lambda c: c + increment, + ) assert agent.llm.max_tokens is None first_response = AIMessage( content="partial", @@ -558,16 +608,22 @@ async def _capture_and_invoke(s): return still_truncated with patch.object(agent, "_invoke_llm", side_effect=_capture_and_invoke): - await agent._retry_on_truncation(first_response, state) + with pytest.raises(RuntimeError, match="LLM output still truncated after 3 retries"): + await agent._retry_on_truncation(first_response, state) - assert observed_max_tokens == [100 + increment, 100 + 2 * increment, 100 + 3 * increment] + assert observed_max_tokens == [356, 612, 868] async def test_retry_on_truncation_increments_from_configured_max_tokens(error_mock_llm, mock_tool): """When max_tokens is already configured on the LLM, it is used as the base instead of usage data.""" increment: int = 512 - agent = _make_agent(error_mock_llm, mock_tool, max_truncation_retries=2, truncation_token_increment=increment) - agent.llm.max_tokens = 50 + error_mock_llm.max_tokens = 50 + agent = _make_agent( + error_mock_llm, + mock_tool, + max_truncation_retries=2, + truncation_scaling_fn=lambda c: c + increment, + ) first_response = AIMessage( content="partial", response_metadata={"finish_reason": "length"}, @@ -585,9 +641,54 @@ async def _capture_and_invoke(s): return still_truncated with patch.object(agent, "_invoke_llm", side_effect=_capture_and_invoke): - await agent._retry_on_truncation(first_response, state) + with pytest.raises(RuntimeError, match="LLM output still truncated after 2 retries"): + await agent._retry_on_truncation(first_response, state) + + assert observed_max_tokens == [562, 1074] + + +async def test_retry_on_truncation_persists_across_calls(error_mock_llm, mock_tool): + """Retries and max_tokens carry forward across multiple truncation events.""" + agent = _make_agent( + error_mock_llm, + mock_tool, + max_truncation_retries=4, + truncation_scaling_fn=lambda c: c + 100, + ) + truncated = AIMessage( + content="partial", + response_metadata={"finish_reason": "length"}, + usage_metadata={ + "output_tokens": 200, "input_tokens": 50, "total_tokens": 250 + }, + ) + good = AIMessage(content="ok", response_metadata={"finish_reason": "stop"}) + state = ToolCallAgentGraphState(messages=[HumanMessage(content="test")]) + + with patch.object(agent, "_invoke_llm", new_callable=AsyncMock, return_value=good): + await agent._retry_on_truncation(truncated, state) + + assert agent._current_max_tokens == 300 + assert agent._truncation_retries_remaining == 3 + + with patch.object(agent, "_invoke_llm", new_callable=AsyncMock, return_value=good): + await agent._retry_on_truncation(truncated, state) + + assert agent._current_max_tokens == 400 + assert agent._truncation_retries_remaining == 2 - assert observed_max_tokens == [50 + increment, 50 + 2 * increment] + still_truncated = AIMessage( + content="partial", + response_metadata={"finish_reason": "length"}, + usage_metadata={ + "output_tokens": 400, "input_tokens": 50, "total_tokens": 450 + }, + ) + with patch.object(agent, "_invoke_llm", new_callable=AsyncMock, return_value=still_truncated): + with pytest.raises(RuntimeError, match="LLM output still truncated after 4 retries"): + await agent._retry_on_truncation(truncated, state) + + assert agent._truncation_retries_remaining == 0 async def test_retry_on_empty_succeeds(error_mock_llm, mock_tool): @@ -601,3 +702,15 @@ async def test_retry_on_empty_succeeds(error_mock_llm, mock_tool): result = await agent._retry_on_empty_response(state, first_meta) assert result.content == "actual answer" + + +async def test_retry_on_empty_exhausted(error_mock_llm, mock_tool): + """All empty-response retries exhausted raises RuntimeError.""" + agent = _make_agent(error_mock_llm, mock_tool, max_empty_response_retries=2) + first_meta: dict = {"finish_reason": "stop"} + still_empty = AIMessage(content="", response_metadata={"finish_reason": "stop"}) + state = ToolCallAgentGraphState(messages=[HumanMessage(content="test")]) + + with patch.object(agent, "_invoke_llm", new_callable=AsyncMock, return_value=still_empty): + with pytest.raises(RuntimeError, match="empty responses after 2 retries"): + await agent._retry_on_empty_response(state, first_meta)