fix(mistralai): handle HTTP errors in async embed documents (#33187)

The async embed function does not properly handle HTTP errors.

For instance with large batches, Mistral AI returns `Too many inputs in
request, split into more batches.` in a 400 error.

This leads to a KeyError in `response.json()["data"]` l.288

This PR fixes the issue by:
- calling `response.raise_for_status()` before returning
- adding a retry similarly to what is done in the synchronous
counterpart `embed_documents`

I also added an integration test, but willing to move it to unit tests
if more relevant.
This commit is contained in:
noeliecherrier
2025-10-01 16:57:47 +02:00
committed by GitHub
parent 7d78ed9b53
commit 08bb74f148
2 changed files with 41 additions and 11 deletions

View File

@@ -1,5 +1,11 @@
"""Test MistralAI Embedding."""
from unittest.mock import patch
import httpx
import pytest
import tenacity
from langchain_mistralai import MistralAIEmbeddings
@@ -29,6 +35,21 @@ async def test_mistralai_embedding_documents_async() -> None:
assert len(output[0]) == 1024
async def test_mistralai_embedding_documents_http_error_async() -> None:
"""Test MistralAI embeddings for documents."""
documents = ["foo bar", "test document"]
embedding = MistralAIEmbeddings(max_retries=0)
mock_response = httpx.Response(
status_code=400,
request=httpx.Request("POST", url=embedding.async_client.base_url),
)
with (
patch.object(embedding.async_client, "post", return_value=mock_response),
pytest.raises(tenacity.RetryError),
):
await embedding.aembed_documents(documents)
async def test_mistralai_embedding_query_async() -> None:
"""Test MistralAI embeddings for query."""
document = "foo bar"