fix(server): honor api_key env vars when overriding main model#1923
fix(server): honor api_key env vars when overriding main model#1923adityasingh2400 wants to merge 4 commits into
Conversation
NVIDIA-NeMo#1895) Signed-off-by: Aditya Singh <adisin650@gmail.com>
Greptile SummaryThis PR fixes a credential-dropping bug (#1895) in the server's request-driven model override path. The old code replaced the entire
|
| Filename | Overview |
|---|---|
| nemoguardrails/server/api.py | Replaces the fragile field-by-field carryover in _update_models_in_config with a generic model_dump(exclude_unset=True) + model_copy approach, so all unset fields on the override (including api_key_env_var, mode, and cache) are transparently preserved from the existing config entry. |
| tests/server/test_api.py | Adds four targeted regression tests: two unit tests for the preserve-from-config and override-wins cases, one for mode/cache preservation, and one end-to-end integration test for env-var credential resolution through _get_rails. |
Sequence Diagram
sequenceDiagram
participant Client
participant API as /v1/chat/completions
participant GR as _get_rails
participant UM as _update_models_in_config
participant RC as RailsConfig
Client->>API: POST with model override
API->>GR: _get_rails(config_ids, model_name)
GR->>RC: RailsConfig.from_path(full_path)
Note over RC: existing entry has api_key_env_var,<br/>mode, cache from config.yml
GR->>GR: Build override Model with only model, type, engine, parameters set
GR->>UM: _update_models_in_config(config, override)
UM->>UM: "model_dump(exclude_unset=True) yields only explicitly-set fields"
Note over UM: api_key_env_var, mode, cache not in override dict
UM->>UM: "existing.model_copy(update=override)"
Note over UM: Unset fields preserved from existing config entry
UM->>GR: updated RailsConfig
GR->>Client: LLMRails with correct credentials
Reviews (4): Last reviewed commit: "refactor(server): carry over all unset m..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| # Preserve fields from the existing entry that the override does not | ||
| # set, so a request-side model swap inherits credential resolution | ||
| # (api_key_env_var) and other model-level settings from config.yml. | ||
| api_key_env_var = main_model.api_key_env_var or existing.api_key_env_var | ||
| models[main_model_index] = main_model | ||
| models[main_model_index].parameters = parameters | ||
| models[main_model_index].api_key_env_var = api_key_env_var |
There was a problem hiding this comment.
api_key_env_var preserved despite broader comment
The docstring says "Fields not explicitly set on the override model … are carried over from the existing entry" and the inline comment mentions "and other model-level settings from config.yml", but in practice only api_key_env_var is restored. Fields like mode (defaults to "chat") and cache from the existing entry are still silently dropped when models[main_model_index] = main_model runs. A config that declares mode: text would lose that after a request-driven model swap, just like api_key_env_var did before this fix.
Prompt To Fix With AI
This is a comment left during a code review.
Path: nemoguardrails/server/api.py
Line: 327-333
Comment:
**Only `api_key_env_var` preserved despite broader comment**
The docstring says "Fields not explicitly set on the override model … are carried over from the existing entry" and the inline comment mentions "and other model-level settings from config.yml", but in practice only `api_key_env_var` is restored. Fields like `mode` (defaults to `"chat"`) and `cache` from the existing entry are still silently dropped when `models[main_model_index] = main_model` runs. A config that declares `mode: text` would lose that after a request-driven model swap, just like `api_key_env_var` did before this fix.
How can I resolve this? If you propose a fix, please make it concise.| # Preserve fields from the existing entry that the override does not | ||
| # set, so a request-side model swap inherits credential resolution | ||
| # (api_key_env_var) and other model-level settings from config.yml. | ||
| api_key_env_var = main_model.api_key_env_var or existing.api_key_env_var |
There was a problem hiding this comment.
or to check whether the override set api_key_env_var conflates None (not set) with "" (explicitly set to empty string). While env-var names are never empty in practice, is not None makes the intent explicit and is marginally safer if the field is ever set programmatically.
| api_key_env_var = main_model.api_key_env_var or existing.api_key_env_var | |
| api_key_env_var = main_model.api_key_env_var if main_model.api_key_env_var is not None else existing.api_key_env_var |
Prompt To Fix With AI
This is a comment left during a code review.
Path: nemoguardrails/server/api.py
Line: 330
Comment:
Using `or` to check whether the override set `api_key_env_var` conflates `None` (not set) with `""` (explicitly set to empty string). While env-var names are never empty in practice, `is not None` makes the intent explicit and is marginally safer if the field is ever set programmatically.
```suggestion
api_key_env_var = main_model.api_key_env_var if main_model.api_key_env_var is not None else existing.api_key_env_var
```
How can I resolve this? If you propose a fix, please make it concise.| api.llm_rails_instances.clear() | ||
|
|
There was a problem hiding this comment.
llm_rails.llm._client is a raw attribute access (not guarded by getattr), so if the OpenAI SDK restructures its client object the test will raise AttributeError rather than a clear assertion failure, obscuring the diagnosis. Meanwhile the getattr(..., "_api_key", None) default only silently swallows changes to the inner attribute, making the fallback behaviour inconsistent: one level of indirection raises, the other returns None. Wrapping both in getattr or using the public api_key surface of the OpenAI client (if available) would give more stable, readable failure messages.
Prompt To Fix With AI
This is a comment left during a code review.
Path: tests/server/test_api.py
Line: 1052-1053
Comment:
**Fragile private-API access in integration test**
`llm_rails.llm._client` is a raw attribute access (not guarded by `getattr`), so if the OpenAI SDK restructures its client object the test will raise `AttributeError` rather than a clear assertion failure, obscuring the diagnosis. Meanwhile the `getattr(..., "_api_key", None)` default only silently swallows changes to the inner attribute, making the fallback behaviour inconsistent: one level of indirection raises, the other returns `None`. Wrapping both in `getattr` or using the public `api_key` surface of the OpenAI client (if available) would give more stable, readable failure messages.
How can I resolve this? If you propose a fix, please make it concise.![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
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@nemoguardrails/server/api.py`:
- Line 330: The current assignment uses `or` which treats falsy explicit
overrides (e.g., `""`) as absent; change the logic so
`main_model.api_key_env_var` wins when it is explicitly provided by checking `is
not None` and otherwise falling back to `existing.api_key_env_var` (i.e.,
implement an explicit None-check for `main_model.api_key_env_var` rather than
using `or` where `api_key_env_var` is set).
🪄 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: Enterprise
Run ID: 349b5b44-7f34-4a2e-bfa7-63e14378fbb0
📒 Files selected for processing (2)
nemoguardrails/server/api.pytests/server/test_api.py
| # Preserve fields from the existing entry that the override does not | ||
| # set, so a request-side model swap inherits credential resolution | ||
| # (api_key_env_var) and other model-level settings from config.yml. | ||
| api_key_env_var = main_model.api_key_env_var or existing.api_key_env_var |
There was a problem hiding this comment.
Use is not None precedence for api_key_env_var override.
Line 330 uses or, so an explicit override like api_key_env_var="" is treated as absent and the existing value is kept. That conflicts with the “explicit override wins” contract.
🔧 Proposed fix
- api_key_env_var = main_model.api_key_env_var or existing.api_key_env_var
+ api_key_env_var = (
+ main_model.api_key_env_var
+ if main_model.api_key_env_var is not None
+ else existing.api_key_env_var
+ )📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| api_key_env_var = main_model.api_key_env_var or existing.api_key_env_var | |
| api_key_env_var = ( | |
| main_model.api_key_env_var | |
| if main_model.api_key_env_var is not None | |
| else existing.api_key_env_var | |
| ) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@nemoguardrails/server/api.py` at line 330, The current assignment uses `or`
which treats falsy explicit overrides (e.g., `""`) as absent; change the logic
so `main_model.api_key_env_var` wins when it is explicitly provided by checking
`is not None` and otherwise falling back to `existing.api_key_env_var` (i.e.,
implement an explicit None-check for `main_model.api_key_env_var` rather than
using `or` where `api_key_env_var` is set).
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
The previous `or` based check treated an explicitly set empty string the same as None, so an override that set api_key_env_var="" would fall back to the existing config entry instead of winning. Use an explicit None check so the override wins whenever the caller sets the field, even to an empty value.
|
Switched the override check from |
📝 WalkthroughWalkthroughUpdated ChangesModel Override Credential and Parameter Merge
Estimated code review effort🎯 2 (Simple) | ⏱️ ~12 minutes 🚥 Pre-merge checks | ✅ 6✅ Passed checks (6 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tests/server/test_api.py (1)
1030-1065: ⚡ Quick winExercise the preserved
api_key_env_varpath end-to-end.This test only covers OpenAI's implicit
OPENAI_API_KEYfallback. It would still pass on the pre-fix implementation, because droppingapi_key_env_varduring the override does not matter when the provider later resolves credentials from its default env var anyway. To lock in the actual regression, make the config declare a non-defaultapi_key_env_varand leaveOPENAI_API_KEYunset.Proposed test adjustment
(config_dir / "config.yml").write_text( - "colang_version: '1.0'\nmodels:\n - type: main\n engine: openai\n model: gpt-3.5-turbo\n" + "colang_version: '1.0'\nmodels:\n - type: main\n engine: openai\n model: gpt-3.5-turbo\n api_key_env_var: MY_CUSTOM_OPENAI_KEY\n" ) - monkeypatch.setenv("OPENAI_API_KEY", "sk-from-env-1895") + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + monkeypatch.setenv("MY_CUSTOM_OPENAI_KEY", "sk-from-env-1895") monkeypatch.setenv("MAIN_MODEL_ENGINE", "openai") monkeypatch.setenv("MAIN_MODEL_BASE_URL", "https://api.openai.com/v1")🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/server/test_api.py` around lines 1030 - 1065, Update the test to exercise the preserved api_key_env_var path end-to-end by having the config declare a non-default env var name and ensuring OPENAI_API_KEY is not used: change the config YAML in test_get_rails_with_model_override_resolves_env_credentials to include api_key_env_var: CUSTOM_API_KEY (or another unique name), remove or monkeypatch.clear the OPENAI_API_KEY env var, set monkeypatch.setenv("CUSTOM_API_KEY","sk-from-env-1895") instead of OPENAI_API_KEY, and keep the rest of the setup and assertion against llm_rails.llm._client._api_key so the test verifies credentials are resolved from the declared api_key_env_var when calling api._get_rails.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@tests/server/test_api.py`:
- Around line 1030-1065: Update the test to exercise the preserved
api_key_env_var path end-to-end by having the config declare a non-default env
var name and ensuring OPENAI_API_KEY is not used: change the config YAML in
test_get_rails_with_model_override_resolves_env_credentials to include
api_key_env_var: CUSTOM_API_KEY (or another unique name), remove or
monkeypatch.clear the OPENAI_API_KEY env var, set
monkeypatch.setenv("CUSTOM_API_KEY","sk-from-env-1895") instead of
OPENAI_API_KEY, and keep the rest of the setup and assertion against
llm_rails.llm._client._api_key so the test verifies credentials are resolved
from the declared api_key_env_var when calling api._get_rails.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 1a9b648b-6022-47ea-9c0e-18db7a9cdddc
📒 Files selected for processing (2)
nemoguardrails/server/api.pytests/server/test_api.py
|
Pushed a formatting-only commit. Lint was failing because ruff-format wanted the api_key_env_var ternary collapsed onto a single line, which fits within the 120 char limit. No logic change. |
|
@adityasingh2400 -- I also encountered this error recently, so thanks for the PR IIUC, your fix patches Perhaps a more robust approach would be to start from the existing model and only override the fields the caller actually changes, using Pydantic's if main_model_index is not None:
existing = models[main_model_index]
models[main_model_index] = existing.model_copy(update={
"model": main_model.model,
"engine": main_model.engine,
"parameters": {**existing.parameters, **main_model.parameters},
})
else:
models.append(main_model)WDYT @adityasingh2400? |
…rride Addresses review feedback on this PR. Instead of replacing the existing main model entry and patching individual fields back (which still dropped mode and cache, and would miss any field added to Model later), start from the existing entry and apply only the fields the override request explicitly set, via existing.model_copy(update=main_model.model_dump(exclude_unset=True)). Parameters keep override-wins merge semantics, and api_key_env_var, mode, cache, and any future Model field are preserved when the override leaves them unset. Adds a regression test covering mode and cache preservation. Signed-off-by: Aditya Singh <adisin650@gmail.com>
|
Thanks, and good catch confirming the broader problem. You are right that patching api_key_env_var back was a band aid. mode and cache were still dropped, and every new Model field would re-hit the same class of bug. I pushed f2eeaae taking your approach. It now starts from the existing entry and applies only the fields the override actually set: override = main_model.model_dump(exclude_unset=True)
if "parameters" in override:
override["parameters"] = {**existing.parameters, **main_model.parameters}
models[main_model_index] = existing.model_copy(update=override)I used model_dump(exclude_unset=True) rather than naming model, engine, and parameters explicitly, so that any field the caller did set (including an explicit api_key_env_var, mode, or cache) still wins, while anything left unset is carried over from config.yml. Parameters keep the override-wins merge, so a request extends the config.yml parameters instead of replacing them. I also added a regression test that sets mode=text and a cache config on the existing entry, overrides only the model name, and asserts mode, cache, api_key_env_var, and the existing parameters all survive. The earlier api_key_env_var tests still pass, and ruff is clean. Happy to switch to an explicit field list instead if you would prefer that, just let me know. |
2 participants
Closes #1895.
When a request to
/v1/chat/completionssupplies amodelfield, the server rebuilds the main model entry by constructing a freshModelfrom the request value plusMAIN_MODEL_ENGINEandMAIN_MODEL_BASE_URL, then hands it to_update_models_in_configto swap into the loadedRailsConfig. That swap replaced the existing entry wholesale, so anyapi_key_env_vardeclared inconfig.ymlwas silently dropped. The downstream initializer then had no env-var hint and the credential resolution path that the CLI and non-server code uses was never taken, surfacing as an invalid-API-key error from the upstream provider.This change carries
api_key_env_varover from the existing main model entry when the override does not set one explicitly, so aconfig.ymlthat points at a custom environment variable continues to resolve credentials correctly when the model name is swapped at request time. Explicitapi_key_env_varon the override still takes precedence. Implicit env-var fallback for engines likeopenaicontinues to work via the default framework, which was already correct on the no-override path.Added regression coverage on
_update_models_in_configfor both the preserve-from-config and override-wins paths, plus an integration-level test that drives_get_railsend to end withOPENAI_API_KEYset and a request-side model override, asserting the resolved HTTP client carries the env-derived key.Summary by CodeRabbit