fix: append DONE message to chat stream (#2221)

* fix: append DONE message to chat stream * fix: update completions endpoint
huggingface · Jul 11, 2024 · d789de3 · d789de3
1 parent cb150eb
commit d789de3
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/router/src/server.rs b/router/src/server.rs
@@ -812,6 +812,10 @@ async fn completions(
             }
         };
 
+        let stream = stream.chain(futures::stream::once(async {
+            Ok(Event::default().data("[DONE]"))
+        }));
+
         let sse = Sse::new(stream).keep_alive(KeepAlive::default());
         Ok((headers, sse).into_response())
     } else {
@@ -1171,6 +1175,11 @@ async fn chat_completions(
             span,
         )
         .await;
+
+        let response_stream = response_stream.chain(futures::stream::once(async {
+            Ok(Event::default().data("[DONE]"))
+        }));
+
         let sse = Sse::new(response_stream).keep_alive(KeepAlive::default());
         Ok((headers, sse).into_response())
     } else {