diff --git a/libs/partners/mistralai/tests/integration_tests/test_embeddings.py b/libs/partners/mistralai/tests/integration_tests/test_embeddings.py index b179f64d80..166bf5c5c1 100644 --- a/libs/partners/mistralai/tests/integration_tests/test_embeddings.py +++ b/libs/partners/mistralai/tests/integration_tests/test_embeddings.py @@ -17,3 +17,37 @@ def test_mistralai_embedding_query() -> None: embedding = MistralAIEmbeddings() output = embedding.embed_query(document) 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