SSE per user

Chat.go

@@ -126,7 +126,16 @@ func generateChatHTML(c *fiber.Ctx) string {
 		panic(err)
 	}
 
-	htmlString := "<div class='columns is-centered' id='chat-container'><div class='column is-12-mobile is-8-tablet is-6-desktop' id='chat-messages'>"
+	var user User
+	err = edgeGlobalClient.WithGlobals(map[string]interface{}{"ext::auth::client_token": c.Cookies("jade-edgedb-auth-token")}).QuerySingle(edgeCtx, `
+	SELECT global currentUser { id } LIMIT 1
+	`, &user)
+	if err != nil {
+		fmt.Println("Error getting user")
+		panic(err)
+	}
+
+	htmlString := "<div class='columns is-centered' id='chat-container' sse-connect='/sse?userID=" + user.ID.String() + "'><div class='column is-12-mobile is-8-tablet is-6-desktop' id='chat-messages'>"
 
 	var templateMessages []TemplateMessage
 
@@ -177,13 +186,6 @@ func generateChatHTML(c *fiber.Ctx) string {
 		}
 	}
 
-	// out, err := messagesPlaceholderTmpl.Execute(pongo2.Context{})
-	// if err != nil {
-	// 	fmt.Println("Error executing message user placeholder template")
-	//	panic(err)
-	// }
-	// htmlString += out
-
 	htmlString += "</div></div>"
 
 	// Render the HTML template with the messages

Request.go

