From d789de329a087301d651ee943e0d76e0dbf5ced5 Mon Sep 17 00:00:00 2001
From: drbh <david.richard.holtz@gmail.com>
Date: Thu, 11 Jul 2024 10:42:58 -0400
Subject: [PATCH] fix: append DONE message to chat stream (#2221)

* fix: append DONE message to chat stream

* fix: update completions endpoint
---
 router/src/server.rs | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/router/src/server.rs b/router/src/server.rs
index 4e5af99c..d3a280ca 100644
--- 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 {