353 lines
7.6 KiB
Go
353 lines
7.6 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
"time"
|
|
|
|
"github.com/edgedb/edgedb-go"
|
|
)
|
|
|
|
var edgeCtx context.Context
|
|
var edgeClient *edgedb.Client
|
|
|
|
type Identity struct {
|
|
ID edgedb.UUID `edgedb:"id"`
|
|
Issuer string `edgedb:"issuer"`
|
|
}
|
|
|
|
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 CompanyInfo `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 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"`
|
|
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"`
|
|
}
|
|
|
|
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 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 {
|
|
panic(err)
|
|
}
|
|
return inserted.id
|
|
}
|
|
|
|
func checkIfLogin() bool {
|
|
var result User
|
|
err := edgeClient.QuerySingle(edgeCtx, "SELECT global currentUser LIMIT 1;", &result)
|
|
return err == nil
|
|
}
|
|
|
|
func addUsage(inputCost float32, outputCost float32, inputToken int32, outputToken int32, modelID string) {
|
|
// Create a new usage
|
|
err := edgeClient.Execute(edgeCtx, `
|
|
INSERT Usage {
|
|
input_cost := <float32>$0,
|
|
output_cost := <float32>$1,
|
|
input_token := <int32>$2,
|
|
output_token := <int32>$3,
|
|
model_id := <str>$4,
|
|
user := global currentUser
|
|
}
|
|
`, inputCost, outputCost, inputToken, outputToken, modelID)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
}
|
|
|
|
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 {
|
|
panic(err)
|
|
}
|
|
return inserted.id
|
|
}
|
|
|
|
func insertArea() (edgedb.UUID, int64) {
|
|
// If the Default conversation doesn't exist, create it.
|
|
var convExists bool
|
|
edgeClient.QuerySingle(edgeCtx, `
|
|
select exists (
|
|
select Conversation
|
|
filter .name = 'Default' AND .user = global currentUser
|
|
);
|
|
`, &convExists)
|
|
if !convExists {
|
|
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 {
|
|
panic(err)
|
|
}
|
|
|
|
var positionSet struct{ position int64 }
|
|
err = edgeClient.QuerySingle(edgeCtx, `
|
|
SELECT Area {
|
|
position
|
|
}
|
|
FILTER .conversation.name = 'Default' AND .conversation.user = global currentUser
|
|
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.
|
|
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
|
|
),
|
|
selected := true,
|
|
llm := (
|
|
SELECT LLM FILTER .id = <uuid>"a32c43ec-12fc-11ef-9dc9-b38e0de8bff0"
|
|
)
|
|
}
|
|
`, &inserted, "user", content, lastAreaID)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
return inserted.id
|
|
}
|
|
|
|
func insertBotMessage(content string, selected bool, llmUUID edgedb.UUID) edgedb.UUID {
|
|
lastAreaID := getLastArea()
|
|
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, lastAreaID)
|
|
if err != nil {
|
|
panic(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,
|
|
selected,
|
|
role,
|
|
content,
|
|
date,
|
|
llm : {
|
|
name,
|
|
modelInfo : {
|
|
modelID,
|
|
name,
|
|
company : {
|
|
icon
|
|
}
|
|
}
|
|
}
|
|
}
|
|
FILTER .conversation.name = 'Default' AND .conversation.user = global currentUser
|
|
ORDER BY .date ASC
|
|
`, &messages)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
|
|
return messages
|
|
}
|
|
|
|
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.name = 'Default' AND .conversation.user = global currentUser AND .selected = true
|
|
ORDER BY .date ASC
|
|
`, &messages)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
|
|
return messages
|
|
}
|
|
|
|
func getCurrentUserKeys() []Key {
|
|
var result []Key
|
|
err := edgeClient.Query(edgeCtx, "SELECT global currentUser.setting.keys", &result)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
return result
|
|
}
|
|
|
|
func checkIfHaveKey() bool {
|
|
keys := getCurrentUserKeys()
|
|
return len(keys) > 0
|
|
}
|