Jade/main.go

package main

import (
	"bufio"
	"fmt"
	"log"
	"sync"

	"github.com/flosch/pongo2"
	"github.com/gofiber/fiber/v2"
	"github.com/gofiber/fiber/v2/middleware/logger"
	"github.com/gofiber/fiber/v2/middleware/recover"
	"github.com/gofiber/template/django/v3"
)

var userTmpl *pongo2.Template
var botTmpl *pongo2.Template
var selectBtnTmpl *pongo2.Template

// SSE event structure
type SSE struct {
	Event string
	Data  string
}

// Global channel and mutex
var (
	clients = make(map[chan SSE]bool)
	mu      sync.Mutex
)

// 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() {
	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"))

	// Import HTML using django engine/template
	engine := django.New("./views", ".html")

	// Create new Fiber instance. Can use any framework. I use fiber for speed and simplicity
	app := fiber.New(fiber.Config{
		Views:             engine,
		AppName:           "JADE 2.0",
		EnablePrintRoutes: true,
	})

	// Add default logger
	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)

	// 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("/loadUsageKPI", LoadUsageKPIHandler)
	app.Get("/loadKeys", LoadKeysHandler)
	app.Get("/loadSettings", LoadSettingsHandler)

	// 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("/test", func(c *fiber.Ctx) error {
		fmt.Println("Hello from test")
		return c.SendString("")
	})

	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
	log.Fatal(app.Listen(":8080"))
}

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) {
			fmt.Println("Invalid OpenAI API Key")
			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 {
			panic(err)
		}

		if Exists {
			fmt.Println("Company key already 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 {
				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 {
				panic(err)
			}
		}
	}

	// Handle Anthropic key
	if anthropicKey != "" {
		if !TestAnthropicKey(anthropicKey) {
			fmt.Println("Invalid Anthropic API Key")
			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 {
			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 {
				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 {
				panic(err)
			}
		}
	}

	// Handle Mistral key
	if mistralKey != "" {
		if !TestMistralKey(mistralKey) {
			fmt.Println("Invalid Mistral API Key")
			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 {
			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 {
				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 {
				panic(err)
			}
		}
	}

	// Handle Groq key
	if groqKey != "" {
		if !TestGroqKey(groqKey) {
			fmt.Println("Invalid Groq API Key")
			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 {
			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 {
				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 {
				panic(err)
			}
		}
	}

	// Handle Gooseai key
	if gooseaiKey != "" {
		if !TestGooseaiKey(gooseaiKey) {
			fmt.Println("Invalid Gooseai API Key")
			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 {
			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 {
				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 {
				panic(err)
			}
		}
	}

	// Handle Google key
	if googleKey != "" {
		if !TestGoogleKey(googleKey) {
			fmt.Println("Invalid Google API Key")
			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 {
			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 {
				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 {
				panic(err)
			}
		}
	}

	return c.SendString("<script>window.location.reload()</script>")
}