Jade/Chat.go

package main

import (
	"context"
	"fmt"
	"log"
	"sort"
	"time"

	"github.com/edgedb/edgedb-go"
	"github.com/flosch/pongo2"
	"github.com/gofiber/fiber/v2"
)

func ChatPageHandler(c *fiber.Ctx) error {
	authCookie := c.Cookies("jade-edgedb-auth-token", "")

	if authCookie != "" && !checkIfLogin() {
		edgeClient = edgeClient.WithGlobals(map[string]interface{}{"ext::auth::client_token": authCookie})
	}

	return c.Render("chat", fiber.Map{"IsLogin": checkIfLogin(), "HaveKey": checkIfHaveKey()}, "layouts/main")
}

func DeleteMessageHandler(c *fiber.Ctx) error {
	messageId := c.FormValue("id")

	messageUUID, err := edgedb.ParseUUID(messageId)
	if err != nil {
		fmt.Println("Error in uuid.FromString: in DeleteMessageHandler")
		fmt.Println(err)
	}

	// Delete all messages
	err = edgeClient.Execute(edgeCtx, `
	WITH
		messageArea := (SELECT Message FILTER .id = <uuid>$0).area
	DELETE Area
		FILTER .position >= messageArea.position AND .conversation = messageArea.conversation;
	`, messageUUID)
	if err != nil {
		fmt.Println("Error in edgeClient.Execute: in DeleteMessageHandler")
		fmt.Println(err)
	}

	return c.SendString(generateChatHTML())
}

func LoadChatHandler(c *fiber.Ctx) error {
	if checkIfLogin() {
		if getCurrentUserKeys() == nil {
			return c.SendString(generateEnterKeyChatHTML())
		}
		return c.SendString(generateChatHTML())
	} else {
		return c.SendString(generateWelcomeChatHTML())
	}
}

type TemplateMessage struct {
	Icon    string
	Content string
	Hidden  bool
	Id      string
	Name    string
	ModelID string
}

func generateChatHTML() string {
	// Get the messages from the database
	// Maybe redo that to be area by area because look like shit rn. It come from early stage of dev. It work tho soooo...
	Messages := getAllMessages()

	htmlString := "<div class='columns is-centered' id='chat-container'><div class='column is-12-mobile is-8-tablet is-6-desktop' id='chat-messages'>"

	var templateMessages []TemplateMessage

	for i, message := range Messages {
		if message.Role == "user" {
			htmlContent := markdownToHTML(message.Content)
			userOut, err := userTmpl.Execute(pongo2.Context{"Content": htmlContent, "ID": message.ID.String()})
			if err != nil {
				panic(err)
			}
			htmlString += userOut
			// Reset NextMessages when a user message is encountered
			templateMessages = []TemplateMessage{}
		} else {
			// For bot messages, add them to NextMessages with only the needed fields
			templateMessage := TemplateMessage{
				Icon:    message.LLM.Model.Company.Icon, // Assuming Icon is a field you want to include from Message
				Content: markdownToHTML(message.Content),
				Hidden:  !message.Selected, // Assuming Hidden is a field you want to include from Message
				Id:      message.ID.String(),
				Name:    message.LLM.Model.Name,
				ModelID: message.LLM.Model.ModelID,
			}
			templateMessages = append(templateMessages, templateMessage)

			// Check if the next message is not a bot or if it's the last message
			if i+1 == len(Messages) || Messages[i+1].Role != "bot" {
				sort.Slice(templateMessages, func(i, j int) bool {
					if !templateMessages[i].Hidden && templateMessages[j].Hidden {
						return true
					}
					if templateMessages[i].Hidden && !templateMessages[j].Hidden {
						return false
					}
					return true
				})

				botOut, err := botTmpl.Execute(pongo2.Context{"Messages": templateMessages, "ConversationAreaId": i})
				if err != nil {
					panic(err)
				}
				htmlString += botOut
				htmlString += "<div style='height: 10px;'></div>"
			}
		}
	}

	htmlString += "</div></div>"

	// Render the HTML template with the messages
	return htmlString
}

