Jade/main.go

package main

import (
	"fmt"
	"log"
	"os"
	"strings"
	"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))
	app.Get("/errorLogs", ErrorCodeLogsHandler)

	// Update models routes
	app.Get("/models/together", UpdateTogetherModelsHandler)

	// 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"),
	}

	requestFunctions := map[string]func(*fiber.Ctx, LLM, []Message, string) string{
		"openai":     RequestOpenai,
		"anthropic":  RequestAnthropic,
		"mistral":    RequestMistral,
		"groq":       RequestGroq,
		"google":     RequestGoogle,
		"nim":        RequestNim,
		"perplexity": RequestPerplexity,
		"fireworks":  RequestFirework,
		"together":   RequestTogether,
		"deepseek":   RequestDeepseek,
	}

	testModels := map[string]string{
		"openai":     "gpt-3.5-turbo",
		"anthropic":  "claude-3-haiku-20240307",
		"mistral":    "open-mistral-7b",
		"groq":       "llama3-8b-8192",
		"google":     "gemini-pro",
		"nim":        "meta/llama3-8b-instruct",
		"perplexity": "llama-3-8b-instruct",
		"fireworks":  "accounts/fireworks/models/llama-v2-7b-chat",
		"together":   "meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo",
		"deepseek":   "deepseek-chat",
	}

	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,
	}

	var messages = []Message{
		{
			Content: "You are useful",
			Role:    "context",
		},
		{
			Content: "hello",
			Role:    "user",
		},
	}

	for company, key := range keys {
		if key != "" {
			var llm LLM = LLM{
				Model:       ModelInfo{ModelID: testModels[company]},
				Temperature: 0,
				MaxToken:    10,
				Context:     "",
			}

			var responseText string = requestFunctions[company](c, llm, messages, key)

			if responseText == "" || strings.Contains(responseText, "JADE internal error") || strings.Contains(responseText, "Provider error") {
				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>")
}