From 825f39d12b1586b48db6e2c2b8f8d0500abb0ec3 Mon Sep 17 00:00:00 2001
From: drbh <david.richard.holtz@gmail.com>
Date: Tue, 22 Oct 2024 11:59:14 -0400
Subject: [PATCH] fix: improve stream api error format and add status code

---
 router/src/server.rs | 23 +++++++++++------------
 1 file changed, 11 insertions(+), 12 deletions(-)

diff --git a/router/src/server.rs b/router/src/server.rs
index 76e1fce1..da63e9b1 100644
--- a/router/src/server.rs
+++ b/router/src/server.rs
@@ -849,10 +849,10 @@ async fn completions(
                                     yield Ok(event);
                                 }
                                 Err(err) => {
-                                    let error_event: ErrorEvent = err.into();
-                                    let event = Event::default().json_data(error_event).unwrap_or_else(|e| {
-                                        InferError::StreamSerializationError(e.to_string()).into()
-                                    });
+                                    let event = Event::default()
+                                        .json_data(ErrorEvent::into_api_error(err, 422))
+                                        .unwrap_or_else(|e| InferError::StreamSerializationError(e.to_string()).into());
+                                    println!("{:?}", event);
                                     yield Ok::<Event, Infallible>(event);
                                     break
                                 }
@@ -1358,10 +1358,9 @@ async fn chat_completions(
                         }
                     }
                     Err(err) => {
-                        let error_event: ErrorEvent = err.into();
-                        let event = Event::default().json_data(error_event).unwrap_or_else(|e| {
-                            InferError::StreamSerializationError(e.to_string()).into()
-                        });
+                        let event = Event::default()
+                            .json_data(ErrorEvent::into_api_error(err, 422))
+                            .unwrap_or_else(|e| InferError::StreamSerializationError(e.to_string()).into());
                         yield Ok::<Event, Infallible>(event);
                         break;
                     }
@@ -2525,7 +2524,7 @@ impl From<InferError> for Event {
 #[derive(serde::Serialize)]
 pub struct APIError {
     message: String,
-    http_status_code: i32,
+    http_status_code: usize,
 }
 
 #[derive(serde::Serialize)]
@@ -2533,12 +2532,12 @@ pub struct ErrorEvent {
     error: APIError,
 }
 
-impl From<InferError> for ErrorEvent {
-    fn from(err: InferError) -> Self {
+impl ErrorEvent {
+    fn into_api_error(err: InferError, http_status_code: usize) -> Self {
         ErrorEvent {
             error: APIError {
                 message: err.to_string(),
-                http_status_code: 500,
+                http_status_code,
             },
         }
     }