Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 15 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Comment on lines +303 to +306

Copilot AI Apr 13, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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).

Suggested change
# 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",

Copilot uses AI. Check for mistakes.
)

# 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?
Expand Down
2 changes: 2 additions & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand All @@ -57,6 +58,7 @@ all = [
"sentence-transformers>=2.6",
"FlagEmbedding>=1.2",
"litellm>=1.30",
"openai>=1.12", # NVIDIA NIM support
]

[project.scripts]
Expand Down
48 changes: 35 additions & 13 deletions tests/unit/test_factory.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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__}"
)
2 changes: 2 additions & 0 deletions vektori/models/factory.py
Original file line number Diff line number Diff line change
Expand Up @@ -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] = {
Expand All @@ -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",
}


Expand Down
Loading
Loading