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

@@ -267,20 +267,29 @@ class MistralAIEmbeddings(BaseModel, Embeddings):
Returns:
List of embeddings, one for each text.
"""
try:
@retry(
retry=retry_if_exception_type(
(httpx.TimeoutException, httpx.HTTPStatusError)
),
wait=wait_fixed(self.wait_time),
stop=stop_after_attempt(self.max_retries),
)
async def _aembed_batch(batch: list[str]) -> Response:
response = await self.async_client.post(
url="/embeddings",
json={
"model": self.model,
"input": batch,
},
)
response.raise_for_status()
return response
batch_responses = await asyncio.gather(
*[
self.async_client.post(
url="/embeddings",
json={
"model": self.model,
"input": batch,
},
)
for batch in self._get_batches(texts)
]
*[_aembed_batch(batch) for batch in self._get_batches(texts)]
)
return [
list(map(float, embedding_obj["embedding"]))