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

	fmt.Println("Current User: ", getCurrentUser(), " - ", checkIfLogin())

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

func LoadModelSelectionHandler(c *fiber.Ctx) error {
	openaiExists, anthropicExists, mistralExists := getExistingKeys()

	var CompanyInfosAvailable []CompanyInfo

	if openaiExists {
		var openaiCompanyInfo CompanyInfo
		for _, info := range CompanyInfos {
			if info.ID == "openai" {
				openaiCompanyInfo = info
				break
			}
		}
		CompanyInfosAvailable = append(CompanyInfosAvailable, openaiCompanyInfo)
	}
	if anthropicExists {
		var anthropicCompanyInfo CompanyInfo
		for _, info := range CompanyInfos {
			if info.ID == "anthropic" {
				anthropicCompanyInfo = info
				break
			}
		}
		CompanyInfosAvailable = append(CompanyInfosAvailable, anthropicCompanyInfo)
	}
	if mistralExists {
		var mistralCompanyInfo CompanyInfo
		for _, info := range CompanyInfos {
			if info.ID == "mistral" {
				mistralCompanyInfo = info
				break
			}
		}
		CompanyInfosAvailable = append(CompanyInfosAvailable, mistralCompanyInfo)
	}

	CheckedModels := []string{"gpt-3.5-turbo"} // Default model
	out, err := pongo2.Must(pongo2.FromFile("views/partials/popover-models.html")).Execute(pongo2.Context{
		"CompanyInfos":  CompanyInfosAvailable,
		"CheckedModels": CheckedModels,
	})
	if err != nil {
		c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{
			"error": "Error rendering template",
		})
	}
	return c.SendString(out)
}

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

	fmt.Println("Deleting message: " + messageId)

	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 NextMessage struct {
	Icon    string
	Content string
	Hidden  bool
	Id      string
	Name    string
}

func generateChatHTML() string {
	// Get the messages from the database
	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 NextMessages []NextMessage // Use the custom NextMessage struct

	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
			NextMessages = []NextMessage{}
		} else {
			modelID, exist := message.ModelID.Get()
			if !exist {
				modelID = "gpt-3.5-turbo"
			}
			selected, exist := message.Selected.Get()
			if !exist {
				selected = false
			}
			// For bot messages, add them to NextMessages with only the needed fields
			nextMsg := NextMessage{
				Icon:    model2Icon(modelID), // Assuming Icon is a field you want to include from Message
				Content: markdownToHTML(message.Content),
				Hidden:  !selected, // Assuming Hidden is a field you want to include from Message
				Id:      message.ID.String(),
				Name:    model2Name(modelID),
			}
			NextMessages = append(NextMessages, nextMsg)

			// 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(NextMessages, func(i, j int) bool {
					if !NextMessages[i].Hidden && NextMessages[j].Hidden {
						return true
					}
					if NextMessages[i].Hidden && !NextMessages[j].Hidden {
						return false
					}
					return true
				})

				botOut, err := botTmpl.Execute(pongo2.Context{"Messages": NextMessages, "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")

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

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

	modelID, _ := selectedMessage.ModelID.Get()

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

	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 := []NextMessage{}
	nextMsg := NextMessage{
		Icon:    "bouvai2", // 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 := []NextMessage{}
	nextMsg := NextMessage{
		Icon:    "bouvai2", // 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
}

// Popover stuff

func LoadUsageKPIHandler(c *fiber.Ctx) error {
	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 LoadSettingsHandler(c *fiber.Ctx) error {
	openaiExists, anthropicExists, mistralExists := getExistingKeys()

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