Jade/Request.go

package main

import (
	"context"
	"fmt"
	"log"
	"sync"
	"time"

	"github.com/edgedb/edgedb-go"
	"github.com/flosch/pongo2"
	"github.com/gofiber/fiber/v2"
)

func GenerateMultipleMessagesHandler(c *fiber.Ctx) error {
	message := c.FormValue("message", "")
	selectedLLMIds := []string{"1e5a07c4-12fe-11ef-8da6-67d29b408c53"} // TODO Hanle in the UI

	var selectedLLMs []LLM
	var selectedLLM LLM

	for _, id := range selectedLLMIds {
		idUUID, _ := edgedb.ParseUUID(id)
		err := edgeClient.QuerySingle(context.Background(), `
			SELECT LLM {
				id,
				name,
				context,
				temperature,
				modelInfo : {
					modelID,
					company : {
						icon,
						name
					}
				}
			}
			FILTER
				.id = <uuid>$0;
		`, &selectedLLM, idUUID)
		if err != nil {
			fmt.Println("Error trying to select the unique LLM")
			log.Fatal(err)
		}
		selectedLLMs = append(selectedLLMs, selectedLLM)
	}

	fmt.Println("Selected LLMs: ", selectedLLMs)

	_, position := insertArea()
	messageID := insertUserMessage(message)

	out := ""
	messageOut, _ := userTmpl.Execute(pongo2.Context{"Content": markdownToHTML(message), "ID": messageID.String()})
	out += messageOut

	messageOut, _ = botTmpl.Execute(pongo2.Context{"IsPlaceholder": true, "SelectedLLMs": selectedLLMs, "ConversationAreaId": position + 1})
	out += messageOut

	go HandleGenerateMultipleMessages(selectedLLMs)

	return c.SendString(out)
}

func HandleGenerateMultipleMessages(selectedLLMs []LLM) {
	insertArea()

	// Create a wait group to synchronize the goroutines
	var wg sync.WaitGroup

	// Add the length of selectedModelIds goroutines to the wait group
	wg.Add(len(selectedLLMs))

	// Create a channel to receive the index of the first completed goroutine
	firstDone := make(chan int, 1)

	for i := range selectedLLMs {
		idx := i
		go func() {
			defer wg.Done()

			// Create a context with a 1-minute timeout
			ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
			defer cancel() // Ensure the context is cancelled to free resources

			// Determine which message function to call based on the model
			var addMessageFunc func(selectedLLM LLM, selected bool) edgedb.UUID
			switch selectedLLMs[idx].Model.Company.Name {
			case "openai":
				addMessageFunc = addOpenaiMessage
			case "anthropic":
				addMessageFunc = addAnthropicMessage
			case "mistral":
				addMessageFunc = addMistralMessage
			case "groq":
				addMessageFunc = addGroqMessage
			}

			var messageID edgedb.UUID
			if addMessageFunc != nil {
				messageID = addMessageFunc(selectedLLMs[idx], idx == 0)
			} else {
				fmt.Println("Invalid model: ", selectedLLMs[idx].Model.Company.Name)
			}

			fmt.Println("Message ID: ", messageID)

			var message Message
			err := edgeClient.QuerySingle(edgeCtx, `
				SELECT Message {
					id,
					content,
					area : {
						id,
						position
					},
					llm : {
						modelInfo : {
							modelID
						}
					}
				}
				FILTER .id = <uuid>$0;
			`, &message, messageID)
			if err != nil {
				fmt.Println("Error in edgedb.QuerySingle: in HandleGenerateMultipleMessages 1")
				log.Fatal(err)
			}

			templateMessage := TemplateMessage{
				Icon:    message.LLM.Model.Company.Icon,
				Content: markdownToHTML(message.Content),
				Hidden:  false,
				Id:      message.ID.String(),
				Name:    message.LLM.Model.Name,
				ModelID: message.LLM.Model.ModelID,
			}

			// Check if the context's deadline is exceeded
			select {
			case <-ctx.Done():
				// The context's deadline was exceeded
				fmt.Printf("Goroutine %d timed out\n", idx)
			default:
				// Send the index of the completed goroutine to the firstDone channel
				select {
				case firstDone <- idx:
					// Generate the HTML content
					out := "<div class='message-header'>"
					out += "<p>"
					out += templateMessage.Name
					out += "	</p>"
					out += "</div>"
					out += "<div class='message-body'>"
					out += "	<ct class='content'>"
					out += templateMessage.Content
					out += "	</ct>"
					out += "</div>"

					// Send Content event
					sseChanel.SendEvent(
						"swapContent-"+fmt.Sprintf("%d", message.Area.Position),
						out,
					)

					fmt.Println(templateMessage)

					out, err := modelSelecBtnTmpl.Execute(map[string]interface{}{
						"message":            templateMessage,
						"ConversationAreaId": message.Area.Position,
					})
					if err != nil {
						fmt.Println("Error in modelSelecBtnTmpl.Execute: in HandleGenerateMultipleMessages 3")
						log.Fatal(err)
					}

					fmt.Println("Sending event: swapSelectionBtn-" + selectedLLMs[idx].ID.String())
					fmt.Println(out)

					// Send Content event
					sseChanel.SendEvent(
						"swapSelectionBtn-"+templateMessage.ModelID,
						out,
					)

					// Send Icon Swap event
					sseChanel.SendEvent(
						"swapIcon-"+fmt.Sprintf("%d", message.Area.Position),
						`<img src="icons/`+selectedLLMs[idx].Model.Company.Name+`.png" alt="User Image">`,
					)
				default:
					out, err := modelSelecBtnTmpl.Execute(map[string]interface{}{
						"message":            templateMessage,
						"ConversationAreaId": message.Area.Position,
					})
					if err != nil {
						fmt.Println("Error in modelSelecBtnTmpl.Execute: in HandleGenerateMultipleMessages 4")
						log.Fatal(err)
					}

					// Send Content event
					sseChanel.SendEvent(
						"swapSelectionBtn-"+templateMessage.ModelID,
						out,
					)
				}
			}
		}()
	}

	// Wait for all goroutines to finish
	wg.Wait()
}