Files
NexQuant/test/oai/test_embedding_and_similarity.py
T
amstrongzyf 9890bb4b00 docs: update explanation for separate config use in litellm (#958)
* docs: update explanation for separate config use in litellm

* docs: update default backend to `rdagent.oai.backend.LiteLLMAPIBackend`

* docs: update .rst format

* Update installation_and_configuration.rst
2025-06-13 18:32:39 +08:00

33 lines
963 B
Python

import unittest
from rdagent.oai.llm_utils import (
APIBackend,
calculate_embedding_distance_between_str_list,
)
class TestEmbedding(unittest.TestCase):
def test_embedding(self) -> None:
emb = APIBackend().create_embedding("hello")
assert emb is not None
assert isinstance(emb, list)
assert len(emb) > 0
def test_embedding_list(self) -> None:
emb = APIBackend().create_embedding(["hello", "hi"])
assert emb is not None
assert isinstance(emb, list)
assert len(emb) == 2
def test_embedding_similarity(self) -> None:
similarity = calculate_embedding_distance_between_str_list(["Hello"], ["Hi"])[0][0]
assert similarity is not None
assert isinstance(similarity, float)
min_similarity_threshold = 0.8
print(f"similarity: {similarity}")
assert similarity >= min_similarity_threshold
if __name__ == "__main__":
unittest.main()