22
33import pytest
44
5- from vektori .models .factory import create_embedder , create_llm
5+ from vektori .models .anthropic import AnthropicLLM
6+ from vektori .models .factory import LLM_REGISTRY , create_embedder , create_llm
7+ from vektori .models .nvidia import DEFAULT_EMBEDDING_MODEL , NvidiaEmbedder , NvidiaLLM
8+ from vektori .models .ollama import OllamaEmbedder , OllamaLLM
9+ from vektori .models .openai import OpenAIEmbedder , OpenAILLM
610
711
812def test_create_openai_embedder ():
9- from vektori .models .openai import OpenAIEmbedder
10-
1113 embedder = create_embedder ("openai:text-embedding-3-small" )
1214 assert isinstance (embedder , OpenAIEmbedder )
1315 assert embedder .model == "text-embedding-3-small"
1416
1517
1618def test_create_openai_embedder_default_model ():
17- from vektori .models .openai import OpenAIEmbedder
18-
1919 embedder = create_embedder ("openai" )
2020 assert isinstance (embedder , OpenAIEmbedder )
2121
2222
2323def test_create_ollama_embedder ():
24- from vektori .models .ollama import OllamaEmbedder
25-
2624 embedder = create_embedder ("ollama:nomic-embed-text" )
2725 assert isinstance (embedder , OllamaEmbedder )
2826 assert embedder .model == "nomic-embed-text"
@@ -41,27 +39,51 @@ def test_unknown_embedding_provider_raises():
4139
4240
4341def test_create_openai_llm ():
44- from vektori .models .openai import OpenAILLM
45-
4642 llm = create_llm ("openai:gpt-4o-mini" )
4743 assert isinstance (llm , OpenAILLM )
4844 assert llm .model == "gpt-4o-mini"
4945
5046
5147def test_create_ollama_llm ():
52- from vektori .models .ollama import OllamaLLM
53-
5448 llm = create_llm ("ollama:llama3" )
5549 assert isinstance (llm , OllamaLLM )
5650
5751
5852def test_create_anthropic_llm ():
59- from vektori .models .anthropic import AnthropicLLM
60-
6153 llm = create_llm ("anthropic:claude-haiku-4-5-20251001" )
6254 assert isinstance (llm , AnthropicLLM )
6355
6456
6557def test_unknown_llm_provider_raises ():
6658 with pytest .raises (ValueError , match = "Unknown LLM provider" ):
6759 create_llm ("nonexistent:model" )
60+
61+
62+ def test_create_nvidia_embedder_default_model ():
63+ embedder = create_embedder ("nvidia" )
64+ assert isinstance (embedder , NvidiaEmbedder )
65+ assert embedder .model == DEFAULT_EMBEDDING_MODEL
66+
67+ def test_create_nvidia_embedder_custom_dimensions ():
68+ embedder = create_embedder ("nvidia:llama-nemotron-embed-1b-v2" , dimensions = 1024 )
69+ assert isinstance (embedder , NvidiaEmbedder )
70+ assert embedder .dimension == 1024 # Matryoshka support
71+
72+ def test_create_nvidia_llm ():
73+ llm = create_llm ("nvidia:llama-3.3-nemotron-super-49b-v1" )
74+ assert isinstance (llm , NvidiaLLM )
75+ assert llm .model == "nvidia/llama-3.3-nemotron-super-49b-v1"
76+
77+ def test_create_nvidia_llm_default_model ():
78+ llm = create_llm ("nvidia" )
79+ assert isinstance (llm , NvidiaLLM )
80+ assert "nvidia/llama-3.3-nemotron-super-49b-v1" == llm .model
81+
82+ def test_nvidia_llm_registered ():
83+ """Verify NVIDIA LLM is registered in factory."""
84+ assert "nvidia" in LLM_REGISTRY , "nvidia should be registered in LLM_REGISTRY"
85+ llm = create_llm ("nvidia" )
86+ assert llm is not None , "create_llm('nvidia') should return a valid instance"
87+ assert "Nvidia" in llm .__class__ .__name__ , (
88+ f"LLM class name should include 'Nvidia', got { llm .__class__ .__name__ } "
89+ )
0 commit comments