[ericevans-nv]

Description

Problem

When the LLM hits its max_tokens limit mid-response, the truncated output lacks valid tool calls or a complete answer, causing the agent to silently loop until it hits a GraphRecursionError with no indication of the root cause. Additionally, other LLM failure modes (empty responses, invalid tool calls, content filtering) had no detection or logging.

Changes

• Error detection — Add _validate_llm_response to inspect finish_reason and response content after every LLM call, classifying truncation (length), empty response, invalid tool calls, and content filter scenarios.
• Truncation retry — Add _retry_on_truncation with configurable max_truncation_retries and truncation_token_increment. Derives the base max_tokens from the LLM config or usage metadata, then incrementally increases it via .bind(max_tokens=...).
• Empty response retry — Add _retry_on_empty_response with configurable max_empty_response_retries.
• Diagnostic logging — Surface model name, token usage (input_tokens, output_tokens, total_tokens), and truncated content in error messages. Add _get_token_usage helper that normalizes both LangChain (usage_metadata) and OpenAI (response_metadata["usage"]) formats.
• Configuration — Add TruncationRetryConfig Pydantic model and wire new parameters through ToolCallAgentWorkflowConfig for YAML-based configuration.
• Shared constants — Add LLMFinishReason StrEnum to nvidia_nat_core for reuse across agents.
• Tests — Add 10 focused unit tests covering all error paths, retry success/exhaustion, and token usage extraction.

By Submitting this PR I confirm:

• I am familiar with the Contributing Guidelines.
• We require that all contributors "sign-off" on their commits. This certifies that the contribution is your original work, or you have rights to submit it under the same license, or a compatible license.
  • Any contribution which contains commits that are not Signed-Off will not be accepted.
• When the PR is ready for review, new or existing tests cover these changes.
• When the PR is ready for review, the documentation is up to date with these changes.

Summary by CodeRabbit

• New Features

  • Tool-calling agent now automatically retries LLM calls on truncated or empty responses with configurable retry limits and token-scaling strategies via workflow config
  • Added standardized LLM finish-reason values for consistent response handling

• Bug Fixes

  • Returns a user-facing retry-exceeded message instead of raising on excessive recursion/retry failures

• Tests

  • Added comprehensive tests for truncation, retry, empty-response, and token-usage scenarios

[coderabbitai]

Walkthrough

Adds a finish-reason enum, refactors LLM streaming into a shared invoker, and implements retry/recovery logic and configuration for truncated or empty LLM outputs; tests expanded to cover truncation scaling, empty-response retries, and token-usage extraction.

Changes

Cohort / File(s)	Summary
LLM Status Enum packages/nvidia_nat_core/src/nat/llm/utils/constants.py	Added LLMFinishReason(StrEnum) with members STOP, LENGTH, CONTENT_FILTER, and TOOL_CALLS.
Agent invocation & recovery packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/tool_calling_agent/agent.py	Extended ToolCallAgentGraph constructor with max_truncation_retries, truncation_scaling_fn, and max_empty_response_retries. Introduced _invoke_llm, _validate_llm_response, _retry_on_truncation, _retry_on_empty_response, and _get_token_usage. Replaced inline streaming with _invoke_llm and added finish-reason based recovery flow.
Workflow config & registration packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/tool_calling_agent/register.py	Added TruncationRetryConfig(BaseModel) and exposed truncation_retry and max_empty_response_retries on ToolCallAgentWorkflowConfig. Threaded truncation/empty-response params into graph construction and changed _response_fn/_stream_fn to return/yield a retry-exceeded message on GraphRecursionError.
Tests: truncation & retries packages/nvidia_nat_langchain/tests/agent/test_tool_calling.py	Expanded tests with error_mock_llm fixture, _make_agent helper, and numerous scenarios verifying truncation retry behavior, empty-response retries, TruncationRetryConfig validation/defaults, and token-usage extraction. Removed some prior streaming/chunk tests in favor of retry-focused coverage.

Sequence Diagram(s)

