File size: 1,866 Bytes
60e3a80 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 |
from chromadb.utils import embedding_functions
from chromadb.api.types import EmbeddingFunction
def test_get_builtins_holds() -> None:
"""
Ensure that `get_builtins` is consistent after the ef migration.
This test is intended to be temporary until the ef migration is complete as
these expected builtins are likely to grow as long as users add new
embedding functions.
REMOVE ME ON THE NEXT EF ADDITION
"""
expected_builtins = {
"AmazonBedrockEmbeddingFunction",
"CohereEmbeddingFunction",
"GoogleGenerativeAiEmbeddingFunction",
"GooglePalmEmbeddingFunction",
"GoogleVertexEmbeddingFunction",
"HuggingFaceEmbeddingFunction",
"HuggingFaceEmbeddingServer",
"InstructorEmbeddingFunction",
"JinaEmbeddingFunction",
"ONNXMiniLM_L6_V2",
"OllamaEmbeddingFunction",
"OpenAIEmbeddingFunction",
"OpenCLIPEmbeddingFunction",
"RoboflowEmbeddingFunction",
"SentenceTransformerEmbeddingFunction",
"Text2VecEmbeddingFunction",
"ChromaLangchainEmbeddingFunction",
}
assert expected_builtins == embedding_functions.get_builtins()
def test_default_ef_exists() -> None:
assert hasattr(embedding_functions, "DefaultEmbeddingFunction")
default_ef = embedding_functions.DefaultEmbeddingFunction()
assert default_ef is not None
assert isinstance(default_ef, EmbeddingFunction)
def test_ef_imports() -> None:
for ef in embedding_functions.get_builtins():
# Langchain embedding function is a special snowflake
if ef == "ChromaLangchainEmbeddingFunction":
continue
assert hasattr(embedding_functions, ef)
assert isinstance(getattr(embedding_functions, ef), type)
assert issubclass(getattr(embedding_functions, ef), EmbeddingFunction)
|