inference errors (#6759)

* fix error handling in inference

* fmt

* Update src/common/inference/service.rs

---------

Co-authored-by: Tim Visée <tim+github@visee.me>

src/common/inference/infer_processing.rs

@@ -45,15 +45,9 @@ impl BatchAccumInferred {
             .map(InferenceInput::from)
             .collect();
 
-        let InferenceResponse {
-            embeddings,
-            usage,
-        } = service
+        let InferenceResponse { embeddings, usage } = service
             .infer(inference_inputs, inference_type, inference_token)
-            .await
-            .map_err(|e| StorageError::service_error(
-                format!("Inference request failed. Check if inference service is running and properly configured: {e}")
-            ))?;
+            .await?;
 
         if embeddings.is_empty() {
             return Err(StorageError::service_error(

src/common/inference/service.rs

@@ -63,6 +63,11 @@ pub enum InferenceData {
     Object(InferenceObject),
 }
 
+#[derive(Debug, Deserialize)]
+struct InferenceError {
+    pub error: String,
+}
+
 impl InferenceData {
     pub(crate) fn type_name(&self) -> &'static str {
         match self {
@@ -193,24 +198,30 @@ impl InferenceService {
             )
         })?;
 
-        let response = self
-            .client
-            .post(url)
-            .json(&request)
-            .send()
-            .await
-            .map_err(|e| {
-                let error_body = e.to_string();
-                StorageError::service_error(format!(
-                    "Failed to send inference request to {url}: {e}, error details: {error_body}",
-                ))
-            })?;
-
-        let status = response.status();
-        let response_body = response.text().await.map_err(|e| {
-            StorageError::service_error(format!("Failed to read inference response body: {e}",))
-        })?;
+        let response = self.client.post(url).json(&request).send().await;
 
+        let (response_body, status) = match response {
+            Ok(response) => {
+                let status = response.status();
+                match response.text().await {
+                    Ok(body) => (body, status),
+                    Err(err) => {
+                        return Err(StorageError::service_error(format!(
+                            "Failed to read inference response body: {err}"
+                        )));
+                    }
+                }
+            }
+            Err(error) => {
+                if let Some(status) = error.status() {
+                    (error.to_string(), status)
+                } else {
+                    return Err(StorageError::service_error(format!(
+                        "Failed to send inference request: {error}"
+                    )));
+                }
+            }
+        };
         Self::handle_inference_response(status, &response_body)
     }
 
@@ -228,20 +239,19 @@ impl InferenceService {
                     })
             }
             reqwest::StatusCode::BAD_REQUEST => {
-                let error_json: Value = serde_json::from_str(response_body).map_err(|e| {
-                    StorageError::service_error(format!(
-                        "Failed to parse error response: {e}. Raw response: {response_body}",
-                    ))
-                })?;
-
-                if let Some(error_message) = error_json["error"].as_str() {
-                    Err(StorageError::bad_request(format!(
-                        "Inference request validation failed: {error_message}",
-                    )))
-                } else {
-                    Err(StorageError::bad_request(format!(
-                        "Invalid inference request: {response_body}",
-                    )))
+                // Try to extract error description from the response body, if it is a valid JSON
+                let parsed_body: Result<InferenceError, _> = serde_json::from_str(response_body);
+                match parsed_body {
+                    Ok(InferenceError { error}) => {
+                        Err(StorageError::bad_request(format!(
+                            "Inference request validation failed: {error}",
+                        )))
+                    }
+                    Err(_) => {
+                        Err(StorageError::bad_request(format!(
+                            "Invalid inference request: {response_body}",
+                        )))
+                    }
                 }
             }
             status @ (reqwest::StatusCode::UNAUTHORIZED | reqwest::StatusCode::FORBIDDEN) => {
@@ -254,9 +264,21 @@ impl InferenceService {
             | reqwest::StatusCode::GATEWAY_TIMEOUT) => Err(StorageError::service_error(format!(
                 "Inference service error ({status}): {response_body}",
             ))),
-            _ => Err(StorageError::service_error(format!(
-                "Unexpected inference service response ({status}): {response_body}"
-            ))),
+            _ => {
+                if status.is_server_error() {
+                    Err(StorageError::service_error(format!(
+                        "Inference service error ({status}): {response_body}",
+                    )))
+                } else if status.is_client_error() {
+                    Err(StorageError::bad_request(format!(
+                        "Inference can't process request ({status}): {response_body}",
+                    )))
+                } else {
+                    Err(StorageError::service_error(format!(
+                        "Unexpected inference error ({status}): {response_body}",
+                    )))
+                }
+            },
         }
     }
 }