sequenceDiagram
    participant Agent as ToolCallAgentGraph
    participant LLM as BaseChatModel
    participant Validator as _validate_llm_response
    participant Retry as RecoveryLogic

    Agent->>LLM: _invoke_llm(prompt, max_tokens)
    LLM-->>Agent: streamed chunks + finish_reason + usage
    Agent->>Validator: validate(response, finish_reason, usage)
    alt finish_reason == LENGTH
        Validator->>Retry: _retry_on_truncation()
        Retry->>Agent: compute increased max_tokens
        Agent->>LLM: _invoke_llm(...) [retry]
    else finish_reason == CONTENT_FILTER
        Validator-->>Agent: surface content-filter condition / abort
    else finish_reason == STOP or TOOL_CALLS
        Validator-->>Agent: proceed to tool invocation / normal flow
    else empty response
        Validator->>Retry: _retry_on_empty_response()
        Retry-->>Agent: retry or fail after max_empty_response_retries
    end

Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 66.67% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title accurately describes the main changes: adding error recovery, retry logic, and diagnostic logging to the Tool Calling Agent. It is specific, directly related to the changeset, and captures the core improvements.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

🧹 Nitpick comments (2)

packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/tool_calling_agent/agent.py (1)

266-321: Consider edge case where output_tokens is 0.

At Line 287, if usage.get("output_tokens", 0) returns 0 (e.g., when usage metadata is unavailable), the fallback to 4096 is applied via or 4096. This is a reasonable default, but logging this fallback would help operators understand why the agent chose a specific starting token count.

💡 Optional: Log when falling back to default token count

         else:
             usage: UsageMetadata = self._get_token_usage(first_response)
-            current_tokens = usage.get("output_tokens", 0) or 4096
+            output_tokens = usage.get("output_tokens", 0)
+            if output_tokens:
+                current_tokens = output_tokens
+            else:
+                current_tokens = 4096
+                logger.debug("%s No output_tokens in usage; defaulting to %d", AGENT_LOG_PREFIX, current_tokens)

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In
`@packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/tool_calling_agent/agent.py`
around lines 266 - 321, In _retry_on_truncation, detect when configured_max is
None and usage.get("output_tokens", 0) yields 0 so the code falls back to 4096;
add a concise log entry (using logger.*) right after computing current_tokens to
record that output_tokens was missing/zero and the default 4096 is being used
(include AGENT_LOG_PREFIX, the usage dict or its output_tokens value, and the
chosen current_tokens) so operators can see why the starting max token count was
selected.

packages/nvidia_nat_langchain/tests/agent/test_tool_calling.py (1)

593-603: Consider adding a test for empty-response retries exhausted.

There's a test for test_retry_on_truncation_exhausted but no equivalent test_retry_on_empty_exhausted to verify behavior when all empty-response retries fail. This would provide symmetric coverage.

