fix: guard nim_langchain against ChatNVIDIA model lookup crash

[bbednarski9]

Guard nim_langchain against ChatNVIDIA model lookup crash

Background: https://forums.developer.nvidia.com/t/nim-api-duplicate-model-issue/365698
Remove after nemo-agent-toolkit uses version with upstreamed fix: langchain-ai/langchain-nvidia#282

Problem

ChatNVIDIA can crash with an AssertionError when a model is not present in the static MODEL_TABLE and appears more than once in the NIM (/v1/models) API response. In _NVIDIAClient.init(), the code asserts that there is at most one matching candidate, but duplicate API entries can produce two:

candidates = [model for model in self.available_models if model.id == self.mdl_name]
assert len(candidates) <= 1

This currently affects nvidia/nemotron-3-super-120b-a12b and could affect any future model that meets both conditions. The error propagates unhandled through nim_langchain(), surfacing as an opaque crash.

Root cause

The NIM API at integrate.api.nvidia.com/v1/models can return duplicate entries for some models. langchain-nvidia-ai-endpoints does not deduplicate this response and instead asserts uniqueness.

An upstream fix has already been prepared in langchain-ai/langchain-nvidia on branch bb/fix-chat-model-dedup, but NAT should not depend on an unreleased upstream change to avoid crashing.

Fix

Pre-register unknown models in ChatNVIDIA’s static MODEL_TABLE before constructing the client. This allows the static determine_model() lookup to succeed and avoids the (/v1/models) API call entirely, bypassing the problematic assertion.

if llm_config.model_name not in MODEL_TABLE:
    MODEL_TABLE[llm_config.model_name] = Model(
        id=llm_config.model_name,
        model_type="chat",
        client="ChatNVIDIA",
    )

Models already present in MODEL_TABLE are unaffected because the guard skips them.

Trade-offs

Pre-registered entries default to supports_tools=False, supports_structured_output=False, and supports_thinking=False.

As a result, bind_tools() on an unknown model may emit a warning such as “not known to support tools,” even though tool calling still works.

This is still strictly better than the current behavior, which is a hard crash.

The fix accesses MODEL_TABLE directly, which is not a public API of langchain-nvidia-ai-endpoints.

The public register_model() helper requires an endpoint parameter that is not available for standard hosted models.

Given that this is a temporary guard until the upstream deduplication fix is released, this trade-off is acceptable.

Testing

Verified that ChatNVIDIA(model="nvidia/nemotron-3-super-120b-a12b") no longer crashes.

Verified that models already present in MODEL_TABLE, such as meta/llama-3.3-70b-instruct, are unaffected by the guard.

Summary by CodeRabbit

• Bug Fixes
  • Improved LLM startup so configured models are more reliably recognized during initialization, reducing cases where a configured model wasn't available.
  • Keeps existing fallback behavior intact to avoid disruptions if advanced model registration isn't available.

[coderabbitai]

Walkthrough

Attempts to pre-register the configured model into langchain_nvidia_ai_endpoints._statics.MODEL_TABLE by importing Model and inserting Model(id=llm_config.model_name, model_type="chat", client="ChatNVIDIA") if missing. The operation is wrapped in a try/except that ignores ImportError and AttributeError, preserving existing discovery behavior.

Changes