func GetMessageContentHandler(c *fiber.Ctx) error {
	messageId := c.FormValue("id")
	onlyContent := c.FormValue("onlyContent") // To init the text area of the edit message form

	messageUUID, _ := edgedb.ParseUUID(messageId)

	var selectedMessage Message
	err := edgeClient.QuerySingle(context.Background(), `
		SELECT Message {
			content,
			llm : {
				modelInfo : {
					modelID,
					name,
				}
			}
		}
		FILTER
			.id = <uuid>$0;
	`, &selectedMessage, messageUUID)
	if err != nil {
		log.Fatal(err)
	}

	if onlyContent == "true" {
		return c.SendString(markdownToHTML(selectedMessage.Content))
	}

	out := "<div class='message-header'>"
	out += "<p>"
	out += selectedMessage.LLM.Model.Name
	out += "	</p>"
	out += "</div>"
	out += "<div class='message-body'>"
	out += "	<ct class='content'>"
	out += markdownToHTML(selectedMessage.Content)
	out += "	</ct>"
	out += "</div>"

	// Update the selected value of messages in the database
	err = edgeClient.Execute(edgeCtx, `
	WITH m := (SELECT Message FILTER .id = <uuid>$0)
	UPDATE Message
	FILTER .area = m.area
	SET {selected := false};
	`, messageUUID)
	if err != nil {
		log.Fatal(err)
	}

	_ = edgeClient.Execute(edgeCtx, `
	UPDATE Message
	FILTER .id = <uuid>$0
	SET {selected := true};
	`, messageUUID)

	return c.SendString(out)
}

func generateWelcomeChatHTML() string {
	welcomeMessage := `To start using JADE, please login.`

	loginButton := `
	<a class="button is-primary is-small" href="/signin">
		Log in
	</a>`

	htmlString := "<div class='columns is-centered' id='chat-container'><div class='column is-12-mobile is-8-tablet is-6-desktop' id='chat-messages'>"

	NextMessages := []TemplateMessage{}
	nextMsg := TemplateMessage{
		Icon:    "icons/bouvai2.png", // Assuming Icon is a field you want to include from Message
		Content: "<br>" + markdownToHTML(welcomeMessage) + loginButton,
		Hidden:  false, // Assuming Hidden is a field you want to include from Message
		Id:      "0",
		Name:    "JADE",
	}
	NextMessages = append(NextMessages, nextMsg)

	botOut, err := botTmpl.Execute(pongo2.Context{"Messages": NextMessages, "ConversationAreaId": 0, "NotClickable": true})
	if err != nil {
		panic(err)
	}
	htmlString += botOut
	htmlString += "<div style='height: 10px;'></div>"
	htmlString += "</div></div>"

	// Render the HTML template with the messages
	return htmlString
}

func generateEnterKeyChatHTML() string {
	welcomeMessage := `To start using JADE, please enter at least one key in the settings.`

	htmlString := "<div class='columns is-centered' id='chat-container'><div class='column is-12-mobile is-8-tablet is-6-desktop' id='chat-messages'>"

	NextMessages := []TemplateMessage{}
	nextMsg := TemplateMessage{
		Icon:    "icons/bouvai2.png", // Assuming Icon is a field you want to include from Message
		Content: "<br>" + markdownToHTML(welcomeMessage),
		Hidden:  false, // Assuming Hidden is a field you want to include from Message
		Id:      "0",
		Name:    "JADE",
	}
	NextMessages = append(NextMessages, nextMsg)

	botOut, err := botTmpl.Execute(pongo2.Context{"Messages": NextMessages, "ConversationAreaId": 0, "NotClickable": true})
	if err != nil {
		panic(err)
	}
	htmlString += botOut
	htmlString += "<div style='height: 10px;'></div>"
	htmlString += "</div></div>"

	// Render the HTML template with the messages
	return htmlString
}

// Buton actions
func GetEditMessageFormHandler(c *fiber.Ctx) error {
	id := c.FormValue("id")
	idUUID, _ := edgedb.ParseUUID(id)

	var message Message
	err := edgeClient.QuerySingle(context.Background(), `
		SELECT Message { content }
		FILTER .id = <uuid>$0;
	`, &message, idUUID)
	if err != nil {
		log.Fatal(err)
	}

	tmpl := pongo2.Must(pongo2.FromFile("views/partials/edit-message-form.html"))
	out, err := tmpl.Execute(pongo2.Context{"Content": message.Content, "ID": id})
	if err != nil {
		panic(err)
	}
	return c.SendString(out)
}

func RedoMessageHandler(c *fiber.Ctx) error {
	messageId := c.FormValue("id")
	messageUUID, _ := edgedb.ParseUUID(messageId)

	// Delete messages
	err := edgeClient.Execute(edgeCtx, `
	WITH
		messageArea := (SELECT Message FILTER .id = <uuid>$0).area
	DELETE Area
		FILTER .position > messageArea.position AND .conversation = messageArea.conversation;
	`, messageUUID)
	if err != nil {
		fmt.Println("Error in edgeClient.Execute: in DeleteMessageHandler")
		fmt.Println(err)
	}

	messageOut, err := botTmpl.Execute(pongo2.Context{"IsPlaceholder": true})
	if err != nil {
		c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{
			"error": "Error rendering template",
		})
	}

	return c.SendString(messageOut)
}