Would you like me to draft a test_retry_on_empty_exhausted test case?

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@packages/nvidia_nat_langchain/tests/agent/test_tool_calling.py` around lines
593 - 603, Add a new test mirroring test_retry_on_truncation_exhausted but for
empty responses: create test_retry_on_empty_exhausted that constructs agent =
_make_agent(error_mock_llm, mock_tool, max_empty_response_retries=2), build
state = ToolCallAgentGraphState(messages=[HumanMessage(content="test")]) and
first_meta = {"finish_reason": "stop"}, patch agent._invoke_llm (AsyncMock) to
always return AIMessage(content="") (or an AIMessage with empty content and
response_metadata indicating "stop"), then call await
agent._retry_on_empty_response(state, first_meta) and assert it raises the same
exception type used by test_retry_on_truncation_exhausted (match its exception
assertion) to verify retries are exhausted. Ensure the test name is
test_retry_on_empty_exhausted and reference _retry_on_empty_response and
_invoke_llm in the setup.

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Nitpick comments:
In
`@packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/tool_calling_agent/agent.py`:
- Around line 266-321: In _retry_on_truncation, detect when configured_max is
None and usage.get("output_tokens", 0) yields 0 so the code falls back to 4096;
add a concise log entry (using logger.*) right after computing current_tokens to
record that output_tokens was missing/zero and the default 4096 is being used
(include AGENT_LOG_PREFIX, the usage dict or its output_tokens value, and the
chosen current_tokens) so operators can see why the starting max token count was
selected.

In `@packages/nvidia_nat_langchain/tests/agent/test_tool_calling.py`:
- Around line 593-603: Add a new test mirroring
test_retry_on_truncation_exhausted but for empty responses: create
test_retry_on_empty_exhausted that constructs agent =
_make_agent(error_mock_llm, mock_tool, max_empty_response_retries=2), build
state = ToolCallAgentGraphState(messages=[HumanMessage(content="test")]) and
first_meta = {"finish_reason": "stop"}, patch agent._invoke_llm (AsyncMock) to
always return AIMessage(content="") (or an AIMessage with empty content and
response_metadata indicating "stop"), then call await
agent._retry_on_empty_response(state, first_meta) and assert it raises the same
exception type used by test_retry_on_truncation_exhausted (match its exception
assertion) to verify retries are exhausted. Ensure the test name is
test_retry_on_empty_exhausted and reference _retry_on_empty_response and
_invoke_llm in the setup.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: b6c5d326-156c-4b47-ba61-450e663e2c5c

📥 Commits

Reviewing files that changed from the base of the PR and between 9d22b17 and 0dae46d.

📒 Files selected for processing (4)

• packages/nvidia_nat_core/src/nat/llm/utils/constants.py
• packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/tool_calling_agent/agent.py
• packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/tool_calling_agent/register.py
• packages/nvidia_nat_langchain/tests/agent/test_tool_calling.py

[willkill07]

I like the robustness improvements. Minor feedback suggested to improve performance and end-user experience.

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (4)

packages/nvidia_nat_langchain/tests/agent/test_tool_calling.py (2)

627-637: Same suggestion: use underscore prefix for unused parameter.

♻️ Suggested fix

-    async def _capture_and_invoke(s):
+    async def _capture_and_invoke(_s):
         observed_max_tokens.append(getattr(agent.bound_llm, "kwargs", {}).get("max_tokens"))
         return still_truncated

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@packages/nvidia_nat_langchain/tests/agent/test_tool_calling.py` around lines
627 - 637, The test's helper _capture_and_invoke(s) declares an unused parameter
s; rename it to _s (or prefix with underscore) so linters/readers recognize it's
intentionally unused—update the function signature in the test (function name
_capture_and_invoke and where it's used in patch.object(agent, "_invoke_llm",
side_effect=_capture_and_invoke)) and any references to that parameter
accordingly.

---

594-604: Consider using underscore prefix for unused parameter.

The s parameter in _capture_and_invoke is unused but required to match the _invoke_llm signature. Using _s or _ would suppress the linter warning and clarify intent.

♻️ Suggested fix

