Jade/Request.go

package main

import (
	"context"
	"encoding/json"
	"fmt"
	"strings"
	"sync"
	"time"

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

type RequestMessage struct {
	Role    string `json:"role"`
	Content string `json:"content"`
}

var lastSelectedLLMs = make(map[string][]LLM)

func GeneratePlaceholderHandler(c *fiber.Ctx) error {
	// Step 1 I create a User message and send it as output with a placeholder
	// that will make a request to GenerateMultipleMessagesHandler when loading
	message := c.FormValue("message", "")
	var selectedLLMIds []string
	err := json.Unmarshal([]byte(c.FormValue("selectedLLMIds")), &selectedLLMIds)
	if err != nil {
		fmt.Println("Error unmarshalling selected LLM IDs")
		panic(err)
	}

	return c.SendString(GeneratePlaceholderHTML(c, message, selectedLLMIds, true))
}

func GeneratePlaceholderHTML(c *fiber.Ctx, message string, selectedLLMIds []string, with_user_message bool) string {
	var selectedLLMs []LLM
	var selectedLLM LLM

	for _, id := range selectedLLMIds {
		idUUID, _ := edgedb.ParseUUID(id)
		err := edgeGlobalClient.WithGlobals(map[string]interface{}{"ext::auth::client_token": c.Cookies("jade-edgedb-auth-token")}).QuerySingle(context.Background(), `
			SELECT LLM {
				id,
				name,
				context,
				temperature,
        max_tokens,
				custom_endpoint : {
					id,
					endpoint,
					key
				},
				modelInfo : {
          id,
					modelID,
					company : {
						icon,
						name
					}
				}
			}
			FILTER
				.id = <uuid>$0;
		`, &selectedLLM, idUUID)
		if err != nil {
			fmt.Println("Error getting LLM")
			panic(err)
		}
		selectedLLMs = append(selectedLLMs, selectedLLM)
	}
	lastSelectedLLMs[c.Cookies("jade-edgedb-auth-token")] = selectedLLMs

	_, position := insertArea(c)

	var user User
	err := edgeGlobalClient.WithGlobals(map[string]interface{}{"ext::auth::client_token": c.Cookies("jade-edgedb-auth-token")}).QuerySingle(context.Background(), `
		SELECT global currentUser { id } LIMIT 1
	`, &user)
	if err != nil {
		fmt.Println("Error getting user")
		panic(err)
	}

	out := ""
	if with_user_message {
		messageID := insertUserMessage(c, message)
		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,
		"userID":             user.ID.String(),
	})
	out += messageOut

	return out
}

