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"
)

type ModelInfo struct {
	ID          string
	Name        string
	Icon        string
	MaxToken    int
	InputPrice  float32
	OutputPrice float32
}

type CompanyInfo struct {
	ID         string
	Name       string
	Icon       string
	ModelInfos []ModelInfo
}

var CompanyInfos []CompanyInfo
var ModelsInfos []ModelInfo

type MultipleModelsCompletionRequest struct {
	ModelIds []string
	Messages []Message
	Message  string
}

type BotContentMessage struct {
	Content string
	Hidden  bool
}

var lastSelectedModelIds []string

func addUsage(inputCost float32, outputCost float32, inputToken int32, outputToken int32, modelID string) {
	// Create a new usage
	err := edgeClient.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 in edgedb.QuerySingle: in addUsage")
		log.Fatal(err)
	}
}

func GenerateMultipleMessages(c *fiber.Ctx) error {
	insertArea()

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

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

	for i := range lastSelectedModelIds {
		idx := i
		go func() {
			// 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

			// Use a channel to signal the completion of addxxxMessage
			done := make(chan struct{}, 1)

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

			// Call the selected addMessageFunc in a goroutine
			go func() {
				defer wg.Done()
				if addMessageFunc != nil {
					addMessageFunc(lastSelectedModelIds[idx], idx == 0)
				}
				done <- struct{}{}
			}()

			// Use select to wait on multiple channel operations
			select {
			case <-ctx.Done(): // Context's deadline is exceeded
				// Insert a bot message indicating a timeout
				insertBotMessage(lastSelectedModelIds[idx]+" too long to answer", idx == 0, lastSelectedModelIds[idx])
			case <-done: // addMessageFunc completed within the deadline
				// No action needed, the function completed successfully
			}
		}()
	}

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

	fmt.Println("Done!")

	return c.SendString(generateChatHTML())
}

func RequestMultipleMessagesHandler(c *fiber.Ctx) error {
	message := c.FormValue("message", "")

	selectedModelIds := []string{}
	for ModelInfo := range ModelsInfos {
		out := c.FormValue("model-check-" + ModelsInfos[ModelInfo].ID)
		if out != "" {
			selectedModelIds = append(selectedModelIds, ModelsInfos[ModelInfo].ID)
		}
	}
	lastSelectedModelIds = selectedModelIds

	out := RequestMultipleMessages(message, selectedModelIds)

	return c.SendString(out)
}

func RequestMultipleMessages(message string, selectedModelIds []string) string {
	// Add an Area with the user message inside
	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})
	out += messageOut

	return out
}