func ClearChatHandler(c *fiber.Ctx) error {
	// Delete the default conversation
	err := edgeClient.Execute(edgeCtx, `
	DELETE Conversation
	FILTER .user = global currentUser AND .name = "Default";
	`)
	if err != nil {
		fmt.Println("Error in edgeClient.Execute: in ClearChatHandler")
		fmt.Println(err)
	}

	return c.SendString(generateChatHTML())
}

// Popover stuff
func LoadUsageKPIHandler(c *fiber.Ctx) error {
	if !checkIfLogin() {
		return c.SendString("")
	}
	var TotalUsage float32
	// Using the database. Get the sum of all usage.inputCost and outputCost
	err := edgeClient.QuerySingle(edgeCtx, `
		WITH
			U := (
				SELECT Usage
				FILTER .user = global currentUser
			)
		SELECT sum(U.input_cost) + sum(U.output_cost)
	`, &TotalUsage)
	if err != nil {
		fmt.Println("Error in edgedb.QuerySingle: in LoadUsageKPIHandler")
		log.Fatal(err)
	}

	now := time.Now()

	var TodayUsage float32
	// Using the database. Get the sum of all usage.inputCost and outputCost
	err = edgeClient.QuerySingle(edgeCtx, `
		WITH
			U := (
				SELECT Usage
				FILTER .user = global currentUser AND .date >= <datetime>$0
			)
		SELECT sum(U.input_cost) + sum(U.output_cost)
	`, &TodayUsage, now.Add(time.Hour*-24))
	if err != nil {
		fmt.Println("Error in edgedb.QuerySingle: in LoadUsageKPIHandler")
		log.Fatal(err)
	}

	var WeekUsage float32
	// Using the database. Get the sum of all usage.inputCost and outputCost
	edgeClient.QuerySingle(edgeCtx, `
		WITH
			U := (
				SELECT Usage
				FILTER .user = global currentUser AND .date >= <datetime>$0
			)
		SELECT sum(U.input_cost) + sum(U.output_cost)
	`, &WeekUsage, now.Add(time.Hour*-24*7))

	var MonthUsage float32
	// Using the database. Get the sum of all usage.inputCost and outputCost
	edgeClient.QuerySingle(edgeCtx, `
		WITH
			U := (
				SELECT Usage
				FILTER .user = global currentUser AND .date >= <datetime>$0
			)
		SELECT sum(U.input_cost) + sum(U.output_cost)
	`, &MonthUsage, now.Add(time.Hour*-24*30))

	out, err := pongo2.Must(pongo2.FromFile("views/partials/popover-usage.html")).Execute(pongo2.Context{
		"TotalUsage": TotalUsage, "TodayUsage": TodayUsage, "WeekUsage": WeekUsage, "MonthUsage": MonthUsage,
	})
	if err != nil {
		c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{
			"error": "Error rendering template",
		})
	}
	return c.SendString(out)
}

func LoadKeysHandler(c *fiber.Ctx) error {
	if !checkIfLogin() {
		return c.SendString("")
	}
	openaiExists, anthropicExists, mistralExists, groqExists, gooseaiExists, googleExists := getExistingKeys()

	out, err := pongo2.Must(pongo2.FromFile("views/partials/popover-keys.html")).Execute(pongo2.Context{
		"IsLogin":         checkIfLogin(),
		"OpenaiExists":    openaiExists,
		"AnthropicExists": anthropicExists,
		"MistralExists":   mistralExists,
		"GroqExists":      groqExists,
		"GooseaiExists":   gooseaiExists,
		"GoogleExists":    googleExists,
		"AnyExists":       openaiExists || anthropicExists || mistralExists || groqExists || gooseaiExists || googleExists,
	})
	if err != nil {
		c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{
			"error": "Error rendering template",
		})
	}
	return c.SendString(out)
}

func LoadModelSelectionHandler(c *fiber.Ctx) error {
	if !checkIfLogin() {
		return c.SendString("")
	}
	// openaiExists, anthropicExists, mistralExists, groqExists := getExistingKeys()
	// TODO: Add model selection
	return c.SendString("")
}

func LoadSettingsHandler(c *fiber.Ctx) error {
	if !checkIfLogin() {
		return c.SendString("")
	}
	out, err := pongo2.Must(pongo2.FromFile("views/partials/popover-settings.html")).Execute(pongo2.Context{"IsLogin": checkIfLogin()})
	if err != nil {
		c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{
			"error": "Error rendering template",
		})
	}
	return c.SendString(out)
}