From 09ae71aa915f25b31b10cb3957d53c42254e467a Mon Sep 17 00:00:00 2001
From: kingbri <bdashore3@proton.me>
Date: Mon, 25 Dec 2023 11:25:38 -0500
Subject: [PATCH] OAI: Add finish to completions

OAI spec requires [DONE] to be sent over SSE to signal that a generation
is completed.

Signed-off-by: kingbri <bdashore3@proton.me>
---
 main.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/main.py b/main.py
index 5095bba..8175c72 100644
--- a/main.py
+++ b/main.py
@@ -366,6 +366,9 @@ async def generate_completion(request: Request, data: CompletionRequest):
                     )
 
                     yield get_sse_packet(response.model_dump_json())
+
+                # Yield a finish response on successful generation
+                yield get_sse_packet("[DONE]")
             except CancelledError:
                 logger.error("Completion request cancelled by user.")
             except Exception as exc: