add NVIDIA NIM support with embedder and LLM providers

README.md

@@ -298,6 +298,21 @@ v = Vektori(embedding_model="bge:BAAI/bge-m3")
 v = Vektori(extraction_model="litellm:groq/llama3-8b-8192")
 ```
 
+## NVIDIA NIM - GPU-optimized models via [NVIDIA NIM](https://build.nvidia.com).
+```python
+# 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 LLM models (nvidia/ prefix auto-added)
+v = Vektori(extraction_model="nvidia:llama-3.3-nemotron-super-49b-v1")
+
+# Third-party models hosted on NVIDIA NIM (use full path)
+v = Vektori(extraction_model="nvidia:z-ai/glm5")
+
+```
 ---
 
 ## Why Not Mem0 / Zep?

pyproject.toml

@@ -40,6 +40,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
 dev = [
     "pytest>=8.0",
     "pytest-asyncio>=0.23",
@@ -57,6 +58,7 @@ all = [
     "sentence-transformers>=2.6",
     "FlagEmbedding>=1.2",
     "litellm>=1.30",
+    "openai>=1.12",  # NVIDIA NIM support
 ]
 
 [project.scripts]

tests/unit/test_factory.py

@@ -65,3 +65,10 @@ def test_create_anthropic_llm():
 def test_unknown_llm_provider_raises():
     with pytest.raises(ValueError, match="Unknown LLM provider"):
         create_llm("nonexistent:model")
+
+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"

vektori/models/factory.py

@@ -20,6 +20,7 @@
     "cloudflare": "vektori.models.cloudflare.CloudflareEmbedder",
     # LiteLLM: 100+ providers — Together AI, Cohere, Azure, Ollama, etc.
     "litellm": "vektori.models.litellm_embedder.LiteLLMEmbedder",
+    "nvidia": "vektori.models.nvidia.NvidiaEmbedder",
 }
 
 LLM_REGISTRY: dict[str, str] = {
@@ -31,6 +32,7 @@
     "gemini": "vektori.models.gemini.GeminiLLM",  # Direct Gemini API
     # LiteLLM: single interface for 100+ providers — recommended for extraction
     "litellm": "vektori.models.litellm_provider.LiteLLMProvider",
+    "nvidia": "vektori.models.nvidia.NvidiaLLM",
 }
 
 

vektori/models/nvidia.py

@@ -0,0 +1,301 @@
+"""NVIDIA NIM embedding and LLM providers.
+
+NVIDIA NIM (NVIDIA Inference Microservices) provides GPU-optimized models
+for embedding and text generation via an OpenAI-compatible API.
+
+Get API key: https://build.nvidia.com
+
+Usage:
+    v = Vektori(
+        embedding_model="nvidia:llama-nemotron-embed-1b-v2",
+        extraction_model="nvidia:llama-3.3-nemotron-super-49b-v1",
+    )
+
+Environment:
+    NVIDIA_API_KEY - Your NVIDIA API key (required)
+
+Models:
+    Embeddings (llama-nemotron-embed-1b-v2 is recommended default):
+    - llama-nemotron-embed-1b-v2: 2048 dim, 8192 tokens, multilingual, Matryoshka
+    - llama-3.2-nv-embedqa-1b-v2: 2048 dim, 8192 tokens, QA-optimized
+    - baai/bge-m3: 1024 dim, 8192 tokens, multilingual, dense+sparse+multi-vector
+    - nv-embed-v1: 4096 dim, 32k tokens, generalist
+    - nv-embedqa-e5-v5: 1024 dim, 512 tokens, fast
+
+    LLMs:
+    - llama-3.3-nemotron-super-49b-v1: High quality extraction (default)
+    - llama-3.1-nemotron-ultra-253b-v1: Largest NVIDIA model
+    - nemotron-4-mini-hindi-4b-instruct: Fast, efficient
+    - z-ai/glm5: GLM-5, multilingual LLM
+    - z-ai/glm4.7: GLM-4.7, multilingual LLM
+    - google/gemma-4-31b-it: Google's Gemma 4 31B instruction tuned
+    - minimaxai/minimax-m2.7: MiniMax M2.7 model
+    - moonshotai/kimi-k2.5: Moonshot Kimi K2.5
+    - moonshotai/kimi-k2-instruct: Moonshot Kimi K2 instruction
+    - deepseek-ai/deepseek-v3.1: DeepSeek V3.1
+    - qwen/qwen3-coder-480b-a35b-instruct: Qwen3 Coder 480B
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+from typing import Any
+
+from vektori.models.base import EmbeddingProvider, LLMProvider
+
+logger = logging.getLogger(__name__)
+
+NVIDIA_BASE_URL = "https://integrate.api.nvidia.com/v1"
+
+# Embedding dimensions for known NVIDIA models
+# Matryoshka models support multiple dimensions; max is shown
+EMBEDDING_DIMENSIONS = {
+    # Nemotron family (Matryoshka: 384, 512, 768, 1024, 2048)
+    "nvidia/llama-nemotron-embed-1b-v2": 2048,
+    "nvidia/llama-nemotron-embed-vl-1b-v2": 2048,
+    # NeMo Retriever family (Matryoshka: 384, 512, 768, 1024, 2048)
+    "nvidia/llama-3.2-nv-embedqa-1b-v2": 2048,
+    "nvidia/llama-3.2-nv-embedqa-1b-v1": 2048,
+    "nvidia/llama-3.2-nemoretriever-300m-embed-v2": 1024,
+    "nvidia/llama-3.2-nemoretriever-300m-embed-v1": 1024,
+    "nvidia/llama-3.2-nemoretriever-1b-vlm-embed-v1": 2048,
+    # BAAI BGE-M3 (hosted on NVIDIA NIM)
+    "baai/bge-m3": 1024,
+    # Other models
+    "nvidia/nv-embed-v1": 4096,
+    "nvidia/nv-embedqa-e5-v5": 1024,
+    "nvidia/nv-embedcode-7b-v1": 4096,
+    "nvidia/nv-embedqa-e5-v4": 1024,
+    "nvidia/nvclip": 512,
+}
+
+# Models that support Matryoshka embeddings (configurable dimensions)
+MATRYOSHKA_MODELS = {
+    "nvidia/llama-nemotron-embed-1b-v2",
+    "nvidia/llama-nemotron-embed-vl-1b-v2",
+    "nvidia/llama-3.2-nv-embedqa-1b-v2",
+    "nvidia/llama-3.2-nv-embedqa-1b-v1",
+}
+
+DEFAULT_EMBEDDING_MODEL = "nvidia/llama-nemotron-embed-1b-v2"
+DEFAULT_LLM_MODEL = "nvidia/llama-3.3-nemotron-super-49b-v1"
+
+
+class NvidiaEmbedder(EmbeddingProvider):
+    """NVIDIA NIM embedding models via OpenAI-compatible API.
+
+    Supports all NVIDIA embedding models including:
+    - llama-nemotron-embed-1b-v2 (recommended): multilingual, 2048-dim, Matryoshka
+    - llama-3.2-nv-embedqa-1b-v2: QA-optimized, 2048-dim, Matryoshka
+    - nv-embed-v1: generalist, 4096-dim, 32k context
+    - nv-embedqa-e5-v5: fast, 1024-dim, 512 context
+
+    Matryoshka models support configurable output dimensions:
+    384, 512, 768, 1024, or 2048. Use the dimensions parameter or
+    embedding_dimension config to set custom dimensions.
+
+    Args:
+        model: Model name from NVIDIA catalog. Defaults to llama-nemotron-embed-1b-v2.
+        api_key: NVIDIA API key. Falls back to NVIDIA_API_KEY env var.
+        dimensions: Optional custom embedding dimensions for Matryoshka models.
+    """
+
+    def __init__(
+        self,
+        model: str | None = None,
+        api_key: str | None = None,
+        dimensions: int | None = None,
+    ) -> None:
+        raw_model = model or DEFAULT_EMBEDDING_MODEL
+        # NVIDIA NIM requires full path with namespace prefix (e.g., "nvidia/" or "baai/")
+        # Add "nvidia/" prefix if model doesn't already have a namespace
+        if "/" not in raw_model:
+            raw_model = f"nvidia/{raw_model}"
+        self.model = raw_model
+        self._api_key = api_key
+        self._custom_dims = dimensions
+        self._client = None
+
+    def _get_client(self):
+        """Lazy initialization of OpenAI client configured for NVIDIA API."""
+        if self._client is None:
+            try:
+                from openai import AsyncOpenAI
+            except ImportError as e:
+                raise ImportError(
+                    "openai package required for NVIDIA NIM: pip install openai>=1.12"
+                ) from e
+
+            api_key = self._api_key or os.environ.get("NVIDIA_API_KEY")
+            if not api_key:
+                raise ValueError(
+                    "NVIDIA API key required. Set NVIDIA_API_KEY environment variable "
+                    "or pass api_key parameter."
+                )
+
+            self._client = AsyncOpenAI(
+                api_key=api_key,
+                base_url=NVIDIA_BASE_URL,
+            )
+        return self._client
+
+    @property
+    def dimension(self) -> int:
+        """Return embedding dimension.
+
+        If custom dimensions were specified and the model supports Matryoshka
+        embeddings, returns the custom dimension. Otherwise returns the
+        model's default dimension.
+        """
+        if self._custom_dims and self.model in MATRYOSHKA_MODELS:
+            return self._custom_dims
+        return EMBEDDING_DIMENSIONS.get(self.model, 2048)
+
+    async def embed(self, text: str) -> list[float]:
+        """Embed a single text string.
+
+        Args:
+            text: Text to embed.
+
+        Returns:
+            Embedding vector as list of floats.
+        """
+        return (await self.embed_batch([text]))[0]
+
+    async def embed_batch(self, texts: list[str]) -> list[list[float]]:
+        """Embed multiple texts in a single API call.
+
+        Args:
+            texts: List of text strings to embed.
+
+        Returns:
+            List of embedding vectors, preserving input order.
+        """
+        if not texts:
+            return []
+
+        client = self._get_client()
+
+        # Build extra_body for NVIDIA-specific parameters
+        extra_body: dict[str, Any] = {}
+
+        # Matryoshka models support custom dimensions
+        if self._custom_dims and self.model in MATRYOSHKA_MODELS:
+            extra_body["dimensions"] = self._custom_dims
+
+        # 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"
+
+        response = await client.embeddings.create(
+            input=texts,
+            model=self.model,
+            extra_body=extra_body if extra_body else None,
+        )
+
+        # Sort by index to preserve order
+        return [item.embedding for item in sorted(response.data, key=lambda x: x.index)]
+
+
+class NvidiaLLM(LLMProvider):
+    """NVIDIA NIM LLM for fact and episode extraction.
+
+    Supports all NVIDIA chat completion models and third-party models hosted
+    on NVIDIA NIM, including:
+    - nvidia/llama-3.3-nemotron-super-49b-v1 (default): High quality extraction
+    - nvidia/llama-3.1-nemotron-ultra-253b-v1: Largest NVIDIA model
+    - nvidia/nemotron-4-mini-hindi-4b-instruct: Fast, efficient
+    - z-ai/glm5: GLM-5 multilingual model
+    - z-ai/glm4.7: GLM-4.7 multilingual model
+    - google/gemma-4-31b-it: Google's Gemma 4 31B instruction tuned
+    - minimaxai/minimax-m2.7: MiniMax M2.7 model
+    - moonshotai/kimi-k2.5: Moonshot Kimi K2.5
+    - moonshotai/kimi-k2-instruct: Moonshot Kimi K2 instruction
+    - deepseek-ai/deepseek-v3.1: DeepSeek V3.1
+    - qwen/qwen3-coder-480b-a35b-instruct: Qwen3 Coder 480B
+
+    Uses OpenAI-compatible API with NVIDIA-specific base URL.
+
+    Args:
+        model: Model name from NVIDIA catalog (e.g., "llama-3.3-nemotron-super-49b-v1"
+            for NVIDIA models, or full path like "z-ai/glm5" for third-party models).
+        api_key: NVIDIA API key. Falls back to NVIDIA_API_KEY env var.
+    """
+
+    def __init__(
+        self,
+        model: str | None = None,
+        api_key: str | None = None,
+    ) -> None:
+        raw_model = model or DEFAULT_LLM_MODEL
+        # NVIDIA NIM requires full path with namespace prefix (e.g., "nvidia/")
+        # Add "nvidia/" prefix if model doesn't already have a namespace
+        if "/" not in raw_model:
+            raw_model = f"nvidia/{raw_model}"
+        self.model = raw_model
+        self._api_key = api_key
+        self._client = None
+
+    def _get_client(self):
+        """Lazy initialization of OpenAI client configured for NVIDIA API."""
+        if self._client is None:
+            try:
+                from openai import AsyncOpenAI
+            except ImportError as e:
+                raise ImportError(
+                    "openai package required for NVIDIA NIM: pip install openai>=1.12"
+                ) from e
+
+            api_key = self._api_key or os.environ.get("NVIDIA_API_KEY")
+            if not api_key:
+                raise ValueError(
+                    "NVIDIA API key required. Set NVIDIA_API_KEY environment variable "
+                    "or pass api_key parameter."
+                )
+
+            self._client = AsyncOpenAI(
+                api_key=api_key,
+                base_url=NVIDIA_BASE_URL,
+            )
+        return self._client
+
+    async def generate(self, prompt: str, max_tokens: int | None = None) -> str:
+        """Generate text completion for the given prompt.
+
+        Attempts to use JSON mode for structured outputs. Falls back to
+        regular chat completion if the model doesn't support JSON mode.
+
+        Args:
+            prompt: The prompt text to send to the model.
+            max_tokens: Optional maximum tokens in the response.
+
+        Returns:
+            Generated text response.
+        """
+        client = self._get_client()
+
+        kwargs: dict[str, Any] = {
+            "model": self.model,
+            "messages": [{"role": "user", "content": prompt}],
+            "temperature": 0.1,
+        }
+
+        if max_tokens is not None:
+            kwargs["max_tokens"] = max_tokens
+
+        # 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)
+
+        return response.choices[0].message.content or ""