func GenerateMultipleMessagesHandler(c *fiber.Ctx) error {
	// Step 2 generate multiple messages
	// And send them one by one using events
	insertArea(c)
	selectedLLMs, _ := lastSelectedLLMs[c.Cookies("jade-edgedb-auth-token")]

	var user User
	err := edgeGlobalClient.WithGlobals(map[string]interface{}{"ext::auth::client_token": c.Cookies("jade-edgedb-auth-token")}).QuerySingle(context.Background(), `
		SELECT global currentUser { id } LIMIT 1
	`, &user)
	if err != nil {
		fmt.Println("Error getting user")
		panic(err)
	}

	authCookie := c.Cookies("jade-edgedb-auth-token")

	var messages []Message = getAllSelectedMessages(c)

	// 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(c *fiber.Ctx, llm LLM, messages []Message, testApiKey string) string
			switch selectedLLMs[idx].Model.Company.Name {
			case "openai":
				addMessageFunc = RequestOpenai
			case "anthropic":
				addMessageFunc = RequestAnthropic
			case "mistral":
				addMessageFunc = RequestMistral
			case "groq":
				addMessageFunc = RequestGroq
			case "huggingface":
				addMessageFunc = RequestCustomEndpoint
			case "google":
				addMessageFunc = RequestGoogle
			case "perplexity":
				addMessageFunc = RequestPerplexity
			case "fireworks":
				addMessageFunc = RequestFirework
			case "nim":
				addMessageFunc = RequestNim
			case "together":
				addMessageFunc = RequestTogether
			case "deepseek":
				addMessageFunc = RequestDeepseek
			}

			var content string = addMessageFunc(c, selectedLLMs[idx], messages, "")
			var messageUUID edgedb.UUID = insertBotMessage(c, content, selectedLLMs[idx].ID)

			var message Message
			err := edgeGlobalClient.WithGlobals(map[string]interface{}{"ext::auth::client_token": c.Cookies("jade-edgedb-auth-token")}).QuerySingle(edgeCtx, `
				SELECT Message {
					id,
					content,
					area : {
						id,
						position
					},
					llm : {
						modelInfo : {
							modelID,
							name,
							company : {
								icon,
							}
						}
					}
				}
				FILTER .id = <uuid>$0;
			`, &message, messageUUID)
			if err != nil {
				fmt.Println("Error getting message for the placeholder. The function addProviderMessage seem to not return any message ID.")
				panic(err)
			}

			select {
			case <-ctx.Done():
				fmt.Printf("Goroutine %d timed out\n", idx)
			default:
				select {
				case firstDone <- idx:
					_ = edgeGlobalClient.WithGlobals(map[string]interface{}{"ext::auth::client_token": authCookie}).Execute(edgeCtx, `
	          UPDATE Message
	          FILTER .id = <uuid>$0
	          SET {selected := true};
	        `, messageUUID)

					outIcon := `<img src="` + selectedLLMs[idx].Model.Company.Icon + `" alt="User Image" id="selectedIcon-` + fmt.Sprintf("%d", message.Area.Position) + `">`

					sendEvent(
						user.ID.String(),
						"swapContent-"+fmt.Sprintf("%d", message.Area.Position),
						GenerateMessageContentHTML(
							authCookie,
							message.ID.String(),
							"false",
							true,
						),
					)
					sendEvent(
						user.ID.String(),
						"swapSelectionBtn-"+selectedLLMs[idx].ID.String(),
						GenerateSelectionBtnHTML(authCookie, message.ID.String()),
					)
					sendEvent(
						user.ID.String(),
						"swapIcon-"+fmt.Sprintf("%d", message.Area.Position),
						outIcon,
					)
				default:
					sendEvent(
						user.ID.String(),
						"swapSelectionBtn-"+selectedLLMs[idx].ID.String(),
						GenerateSelectionBtnHTML(authCookie, message.ID.String()),
					)
				}
			}
		}()
	}

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

	delete(lastSelectedLLMs, c.Cookies("jade-edgedb-auth-token"))

	return c.SendString("")
}

func GenerateSelectionBtnHTML(authCookie string, messageID string) string {
	messageUUID, err := edgedb.ParseUUID(messageID)
	if err != nil {
		fmt.Println("Error parsing UUID")
		panic(err)
	}

	var message Message
	err = edgeGlobalClient.WithGlobals(map[string]interface{}{"ext::auth::client_token": authCookie}).QuerySingle(edgeCtx, `
		SELECT Message {
			id,
			content,
			area : {
				id,
				position
			},
			llm : {
        id,
				modelInfo : {
					modelID,
					name,
					company : {
						icon,
					}
				}
			}
		}
		FILTER .id = <uuid>$0;
	`, &message, messageUUID)
	if err != nil {
		fmt.Println("Error getting message")
		panic(err)
	}

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

	outBtn, err := selectBtnTmpl.Execute(map[string]interface{}{
		"message":            templateMessage,
		"ConversationAreaId": message.Area.Position,
	})
	if err != nil {
		fmt.Println("Error generating HTML content")
		panic(err)
	}
	outBtn = strings.ReplaceAll(outBtn, "\n", "")

	return outBtn
}

func addUsage(c *fiber.Ctx, inputCost float32, outputCost float32, inputToken int32, outputToken int32, modelID string) {
	// Create a new usage
	err := edgeGlobalClient.WithGlobals(map[string]interface{}{"ext::auth::client_token": c.Cookies("jade-edgedb-auth-token")}).Execute(edgeCtx, `
		INSERT Usage {
			input_cost := <float32>$0,
			output_cost := <float32>$1,
			input_token := <int32>$2,
			output_token := <int32>$3,
			model_id := <str>$4,
			user := global currentUser
		}
	`, inputCost, outputCost, inputToken, outputToken, modelID)
	if err != nil {
		fmt.Println("Error adding usage")
		panic(err)
	}
}