add NVIDIA NIM support with embedder and LLM providers#53
Conversation
📝 WalkthroughWalkthroughAdded NVIDIA NIM support: new Changes
Sequence DiagramsequenceDiagram
actor User
participant Factory as Factory
participant NvidiaEmbedder as NvidiaEmbedder
participant OpenAI as AsyncOpenAI Client
participant NVIDIA as NVIDIA NIM API
User->>Factory: create_embedder("nvidia")
Factory->>Factory: lookup "nvidia" in EMBEDDING_REGISTRY
Factory->>NvidiaEmbedder: instantiate (model, dimensions)
Factory-->>User: NvidiaEmbedder instance
User->>NvidiaEmbedder: embed_batch(texts)
NvidiaEmbedder->>NvidiaEmbedder: lazy init AsyncOpenAI client (api_key)
NvidiaEmbedder->>OpenAI: client.embeddings.create(model=..., input=[...], extra_body={...})
OpenAI->>NVIDIA: POST /v1/embeddings
NVIDIA-->>OpenAI: embedding vectors (with indices)
OpenAI-->>NvidiaEmbedder: response
NvidiaEmbedder->>NvidiaEmbedder: sort by index, preserve original order
NvidiaEmbedder-->>User: list[list[float]]
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 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.
Actionable comments posted: 2
🧹 Nitpick comments (1)
tests/unit/test_factory.py (1)
69-74: Add NVIDIA LLM factory coverage alongside the embedder test.Line 69 validates the embedder registration path, but the PR also registers NVIDIA in
LLM_REGISTRY(vektori/models/factory.pyLine 35). Add a matching LLM factory assertion to guard both provider mappings.Proposed test addition
def test_create_nvidia_embedder_default_model(): from vektori.models.nvidia import NvidiaEmbedder, DEFAULT_EMBEDDING_MODEL embedder = create_embedder("nvidia") assert isinstance(embedder, NvidiaEmbedder) assert embedder.model == DEFAULT_EMBEDDING_MODEL # "nvidia/llama-nemotron-embed-1b-v2" + + +def test_create_nvidia_llm_default_model(): + from vektori.models.nvidia import NvidiaLLM, DEFAULT_LLM_MODEL + + llm = create_llm("nvidia") + assert isinstance(llm, NvidiaLLM) + assert llm.model == DEFAULT_LLM_MODEL🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/unit/test_factory.py` around lines 69 - 74, Add a parallel LLM factory assertion to the embedder test: import create_llm and LLM_REGISTRY from vektori.models.factory, assert "nvidia" is present in LLM_REGISTRY, call create_llm("nvidia") and assert the result is not None and that its class name includes "Nvidia" (e.g., result.__class__.__name__.contains("Nvidia")) to ensure the LLM mapping for the "nvidia" provider is registered and returns an instance.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@vektori/models/nvidia.py`:
- Around line 104-118: Validate the dimensions parameter in __init__: ensure if
dimensions is not None it is an int and >0 (and optionally belongs to the set of
supported sizes for the selected model), otherwise raise a ValueError; store
only validated values in self._custom_dims so downstream uses (places that read
self._custom_dims or call methods that rely on dimensions) cannot proceed with
invalid config. Also add the same validation/check where dimensions are consumed
(any methods that read self._custom_dims) to fail fast if __init__ was bypassed.
- Around line 292-299: The current try/except around
client.chat.completions.create unconditionally retries without response_format
on any exception; change it so you only fallback when the error clearly
indicates the model doesn't support the JSON response_format (e.g., inspect the
caught exception from client.chat.completions.create and check its type or
message for keywords like "response_format", "json", "json_object", or an
SDK/model error code that denotes unsupported response format), and for all
other exceptions (auth, network, rate-limit, etc.) re-raise the exception
instead of retrying; keep references to kwargs["response_format"] and the same
client.chat.completions.create call, remove the unconditional del/retry, and
perform the conditional delete-and-retry only when the inspected error matches
the JSON-mode incompatibility signature.
---
Nitpick comments:
In `@tests/unit/test_factory.py`:
- Around line 69-74: Add a parallel LLM factory assertion to the embedder test:
import create_llm and LLM_REGISTRY from vektori.models.factory, assert "nvidia"
is present in LLM_REGISTRY, call create_llm("nvidia") and assert the result is
not None and that its class name includes "Nvidia" (e.g.,
result.__class__.__name__.contains("Nvidia")) to ensure the LLM mapping for the
"nvidia" provider is registered and returns an instance.
🪄 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: defaults
Review profile: CHILL
Plan: Pro
Run ID: 9eeb607d-c323-4bf0-9b28-0931cf8b0462
📒 Files selected for processing (5)
README.mdpyproject.tomltests/unit/test_factory.pyvektori/models/factory.pyvektori/models/nvidia.py
There was a problem hiding this comment.
Pull request overview
Adds a new NVIDIA NIM model provider to Vektori so users can run embeddings and extraction via NVIDIA’s OpenAI-compatible inference endpoint.
Changes:
- Introduces
NvidiaEmbedderandNvidiaLLMproviders (OpenAI SDK with NVIDIA base URL) with model normalization and Matryoshka dimension support. - Registers the new provider in the model factory for both embedding and LLM creation.
- Updates docs/tests and project metadata to expose the new provider.
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated 6 comments.
Show a summary per file
| File | Description |
|---|---|
vektori/models/nvidia.py |
New NVIDIA NIM embedding + LLM provider implementations. |
vektori/models/factory.py |
Adds nvidia to embedder/LLM registries. |
tests/unit/test_factory.py |
Adds a basic factory creation test for the NVIDIA embedder. |
README.md |
Documents how to configure NVIDIA NIM models in Vektori. |
pyproject.toml |
Adds an (currently redundant) nvidia optional dependency entry and updates all. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| 384, 512, 768, 1024, or 2048. Use the dimensions parameter or | ||
| embedding_dimension config to set custom dimensions. |
There was a problem hiding this comment.
The docstring says Matryoshka dimensions can be set via the embedding_dimension config, but the embedder only accepts dimensions and nothing in Vektori/create_embedder() passes embedding_dimension through. Either plumb VektoriConfig.embedding_dimension into this provider (or factory) or adjust the docstring to avoid implying it works today.
| 384, 512, 768, 1024, or 2048. Use the dimensions parameter or | |
| embedding_dimension config to set custom dimensions. | |
| 384, 512, 768, 1024, or 2048. Use the dimensions parameter to | |
| set custom dimensions. |
| # Some NVIDIA embedding models require input_type parameter | ||
| # "passage" is used for indexing, "query" for querying | ||
| # We default to "passage" for general embedding use | ||
| if self.model in { | ||
| "nvidia/llama-nemotron-embed-1b-v2", | ||
| "nvidia/llama-nemotron-embed-vl-1b-v2", | ||
| }: | ||
| extra_body["input_type"] = "passage" | ||
|
|
There was a problem hiding this comment.
input_type is hard-coded to "passage" for models that distinguish between indexing (passage) vs querying (query). Vektori uses the same embedder for both ingestion and retrieval queries, so this will also embed user queries as "passage" and can degrade retrieval quality for asymmetric models. Consider making input_type configurable (e.g., constructor arg) and/or exposing separate methods/instances for passage vs query embeddings.
| # Try JSON mode first (for structured extraction) | ||
| try: | ||
| kwargs["response_format"] = {"type": "json_object"} | ||
| response = await client.chat.completions.create(**kwargs) | ||
| except Exception: | ||
| # Fallback if model doesn't support JSON mode | ||
| del kwargs["response_format"] | ||
| response = await client.chat.completions.create(**kwargs) |
There was a problem hiding this comment.
generate() retries without response_format on any exception. This can mask real failures (auth/rate limit/network) and can double-request/cost on transient errors. Prefer only falling back for the specific "response_format not supported"/400-type error (similar to the GitHub provider’s pattern) and re-raise unexpected exceptions.
| # NVIDIA embedding models (Matryoshka: 384-2048 dimensions) | ||
| v = Vektori( | ||
| embedding_model="nvidia:llama-nemotron-embed-1b-v2", | ||
| embedding_dimension=1024, # Optional: 384, 512, 768, 1024, or 2048 |
There was a problem hiding this comment.
The README suggests embedding_dimension=1024 will control NVIDIA Matryoshka output size, but the current code path (Vektori -> create_embedder) doesn’t pass embedding_dimension into NvidiaEmbedder(dimensions=...). Either update the wiring so this example works as written, or adjust the docs to show how to set dimensions (e.g., via provider kwargs/config).
| # NVIDIA embedding models (Matryoshka: 384-2048 dimensions) | |
| v = Vektori( | |
| embedding_model="nvidia:llama-nemotron-embed-1b-v2", | |
| embedding_dimension=1024, # Optional: 384, 512, 768, 1024, or 2048 | |
| # NVIDIA embedding models | |
| # Note: llama-nemotron-embed-1b-v2 supports Matryoshka output sizes, | |
| # but `embedding_dimension=` is not currently forwarded by this Vektori code path. | |
| v = Vektori( | |
| embedding_model="nvidia:llama-nemotron-embed-1b-v2", |
| def test_create_nvidia_embedder_default_model(): | ||
| from vektori.models.nvidia import NvidiaEmbedder, DEFAULT_EMBEDDING_MODEL | ||
|
|
||
| embedder = create_embedder("nvidia") | ||
| assert isinstance(embedder, NvidiaEmbedder) | ||
| assert embedder.model == DEFAULT_EMBEDDING_MODEL # "nvidia/llama-nemotron-embed-1b-v2" |
There was a problem hiding this comment.
The PR description claims “comprehensive NVIDIA tests” (including integration tests and Matryoshka dimension verification), but in this PR’s changes the only NVIDIA test added is a factory default-model check. Either add the missing unit/integration tests that exercise the new provider behavior, or update the PR description/testing notes to match what’s actually included.
| [project.optional-dependencies] | ||
| postgres = ["asyncpg>=0.29", "pgvector>=0.2"] | ||
| neo4j = ["neo4j>=5.14"] | ||
| qdrant = ["qdrant-client>=1.7"] | ||
| anthropic = ["anthropic>=0.20", "voyageai>=0.2"] | ||
| sentence-transformers = ["sentence-transformers>=2.6"] | ||
| bge = ["FlagEmbedding>=1.2"] | ||
| litellm = ["litellm>=1.30"] | ||
| nvidia = ["openai>=1.12"] # NVIDIA NIM uses OpenAI-compatible API |
There was a problem hiding this comment.
openai>=1.12 is already a required base dependency, so adding it again under the nvidia extra (and again in the all extra) is redundant and can mislead readers into thinking OpenAI is optional unless .[nvidia] is installed. Consider removing the nvidia extra entirely (or, if the goal is to make OpenAI optional, moving it out of base deps and updating providers accordingly).
…s in NvidiaEmbedder
sankalphs
left a comment
sankalphs
left a comment
There was a problem hiding this comment.
Made changes suggested by coderabbit
|
Hey @sankalphs , thanks so much for this contribution! The code looks really solid and the integration is clean |
|
@laxmanclo fixed the CI error |
There was a problem hiding this comment.
🧹 Nitpick comments (9)
tests/unit/test_factory.py (3)
77-80: Use consistent assertion style.Line 80 uses Yoda-style assertion (
"literal" == variable) which is inconsistent with line 75 and other tests in this file that usevariable == "literal".🔧 Proposed fix
def test_create_nvidia_llm_default_model(): llm = create_llm("nvidia") assert isinstance(llm, NvidiaLLM) - assert "nvidia/llama-3.3-nemotron-super-49b-v1" == llm.model + assert llm.model == "nvidia/llama-3.3-nemotron-super-49b-v1"🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/unit/test_factory.py` around lines 77 - 80, The test test_create_nvidia_llm_default_model uses a Yoda-style assertion; change the equality check to match the file's consistent style by asserting llm.model == "nvidia/llama-3.3-nemotron-super-49b-v1" (keep the other assertions as-is: llm = create_llm("nvidia") and isinstance(llm, NvidiaLLM)).
82-89: Consider consolidating redundant test logic.This test overlaps significantly with
test_create_nvidia_llm_default_model— both verify thatcreate_llm("nvidia")returns a valid instance. The class name substring check (line 87-88) is also weaker than theisinstancecheck in line 79. Consider whether this test adds sufficient value or could be merged.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/unit/test_factory.py` around lines 82 - 89, The test test_nvidia_llm_registered duplicates assertions already covered in test_create_nvidia_llm_default_model; remove redundancy by either deleting test_nvidia_llm_registered or merging its unique check into test_create_nvidia_llm_default_model: ensure LLM_REGISTRY contains "nvidia" and that create_llm("nvidia") returns the expected type using the stronger isinstance assertion (matching the class used in the other test) rather than the weaker substring check of llm.__class__.__name__.
67-70: Missing blank line between test functions.PEP8 requires two blank lines between top-level function definitions. There should be a blank line before
test_create_nvidia_embedder_custom_dimensions.🔧 Proposed fix
assert embedder.model == DEFAULT_EMBEDDING_MODEL + def test_create_nvidia_embedder_custom_dimensions(): embedder = create_embedder("nvidia:llama-nemotron-embed-1b-v2", dimensions=1024)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/unit/test_factory.py` around lines 67 - 70, Add a single blank line before the top-level test function test_create_nvidia_embedder_custom_dimensions to satisfy PEP8's requirement of two blank lines between top-level function definitions; locate the test function definition (def test_create_nvidia_embedder_custom_dimensions():) near other tests that call create_embedder and instantiate NvidiaEmbedder, and insert the blank line immediately above it.vektori/models/nvidia.py (6)
164-176: Consider extracting dimension validation to a helper method.The same validation logic for
_custom_dimsis repeated in__init__(lines 119-130), here in thedimensionproperty (lines 165-176), and inembed_batch(lines 211-222). While this defensive approach ensures fail-fast behavior, extracting to a private_validate_custom_dims()method would reduce duplication and ensure consistent validation.🔧 Proposed refactor
+ def _validate_custom_dims(self) -> None: + """Validate _custom_dims if set.""" + if self._custom_dims is None: + return + if not isinstance(self._custom_dims, int): + raise ValueError( + f"_custom_dims must be an integer, got {type(self._custom_dims).__name__}" + ) + if self._custom_dims <= 0: + raise ValueError(f"_custom_dims must be positive, got {self._custom_dims}") + supported_dims = {384, 512, 768, 1024, 2048} + if self._custom_dims not in supported_dims: + raise ValueError( + f"_custom_dims must be one of {sorted(supported_dims)}, got {self._custom_dims}" + ) + `@property` def dimension(self) -> int: - # Validate _custom_dims in case __init__ was bypassed - if self._custom_dims is not None: - if not isinstance(self._custom_dims, int): - raise ValueError( - f"_custom_dims must be an integer, got {type(self._custom_dims).__name__}" - ) - ... + self._validate_custom_dims() if self._custom_dims is not None and self.model in MATRYOSHKA_MODELS: return self._custom_dims return EMBEDDING_DIMENSIONS.get(self.model, 2048)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@vektori/models/nvidia.py` around lines 164 - 176, Extract the repeated validation logic for _custom_dims into a single private helper (e.g. _validate_custom_dims()) and call it from __init__, the dimension property, and embed_batch to remove duplication; _validate_custom_dims() should check type is int, >0, and membership in supported_dims = {384,512,768,1024,2048} and raise the same ValueError messages so behavior remains identical, and replace the in-place checks in __init__, dimension, and embed_batch with calls to this helper.
356-356: Add defensive check for empty response choices.If the API returns an empty
choiceslist (edge case, but possible with certain errors or rate limits), this line will raise anIndexError. Consider adding a guard.🛡️ Proposed fix
- return response.choices[0].message.content or "" + if not response.choices: + logger.warning("NVIDIA API returned empty choices") + return "" + return response.choices[0].message.content or ""🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@vektori/models/nvidia.py` at line 356, Add a defensive guard before accessing response.choices[0] to avoid IndexError when the API returns an empty choices list: check that response.choices is truthy and has at least one element and that response.choices[0].message exists, then return response.choices[0].message.content or an empty string; otherwise return an empty string (or raise a clear exception). Update the return site that currently does "return response.choices[0].message.content or \"\"" to perform this guarded check (referencing response.choices and response.choices[0].message).
226-233: Hardcodedinput_type="passage"may affect retrieval quality.For asymmetric embedding models, using
"passage"for both indexing and querying can degrade retrieval quality. When embedding queries for search,input_type="query"is typically preferred. Consider:
- Adding an
input_typeparameter toembed_batch- Or providing a separate
embed_querymethod for query-time embeddingsThis is noted in the docstring (lines 227-228) but the current implementation always defaults to
"passage".🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@vektori/models/nvidia.py` around lines 226 - 233, The NVIDIA model branch hardcodes extra_body["input_type"]="passage" which forces passage-style embeddings for both indexing and queries; modify the embedding API by adding an input_type parameter to embed_batch (default "passage") and pass that value into extra_body for the NVIDIA models (referencing embed_batch, extra_body, and the model name checks), and/or add a convenience embed_query method that calls embed_batch(input_type="query") so callers can request query-style embeddings; update any callers/tests and the docstring to reflect the new parameter or new method.
284-305: Consider extracting shared client initialization.The
_get_clientmethod is nearly identical betweenNvidiaEmbedderandNvidiaLLM. Consider extracting to a module-level helper function or a shared mixin to reduce duplication.🔧 Proposed helper function
def _create_nvidia_client(api_key: str | None): """Create AsyncOpenAI client configured for NVIDIA API.""" try: from openai import AsyncOpenAI except ImportError as e: raise ImportError( "openai package required for NVIDIA NIM: pip install openai>=1.12" ) from e resolved_key = api_key or os.environ.get("NVIDIA_API_KEY") if not resolved_key: raise ValueError( "NVIDIA API key required. Set NVIDIA_API_KEY environment variable " "or pass api_key parameter." ) return AsyncOpenAI(api_key=resolved_key, base_url=NVIDIA_BASE_URL)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@vektori/models/nvidia.py` around lines 284 - 305, Extract the duplicated client-init logic used in NvidiaEmbedder._get_client and NvidiaLLM._get_client into a single module-level helper (e.g. _create_nvidia_client(api_key: str | None)), which should perform the AsyncOpenAI import, resolve api_key from the parameter or os.environ["NVIDIA_API_KEY"], raise the same ImportError/ValueError messages, and return AsyncOpenAI(api_key=resolved_key, base_url=NVIDIA_BASE_URL); then update both NvidiaEmbedder._get_client and NvidiaLLM._get_client to call this helper and assign/return its result (preserve lazy initialization of self._client and all existing error messages).
322-326: Consider making temperature configurable.The
temperatureis hardcoded to0.1, which is reasonable for extraction tasks but limits flexibility. Other LLM providers in the codebase may allow temperature configuration. Consider adding atemperatureparameter to__init__orgenerate()for consistency.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@vektori/models/nvidia.py` around lines 322 - 326, The temperature is hardcoded in the kwargs, so make it configurable by adding a temperature parameter (default 0.1) to the class constructor (__init__) or to the generate() method, persist it as self.temperature when added to __init__, and replace the hardcoded value in the kwargs block ("temperature": 0.1) with self.temperature (or the generate() arg). Update the relevant function signature(s) and any docstring/comments for the class/method (e.g., __init__, generate, kwargs, self.model) so callers can override temperature consistently with other providers.
119-131: Consider validating model compatibility whendimensionsis specified.The validation checks if
dimensionsis a supported Matryoshka size, but doesn't verify whether the specified model actually supports Matryoshka embeddings. A user could specifydimensions=1024for a non-Matryoshka model (e.g.,nvidia/nv-embed-v1), and the parameter would be silently ignored at embedding time (lines 223-224). Failing fast here would prevent confusion.🔧 Proposed fix
if dimensions not in supported_dims: raise ValueError( f"dimensions must be one of {sorted(supported_dims)}, got {dimensions}" ) + # Check model supports Matryoshka dimensions + if self.model not in MATRYOSHKA_MODELS: + raise ValueError( + f"Model '{self.model}' does not support Matryoshka dimensions. " + f"Supported models: {sorted(MATRYOSHKA_MODELS)}" + ) self._custom_dims = dimensions🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@vektori/models/nvidia.py` around lines 119 - 131, When dimensions is provided you must also validate that the selected model supports Matryoshka embeddings rather than only checking supported_dims: after the existing type/size checks (dimensions, supported_dims) add a compatibility check against the current model (e.g., self._model) and/or a known set of Matryoshka-capable model names (or a helper like _is_matryoshka_model(model)); if the model does not support Matryoshka raise a ValueError explaining that dimensions is only valid for Matryoshka models, and only set self._custom_dims if the compatibility check passes.
🤖 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/unit/test_factory.py`:
- Around line 77-80: The test test_create_nvidia_llm_default_model uses a
Yoda-style assertion; change the equality check to match the file's consistent
style by asserting llm.model == "nvidia/llama-3.3-nemotron-super-49b-v1" (keep
the other assertions as-is: llm = create_llm("nvidia") and isinstance(llm,
NvidiaLLM)).
- Around line 82-89: The test test_nvidia_llm_registered duplicates assertions
already covered in test_create_nvidia_llm_default_model; remove redundancy by
either deleting test_nvidia_llm_registered or merging its unique check into
test_create_nvidia_llm_default_model: ensure LLM_REGISTRY contains "nvidia" and
that create_llm("nvidia") returns the expected type using the stronger
isinstance assertion (matching the class used in the other test) rather than the
weaker substring check of llm.__class__.__name__.
- Around line 67-70: Add a single blank line before the top-level test function
test_create_nvidia_embedder_custom_dimensions to satisfy PEP8's requirement of
two blank lines between top-level function definitions; locate the test function
definition (def test_create_nvidia_embedder_custom_dimensions():) near other
tests that call create_embedder and instantiate NvidiaEmbedder, and insert the
blank line immediately above it.
In `@vektori/models/nvidia.py`:
- Around line 164-176: Extract the repeated validation logic for _custom_dims
into a single private helper (e.g. _validate_custom_dims()) and call it from
__init__, the dimension property, and embed_batch to remove duplication;
_validate_custom_dims() should check type is int, >0, and membership in
supported_dims = {384,512,768,1024,2048} and raise the same ValueError messages
so behavior remains identical, and replace the in-place checks in __init__,
dimension, and embed_batch with calls to this helper.
- Line 356: Add a defensive guard before accessing response.choices[0] to avoid
IndexError when the API returns an empty choices list: check that
response.choices is truthy and has at least one element and that
response.choices[0].message exists, then return
response.choices[0].message.content or an empty string; otherwise return an
empty string (or raise a clear exception). Update the return site that currently
does "return response.choices[0].message.content or \"\"" to perform this
guarded check (referencing response.choices and response.choices[0].message).
- Around line 226-233: The NVIDIA model branch hardcodes
extra_body["input_type"]="passage" which forces passage-style embeddings for
both indexing and queries; modify the embedding API by adding an input_type
parameter to embed_batch (default "passage") and pass that value into extra_body
for the NVIDIA models (referencing embed_batch, extra_body, and the model name
checks), and/or add a convenience embed_query method that calls
embed_batch(input_type="query") so callers can request query-style embeddings;
update any callers/tests and the docstring to reflect the new parameter or new
method.
- Around line 284-305: Extract the duplicated client-init logic used in
NvidiaEmbedder._get_client and NvidiaLLM._get_client into a single module-level
helper (e.g. _create_nvidia_client(api_key: str | None)), which should perform
the AsyncOpenAI import, resolve api_key from the parameter or
os.environ["NVIDIA_API_KEY"], raise the same ImportError/ValueError messages,
and return AsyncOpenAI(api_key=resolved_key, base_url=NVIDIA_BASE_URL); then
update both NvidiaEmbedder._get_client and NvidiaLLM._get_client to call this
helper and assign/return its result (preserve lazy initialization of
self._client and all existing error messages).
- Around line 322-326: The temperature is hardcoded in the kwargs, so make it
configurable by adding a temperature parameter (default 0.1) to the class
constructor (__init__) or to the generate() method, persist it as
self.temperature when added to __init__, and replace the hardcoded value in the
kwargs block ("temperature": 0.1) with self.temperature (or the generate() arg).
Update the relevant function signature(s) and any docstring/comments for the
class/method (e.g., __init__, generate, kwargs, self.model) so callers can
override temperature consistently with other providers.
- Around line 119-131: When dimensions is provided you must also validate that
the selected model supports Matryoshka embeddings rather than only checking
supported_dims: after the existing type/size checks (dimensions, supported_dims)
add a compatibility check against the current model (e.g., self._model) and/or a
known set of Matryoshka-capable model names (or a helper like
_is_matryoshka_model(model)); if the model does not support Matryoshka raise a
ValueError explaining that dimensions is only valid for Matryoshka models, and
only set self._custom_dims if the compatibility check passes.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: f5a835c9-e1b6-4e48-a6f5-d6f0a4d6c8bc
📒 Files selected for processing (2)
tests/unit/test_factory.pyvektori/models/nvidia.py
Added support for NVIDIA NIM (NVIDIA Inference Microservices) as a model provider for both embeddings and LLM-based fact extraction. This enables users to use GPU-optimized models hosted on NVIDIA's platform.
Changes Made
NvidiaEmbedder class with Matryoshka embedding support (384-2048 dimensions)
NvidiaLLM class for fact/episode extraction
Automatic namespace handling (nvidia/ prefix for NVIDIA models, full path for third-party models)
Support for input_type parameter on asymmetric embedding models
Testing
✅ All 53 unit tests pass
✅ Integration tests verify actual NVIDIA API calls work correctly
✅ Matryoshka dimensions verified (384, 512, 768, 1024, 2048)
✅ Factory creation tests for all model types
Notes
Summary by CodeRabbit
New Features
Documentation
Chores
Tests