Jade/main.go
2024-05-29 13:01:20 +02:00

510 lines
13 KiB
Go

// Welcome welcome, I guess you started here
// Thank you for reviewing my code.
// As for context. This is the absolutely first time I use Golang or EdgeDB.
// So far I only did python. I have a master but in thermal engineering, not CS.
// My philosophy in coding is simplicity. A "clean code" is a simple code.
// For me you should need few comments. Variables, type and functions name should be enough to understand.
// And what's work, work.
// To build my app I took a simple approach.
// All data are stored in database. If I need it, I ask the DB.
// There is no state in the server nor in the client (or very little for the client).
// For example if I want to ask a question. Here a simplify version:
// 1. Add the user message to the DB.
// 2. Read all messages from the DB to send a request.
// 3. Add the bot message to the DB.
// 4. Read all messages from the DB to generate the chat HTML.
// So all query are simple query. That is my goal. In the future I may optimize it, but for dev, that is the best approach I think.
// Also if I want to change the order of a list (you will see later).
// I first send the new positions -> update the db -> read the db to generate the HTML.
// So Edgedb is the heart of the app. Everything pass by it.
// It need A LOT of optimization, I first do something that work then optimize it.
// I hope you will like it.
package main
import (
"bufio"
"fmt"
"log"
"os"
"sync"
"github.com/flosch/pongo2"
"github.com/gofiber/fiber/v2"
"github.com/gofiber/fiber/v2/middleware/favicon"
"github.com/gofiber/fiber/v2/middleware/logger"
"github.com/gofiber/fiber/v2/middleware/recover"
"github.com/gofiber/template/django/v3"
"github.com/stripe/stripe-go"
)
var (
userTmpl *pongo2.Template
botTmpl *pongo2.Template
selectBtnTmpl *pongo2.Template
modelPopoverTmpl *pongo2.Template
usagePopoverTmpl *pongo2.Template
settingPopoverTmpl *pongo2.Template
messageEditTmpl *pongo2.Template
conversationPopoverTmpl *pongo2.Template
clients = make(map[chan SSE]bool)
mu sync.Mutex
)
// SSE event structure
type SSE struct {
Event string
Data string
}
// Function to send events to all clients
func sendEvent(event, data string) {
mu.Lock()
defer mu.Unlock()
for client := range clients {
client <- SSE{Event: event, Data: data}
}
}
func main() {
// Use STRIPE_KEY environment variable
stripe.Key = os.Getenv("STRIPE_KEY")
botTmpl = pongo2.Must(pongo2.FromFile("views/partials/message-bot.html"))
userTmpl = pongo2.Must(pongo2.FromFile("views/partials/message-user.html"))
selectBtnTmpl = pongo2.Must(pongo2.FromFile("views/partials/model-selection-btn.html"))
modelPopoverTmpl = pongo2.Must(pongo2.FromFile("views/partials/popover-models.html"))
conversationPopoverTmpl = pongo2.Must(pongo2.FromFile("views/partials/popover-conversation.html"))
usagePopoverTmpl = pongo2.Must(pongo2.FromFile("views/partials/popover-usage.html"))
settingPopoverTmpl = pongo2.Must(pongo2.FromFile("views/partials/popover-settings.html"))
messageEditTmpl = pongo2.Must(pongo2.FromFile("views/partials/message-edit-form.html"))
// Import HTML using django engine/template
engine := django.New("./views", ".html")
// Create new Fiber instance
app := fiber.New(fiber.Config{
Views: engine,
AppName: "JADE",
EnablePrintRoutes: true,
})
defer app.Shutdown()
// Add default logger
app.Use(favicon.New())
app.Use(logger.New())
app.Use(recover.New())
// Add static files
app.Static("/", "./static")
// Main routes
app.Get("/", ChatPageHandler)
app.Get("/loadChat", LoadChatHandler)
app.Get("/pricingTable", PricingTableHandler)
app.Get("/generateTermAndService", generateTermAndServiceHandler)
// Chat routes
app.Post("/deleteMessage", DeleteMessageHandler)
app.Post("/generatePlaceholder", GeneratePlaceholderHandler)
app.Get("/generateMultipleMessages", GenerateMultipleMessagesHandler)
app.Get("/messageContent", GetMessageContentHandler)
app.Get("/editMessageForm", GetEditMessageFormHandler)
app.Post("/redoMessage", RedoMessageHandler)
app.Post("/clearChat", ClearChatHandler)
app.Get("/userMessage", GetUserMessageHandler)
app.Post("/editMessage", EditMessageHandler)
// Settings routes
app.Post("/addKeys", addKeys)
// Popovers
app.Get("/loadModelSelection", LoadModelSelectionHandler)
app.Get("/loadConversationSelection", LoadConversationSelectionHandler)
app.Get("/refreshConversationSelection", RefreshConversationSelectionHandler)
app.Get("/loadUsageKPI", LoadUsageKPIHandler)
app.Get("/loadSettings", LoadSettingsHandler)
app.Post("/updateLLMPositionBatch", updateLLMPositionBatch)
app.Get("/createConversation", CreateConversationHandler)
app.Get("/deleteConversation", DeleteConversationHandler)
app.Get("/selectConversation", SelectConversationHandler)
app.Post("/updateConversationPositionBatch", updateConversationPositionBatch)
// Authentication
app.Get("/signin", handleUiSignIn)
app.Get("/signout", handleSignOut)
app.Get("/callback", handleCallback)
app.Get("/callbackSignup", handleCallbackSignup)
// LLM
app.Get("deleteLLM", deleteLLM)
app.Post("/createLLM", createLLM)
app.Get("/empty", func(c *fiber.Ctx) error {
return c.SendString("")
})
app.Get("/sse", func(c *fiber.Ctx) error {
c.Set("Content-Type", "text/event-stream")
c.Set("Cache-Control", "no-cache")
c.Set("Connection", "keep-alive")
events := make(chan SSE, 100)
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 {
log.Fatal(err)
}
}
// The route to add keys, idk where to put it so it's here
// Deserve a good REDO lol
// I mean, look at this shit...
// 300 lines, I'm sure it can be done in 50
// TODO REDO and maybe find it a better place
func addKeys(c *fiber.Ctx) error {
openaiKey := c.FormValue("openai_key")
anthropicKey := c.FormValue("anthropic_key")
mistralKey := c.FormValue("mistral_key")
groqKey := c.FormValue("groq_key")
gooseaiKey := c.FormValue("goose_key")
googleKey := c.FormValue("google_key")
var Exists bool
// Handle OpenAI key
if openaiKey != "" {
if !TestOpenaiKey(openaiKey) {
return c.SendString("Invalid OpenAI API Key\n")
}
// Check if the company key already exists
err := edgeClient.QuerySingle(edgeCtx, `
select exists (
select global currentUser.setting.keys
filter .company.name = "openai"
);
`, &Exists)
if err != nil {
fmt.Println("Error checking if OpenAI key exists")
panic(err)
}
if Exists {
err = edgeClient.Execute(edgeCtx, `
UPDATE Key filter .company.name = <str>$0 AND .key = <str>$1
SET {
key := <str>$1,
}
`, "openai", openaiKey)
if err != nil {
fmt.Println("Error updating OpenAI key")
panic(err)
}
} else {
err = edgeClient.Execute(edgeCtx, `
WITH
c := (SELECT Company FILTER .name = "openai" LIMIT 1)
UPDATE global currentUser.setting
SET {
keys += (
INSERT Key {
company := c,
key := <str>$0,
name := "OpenAI API Key",
}
)
}`, openaiKey)
if err != nil {
fmt.Println("Error adding OpenAI key")
panic(err)
}
}
}
// Handle Anthropic key
if anthropicKey != "" {
if !TestAnthropicKey(anthropicKey) {
return c.SendString("Invalid Anthropic API Key\n")
}
// Check if the company key already exists
err := edgeClient.QuerySingle(edgeCtx, `
select exists (
select global currentUser.setting.keys
filter .company.name = "anthropic"
);
`, &Exists)
if err != nil {
fmt.Println("Error checking if Anthropic key exists")
panic(err)
}
if Exists {
err = edgeClient.Execute(edgeCtx, `
UPDATE Key filter .company.name = "anthropic" AND .key = <str>$0
SET {
key := <str>$0,
}
`, anthropicKey)
if err != nil {
fmt.Println("Error updating Anthropic key")
panic(err)
}
} else {
err = edgeClient.Execute(edgeCtx, `
WITH
c := (SELECT Company FILTER .name = "anthropic" LIMIT 1)
UPDATE global currentUser.setting
SET {
keys += (
INSERT Key {
company := c,
key := <str>$0,
name := "Anthropic API Key",
}
)
}`, anthropicKey)
if err != nil {
fmt.Println("Error adding Anthropic key")
panic(err)
}
}
}
// Handle Mistral key
if mistralKey != "" {
if !TestMistralKey(mistralKey) {
return c.SendString("Invalid Mistral API Key\n")
}
// Check if the company key already exists
err := edgeClient.QuerySingle(edgeCtx, `
select exists (
select global currentUser.setting.keys
filter .company.name = "mistral"
);
`, &Exists)
if err != nil {
fmt.Println("Error checking if Mistral key exists")
panic(err)
}
if Exists {
err = edgeClient.Execute(edgeCtx, `
UPDATE Key filter .company.name = "mistral" AND .key = <str>$0
SET {
key := <str>$0,
}
`, mistralKey)
if err != nil {
fmt.Println("Error updating Mistral key")
panic(err)
}
} else {
err = edgeClient.Execute(edgeCtx, `
WITH
c := (SELECT Company FILTER .name = "mistral" LIMIT 1)
UPDATE global currentUser.setting
SET {
keys += (
INSERT Key {
company := c,
key := <str>$0,
name := "Mistral API Key",
}
)
}`, mistralKey)
if err != nil {
fmt.Println("Error adding Mistral key")
panic(err)
}
}
}
// Handle Groq key
if groqKey != "" {
if !TestGroqKey(groqKey) {
return c.SendString("Invalid Groq API Key\n")
}
// Check if the company key already exists
err := edgeClient.QuerySingle(edgeCtx, `
select exists (
select global currentUser.setting.keys
filter .company.name = "groq"
);
`, &Exists)
if err != nil {
fmt.Println("Error checking if Groq key exists")
panic(err)
}
if Exists {
err = edgeClient.Execute(edgeCtx, `
UPDATE Key filter .company.name = "groq" AND .key = <str>$0
SET {
key := <str>$0,
}
`, groqKey)
if err != nil {
fmt.Println("Error updating Groq key")
panic(err)
}
} else {
err = edgeClient.Execute(edgeCtx, `
WITH
c := (SELECT Company FILTER .name = "groq" LIMIT 1)
UPDATE global currentUser.setting
SET {
keys += (
INSERT Key {
company := c,
key := <str>$0,
name := "Groq API Key",
}
)
}`, groqKey)
if err != nil {
fmt.Println("Error adding Groq key")
panic(err)
}
}
}
// Handle Gooseai key
if gooseaiKey != "" {
if !TestGooseaiKey(gooseaiKey) {
return c.SendString("Invalid Gooseai API Key\n")
}
// Check if the company key already exists
err := edgeClient.QuerySingle(edgeCtx, `
select exists (
select global currentUser.setting.keys
filter .company.name = "gooseai"
);
`, &Exists)
if err != nil {
fmt.Println("Error checking if Gooseai key exists")
panic(err)
}
if Exists {
err = edgeClient.Execute(edgeCtx, `
UPDATE Key filter .company.name = "gooseai" AND .key = <str>$0
SET {
key := <str>$0,
}
`, gooseaiKey)
if err != nil {
fmt.Println("Error updating Gooseai key")
panic(err)
}
} else {
err = edgeClient.Execute(edgeCtx, `
WITH
c := (SELECT Company FILTER .name = "gooseai" LIMIT 1)
UPDATE global currentUser.setting
SET {
keys += (
INSERT Key {
company := c,
key := <str>$0,
name := "Gooseai API Key",
}
)
}`, gooseaiKey)
if err != nil {
fmt.Println("Error adding Gooseai key")
panic(err)
}
}
}
// Handle Google key
if googleKey != "" {
if !TestGoogleKey(googleKey) {
return c.SendString("Invalid Google API Key\n")
}
// Check if the company key already exists
err := edgeClient.QuerySingle(edgeCtx, `
select exists (
select global currentUser.setting.keys
filter .company.name = "google"
);
`, &Exists)
if err != nil {
fmt.Println("Error checking if Google key exists")
panic(err)
}
if Exists {
err = edgeClient.Execute(edgeCtx, `
UPDATE Key filter .company.name = "google" AND .key = <str>$0
SET {
key := <str>$0,
}
`, googleKey)
if err != nil {
fmt.Println("Error updating Google key")
panic(err)
}
} else {
err = edgeClient.Execute(edgeCtx, `
WITH
c := (SELECT Company FILTER .name = "google" LIMIT 1)
UPDATE global currentUser.setting
SET {
keys += (
INSERT Key {
company := c,
key := <str>$0,
name := "Google API Key",
}
)
}`, googleKey)
if err != nil {
fmt.Println("Error adding Google key")
panic(err)
}
}
}
return c.SendString("<script>window.location.reload()</script>")
}