@@ -230,18 +230,20 @@ func GenerateMultipleMessagesHandler(c *fiber.Ctx) error {
 					outIcon := `<img src="` + selectedLLMs[idx].Model.Company.Icon + `" alt="User Image" id="selectedIcon-` + fmt.Sprintf("%d", message.Area.Position) + `">`
 
 					go func() {
-						// I do a ping because of sse size limit
+						// I do a ping because of sse size limit. Do see if it's possible to do without it TODO
-						fmt.Println("Sending event: ", "swapContent-"+fmt.Sprintf("%d", message.Area.Position)+"-"+user.ID.String())
 						sendEvent(
-							"swapContent-"+fmt.Sprintf("%d", message.Area.Position)+"-"+user.ID.String(),
+							user.ID.String(),
+							"swapContent-"+fmt.Sprintf("%d", message.Area.Position),
 							`<hx hx-get="/messageContent?id=`+message.ID.String()+`" hx-trigger="load" hx-swap="outerHTML"></hx>`,
 						)
 						sendEvent(
-							"swapSelectionBtn-"+selectedLLMs[idx].ID.String()+"-"+user.ID.String(),
+							user.ID.String(),
+							"swapSelectionBtn-"+selectedLLMs[idx].ID.String(),
 							outBtn,
 						)
 						sendEvent(
-							"swapIcon-"+fmt.Sprintf("%d", message.Area.Position)+"-"+user.ID.String(),
+							user.ID.String(),
+							"swapIcon-"+fmt.Sprintf("%d", message.Area.Position),
 							outIcon,
 						)
 					}()
@@ -261,7 +263,8 @@ func GenerateMultipleMessagesHandler(c *fiber.Ctx) error {
 					// Send Content event
 					go func() {
 						sendEvent(
-							"swapSelectionBtn-"+selectedLLMs[idx].ID.String()+"-"+user.ID.String(),
+							user.ID.String(),
+							"swapSelectionBtn-"+selectedLLMs[idx].ID.String(),
 							outBtn,
 						)
 					}()

RequestGoogle.go

@@ -154,12 +154,12 @@ func RequestGoogle(c *fiber.Ctx, model string, messages []Message, temperature f
 		if message.Role == "user" {
 			googleMessages = append(googleMessages, GoogleRequestMessage{
 				Role:  "user",
-				Parts: []GooglePart{GooglePart{Text: message.Content}},
+				Parts: []GooglePart{{Text: message.Content}}, // Changed something here, to test
 			})
 		} else {
 			googleMessages = append(googleMessages, GoogleRequestMessage{
 				Role:  "model",
-				Parts: []GooglePart{GooglePart{Text: message.Content}},
+				Parts: []GooglePart{{Text: message.Content}},
 			})
 		}
 	}

main.go

@@ -27,9 +27,9 @@ var (
 	welcomeChatTmpl         *pongo2.Template
 	chatInputTmpl           *pongo2.Template
 	explainLLMconvChatTmpl  *pongo2.Template
-	messagesPlaceholderTmpl *pongo2.Template
-	clients                 = make(map[chan SSE]bool)
 	mu                      sync.Mutex
+	app                     *fiber.App
+	userSSEChannels         = make(map[string]chan SSE)
 )
 
 // SSE event structure
@@ -39,13 +39,16 @@ type SSE struct {
 }
 
 // Function to send events to all clients
-func sendEvent(event, data string) {
+func sendEvent(userID string, event string, data string) {
 	mu.Lock()
 	defer mu.Unlock()
 
-	for client := range clients {
+	userEvents, ok := userSSEChannels[userID]
-		client <- SSE{Event: event, Data: data}
+	if !ok {
+		return
 	}
+
+	userEvents <- SSE{Event: event, Data: data}
 }
 
 func main() {
@@ -63,13 +66,12 @@ func main() {
 	welcomeChatTmpl = pongo2.Must(pongo2.FromFile("views/partials/welcome-chat.html"))
 	chatInputTmpl = pongo2.Must(pongo2.FromFile("views/partials/chat-input.html"))
 	explainLLMconvChatTmpl = pongo2.Must(pongo2.FromFile("views/partials/explain-llm-conv-chat.html"))
-	messagesPlaceholderTmpl = pongo2.Must(pongo2.FromFile("views/partials/messages-placeholder.html"))
 
 	// Import HTML using django engine/template
 	engine := django.New("./views", ".html")
 
 	// Create new Fiber instance
-	app := fiber.New(fiber.Config{
+	app = fiber.New(fiber.Config{
 		Views:   engine,
 		AppName: "JADE",
 	})
@@ -133,39 +135,7 @@ func main() {
 		return c.SendString("")
 	})
 
-	app.Get("/sse", func(c *fiber.Ctx) error {
+	app.Get("/sse", handleSSE)
-		c.Set("Content-Type", "text/event-stream")
-		c.Set("Cache-Control", "no-cache")
-		c.Set("Connection", "keep-alive")
-
-		events := make(chan SSE, 500)
-		mu.Lock()
-		clients[events] = true
-		mu.Unlock()
-
-		// Create a context copy to use in the goroutine
-		ctx := c.Context()
-
-		go func() {
-			<-ctx.Done()
-			mu.Lock()
-			delete(clients, events)
-			mu.Unlock()
-			close(events)
-		}()
-
-		c.Context().SetBodyStreamWriter(func(w *bufio.Writer) {
-			for event := range events {
-				if _, err := fmt.Fprintf(w, "event: %s\ndata: %s\n\n", event.Event, event.Data); err != nil {
-					fmt.Println(err)
-					return
-				}
-				w.Flush()
-			}
-		})
-
-		return nil
-	})
 
 	// Start server
 	if err := app.Listen(":8080"); err != nil {
@@ -173,6 +143,45 @@ func main() {
 	}
 }
 
+func handleSSE(c *fiber.Ctx) error {
+	userID := c.Query("userID") // Get userID from query parameter
+	if userID == "" {
+		return c.Status(fiber.StatusBadRequest).SendString("Missing userID")
+	}
+
+	events := make(chan SSE, 500)
+	mu.Lock()
+	userSSEChannels[userID] = events
+	mu.Unlock()
+
+	// Create a context copy to use in the goroutine
+	ctx := c.Context()
+
+	go func() {
+		<-ctx.Done()
+		mu.Lock()
+		delete(userSSEChannels, userID)
+		mu.Unlock()
+		close(events)
+	}()
+
+	c.Set("Content-Type", "text/event-stream")
+	c.Set("Cache-Control", "no-cache")
+	c.Set("Connection", "keep-alive")
+
+	c.Context().SetBodyStreamWriter(func(w *bufio.Writer) {
+		for event := range events {
+			if _, err := fmt.Fprintf(w, "event: %s\ndata: %s\n\n", event.Event, event.Data); err != nil {
+				fmt.Println(err)
+				return
+			}
+			w.Flush()
+		}
+	})
+
+	return nil
+}
+
 func addKeys(c *fiber.Ctx) error {
 	keys := map[string]string{
 		"openai":     c.FormValue("openai_key"),

views/chat.html

@@ -1,5 +1,4 @@
-<div class="chat-container mt-5" style="padding-bottom: 155px;" hx-indicator="#textarea-control">
+<div class="chat-container mt-5" style="padding-bottom: 155px;" hx-indicator="#textarea-control" hx-ext="sse">
     <hx hx-get="/loadChat" hx-trigger="load once" hx-swap="outerHTML"></hx>
-
     <hx hx-get="/loadChatInput" hx-trigger="load once" hx-swap="outerHTML"></hx>
 </div>

views/partials/message-bot.html

@@ -6,7 +6,7 @@
             {% if IsPlaceholder %}
 
             <figure class="image is-48x48 message-icon" style="flex-shrink: 0;"
-                sse-swap="swapIcon-{{ ConversationAreaId }}-{{ userID }}">
+                sse-swap="swapIcon-{{ ConversationAreaId }}">
                 <img src="icons/bouvai2.png" alt="User Image" id="selectedIcon-{{ ConversationAreaId }}">
             </figure>
 
@@ -72,7 +72,7 @@
             {% elif IsPlaceholder %}
             <div class="is-flex is-align-items-start">
                 <div class="message-content" id="content-{{ ConversationAreaId }}"
-                    sse-swap="swapContent-{{ ConversationAreaId }}-{{ userID }}">
+                    sse-swap="swapContent-{{ ConversationAreaId }}">
                     <hx hx-trigger="load" hx-get="/generateMultipleMessages" id="generate-multiple-messages"></hx>
                     <div class='message-header'>
                         <p>
@@ -98,7 +98,7 @@
 
                 {% for selectedLLM in SelectedLLMs %}
                 <button disable class="button is-small is-primary message-button is-outlined mr-1"
-                    sse-swap="swapSelectionBtn-{{ selectedLLM.ID.String() }}-{{ userID }}" hx-swap="outerHTML"
+                    sse-swap="swapSelectionBtn-{{ selectedLLM.ID.String() }}" hx-swap="outerHTML"
                     hx-target="this">
                     <span class="icon is-small">
                         <!--img src="icons/{{ selectedLLM.Company }}.png" alt="{{ selectedLLM.Name }}"