297 lines
7.4 KiB
Go
297 lines
7.4 KiB
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"
|
|
"fmt"
|
|
"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 {
|
|
fmt.Println("Error connecting to edgedb")
|
|
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 {
|
|
fmt.Println("Error inserting area")
|
|
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 {
|
|
fmt.Println("Error getting last area")
|
|
panic(err)
|
|
}
|
|
|
|
var inserted struct{ id edgedb.UUID } // TODO: Maybe remove that hard code LLM
|
|
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>"32cb4f0c-1c3f-11ef-9247-2fac0086a7d1"
|
|
)
|
|
}
|
|
`, &inserted, "user", content, lastArea.ID)
|
|
if err != nil {
|
|
fmt.Println("Error inserting user message")
|
|
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 {
|
|
fmt.Println("Error getting last area")
|
|
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 {
|
|
fmt.Println("Error inserting bot message")
|
|
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 {
|
|
fmt.Println("Error getting messages")
|
|
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
|
|
}
|