Jade/main.go

package main

import (
	"fmt"
	"log"
	"os"
	"sync"

	"github.com/flosch/pongo2"
	"github.com/gofiber/contrib/websocket"
	"github.com/gofiber/fiber/v2"
	"github.com/gofiber/fiber/v2/middleware/logger"
	"github.com/gofiber/fiber/v2/middleware/monitor"
	"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
	welcomeChatTmpl         *pongo2.Template
	chatInputTmpl           *pongo2.Template
	explainLLMconvChatTmpl  *pongo2.Template
	mu                      sync.Mutex
	app                     *fiber.App
)

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"))
	welcomeChatTmpl = pongo2.Must(pongo2.FromFile("views/partials/welcome-chat.html"))
	chatInputTmpl = pongo2.Must(pongo2.FromFile("views/partials/chat-input.html"))
	explainLLMconvChatTmpl = pongo2.Must(pongo2.FromFile("views/partials/explain-llm-conv-chat.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",
	})
	defer app.Shutdown()

	// Add default logger
	app.Use(logger.New())
	app.Use(recover.New())

	// Main routes
	app.Get("/", ChatPageHandler)
	app.Get("/loadChat", LoadChatHandler)
	app.Get("/loadChatInput", LoadChatInputHandler)
	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)
	app.Get("/help", generateHelpChatHandler)

	// Settings routes
	app.Post("/addKeys", addKeys)

	// Popovers
	app.Get("/loadModelSelection", LoadModelSelectionHandler)
	app.Get("/loadConversationSelection", LoadConversationSelectionHandler)
	app.Get("/loadUsageKPI", LoadUsageKPIHandler)
	app.Get("/loadSettings", LoadSettingsHandler)
	app.Get("/refreshConversationSelection", RefreshConversationSelectionHandler)

	// Conversation routes
	app.Get("/createConversation", CreateConversationHandler)
	app.Get("/deleteConversation", DeleteConversationHandler)
	app.Get("/selectConversation", SelectConversationHandler)
	app.Post("/updateConversationPositionBatch", updateConversationPositionBatch)
	app.Post("/archiveDefaultConversation", ArchiveDefaultConversationHandler)

	// 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.Post("/updateLLMPositionBatch", updateLLMPositionBatch)

	//Admin monitoring
	app.Get("/metrics", monitor.New())

	// Add static files
	app.Static("/", "./static")

	app.Get("/empty", func(c *fiber.Ctx) error {
		return c.SendString("")
	})

	app.Get("/ws", websocket.New(handleWebSocket))

	// Start server
	if err := app.Listen(":8080"); err != nil {
		log.Fatal(err)
	}
}

func addKeys(c *fiber.Ctx) error {
	keys := map[string]string{
		"openai":     c.FormValue("openai_key"),
		"anthropic":  c.FormValue("anthropic_key"),
		"mistral":    c.FormValue("mistral_key"),
		"groq":       c.FormValue("groq_key"),
		"gooseai":    c.FormValue("goose_key"),
		"google":     c.FormValue("google_key"),
		"nim":        c.FormValue("nim_key"),
		"perplexity": c.FormValue("perplexity_key"),
		"fireworks":  c.FormValue("fireworks_key"),
		"together":   c.FormValue("together_key"),
		"deepseek":   c.FormValue("deepseek_key"),
	}

	testFunctions := map[string]func(string) bool{
		"openai":     TestOpenaiKey,
		"anthropic":  TestAnthropicKey,
		"mistral":    TestMistralKey,
		"groq":       TestGroqKey,
		"google":     TestGoogleKey,
		"nim":        TestNimKey,
		"perplexity": TestPerplexityKey,
		"fireworks":  TestFireworkKey,
		"together":   TestTogetherKey,
		"deepseek":   TestDeepseekKey,
	}

	openaiExists, anthropicExists, mistralExists, groqExists, googleExists, perplexityExists, fireworksExists, nimExists, togetherExists, deepseekExists := getExistingKeysNew(c)

	keyExists := map[string]bool{
		"openai":     openaiExists,
		"anthropic":  anthropicExists,
		"mistral":    mistralExists,
		"groq":       groqExists,
		"google":     googleExists,
		"nim":        nimExists,
		"perplexity": perplexityExists,
		"fireworks":  fireworksExists,
		"together":   togetherExists,
		"deepseek":   deepseekExists,
	}

	for company, key := range keys {
		if key != "" {
			if !testFunctions[company](key) {
				return c.SendString(fmt.Sprintf("Invalid %s API Key\n", company))
			}

			if keyExists[company] {
				err := edgeGlobalClient.WithGlobals(map[string]interface{}{"ext::auth::client_token": c.Cookies("jade-edgedb-auth-token")}).Execute(edgeCtx, `
          UPDATE Key
          FILTER .company.name = <str>$0 AND .<keys[IS Setting].<setting[IS User] = global currentUser
          SET {
              key := <str>$1,
          }
          `, company, key)
				if err != nil {
					fmt.Println("Error updating key")
					panic(err)
				}
			} else {
				err := edgeGlobalClient.WithGlobals(map[string]interface{}{"ext::auth::client_token": c.Cookies("jade-edgedb-auth-token")}).Execute(edgeCtx, `
                WITH
                    c := (SELECT Company FILTER .name = <str>$0 LIMIT 1)
                UPDATE global currentUser.setting
                SET {
                    keys += (
                        INSERT Key {
                            company := c,
                            key := <str>$1,
                            name := <str>$2 ++ " API Key",
                        }
                    )
                }`, company, key, company)
				if err != nil {
					fmt.Println("Error adding key")
					panic(err)
				}
			}
		}
	}

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