codekingpro/portable-devtools

1import pytest
2
3from chromadb.utils.embedding_functions.ollama_embedding_function import (
4    OllamaEmbeddingFunction,
5)
6
7
8def test_ollama_default_model() -> None:
9    pytest.importorskip("ollama", reason="ollama not installed")
10    ef = OllamaEmbeddingFunction()
11    embeddings = ef(["Here is an article about llamas...", "this is another article"])
12    assert embeddings is not None
13    assert len(embeddings) == 2
14    assert all(len(e) == 384 for e in embeddings)
15
16
17def test_ollama_unknown_model() -> None:
18    pytest.importorskip("ollama", reason="ollama not installed")
19    model_name = "unknown-model"
20    ef = OllamaEmbeddingFunction(model_name=model_name)
21    with pytest.raises(Exception) as e:
22        ef(["Here is an article about llamas...", "this is another article"])
23    assert f'model "{model_name}" not found' in str(e.value)
24
25
26def test_ollama_backward_compat() -> None:
27    pytest.importorskip("ollama", reason="ollama not installed")
28    ef = OllamaEmbeddingFunction(url="http://localhost:11434/api/embeddings")
29    embeddings = ef(["Here is an article about llamas...", "this is another article"])
30    assert embeddings is not None
31
32
33def test_wrong_url() -> None:
34    pytest.importorskip("ollama", reason="ollama not installed")
35    ef = OllamaEmbeddingFunction(url="http://localhost:11434/this_is_wrong")
36    with pytest.raises(Exception) as e:
37        ef(["Here is an article about llamas...", "this is another article"])
38    assert "404" in str(e.value)
39
40
41def test_ollama_ask_user_to_install() -> None:
42    try:
43        from ollama import Client  # noqa: F401
44    except ImportError:
45        pass
46    else:
47        pytest.skip("ollama python package is installed")
48    with pytest.raises(ValueError) as e:
49        OllamaEmbeddingFunction()
50    assert "The ollama python package is not installed" in str(e.value)
51