Conversation
Greptile SummaryThis PR fixes a cache-invalidation bug in the KB annoy index where switching the embedding model via the The fix backfills Key implementation details:
|
| Filename | Overview |
|---|---|
| nemoguardrails/rails/llm/llmrails.py | Adds backfill logic for embedding_model/embedding_engine into embedding_search_provider.parameters for both core and knowledge_base configs when a models[type=embeddings] entry is found; correctly guards against None model.model and only targets the default provider. |
| tests/test_llmrails.py | Adds five targeted test cases covering: backfill of both providers, no overwrite of explicit params, partial backfill, no backfill for custom providers, and no backfill when embeddings model is absent; coverage is comprehensive. |
Sequence Diagram
sequenceDiagram
participant User as User Config (YAML)
participant LLMRails as LLMRails.__init__
participant ESP_Core as core.embedding_search_provider
participant ESP_KB as knowledge_base.embedding_search_provider
participant KB as kb.py build()
User->>LLMRails: models[type=embeddings] entry
LLMRails->>LLMRails: set default_embedding_model/engine
LLMRails->>ESP_Core: if name=="default" and key missing → backfill embedding_model/engine
LLMRails->>ESP_KB: if name=="default" and key missing → backfill embedding_model/engine
KB->>ESP_KB: parameters.get("embedding_engine","") + parameters.get("embedding_model","")
KB->>KB: hash_prefix + content → cache hash
KB->>KB: load or build annoy index with correct model
Reviews (3): Last reviewed commit: "assert KB provider gets backfill when co..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
📝 WalkthroughWalkthroughThe code now populates embedding search provider parameters with embedding model and engine values during Changes
Estimated code review effort🎯 2 (Simple) | ⏱️ ~15 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 docstrings
🧪 Generate unit tests (beta)
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.
🧹 Nitpick comments (1)
tests/test_llmrails.py (1)
1439-1456: Consider also asserting knowledge_base backfill in this test.This test verifies that custom providers are skipped, but only checks
core.embedding_search_provider. Theknowledge_base.embedding_search_providerdefaults toname: "default"and should receive the backfill. Adding assertions for the knowledge_base would strengthen this test by verifying the selectivity of the backfill logic—ensuring it correctly applies to default providers while skipping custom ones in the same initialization.💡 Optional enhancement to verify selective backfill
def test_embedding_model_no_backfill_for_custom_provider(): config = RailsConfig.parse_object( { **EMBEDDING_MODEL_CONFIG_BASE, "core": { "embedding_search_provider": { "name": "custom", "parameters": {"some_param": "value"}, } }, } ) rails = LLMRails(config=config, llm=FakeLLM(responses=[" express greeting"])) assert "embedding_model" not in rails.config.core.embedding_search_provider.parameters assert "embedding_engine" not in rails.config.core.embedding_search_provider.parameters assert rails.config.core.embedding_search_provider.parameters["some_param"] == "value" + # knowledge_base remains default and should still be backfilled + assert rails.config.knowledge_base.embedding_search_provider.parameters["embedding_model"] == "intfloat/e5-large-v2" + assert rails.config.knowledge_base.embedding_search_provider.parameters["embedding_engine"] == "SentenceTransformers"🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/test_llmrails.py` around lines 1439 - 1456, Extend test_embedding_model_no_backfill_for_custom_provider to also assert that the knowledge_base.embedding_search_provider received the backfill: after creating rails = LLMRails(...), add assertions that rails.config.knowledge_base.embedding_search_provider.parameters contains the backfilled keys ("embedding_model" and/or "embedding_engine") and that their values match the expected values from EMBEDDING_MODEL_CONFIG_BASE (so the custom provider remains untouched while the default knowledge_base provider is populated).
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@tests/test_llmrails.py`:
- Around line 1439-1456: Extend
test_embedding_model_no_backfill_for_custom_provider to also assert that the
knowledge_base.embedding_search_provider received the backfill: after creating
rails = LLMRails(...), add assertions that
rails.config.knowledge_base.embedding_search_provider.parameters contains the
backfilled keys ("embedding_model" and/or "embedding_engine") and that their
values match the expected values from EMBEDDING_MODEL_CONFIG_BASE (so the custom
provider remains untouched while the default knowledge_base provider is
populated).
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: dfaf86f1-bea6-45cf-a0d1-0b1c2dcef416
📒 Files selected for processing (2)
nemoguardrails/rails/llm/llmrails.pytests/test_llmrails.py
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
Pouyanpi
changed the title
Pouyanpi
changed the title
tgasser-nv
left a comment
tgasser-nv
left a comment
There was a problem hiding this comment.
Looks good, approving! Not blocking for this PR, but items to consider as follow-ons:
- It feels a bit wrong for an object to mutate its data object passed in during initialization. I'd prefer to handle this at the Pydantic validation stage.
- Some embedding models (Gemini) have a
dimensionalityvalue inembedding_parameterswhich we'd need to add to item hash. Without this if a users changes the dimensionality we won't invalidate the cache and cause errors.
| ]: | ||
| if esp.name != "default": | ||
| continue | ||
| if "embedding_model" not in esp.parameters and model.model is not None: |
There was a problem hiding this comment.
nit: The Model model and engine fields can't be None. The model fields has its own validator, and engine isn't optional so it can't be None. I think the and model.{model,engine} is not None checks can be removed
2 participants
## Description
Resolves #1751
models: [{type: embeddings, ...}]config sectionmodels[type=embeddings]entry is found, backfillembedding_modelandembedding_engineintoembedding_search_provider.parametersfor bothcoreandknowledge_baseconfigs so the cache hash includes the actual model nameIndexError: Vector has wrong lengthat runtime with no clear error messageTest plan
poetry run nemoguardrails chat --config=./examples/bots/abc)modelssection (addpytest tests/ -k embeddingor just run (ensure .cache dir is deleted prior running it):
Expected Output:
After the Fix:
Prior to the Fix:
Summary by CodeRabbit
Release Notes
New Features
Tests