Cohort / File(s)	Summary
Model Pre-registration packages/nvidia_nat_langchain/src/nat/plugins/langchain/llm.py	Adds import of Model and a guarded attempt to access langchain_nvidia_ai_endpoints._statics.MODEL_TABLE. If present and missing the configured llm_config.model_name, inserts a Model(...) entry keyed by that name. The insertion is wrapped to silently ignore ImportError/AttributeError; main client-construction flow remains unchanged.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% 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 clearly and concisely describes the main change: guarding nim_langchain against a ChatNVIDIA model lookup crash. It uses imperative mood and is within the 72-character limit at 62 characters.

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

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{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/llm.py (2)

176-184: Consider guarding against empty model_name.

The fix correctly addresses the upstream duplicate-entry crash. However, llm_config.model_name is typed as str without non-empty validation at the Pydantic level (per NIMModelConfig). If an empty string slips through, it would create an invalid MODEL_TABLE entry with id="".

🛡️ Proposed defensive check

     # Pre-register unknown models so ChatNVIDIA skips the /v1/models API
     # call. This guards against upstream issues such as duplicate entries
     # in the API response that cause ChatNVIDIA to crash with AssertionError.
-    if llm_config.model_name not in MODEL_TABLE:
+    if llm_config.model_name and llm_config.model_name not in MODEL_TABLE:
         MODEL_TABLE[llm_config.model_name] = Model(
             id=llm_config.model_name,
             model_type="chat",
             client="ChatNVIDIA",
         )

🤖 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/llm.py` around lines
176 - 184, The code pre-registers models using llm_config.model_name without
checking for emptiness; update the block that touches MODEL_TABLE to first
validate llm_config.model_name (e.g., if not llm_config.model_name or
llm_config.model_name.strip() == "") and handle it by either raising a clear
error (ValueError) or skipping registration with a warning logged via the
existing logger, so you never insert Model(id="") into MODEL_TABLE; reference
llm_config.model_name, MODEL_TABLE, Model and NIMModelConfig when making the
change.

---

172-172: Version pinning is already in place; consider adding an upstream issue reference to the workaround comment.

The pyproject.toml already constrains langchain-nvidia-ai-endpoints>=1.0.2,<2.0.0, so accidental breaking upgrades are prevented. The code comment at line 176 adequately explains the workaround for the upstream AssertionError issue. To improve maintainability, consider appending the upstream issue or PR number (e.g., langchain-ai/langchain-nvidia#XYZ) to the comment so it can be easily removed once that fix is released.

🤖 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/llm.py` at line 172,
Update the inline workaround comment next to the import of MODEL_TABLE and Model
from langchain_nvidia_ai_endpoints in nat/plugins/langchain/llm.py to include
the upstream issue/PR reference (e.g., "langchain-ai/langchain-nvidia#<number>"
or URL) so it's clear when to remove the hack; also mention the existing
pyproject.toml version pin (langchain-nvidia-ai-endpoints>=1.0.2,<2.0.0) in that
comment for context. Ensure the comment directly follows the import of
MODEL_TABLE and Model and is concise, referencing the AssertionError workaround
and the specific upstream issue/PR identifier.

🤖 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/llm.py`:
- Around line 176-184: The code pre-registers models using llm_config.model_name
without checking for emptiness; update the block that touches MODEL_TABLE to
first validate llm_config.model_name (e.g., if not llm_config.model_name or
llm_config.model_name.strip() == "") and handle it by either raising a clear
error (ValueError) or skipping registration with a warning logged via the
existing logger, so you never insert Model(id="") into MODEL_TABLE; reference
llm_config.model_name, MODEL_TABLE, Model and NIMModelConfig when making the
change.
- Line 172: Update the inline workaround comment next to the import of
MODEL_TABLE and Model from langchain_nvidia_ai_endpoints in
nat/plugins/langchain/llm.py to include the upstream issue/PR reference (e.g.,
"langchain-ai/langchain-nvidia#<number>" or URL) so it's clear when to remove
the hack; also mention the existing pyproject.toml version pin
(langchain-nvidia-ai-endpoints>=1.0.2,<2.0.0) in that comment for context.
Ensure the comment directly follows the import of MODEL_TABLE and Model and is
concise, referencing the AssertionError workaround and the specific upstream
issue/PR identifier.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 5ee12f50-5473-406e-9ead-70a40d2beac3

📥 Commits

Reviewing files that changed from the base of the PR and between 7c46d71 and 63ce3f0.

📒 Files selected for processing (1)

• packages/nvidia_nat_langchain/src/nat/plugins/langchain/llm.py

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (1)

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

172-195: Move Model import inside the try block to align with the stated fallback intent.

The comment states the code should fall back gracefully "if the private module changes between langchain-nvidia-ai-endpoints versions", but the Model import at line 172 is outside the try/except. If Model is renamed or removed in a future version while ChatNVIDIA remains, the function will fail entirely instead of gracefully skipping pre-registration.

Proposed fix

     from langchain_nvidia_ai_endpoints import ChatNVIDIA
-    from langchain_nvidia_ai_endpoints import Model

     validate_no_responses_api(llm_config, LLMFrameworkEnum.LANGCHAIN)

-    # TODO: Remove after upgrading to a langchain-nvidia-ai-endpoints release
+    # Temporary guard: remove after upgrading to a langchain-nvidia-ai-endpoints release
     # that includes https://github.com/langchain-ai/langchain-nvidia/pull/282.
     #
     # Pre-register unknown models so ChatNVIDIA skips the /v1/models API
     # call. This guards against upstream issues such as duplicate entries
     # in the API response that cause ChatNVIDIA to crash with AssertionError.
     # Uses internal MODEL_TABLE with fallback — if the private module
     # changes between langchain-nvidia-ai-endpoints versions, we skip
     # pre-registration and let ChatNVIDIA discover the model via /v1/models.
     try:
+        from langchain_nvidia_ai_endpoints import Model
         from langchain_nvidia_ai_endpoints._statics import MODEL_TABLE

         if llm_config.model_name not in MODEL_TABLE:
             MODEL_TABLE[llm_config.model_name] = Model(
                 id=llm_config.model_name,
                 model_type="chat",
                 client="ChatNVIDIA",
             )
     except (ImportError, AttributeError):
         pass

🤖 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/llm.py` around lines
172 - 195, The import of Model is outside the try/except which breaks the
intended fallback; move the "from langchain_nvidia_ai_endpoints import Model"
into the existing try block that imports MODEL_TABLE so both imports are
guarded, then only attempt to pre-register into MODEL_TABLE when
llm_config.model_name is missing; ensure the rest of the logic using
MODEL_TABLE, Model, llm_config, validate_no_responses_api and
LLMFrameworkEnum.LANGCHAIN remains unchanged so ChatNVIDIA discovery is used as
the fallback if either import fails.

🤖 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/llm.py`:
- Line 176: Replace the forbidden "TODO" comment string "# TODO: Remove after
upgrading to a langchain-nvidia-ai-endpoints release" with a non-TODO phrasing
that preserves the intent (for example "# NOTE: Remove after upgrading to a
langchain-nvidia-ai-endpoints release" or "# To be removed after upgrading to a
langchain-nvidia-ai-endpoints release"); update the comment in the
nat/plugins/langchain/llm.py module where that exact comment occurs so the
intent remains documented without using the "TODO" keyword.

---

Nitpick comments:
In `@packages/nvidia_nat_langchain/src/nat/plugins/langchain/llm.py`:
- Around line 172-195: The import of Model is outside the try/except which
breaks the intended fallback; move the "from langchain_nvidia_ai_endpoints
import Model" into the existing try block that imports MODEL_TABLE so both
imports are guarded, then only attempt to pre-register into MODEL_TABLE when
llm_config.model_name is missing; ensure the rest of the logic using
MODEL_TABLE, Model, llm_config, validate_no_responses_api and
LLMFrameworkEnum.LANGCHAIN remains unchanged so ChatNVIDIA discovery is used as
the fallback if either import fails.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 37f33ad1-bc1a-4db1-929c-a90ffd22fdfe

📥 Commits

Reviewing files that changed from the base of the PR and between c173ddd and 9cef1f1.

📒 Files selected for processing (1)

• packages/nvidia_nat_langchain/src/nat/plugins/langchain/llm.py

[bbednarski9]

/ok to test 9cef1f1

[bbednarski9]

/merge