feat: handle error status (#338)

* handle non-success status * fmt
atoma-network · Jan 14, 2025 · 7df38e6 · 7df38e6
1 parent d385f60
commit 7df38e6
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 0 deletions.
diff --git a/atoma-service/src/handlers/chat_completions.rs b/atoma-service/src/handlers/chat_completions.rs
@@ -1207,6 +1207,17 @@ pub(crate) mod utils {
                 endpoint: endpoint.to_string(),
             }
         })?;
+
+        if !response.status().is_success() {
+            return Err(AtomaServiceError::InternalError {
+                message: format!(
+                    "Inference service returned non-success status code: {}",
+                    response.status()
+                ),
+                endpoint: endpoint.to_string(),
+            });
+        }
+
         response.json::<Value>().await.map_err(|e| {
             AtomaServiceError::InternalError {
                 message: format!(

diff --git a/atoma-service/src/handlers/embeddings.rs b/atoma-service/src/handlers/embeddings.rs
@@ -307,6 +307,17 @@ async fn handle_embeddings_response(
             message: format!("Error sending request to embeddings service: {}", e),
             endpoint: endpoint.to_string(),
         })?;
+
+    if !response.status().is_success() {
+        return Err(AtomaServiceError::InternalError {
+            message: format!(
+                "Inference service returned non-success status code: {}",
+                response.status()
+            ),
+            endpoint: endpoint.to_string(),
+        });
+    }
+
     let mut response_body =
         response
             .json::<Value>()

diff --git a/atoma-service/src/handlers/image_generations.rs b/atoma-service/src/handlers/image_generations.rs
@@ -293,6 +293,17 @@ async fn handle_image_generations_response(
             message: format!("Error sending request to image generations service: {}", e),
             endpoint: endpoint.to_string(),
         })?;
+
+    if !response.status().is_success() {
+        return Err(AtomaServiceError::InternalError {
+            message: format!(
+                "Inference service returned non-success status code: {}",
+                response.status()
+            ),
+            endpoint: endpoint.to_string(),
+        });
+    }
+
     let mut response_body =
         response
             .json::<Value>()