codekingpro/portable-devtools

1import os
2
3import pytest
4
5from chromadb.utils.embedding_functions.openai_embedding_function import (
6    OpenAIEmbeddingFunction,
7)
8
9
10def test_with_embedding_dimensions() -> None:
11    if os.environ.get("OPENAI_API_KEY") is None:
12        pytest.skip("OPENAI_API_KEY not set")
13    ef = OpenAIEmbeddingFunction(
14        api_key=os.environ["OPENAI_API_KEY"],
15        model_name="text-embedding-3-small",
16        dimensions=64,
17    )
18    embeddings = ef(["hello world"])
19    assert embeddings is not None
20    assert len(embeddings) == 1
21    assert len(embeddings[0]) == 64
22
23
24def test_with_embedding_dimensions_not_working_with_old_model() -> None:
25    if os.environ.get("OPENAI_API_KEY") is None:
26        pytest.skip("OPENAI_API_KEY not set")
27    ef = OpenAIEmbeddingFunction(api_key=os.environ["OPENAI_API_KEY"], dimensions=64)
28    with pytest.raises(
29        Exception, match="This model does not support specifying dimensions"
30    ):
31        ef(["hello world"])
32
33
34def test_with_incorrect_api_key() -> None:
35    pytest.importorskip("openai", reason="openai not installed")
36    ef = OpenAIEmbeddingFunction(api_key="incorrect_api_key", dimensions=64)
37    with pytest.raises(Exception, match="Incorrect API key provided"):
38        ef(["hello world"])
39