-    async def _capture_and_invoke(s):
+    async def _capture_and_invoke(_s):
         observed_max_tokens.append(getattr(agent.bound_llm, "kwargs", {}).get("max_tokens"))
         return still_truncated

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@packages/nvidia_nat_langchain/tests/agent/test_tool_calling.py` around lines
594 - 604, The test helper function _capture_and_invoke defines an unused
parameter s that exists only to match agent._invoke_llm's signature; rename it
to _s or _ to silence linters and signal it's intentionally unused (e.g., change
def _capture_and_invoke(s): to def _capture_and_invoke(_s):) and keep the body
intact so the patched call behavior and observed_max_tokens assertions remain
unchanged.

packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/tool_calling_agent/agent.py (2)

126-130: Design consideration: Retry budget is shared across all truncation events.

The retry counter (_truncation_retries_remaining) persists across multiple truncation events during the agent's lifetime. Once exhausted, no further truncation retries are possible for any subsequent requests handled by this agent instance.

This matches the test behavior in test_retry_on_truncation_persists_across_calls, but it may be worth documenting this in the config or class docstring so users understand the retry budget is not per-request.

Also applies to: 295-296

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In
`@packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/tool_calling_agent/agent.py`
around lines 126 - 130, The retry counter _truncation_retries_remaining is
shared across the agent instance (initialized from _max_truncation_retries) so
truncation retries are consumed globally across multiple handle calls; update
the class docstring (or public config docs) for the Agent implementing these
fields to explicitly state that _truncation_retries_remaining is a per-agent
budget (not per-request) and note the related fields (_max_truncation_retries,
_truncation_scaling_fn, _current_max_tokens) and the same behavior at the other
occurrence (lines referenced near _truncation_retries_remaining at the second
location) so users know retries persist across calls and how to configure/reset
them if needed.

---

367-372: Inconsistency: Empty response retry returns empty response on exhaustion, truncation raises.

When truncation retries are exhausted, _retry_on_truncation raises RuntimeError. When empty response retries are exhausted, _retry_on_empty_response logs a warning and returns the last empty response. This asymmetry may be intentional (empty responses are less harmful than truncated tool calls), but consider documenting or aligning the behavior.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In
`@packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/tool_calling_agent/agent.py`
around lines 367 - 372, The behavior is inconsistent: _retry_on_truncation
raises RuntimeError on exhaustion while _retry_on_empty_response logs a warning
and returns the last empty response; make them consistent by updating
_retry_on_empty_response to raise a RuntimeError (including context like
AGENT_LOG_PREFIX, self._max_empty_response_retries and the last response
content) when retries are exhausted, and update the _retry_on_empty_response
docstring/tests to reflect the raising behavior so callers can handle both cases
uniformly.

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In
`@packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/tool_calling_agent/agent.py`:
- Around line 234-247: The current empty-response branch in the agent (inside
the method handling LLM responses) only logs an error when
self._max_empty_response_retries == 0 but the documented behavior says recovery
is disabled and a RuntimeError should be raised; change the branch where (not
response.tool_calls ... and not (response.content ...)) to raise a RuntimeError
with a clear message (including finish_reason and metadata) when
self._max_empty_response_retries == 0 instead of just logging, and keep the
existing call to await self._retry_on_empty_response(state, metadata) when
_max_empty_response_retries > 0 so behavior matches the register.py field
description.

---

Nitpick comments:
In
`@packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/tool_calling_agent/agent.py`:
- Around line 126-130: The retry counter _truncation_retries_remaining is shared
across the agent instance (initialized from _max_truncation_retries) so
truncation retries are consumed globally across multiple handle calls; update
the class docstring (or public config docs) for the Agent implementing these
fields to explicitly state that _truncation_retries_remaining is a per-agent
budget (not per-request) and note the related fields (_max_truncation_retries,
_truncation_scaling_fn, _current_max_tokens) and the same behavior at the other
occurrence (lines referenced near _truncation_retries_remaining at the second
location) so users know retries persist across calls and how to configure/reset
them if needed.
- Around line 367-372: The behavior is inconsistent: _retry_on_truncation raises
RuntimeError on exhaustion while _retry_on_empty_response logs a warning and
returns the last empty response; make them consistent by updating
_retry_on_empty_response to raise a RuntimeError (including context like
AGENT_LOG_PREFIX, self._max_empty_response_retries and the last response
content) when retries are exhausted, and update the _retry_on_empty_response
docstring/tests to reflect the raising behavior so callers can handle both cases
uniformly.

In `@packages/nvidia_nat_langchain/tests/agent/test_tool_calling.py`:
- Around line 627-637: The test's helper _capture_and_invoke(s) declares an
unused parameter s; rename it to _s (or prefix with underscore) so
linters/readers recognize it's intentionally unused—update the function
signature in the test (function name _capture_and_invoke and where it's used in
patch.object(agent, "_invoke_llm", side_effect=_capture_and_invoke)) and any
references to that parameter accordingly.
- Around line 594-604: The test helper function _capture_and_invoke defines an
unused parameter s that exists only to match agent._invoke_llm's signature;
rename it to _s or _ to silence linters and signal it's intentionally unused
(e.g., change def _capture_and_invoke(s): to def _capture_and_invoke(_s):) and
keep the body intact so the patched call behavior and observed_max_tokens
assertions remain unchanged.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 52cc7cc5-f2c4-410c-bb69-6641b8bfaf85

📥 Commits

Reviewing files that changed from the base of the PR and between 0dae46d and 43302df.

📒 Files selected for processing (3)

• packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/tool_calling_agent/agent.py
• packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/tool_calling_agent/register.py
• packages/nvidia_nat_langchain/tests/agent/test_tool_calling.py

[coderabbitai]

Actionable comments posted: 2

🧹 Nitpick comments (2)

packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/tool_calling_agent/agent.py (1)

134-156: Add return type hint to _invoke_llm.

The method is missing a return type annotation. Per coding guidelines, all methods should use type hints for parameters and return values.

✏️ Proposed fix

-    async def _invoke_llm(self, state: ToolCallAgentGraphState):
+    async def _invoke_llm(self, state: ToolCallAgentGraphState) -> AIMessageChunk:

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In
`@packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/tool_calling_agent/agent.py`
around lines 134 - 156, The method _invoke_llm lacks a return type annotation;
update its signature in agent.py to include the appropriate return type (e.g.,
-> AIMessageChunk or -> AIMessage depending on actual expected type) so it
complies with the project's type-hinting rules; modify the async def
_invoke_llm(self, state: ToolCallAgentGraphState): declaration to add the
correct return type and adjust any imports/types if needed to ensure
AIMessageChunk (or chosen type) is available in the module.

packages/nvidia_nat_langchain/tests/agent/test_tool_calling.py (1)

650-691: Consider: Test documents the instance-state persistence issue.

test_retry_on_truncation_persists_across_calls explicitly tests that retry state persists across multiple calls to _retry_on_truncation. While this validates the current implementation, it also documents the problematic behavior flagged in the agent code review—retry budget exhaustion affects subsequent unrelated requests.

If the agent code is fixed to reset retry state per-invocation, this test will need to be updated accordingly.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@packages/nvidia_nat_langchain/tests/agent/test_tool_calling.py` around lines
650 - 691, The test test_retry_on_truncation_persists_across_calls documents
current behavior where agent instance state (_current_max_tokens and
_truncation_retries_remaining) persists across multiple _retry_on_truncation
calls; if you change the agent to reset retry state per invocation, update this
test to reflect that by either re-creating the agent via _make_agent before the
second call or by changing assertions to expect restarted values (initial
_current_max_tokens and full _truncation_retries_remaining) after each
independent invocation of _retry_on_truncation; update references to
_current_max_tokens, _truncation_retries_remaining, _retry_on_truncation, and
_make_agent accordingly.

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In
`@packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/tool_calling_agent/agent.py`:
- Around line 125-131: The instance-level retry counters
_truncation_retries_remaining and _current_max_tokens (set in __init__) are
mutated by _retry_on_truncation causing cross-invocation state bleed; modify
agent_node to initialize or reset these per invocation (e.g., set
_truncation_retries_remaining = self._max_truncation_retries and
_current_max_tokens = getattr(llm, "max_tokens", None) at the start of
agent_node) or refactor _retry_on_truncation to accept and return per-call retry
state so the retry budget is local to each invocation rather than stored on the
ToolCallAgentGraph instance.
- Around line 306-307: The code mutates instance state by reassigning
self.bound_llm and self.agent during retries; instead, inside the
retry/truncation path create local variables (e.g., local_bound_llm =
self.llm.bind_tools(self.tools).bind(max_tokens=new_limit) and local_agent =
self.agent.first | local_bound_llm) and use those locals for the retried call so
you never overwrite self.bound_llm or self.agent; ensure all retry logic calls
the local_agent/local_bound_llm and leave instance attributes unchanged to avoid
cross-request state leaks and races.

