Jade/EdgeDatabase.go

package main

import (
	"context"
	"fmt"
	"time"

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

type Identity struct {
	ID     edgedb.UUID `edgedb:"id"`
	Issuer string      `edgedb:"issuer"`
}

type User struct {
	ID       edgedb.UUID `edgedb:"id"`
	Setting  Setting     `edgedb:"setting"`
	StripeID string      `edgedb:"stripe_id"`
	Email    string      `edgedb:"email"`
	Name     string      `edgedb:"name"`
	Avatar   string      `edgedb:"avatar"`
}

type Key struct { // API key
	ID      edgedb.UUID `edgedb:"id"`
	Name    string      `edgedb:"name"`
	Company CompanyInfo `edgedb:"company"`
	Key     string      `edgedb:"key"`
	Date    time.Time   `edgedb:"date"`
}

type Setting struct { // Per user
	ID           edgedb.UUID        `edgedb:"id"`
	Keys         []Key              `edgedb:"keys"`
	DefaultModel edgedb.OptionalStr `edgedb:"default_model"`
}

type Conversation struct {
	ID       edgedb.UUID `edgedb:"id"`
	Name     string      `edgedb:"name"`
	Selected bool        `edgedb:"selected"`
	Position int32       `edgedb:"position"`
	Date     time.Time   `edgedb:"date"`
	User     User        `edgedb:"user"`
}

type Area struct {
	ID       edgedb.UUID  `edgedb:"id"`
	Position int64        `edgedb:"position"`
	Conv     Conversation `edgedb:"conversation"`
}

type Message struct {
	ID       edgedb.UUID  `edgedb:"id"`
	Content  string       `edgedb:"content"`
	Role     string       `edgedb:"role"`
	Selected bool         `edgedb:"selected"`
	Date     time.Time    `edgedb:"date"`
	Area     Area         `edgedb:"area"`
	Conv     Conversation `edgedb:"conversation"`
	LLM      LLM          `edgedb:"llm"`
}

type Usage struct {
	ID          edgedb.UUID `edgedb:"id"`
	ModelID     string      `edgedb:"model_id"`
	Date        time.Time   `edgedb:"date"`
	InputCost   float32     `edgedb:"input_cost"`
	OutputCost  float32     `edgedb:"output_cost"`
	InputToken  int32       `edgedb:"input_token"`
	OutputToken int32       `edgedb:"output_token"`
}

type LLM struct {
	ID          edgedb.UUID    `edgedb:"id"`
	Name        string         `edgedb:"name"`
	Context     string         `edgedb:"context"`
	MaxToken    int32          `edgedb:"max_tokens"`
	Temperature float32        `edgedb:"temperature"`
	Model       ModelInfo      `edgedb:"modelInfo"`
	Endpoint    CustomEndpoint `edgedb:"custom_endpoint"`
}

type CustomEndpoint struct {
	edgedb.Optional
	ID       edgedb.UUID `edgedb:"id"`
	Endpoint string      `edgedb:"endpoint"`
	Key      string      `edgedb:"key"`
}

type ModelInfo struct {
	ID          edgedb.UUID `edgedb:"id"`
	Name        string      `edgedb:"name"`
	InputPrice  float32     `edgedb:"inputPrice"`
	OutputPrice float32     `edgedb:"outputPrice"`
	ModelID     string      `edgedb:"modelID"`
	Company     CompanyInfo `edgedb:"company"`
}

type CompanyInfo struct {
	ID   edgedb.UUID `edgedb:"id"`
	Name string      `edgedb:"name"`
	Icon string      `edgedb:"icon"`
}

var edgeCtx context.Context
var edgeGlobalClient *edgedb.Client

func init() {
	var ctx = context.Background()
	client, err := edgedb.CreateClient(ctx, edgedb.Options{})
	if err != nil {
		fmt.Println("Error connecting to edgedb")
		panic(err)
	}

	edgeCtx = ctx
	edgeGlobalClient = client
}

func checkIfLogin(c *fiber.Ctx) bool {
	return c.Cookies("jade-edgedb-auth-token") != "" && c.Cookies("jade-edgedb-auth-token") != "null"
}

func insertArea(c *fiber.Ctx) (edgedb.UUID, int64) {
	// Insert a new area.
	var inserted struct{ id edgedb.UUID }
	err := edgeGlobalClient.WithGlobals(map[string]interface{}{"ext::auth::client_token": c.Cookies("jade-edgedb-auth-token")}).QuerySingle(edgeCtx, `
	WITH
		positionVar := count((SELECT Area FILTER .conversation = global currentConversation)) + 1
	INSERT Area {
			position := positionVar,
			conversation := global currentConversation
		}
	`, &inserted)
	if err != nil {
		fmt.Println("Error inserting area")
		panic(err)
	}

	var positionSet struct{ position int64 }
	err = edgeGlobalClient.WithGlobals(map[string]interface{}{"ext::auth::client_token": c.Cookies("jade-edgedb-auth-token")}).QuerySingle(edgeCtx, `
		SELECT Area {
			position,
		}
		FILTER .conversation = global currentConversation
		ORDER BY .position desc
		LIMIT 1
	`, &positionSet)
	if err != nil {
		panic(err)
	}

	return inserted.id, positionSet.position
}

func insertUserMessage(c *fiber.Ctx, content string) edgedb.UUID {
	// Insert a new user.
	var lastArea Area
	err := edgeGlobalClient.WithGlobals(map[string]interface{}{"ext::auth::client_token": c.Cookies("jade-edgedb-auth-token")}).QuerySingle(edgeCtx, `
		SELECT Area {
			id
		}
		FILTER .conversation = global currentConversation
		ORDER BY .position DESC
		LIMIT 1
	`, &lastArea)
	if err != nil {
		fmt.Println("Error getting last area")
		panic(err)
	}

	var inserted struct{ id edgedb.UUID } // TODO: Maybe remove that hard code LLM
	err = edgeGlobalClient.WithGlobals(map[string]interface{}{"ext::auth::client_token": c.Cookies("jade-edgedb-auth-token")}).QuerySingle(edgeCtx, `
        INSERT Message {
			role := <str>$0,
			content := <str>$1,
			area := (
				SELECT Area
				FILTER .id = <uuid>$2
			),
			conversation := global currentConversation,
			selected := true,
			llm := (
				SELECT LLM FILTER .id = <uuid>"f777f024-1da9-11ef-afd4-bf1c8b177a62"
			)
        }
    `, &inserted, "user", content, lastArea.ID)
	if err != nil {
		fmt.Println("Error inserting user message")
		panic(err)
	}
	return inserted.id
}

func insertBotMessage(c *fiber.Ctx, content string, llmUUID edgedb.UUID) edgedb.UUID {
	var lastArea Area
	err := edgeGlobalClient.WithGlobals(map[string]interface{}{"ext::auth::client_token": c.Cookies("jade-edgedb-auth-token")}).QuerySingle(edgeCtx, `
		SELECT Area {
			id
		}
		FILTER .conversation = global currentConversation AND .conversation.user = global currentUser
		ORDER BY .position desc
		LIMIT 1
	`, &lastArea)
	if err != nil {
		fmt.Println("Error getting last area")
		panic(err)
	}

	var inserted struct{ id edgedb.UUID }
	err = edgeGlobalClient.WithGlobals(map[string]interface{}{"ext::auth::client_token": c.Cookies("jade-edgedb-auth-token")}).QuerySingle(edgeCtx, `
		INSERT Message {
			role := <str>$0,
			content := <str>$2,
			selected := false,
			conversation := (
				SELECT Area
				FILTER .id = <uuid>$3
				LIMIT 1
			).conversation,
			area := (
				SELECT Area
				FILTER .id = <uuid>$3
				LIMIT 1
			),
			llm := (
				SELECT LLM
				FILTER .id = <uuid>$1
				LIMIT 1
			)
		}
	`, &inserted, "bot", llmUUID, content, lastArea.ID)
	if err != nil {
		fmt.Println("Error inserting bot message")
		panic(err)
	}
	return inserted.id
}

func getAllSelectedMessages(c *fiber.Ctx) []Message {
	// If no CurrentUser, return an empty array
	if !checkIfLogin(c) {
		return []Message{}
	}

	var messages []Message

	err := edgeGlobalClient.WithGlobals(map[string]interface{}{"ext::auth::client_token": c.Cookies("jade-edgedb-auth-token")}).Query(edgeCtx, `
	SELECT Message {
		id,
		selected,
		role,
		content,
		date,
		llm : {
			modelInfo : {
				modelID,
				company : {
					icon
				}
			}
		}
	  } FILTER .conversation = global currentConversation AND .conversation.user = global currentUser AND .selected = true
	  ORDER BY .date ASC
    `, &messages)
	if err != nil {
		fmt.Println("Error getting messages")
		panic(err)
	}

	return messages
}

func checkIfHaveKey(c *fiber.Ctx) bool {
	var keys []Key
	err := edgeGlobalClient.WithGlobals(map[string]interface{}{"ext::auth::client_token": c.Cookies("jade-edgedb-auth-token")}).Query(edgeCtx,
		"SELECT global currentUser.setting.keys", &keys)
	if err != nil {
		panic(err)
	}
	return len(keys) > 0
}