feat(backend): add Nvidia NIM as a first-class backend provider#962
feat(backend): add Nvidia NIM as a first-class backend provider#962Henry-811 merged 1 commit intodev/v0.1.58from
Conversation
📝 WalkthroughWalkthroughThe PR adds support for Nvidia NIM as a new chat completion backend provider across the codebase. Changes include backend capability registration, provider detection logic, CLI configuration, model catalog integration, model matcher updates, and documentation additions. Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes 🚥 Pre-merge checks | ✅ 2 | ❌ 4❌ Failed checks (3 warnings, 1 inconclusive)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches
🧪 Generate unit tests (beta)
Tip Try Coding Plans. Let us write the prompt for your AI agent so you can ship faster (with fewer bugs). Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (3)
massgen/backend/capabilities.py (1)
575-588: Addmodel_release_datesfor the new Nvidia NIM model entry.Line 583 introduces an explicit model list, but this new backend entry does not include
model_release_dates. Please add release-date metadata for the listed concrete model(s) to keep registry metadata consistent.Based on learnings: When adding new models, always update
massgen/backend/capabilities.pyby adding to themodelslist (newest first), adding tomodel_release_dates, and updatingsupported_capabilitiesif new features are introduced.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@massgen/backend/capabilities.py` around lines 575 - 588, The new BackendCapabilities entry for "nvidia_nim" is missing model_release_dates metadata; update the same backend object by adding a model_release_dates mapping that includes release dates for the concrete models listed in models (e.g., "moonshotai/kimi-k2.5" and "custom") so registrations remain consistent; locate the BackendCapabilities instantiation for backend_type "nvidia_nim" and add a model_release_dates field (following the format used elsewhere in this file) with the newest model first and appropriate ISO date strings for each model name.massgen/cli.py (2)
1662-1688: Use Google-style sections in the updatedcreate_backenddocstring.
create_backend()was changed, but its docstring is still free-form bullets. Please convert to Google-style (Args,Returns,Raises) for repository consistency.As per coding guidelines
**/*.py: "For new or changed functions, include Google-style docstrings".🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@massgen/cli.py` around lines 1662 - 1688, Update the create_backend() docstring to Google-style sections: add an "Args" section describing the parameters (e.g., backend type and params dict), a "Returns" section describing the created backend instance, and a "Raises" section for possible exceptions (e.g., missing API key or unsupported backend). Move the existing supported-backend bullets into a "Raises" or "Args" note as appropriate (or an "Raises/Notes" subsection) and include the auto-detected chatcompletion providers under "Notes" or "Raises" as contextual info; keep the descriptive list of backend types and external dependencies but restructure them into the Google-style sections to match repo conventions.
1811-1819: Use_api_key_error_message()for the Nvidia NIM branch.Line 1814 hardcodes a long error string while this function already has
_api_key_error_message(). Reusing it improves consistency and keeps linter warnings down.Suggested patch
elif base_url and "nvidia.com" in base_url: api_key = os.getenv("NGC_API_KEY") if not api_key: - raise ConfigurationError( - "Nvidia NIM API key not found. Set NGC_API_KEY environment variable.\n" - "You can add it to a .env file in:\n" - " - Current directory: .env\n" - " - Global config: ~/.massgen/.env", - ) + raise ConfigurationError( + _api_key_error_message("Nvidia NIM", "NGC_API_KEY", config_path), + )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@massgen/cli.py` around lines 1811 - 1819, Replace the hardcoded Nvidia NIM API key error string in the branch that checks `elif base_url and "nvidia.com" in base_url` with a call to the existing helper `_api_key_error_message()`; specifically, when `api_key = os.getenv("NGC_API_KEY")` is falsy, raise `ConfigurationError(_api_key_error_message())` instead of the long literal so the `api_key` check in that branch reuses the shared message logic.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@massgen/cli.py`:
- Around line 1915-1924: The CLI currently implements backend_type "nvidia_nim"
in the factory branch (the elif handling that returns ChatCompletionsBackend)
but the CLI argument parser's allowed backend choices still omits "nvidia_nim",
causing argparse validation to reject it; locate where the backend choices are
defined (e.g., the --backend add_argument call or a BACKEND_CHOICES /
SUPPORTED_BACKENDS list used by parser.add_argument) and add the string
"nvidia_nim" to that choices list and any corresponding help text or
documentation strings so the CLI accepts --backend nvidia_nim and maps to the
existing backend_type handling (keep symbols: backend_type,
ChatCompletionsBackend, and the --backend argument).
---
Nitpick comments:
In `@massgen/backend/capabilities.py`:
- Around line 575-588: The new BackendCapabilities entry for "nvidia_nim" is
missing model_release_dates metadata; update the same backend object by adding a
model_release_dates mapping that includes release dates for the concrete models
listed in models (e.g., "moonshotai/kimi-k2.5" and "custom") so registrations
remain consistent; locate the BackendCapabilities instantiation for backend_type
"nvidia_nim" and add a model_release_dates field (following the format used
elsewhere in this file) with the newest model first and appropriate ISO date
strings for each model name.
In `@massgen/cli.py`:
- Around line 1662-1688: Update the create_backend() docstring to Google-style
sections: add an "Args" section describing the parameters (e.g., backend type
and params dict), a "Returns" section describing the created backend instance,
and a "Raises" section for possible exceptions (e.g., missing API key or
unsupported backend). Move the existing supported-backend bullets into a
"Raises" or "Args" note as appropriate (or an "Raises/Notes" subsection) and
include the auto-detected chatcompletion providers under "Notes" or "Raises" as
contextual info; keep the descriptive list of backend types and external
dependencies but restructure them into the Google-style sections to match repo
conventions.
- Around line 1811-1819: Replace the hardcoded Nvidia NIM API key error string
in the branch that checks `elif base_url and "nvidia.com" in base_url` with a
call to the existing helper `_api_key_error_message()`; specifically, when
`api_key = os.getenv("NGC_API_KEY")` is falsy, raise
`ConfigurationError(_api_key_error_message())` instead of the long literal so
the `api_key` check in that branch reuses the shared message logic.
ℹ️ Review info
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (7)
massgen/backend/capabilities.pymassgen/backend/chat_completions.pymassgen/cli.pymassgen/config_builder.pymassgen/configs/BACKEND_CONFIGURATION.mdmassgen/utils/model_catalog.pymassgen/utils/model_matcher.py
| elif backend_type == "nvidia_nim": | ||
| # Nvidia NIM uses OpenAI-compatible Chat Completions API | ||
| api_key = kwargs.get("api_key") or os.getenv("NGC_API_KEY") | ||
| if not api_key: | ||
| raise ConfigurationError( | ||
| _api_key_error_message("Nvidia NIM", "NGC_API_KEY", config_path), | ||
| ) | ||
| if "base_url" not in kwargs: | ||
| kwargs["base_url"] = "https://integrate.api.nvidia.com/v1" | ||
| return ChatCompletionsBackend(api_key=api_key, **kwargs) |
There was a problem hiding this comment.
nvidia_nim is implemented here but still not selectable via --backend.
Line 1915 adds first-class handling, but CLI parser choices (later in this file) still omit "nvidia_nim", so massgen --backend nvidia_nim ... fails argument validation.
Suggested patch
@@
config_group.add_argument(
"--backend",
type=str,
choices=[
"chatcompletion",
"claude",
"gemini",
"grok",
"openai",
"azure_openai",
"copilot",
"claude_code",
"zai",
+ "nvidia_nim",
"lmstudio",
"vllm",
"sglang",
],
help="Backend type for quick setup",
)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@massgen/cli.py` around lines 1915 - 1924, The CLI currently implements
backend_type "nvidia_nim" in the factory branch (the elif handling that returns
ChatCompletionsBackend) but the CLI argument parser's allowed backend choices
still omits "nvidia_nim", causing argparse validation to reject it; locate where
the backend choices are defined (e.g., the --backend add_argument call or a
BACKEND_CHOICES / SUPPORTED_BACKENDS list used by parser.add_argument) and add
the string "nvidia_nim" to that choices list and any corresponding help text or
documentation strings so the CLI accepts --backend nvidia_nim and maps to the
existing backend_type handling (keep symbols: backend_type,
ChatCompletionsBackend, and the --backend argument).
There was a problem hiding this comment.
Pull request overview
Adds Nvidia NIM as a supported OpenAI-compatible provider across MassGen’s provider registry, CLI/backend factory, model discovery, and documentation, making it selectable/configurable like other first-class backends.
Changes:
- Register
nvidia_nimin backend capabilities (models, env var, base URL) and provider-name inference. - Add CLI/config-builder support for selecting
nvidia_nimand guiding users toNGC_API_KEY. - Enable model listing via the OpenAI-compatible
/modelsendpoint and expand fuzzy model matching + docs.
Reviewed changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
massgen/utils/model_matcher.py |
Adds curated common model IDs and includes nvidia_nim in chatcompletion providers for model fetching/fuzzy matching. |
massgen/utils/model_catalog.py |
Adds nvidia_nim branch to fetch models from Nvidia’s OpenAI-compatible endpoint. |
massgen/configs/BACKEND_CONFIGURATION.md |
Documents Nvidia NIM base URL and NGC_API_KEY. |
massgen/config_builder.py |
Allows nvidia_nim in smart model selection and interactive API key setup. |
massgen/cli.py |
Adds nvidia_nim backend type and auto-detection of NGC_API_KEY for chatcompletion base URLs. |
massgen/backend/chat_completions.py |
Adds provider-name inference for Nvidia NIM based on base URL. |
massgen/backend/capabilities.py |
Registers nvidia_nim in BACKEND_CAPABILITIES with base URL, env var, models, and notes. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| models=["moonshotai/kimi-k2.5", "custom"], | ||
| default_model="moonshotai/kimi-k2.5", |
There was a problem hiding this comment.
ConfigValidator treats backends with default_model != "custom" as not requiring an explicit model, but the Chat Completions request builder here does not inject a default model. That means configs using type: nvidia_nim without a model will pass validation yet fail at runtime (OpenAI-compatible APIs require model). Consider setting default_model to "custom" (and/or only listing "custom" in models) or adding logic elsewhere to auto-fill model from capabilities when it’s missing.
| models=["moonshotai/kimi-k2.5", "custom"], | |
| default_model="moonshotai/kimi-k2.5", | |
| models=["custom"], | |
| default_model="custom", |
There was a problem hiding this comment.
No default model should be moonshotai/kimi-k2.5
| elif "nebius.com" in base_url: | ||
| return "Nebius AI Studio" |
There was a problem hiding this comment.
Nebius provider detection checks only for nebius.com, but the default Nebius base URL used elsewhere is api.studio.nebius.ai. With the current condition, get_provider_name() will not recognize Nebius when using the default .ai endpoint, which can affect logging/routing that relies on the inferred provider name. Update the check to also match nebius.ai (and/or studio.nebius.ai).
| "Nvidia NIM API key not found. Set NGC_API_KEY environment variable.\n" | ||
| "You can add it to a .env file in:\n" | ||
| " - Current directory: .env\n" | ||
| " - Global config: ~/.massgen/.env", |
There was a problem hiding this comment.
This branch adds a new provider-specific API-key-missing message, but it duplicates logic and diverges from the standard _api_key_error_message(...) helper used elsewhere in create_backend (different listed .env locations, no --setup hint, etc.). Consider reusing _api_key_error_message("Nvidia NIM", "NGC_API_KEY", config_path) here for consistency and easier future maintenance.
| "Nvidia NIM API key not found. Set NGC_API_KEY environment variable.\n" | |
| "You can add it to a .env file in:\n" | |
| " - Current directory: .env\n" | |
| " - Global config: ~/.massgen/.env", | |
| _api_key_error_message("Nvidia NIM", "NGC_API_KEY", config_path), |
PR Title Format
Your PR title must follow the format:
<type>: <brief description>Valid types:
fix:- Bug fixesfeat:- New featuresbreaking:- Breaking changesdocs:- Documentation updatesrefactor:- Code refactoringtest:- Test additions/modificationschore:- Maintenance tasksperf:- Performance improvementsstyle:- Code style changesci:- CI/CD configuration changesExamples:
fix: resolve memory leak in data processingfeat: add export to CSV functionalitybreaking: change API response formatdocs: update installation guideDescription
Brief description of the changes in this PR
Type of change
fix:) - Non-breaking change which fixes an issuefeat:) - Non-breaking change which adds functionalitybreaking:) - Fix or feature that would cause existing functionality to not work as expecteddocs:) - Documentation updatesrefactor:) - Code changes that neither fix a bug nor add a featuretest:) - Adding missing tests or correcting existing testschore:) - Maintenance tasks, dependency updates, etc.perf:) - Code changes that improve performancestyle:) - Changes that do not affect the meaning of the code (formatting, missing semi-colons, etc.)ci:) - Changes to CI/CD configuration files and scriptsChecklist
Pre-commit status
How to Test
Add test method for this PR.
Test CLI Command
Write down the test bash command. If there is pre-requests, please emphasize.
Expected Results
Description/screenshots of expected results.
Additional context
Add any other context about the PR here.
Summary by CodeRabbit
Release Notes