131 lines
3.2 KiB
Go
131 lines
3.2 KiB
Go
package main
|
|
|
|
import (
|
|
"fmt"
|
|
"log"
|
|
"sync"
|
|
|
|
"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 {
|
|
// Create a wait group to synchronize the goroutines
|
|
var wg sync.WaitGroup
|
|
var InsertedIDs []edgedb.UUID
|
|
|
|
// Add the length of lastSelectedModelIds goroutines to the wait group
|
|
wg.Add(len(lastSelectedModelIds))
|
|
|
|
for i := range lastSelectedModelIds {
|
|
idx := i
|
|
if model2Icon(lastSelectedModelIds[i]) == "openai" {
|
|
go func() {
|
|
defer wg.Done()
|
|
response := addOpenaiMessage(lastSelectedModelIds[idx], idx == 0)
|
|
InsertedIDs = append(InsertedIDs, response)
|
|
}()
|
|
} else if model2Icon(lastSelectedModelIds[i]) == "anthropic" {
|
|
go func() {
|
|
defer wg.Done()
|
|
response := addAnthropicMessage(lastSelectedModelIds[idx], idx == 0)
|
|
InsertedIDs = append(InsertedIDs, response)
|
|
}()
|
|
} else if model2Icon(lastSelectedModelIds[i]) == "mistral" {
|
|
go func() {
|
|
defer wg.Done()
|
|
response := addMistralMessage(lastSelectedModelIds[idx], idx == 0)
|
|
InsertedIDs = append(InsertedIDs, response)
|
|
}()
|
|
}
|
|
}
|
|
|
|
// Wait for both goroutines to finish
|
|
wg.Wait()
|
|
|
|
return c.SendString(generateChatHTML())
|
|
}
|
|
|
|
func RequestMultipleMessages(c *fiber.Ctx) error {
|
|
message := c.FormValue("message")
|
|
if chatString, commandExecuted := DetectCommand(message, c); commandExecuted {
|
|
return c.SendString(chatString)
|
|
}
|
|
|
|
// Add an Area with the user message inside
|
|
insertArea()
|
|
messageID := insertUserMessage(message)
|
|
|
|
selectedModelIds := []string{}
|
|
for CompanyInfo := range CompanyInfos {
|
|
for ModelInfo := range CompanyInfos[CompanyInfo].ModelInfos {
|
|
out := c.FormValue("model-check-" + CompanyInfos[CompanyInfo].ModelInfos[ModelInfo].ID)
|
|
if out != "" {
|
|
selectedModelIds = append(selectedModelIds, CompanyInfos[CompanyInfo].ModelInfos[ModelInfo].ID)
|
|
}
|
|
}
|
|
}
|
|
lastSelectedModelIds = selectedModelIds
|
|
|
|
out := ""
|
|
|
|
messageOut, _ := userTmpl.Execute(pongo2.Context{"Content": markdownToHTML(message), "ID": messageID.String()})
|
|
out += messageOut
|
|
|
|
messageOut, _ = botTmpl.Execute(pongo2.Context{"IsPlaceholder": true, "SelectedModelIds": selectedModelIds, "Message": message})
|
|
out += messageOut
|
|
|
|
return c.SendString(out)
|
|
}
|