457 lines
10 KiB
Go
457 lines
10 KiB
Go
package main
|
|
|
|
import (
|
|
"bufio"
|
|
"fmt"
|
|
"sync"
|
|
|
|
"github.com/flosch/pongo2"
|
|
"github.com/gofiber/fiber/v2"
|
|
"github.com/gofiber/fiber/v2/middleware/logger"
|
|
"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())
|
|
|
|
// Add static files
|
|
app.Static("/", "./static")
|
|
|
|
// Main routes
|
|
app.Get("/", ChatPageHandler)
|
|
app.Get("/loadChat", LoadChatHandler)
|
|
|
|
// Chat routes
|
|
app.Post("/deleteMessage", DeleteMessageHandler)
|
|
app.Get("/generatePlaceholder", GeneratePlaceholderHandler)
|
|
app.Get("/generateMultipleMessages", GenerateMultipleMessagesHandler)
|
|
app.Get("/messageContent", GetMessageContentHandler)
|
|
app.Get("/editMessageForm", GetEditMessageFormHandler)
|
|
app.Post("/redoMessage", RedoMessageHandler)
|
|
app.Post("/clearChat", ClearChatHandler)
|
|
|
|
// Settings routes
|
|
app.Post("/addKeys", addKeys)
|
|
app.Get("/modal", openLlmModal)
|
|
|
|
// 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.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)
|
|
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 {
|
|
return
|
|
}
|
|
w.Flush()
|
|
}
|
|
})
|
|
|
|
return nil
|
|
})
|
|
|
|
// Start server
|
|
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 {
|
|
fmt.Println("Error in edgedb.QuerySingle: in addOpenaiKey")
|
|
fmt.Println(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 {
|
|
fmt.Println("Error in edgedb.QuerySingle: in addOpenaiKey")
|
|
fmt.Println(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 in edgedb.QuerySingle: in addOpenaiKey")
|
|
fmt.Println(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 {
|
|
fmt.Println("Error in edgedb.QuerySingle: in addAnthropicKey")
|
|
fmt.Println(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 in edgedb.QuerySingle: in addAnthropicKey")
|
|
fmt.Println(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 in edgedb.QuerySingle: in addAnthropicKey")
|
|
fmt.Println(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 {
|
|
fmt.Println("Error in edgedb.QuerySingle: in addMistralKey")
|
|
fmt.Println(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 in edgedb.QuerySingle: in addMistralKey")
|
|
fmt.Println(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 in edgedb.QuerySingle: in addMistralKey")
|
|
fmt.Println(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 {
|
|
fmt.Println("Error in edgedb.QuerySingle: in addGroqKey")
|
|
fmt.Println(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 in edgedb.QuerySingle: in addGroqKey")
|
|
fmt.Println(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 in edgedb.QuerySingle: in addGroqKey")
|
|
fmt.Println(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 {
|
|
fmt.Println("Error in edgedb.QuerySingle: in addGooseaiKey")
|
|
fmt.Println(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 in edgedb.QuerySingle: in addGooseaiKey")
|
|
fmt.Println(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 in edgedb.QuerySingle: in addGooseaiKey")
|
|
fmt.Println(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 {
|
|
fmt.Println("Error in edgedb.QuerySingle: in addGoogleKey")
|
|
fmt.Println(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 in edgedb.QuerySingle: in addGoogleKey")
|
|
fmt.Println(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 in edgedb.QuerySingle: in addGoogleKey")
|
|
fmt.Println(err)
|
|
}
|
|
}
|
|
}
|
|
|
|
return c.SendString("<script>window.location.reload()</script>")
|
|
}
|