From a5e9da6bf054c53d088a5d06f0914fd32686c414 Mon Sep 17 00:00:00 2001
From: krishnaduttPanchagnula <krishnadutt123@gmail.com>
Date: Tue, 30 May 2023 01:37:13 +0530
Subject: [PATCH] added Enable token stream on /completions

---
 api/openai.go | 31 +++++++++++++++++++++++++++++++
 1 file changed, 31 insertions(+)

diff --git a/api/openai.go b/api/openai.go
index b97b4e56..3db80d8f 100644
--- a/api/openai.go
+++ b/api/openai.go
@@ -192,6 +192,37 @@ func completionEndpoint(cm *ConfigMerger, o *Option) func(c *fiber.Ctx) error {
 
 		jsonResult, _ := json.Marshal(resp)
 		log.Debug().Msgf("Response: %s", jsonResult)
+		if input.Stream {
+			responses := make(chan OpenAIResponse)
+
+			go func() {
+				defer close(responses)
+				for _, r := range result {
+					responses <- OpenAIResponse{
+						Model:   input.Model,
+						Choices: []Choice{r},
+						Object:  "text_completion",
+					}
+				}
+			}()
+
+			c.Context().SetContentType("text/event-stream")
+			c.Set("Cache-Control", "no-cache")
+			c.Set("Connection", "keep-alive")
+			c.Set("Transfer-Encoding", "chunked")
+
+			c.Context().SetBodyStreamWriter(fasthttp.StreamWriter(func(w *bufio.Writer) {
+				for ev := range responses {
+					var buf bytes.Buffer
+					enc := json.NewEncoder(&buf)
+					enc.Encode(ev)
+
+					log.Debug().Msgf("Sending chunk: %s", buf.String())
+					fmt.Fprintf(w, "data: %v\n", buf.String())
+					w.Flush()
+				}
+			}))
+		}
 
 		// Return the prediction in the response body
 		return c.JSON(resp)