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"`
	Email         string         `edgedb:"email"`
	Name          string         `edgedb:"name"`
	Setting       Setting        `edgedb:"setting"`
	Conversations []Conversation `edgedb:"conversations"`
	Usages        []Usage        `edgedb:"usages"`
}

type Key struct {
	ID   edgedb.UUID `edgedb:"id"`
	Name string      `edgedb:"name"`
	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"`
	Areas []Area      `edgedb:"areas"`
}

type Area struct {
	ID       edgedb.UUID `edgedb:"id"`
	Position int         `edgedb:"position"`
	Messages []Message   `edgedb:"messages"`
}

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

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

	// TODO Change
	edgeCtx = ctx
	var clientUUID edgedb.UUID
	clientUUID, err = edgedb.ParseUUID("9323365e-0b09-11ef-8f41-c3575d386283")
	if err != nil {
		fmt.Println("Error in edgedb.ParseUUID: in init")
		log.Fatal(err)
	}
	edgeClient = client.WithGlobals(map[string]interface{}{"current_user_id": clientUUID})
}

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 checkIfLogin() bool {
	var result User
	err := edgeClient.QuerySingle(edgeCtx, "SELECT global currentUser LIMIT 1;", &result)
	if err != nil {
		fmt.Println("Error in edgedb.QuerySingle: in checkIfLogin")
		fmt.Println(err)
		return false
	}
	return true
}

func insertArea() edgedb.UUID {
	// 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 Conversation
				FILTER .name = 'Default' AND .user = global currentUser
				LIMIT 1
			),
			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
}