From 0c2e53a83f63e6c48ae7ac665f4d56e92fb1cbad Mon Sep 17 00:00:00 2001
From: Georges Haidar <ghaidar0@gmail.com>
Date: Mon, 15 Jul 2024 22:25:44 +0100
Subject: [PATCH] feat: add sse chat endpoint with arbitrary chunks (#12)

This change adds a new endpoint which streams server-sent events in
arbitrarily sized chunks to exercise the buffering and scanning
functionality in SDKs.
---
 cmd/server/main.go               |  1 +
 internal/eventstreams/service.go | 24 ++++++++++++++++++++++++
 2 files changed, 25 insertions(+)

diff --git a/cmd/server/main.go b/cmd/server/main.go
index 81488a3..e3d883e 100644
--- a/cmd/server/main.go
+++ b/cmd/server/main.go
@@ -51,6 +51,7 @@ func main() {
 	r.HandleFunc("/eventstreams/multiline", eventstreams.HandleEventStreamMultiLine).Methods(http.MethodPost)
 	r.HandleFunc("/eventstreams/rich", eventstreams.HandleEventStreamRich).Methods(http.MethodPost)
 	r.HandleFunc("/eventstreams/chat", eventstreams.HandleEventStreamChat).Methods(http.MethodPost)
+	r.HandleFunc("/eventstreams/chat-chunked", eventstreams.HandleEventStreamChat).Methods(http.MethodPost)
 	r.HandleFunc("/eventstreams/differentdataschemas", eventstreams.HandleEventStreamDifferentDataSchemas).Methods(http.MethodPost)
 	r.HandleFunc("/clientcredentials/token", clientcredentials.HandleTokenRequest).Methods(http.MethodPost)
 	r.HandleFunc("/clientcredentials/authenticatedrequest", clientcredentials.HandleAuthenticatedRequest).Methods(http.MethodPost)
diff --git a/internal/eventstreams/service.go b/internal/eventstreams/service.go
index d911687..d760135 100644
--- a/internal/eventstreams/service.go
+++ b/internal/eventstreams/service.go
@@ -6,6 +6,18 @@ import (
 	"time"
 )
 
+func pushChunks(rw http.ResponseWriter, chunks []string) {
+	for _, chunk := range chunks {
+		fmt.Fprintln(rw, chunk)
+
+		if f, ok := rw.(http.Flusher); ok {
+			f.Flush()
+		}
+
+		time.Sleep(100 * time.Millisecond)
+	}
+}
+
 func pushEvents(rw http.ResponseWriter, events [][]string) {
 	for _, event := range events {
 		for _, line := range event {
@@ -127,6 +139,18 @@ func HandleEventStreamChat(rw http.ResponseWriter, _ *http.Request) {
 	})
 }
 
+func HandleEventStreamChatChunked(rw http.ResponseWriter, _ *http.Request) {
+	rw.Header().Add("Content-Type", "text/event-stream")
+
+	pushChunks(rw, []string{
+		"data: {\"content\": ",
+		"\"Hello\"}\n\ndata: {\"content\": \" \"}",
+		"data: {\"content\": \"world\"}",
+		"data: {\"content\": \"!\"}\n\ndata: [DONE]\n",
+		"\ndata: {\"content\": \"Post sentinel data\"}\n\n",
+	})
+}
+
 func HandleEventStreamDifferentDataSchemas(rw http.ResponseWriter, _ *http.Request) {
 	rw.Header().Add("Content-Type", "text/event-stream")