Jade/database.go

// Here are all type that I use with EdgeDB. I really love how this works together.
// The functions are from the very beginning. Of the app, I will remove it I think.
// I need to think of a better way to organize that.
package main

import (
	"context"
	"time"

	"github.com/edgedb/edgedb-go"
)

// So I have one client and one context for the database. All query use the same and it work well so far.
var edgeCtx context.Context
var edgeClient *edgedb.Client

// I will not put a comment on all type, I think they are self-explaining.
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"`
	Position int32       `edgedb:"position"`
	Date     time.Time   `edgedb:"date"`
	User     User        `edgedb:"user"`
}

// An area is in between messages and conversation.
// In a normal chat, you have a list of message, easy. By here you need to add, kind of a new dimension.
// All "message" can have multiple messages. So I created a new type named Area.
// A conversation is a list of Area and an Area is a list of Message. Easy enough.
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"` // Selected can also be seen as "Active". This is the message that will be use for the request.
	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"`
}

// A LLM is a bad name but I like it.
// It is more one instance of a model with it parameters.
// Maybe I will add more options later.
type LLM struct {
	ID          edgedb.UUID    `edgedb:"id"`
	Name        string         `edgedb:"name"`
	Context     string         `edgedb:"context"`
	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"`
	MaxToken    int32       `edgedb:"maxToken"`
	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"`
}

func init() {
	var ctx = context.Background()
	client, err := edgedb.CreateClient(ctx, edgedb.Options{})
	if err != nil {
		panic(err)
	}

	edgeCtx = ctx
	edgeClient = client
}

func checkIfLogin() bool {
	var result User
	err := edgeClient.QuerySingle(edgeCtx, "SELECT global currentUser LIMIT 1;", &result)
	return err == nil
}

func insertArea() (edgedb.UUID, int64) {
	// Insert a new area.
	var inserted struct{ id edgedb.UUID }
	err := edgeClient.QuerySingle(edgeCtx, `
	WITH
		positionVar := count((SELECT Area FILTER .conversation = global currentConversation)) + 1
	INSERT Area {
			position := positionVar,
			conversation := global currentConversation
		}
	`, &inserted)
	if err != nil {
		panic(err)
	}

	var positionSet struct{ position int64 }
	err = edgeClient.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(content string) edgedb.UUID {
	// Insert a new user.
	var lastArea Area
	err := edgeClient.QuerySingle(edgeCtx, `
		SELECT Area {
			id
		}
		FILTER .conversation = global currentConversation
		ORDER BY .position DESC
		LIMIT 1
	`, &lastArea)
	if err != nil {
		panic(err)
	}

	var inserted struct{ id edgedb.UUID }
	err = edgeClient.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>"a32c43ec-12fc-11ef-9dc9-b38e0de8bff0"
			)
        }
    `, &inserted, "user", content, lastArea.ID)
	if err != nil {
		panic(err)
	}
	return inserted.id
}

func insertBotMessage(content string, selected bool, llmUUID edgedb.UUID) edgedb.UUID {
	var lastArea Area
	err := edgeClient.QuerySingle(edgeCtx, `
		SELECT Area {
			id
		}
		FILTER .conversation = global currentConversation AND .conversation.user = global currentUser
		ORDER BY .position desc
		LIMIT 1
	`, &lastArea)
	if err != nil {
		panic(err)
	}

	var inserted struct{ id edgedb.UUID }
	err = edgeClient.QuerySingle(edgeCtx, `
		INSERT Message {
			role := <str>$0,
			content := <str>$2,
			selected := <bool>$3,
			conversation := (
				SELECT Area
				FILTER .id = <uuid>$4
				LIMIT 1
			).conversation,
			area := (
				SELECT Area
				FILTER .id = <uuid>$4
				LIMIT 1
			),
			llm := (
				SELECT LLM
				FILTER .id = <uuid>$1
				LIMIT 1
			)
		}
	`, &inserted, "bot", llmUUID, content, selected, lastArea.ID)
	if err != nil {
		panic(err)
	}
	return inserted.id
}

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

	var messages []Message

	err := edgeClient.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 {
		panic(err)
	}

	return messages
}

func checkIfHaveKey() bool {
	var keys []Key
	err := edgeClient.Query(edgeCtx, "SELECT global currentUser.setting.keys", &keys)
	if err != nil {
		panic(err)
	}
	return len(keys) > 0
}