mistralai[patch]: async integration tests (#16214)

This commit is contained in:
Erick Friis 2024-01-18 09:45:44 -08:00 committed by GitHub
parent ed118950fe
commit 65b231d40b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -17,3 +17,37 @@ def test_mistralai_embedding_query() -> None:
embedding = MistralAIEmbeddings() embedding = MistralAIEmbeddings()
output = embedding.embed_query(document) output = embedding.embed_query(document)
assert len(output) == 1024 assert len(output) == 1024
async def test_mistralai_embedding_documents_async() -> None:
"""Test MistralAI embeddings for documents."""
documents = ["foo bar", "test document"]
embedding = MistralAIEmbeddings()
output = await embedding.aembed_documents(documents)
assert len(output) == 2
assert len(output[0]) == 1024
async def test_mistralai_embedding_query_async() -> None:
"""Test MistralAI embeddings for query."""
document = "foo bar"
embedding = MistralAIEmbeddings()
output = await embedding.aembed_query(document)
assert len(output) == 1024
def test_mistralai_embedding_documents_long() -> None:
"""Test MistralAI embeddings for documents."""
documents = ["foo bar " * 1000, "test document " * 1000]
embedding = MistralAIEmbeddings()
output = embedding.embed_documents(documents)
assert len(output) == 2
assert len(output[0]) == 1024
def test_mistralai_embed_query_character() -> None:
"""Test MistralAI embeddings for query."""
document = "😳"
embedding = MistralAIEmbeddings()
output = embedding.embed_query(document)
assert len(output) == 1024