Merge branch 'main' into feat/pipecat-voice-integration

Author: laxmanclo

README.md

@@ -321,6 +321,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

@@ -42,6 +42,7 @@ sentence-transformers = ["sentence-transformers>=2.6"]
 bge = ["FlagEmbedding>=1.2"]
 litellm = ["litellm>=1.30"]
 pipecat = ["pipecat-ai>=0.0.50", "fastapi>=0.110", "uvicorn[standard]>=0.29"]
+nvidia = ["openai>=1.12"]  # NVIDIA NIM uses OpenAI-compatible API
 dev = [
     "pytest>=8.0",
     "pytest-asyncio>=0.23",
@@ -60,6 +61,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

@@ -2,27 +2,25 @@
 
 import pytest
 
-from vektori.models.factory import create_embedder, create_llm
+from vektori.models.anthropic import AnthropicLLM
+from vektori.models.factory import LLM_REGISTRY, create_embedder, create_llm
+from vektori.models.nvidia import DEFAULT_EMBEDDING_MODEL, NvidiaEmbedder, NvidiaLLM
+from vektori.models.ollama import OllamaEmbedder, OllamaLLM
+from vektori.models.openai import OpenAIEmbedder, OpenAILLM
 
 
 def test_create_openai_embedder():
-    from vektori.models.openai import OpenAIEmbedder
-
     embedder = create_embedder("openai:text-embedding-3-small")
     assert isinstance(embedder, OpenAIEmbedder)
     assert embedder.model == "text-embedding-3-small"
 
 
 def test_create_openai_embedder_default_model():
-    from vektori.models.openai import OpenAIEmbedder
-
     embedder = create_embedder("openai")
     assert isinstance(embedder, OpenAIEmbedder)
 
 
 def test_create_ollama_embedder():
-    from vektori.models.ollama import OllamaEmbedder
-
     embedder = create_embedder("ollama:nomic-embed-text")
     assert isinstance(embedder, OllamaEmbedder)
     assert embedder.model == "nomic-embed-text"
@@ -41,27 +39,51 @@ def test_unknown_embedding_provider_raises():
 
 
 def test_create_openai_llm():
-    from vektori.models.openai import OpenAILLM
-
     llm = create_llm("openai:gpt-4o-mini")
     assert isinstance(llm, OpenAILLM)
     assert llm.model == "gpt-4o-mini"
 
 
 def test_create_ollama_llm():
-    from vektori.models.ollama import OllamaLLM
-
     llm = create_llm("ollama:llama3")
     assert isinstance(llm, OllamaLLM)
 
 
 def test_create_anthropic_llm():
-    from vektori.models.anthropic import AnthropicLLM
-
     llm = create_llm("anthropic:claude-haiku-4-5-20251001")
     assert isinstance(llm, AnthropicLLM)
 
 
 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():
+    embedder = create_embedder("nvidia")
+    assert isinstance(embedder, NvidiaEmbedder)
+    assert embedder.model == DEFAULT_EMBEDDING_MODEL
+
+def test_create_nvidia_embedder_custom_dimensions():
+    embedder = create_embedder("nvidia:llama-nemotron-embed-1b-v2", dimensions=1024)
+    assert isinstance(embedder, NvidiaEmbedder)
+    assert embedder.dimension == 1024  # Matryoshka support
+
+def test_create_nvidia_llm():
+    llm = create_llm("nvidia:llama-3.3-nemotron-super-49b-v1")
+    assert isinstance(llm, NvidiaLLM)
+    assert llm.model == "nvidia/llama-3.3-nemotron-super-49b-v1"
+
+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
+
+def test_nvidia_llm_registered():
+    """Verify NVIDIA LLM is registered in factory."""
+    assert "nvidia" in LLM_REGISTRY, "nvidia should be registered in LLM_REGISTRY"
+    llm = create_llm("nvidia")
+    assert llm is not None, "create_llm('nvidia') should return a valid instance"
+    assert "Nvidia" in llm.__class__.__name__, (
+        f"LLM class name should include 'Nvidia', got {llm.__class__.__name__}"
+    )

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",
 }