Jade/database.go

package main

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

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

var edgeCtx context.Context
var edgeClient *edgedb.Client
var CurrentUser User

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

	// Get the user test@example.com
	var user User
	err = client.QuerySingle(ctx, `
		SELECT User { id, email, name, setting } FILTER .email = <str>$0 LIMIT 1
	`, &user, "test@example.com")
	if err != nil {
		fmt.Println("Error in edgedb.QuerySingle: in init")
		log.Fatal(err)
	}
	CurrentUser = user

	fmt.Print("Current User: ")
	fmt.Println(CurrentUser)

	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.id = <uuid>$0
		order by .position desc
		limit 1
	`, &inserted, CurrentUser.ID)
	if err != nil {
		fmt.Println("Error in edgedb.QuerySingle: in getLastArea")
		log.Fatal(err)
	}
	return inserted.id
}

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.id = <uuid>$0)) + 1
	INSERT Area {
			position := positionVar,
			conversation := (
				SELECT Conversation
				FILTER .name = 'Default' AND .user.id = <uuid>$0
				LIMIT 1
			)
		}
	`, &inserted, CurrentUser.ID)
	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.id = <uuid>$3
				LIMIT 1
			)
        }
    `, &inserted, "user", content, lastAreaID, CurrentUser.ID)
	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.id = <uuid>$4
				LIMIT 1
			),
			area := (
				SELECT Area
				FILTER .id = <uuid>$5
				LIMIT 1
			)
		}
	`, &inserted, "bot", model, content, selected, CurrentUser.ID, lastAreaID)
	if err != nil {
		fmt.Println("Error in edgedb.QuerySingle: in insertBotMessage")
		log.Fatal(err)
	}
	return inserted.id
}

func getAllMessages() []Message {
	var messages []Message

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

	return messages
}