---

Nitpick comments:
In
`@packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/tool_calling_agent/agent.py`:
- Around line 134-156: The method _invoke_llm lacks a return type annotation;
update its signature in agent.py to include the appropriate return type (e.g.,
-> AIMessageChunk or -> AIMessage depending on actual expected type) so it
complies with the project's type-hinting rules; modify the async def
_invoke_llm(self, state: ToolCallAgentGraphState): declaration to add the
correct return type and adjust any imports/types if needed to ensure
AIMessageChunk (or chosen type) is available in the module.

In `@packages/nvidia_nat_langchain/tests/agent/test_tool_calling.py`:
- Around line 650-691: The test test_retry_on_truncation_persists_across_calls
documents current behavior where agent instance state (_current_max_tokens and
_truncation_retries_remaining) persists across multiple _retry_on_truncation
calls; if you change the agent to reset retry state per invocation, update this
test to reflect that by either re-creating the agent via _make_agent before the
second call or by changing assertions to expect restarted values (initial
_current_max_tokens and full _truncation_retries_remaining) after each
independent invocation of _retry_on_truncation; update references to
_current_max_tokens, _truncation_retries_remaining, _retry_on_truncation, and
_make_agent accordingly.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 50ff3ccd-1aea-4733-8f01-202200087e7a

📥 Commits

Reviewing files that changed from the base of the PR and between 43302df and 40f2d9c.

📒 Files selected for processing (4)

• packages/nvidia_nat_core/src/nat/llm/utils/constants.py
• packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/tool_calling_agent/agent.py
• packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/tool_calling_agent/register.py
• packages/nvidia_nat_langchain/tests/agent/test_tool_calling.py

🚧 Files skipped from review as they are similar to previous changes (1)

• packages/nvidia_nat_core/src/nat/llm/utils/constants.py

[ericevans-nv]

/merge