From 5d8726d92d215c72a7938166ec6eb3145a7fe919 Mon Sep 17 00:00:00 2001
From: Louis Dureuil <louis@meilisearch.com>
Date: Thu, 14 Nov 2024 15:28:44 +0100
Subject: [PATCH] Retry in case where the JSON deserialization fails

---
 milli/src/vector/error.rs | 2 +-
 milli/src/vector/rest.rs  | 9 +++++----
 2 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/milli/src/vector/error.rs b/milli/src/vector/error.rs
index 3c8cb4b06..45ed5f3a5 100644
--- a/milli/src/vector/error.rs
+++ b/milli/src/vector/error.rs
@@ -58,7 +58,7 @@ pub enum EmbedErrorKind {
     ManualEmbed(String),
     #[error("model not found. Meilisearch will not automatically download models from the Ollama library, please pull the model manually{}", option_info(.0.as_deref(), "server replied with "))]
     OllamaModelNotFoundError(Option<String>),
-    #[error("error deserialization the response body as JSON:\n  - {0}")]
+    #[error("error deserializing the response body as JSON:\n  - {0}")]
     RestResponseDeserialization(std::io::Error),
     #[error("expected a response containing {0} embeddings, got only {1}")]
     RestResponseEmbeddingCount(usize, usize),
diff --git a/milli/src/vector/rest.rs b/milli/src/vector/rest.rs
index 05578ef1f..02ade88e8 100644
--- a/milli/src/vector/rest.rs
+++ b/milli/src/vector/rest.rs
@@ -252,12 +252,13 @@ where
 
     for attempt in 0..10 {
         let response = request.clone().send_json(&body);
-        let result = check_response(response, data.configuration_source);
+        let result = check_response(response, data.configuration_source).and_then(|response| {
+            response_to_embedding(response, data, expected_count, expected_dimension)
+                .map_err(Retry::retry_later)
+        });
 
         let retry_duration = match result {
-            Ok(response) => {
-                return response_to_embedding(response, data, expected_count, expected_dimension)
-            }
+            Ok(response) => return Ok(response),
             Err(retry) => {
                 tracing::warn!("Failed: {}", retry.error);
                 if let Some(deadline) = deadline {