Jade/database.go

package main

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

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

var edgeCtx context.Context
var edgeClient *edgedb.Client

type User struct {
	ID      edgedb.UUID `edgedb:"id"`
	Setting Setting     `edgedb:"setting"`
}

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

type Setting struct {
	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"`
	Date time.Time   `edgedb:"date"`
	User User        `edgedb:"user"`
}

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

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

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

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

	edgeCtx = ctx
	edgeClient = client
}

func getLastArea() edgedb.UUID {
	var inserted struct{ id edgedb.UUID }
	err := edgeClient.QuerySingle(edgeCtx, `
		select Area
		filter .conversation.name = 'Default' AND .conversation.user = global currentUser
		order by .position desc
		limit 1
	`, &inserted)
	if err != nil {
		fmt.Println("Error in edgedb.QuerySingle: in getLastArea")
		log.Fatal(err)
	}
	return inserted.id
}

func getCurrentUser() User {
	var result User
	err := edgeClient.QuerySingle(edgeCtx, "SELECT global currentUser LIMIT 1;", &result)
	if err != nil {
		fmt.Println("Error in edgedb.QuerySingle: in getCurrentUser")
		fmt.Println(err)
		return User{}
	}
	return result
}

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

func insertNewConversation() edgedb.UUID {
	var inserted struct{ id edgedb.UUID }
	err := edgeClient.QuerySingle(edgeCtx, `
		INSERT Conversation {
			name := 'Default',
			user := global currentUser
		}
	`, &inserted)
	if err != nil {
		fmt.Println("Error in edgedb.QuerySingle: in insertNewConversation")
		log.Fatal(err)
	}
	return inserted.id
}

func insertArea() edgedb.UUID {
	// If the Default conversation doesn't exist, create it.
	err := edgeClient.QuerySingle(edgeCtx, `
		SELECT Conversation
		FILTER .name = 'Default' AND .user = global currentUser
		LIMIT 1
	`, nil)
	if err != nil {
		insertNewConversation()
	}
	// Insert a new area.
	var inserted struct{ id edgedb.UUID }
	err = edgeClient.QuerySingle(edgeCtx, `
	WITH
		positionVar := count((SELECT Area FILTER .conversation.name = 'Default' AND .conversation.user = global currentUser)) + 1
	INSERT Area {
			position := positionVar,
			conversation := (
				SELECT Conversation
				FILTER .name = 'Default' AND .user = global currentUser
				LIMIT 1
			)
		}
	`, &inserted)
	if err != nil {
		fmt.Println("Error in edgedb.QuerySingle: in insertArea")
		log.Fatal(err)
	}

	return inserted.id
}

func insertUserMessage(content string) edgedb.UUID {
	// Insert a new user.
	lastAreaID := getLastArea()
	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 := (
				SELECT Conversation
				FILTER .name = 'Default' AND .user = global currentUser
				LIMIT 1
			)
        }
    `, &inserted, "user", content, lastAreaID)
	if err != nil {
		fmt.Println("Error in edgedb.QuerySingle: in insertUserMessage")
		log.Fatal(err)
	}
	return inserted.id
}

func insertBotMessage(content string, selected bool, model string) edgedb.UUID {
	lastAreaID := getLastArea()
	var inserted struct{ id edgedb.UUID }
	err := edgeClient.QuerySingle(edgeCtx, `
		INSERT Message {
			role := <str>$0,
			model_id := <str>$1,
			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
			)
		}
	`, &inserted, "bot", model, content, selected, lastAreaID)
	if err != nil {
		fmt.Println("Error in edgedb.QuerySingle: in insertBotMessage")
		log.Fatal(err)
	}
	return inserted.id
}

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

	var messages []Message

	err := edgeClient.Query(edgeCtx, `
	SELECT Message {
		id,
		model_id,
		selected,
		role,
		content,
		date
	  } FILTER .conversation.name = 'Default' AND .conversation.user = global currentUser
	  ORDER BY .date ASC
    `, &messages)
	if err != nil {
		fmt.Println("Error in edgedb.Query: in getAllMessages")
		fmt.Println(err)
	}

	return messages
}

func getCurrentUserKeys() []Key {
	var result []Key
	err := edgeClient.Query(edgeCtx, "SELECT global currentUser.setting.keys", &result)
	if err != nil {
		fmt.Println("Error in edgedb.Query: in getCurrentUserKeys")
		fmt.Println(err)
	}
	return result
}

func checkIfHaveKey() bool {
	keys := getCurrentUserKeys()
	return keys != nil && len(keys) > 0
}