Adding usage to db

This commit is contained in:
Adrien Bouvais 2024-05-05 21:33:46 +02:00
parent 8f28d07051
commit c363f12664
4 changed files with 49 additions and 20 deletions

View File

@ -1,6 +1,8 @@
package main package main
import ( import (
"fmt"
"log"
"sync" "sync"
"github.com/edgedb/edgedb-go" "github.com/edgedb/edgedb-go"
@ -13,8 +15,8 @@ type ModelInfo struct {
Name string Name string
Icon string Icon string
MaxToken int MaxToken int
InputPrice float64 InputPrice float32
OutputPrice float64 OutputPrice float32
} }
type CompanyInfo struct { type CompanyInfo struct {
@ -40,6 +42,24 @@ type BotContentMessage struct {
var lastSelectedModelIds []string var lastSelectedModelIds []string
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 {
fmt.Println("Error in edgedb.QuerySingle: in addUsage")
log.Fatal(err)
}
}
func GenerateMultipleMessages(c *fiber.Ctx) error { func GenerateMultipleMessages(c *fiber.Ctx) error {
// Create a wait group to synchronize the goroutines // Create a wait group to synchronize the goroutines
var wg sync.WaitGroup var wg sync.WaitGroup

View File

@ -36,8 +36,8 @@ type AnthropicContentItem struct {
} }
type AnthropicUsage struct { type AnthropicUsage struct {
InputTokens int `json:"input_tokens"` InputTokens int32 `json:"input_tokens"`
OutputTokens int `json:"output_tokens"` OutputTokens int32 `json:"output_tokens"`
} }
func init() { func init() {
@ -169,5 +169,15 @@ func RequestAnthropic(model string, messages []Message, maxTokens int, temperatu
return AnthropicChatCompletionResponse{}, fmt.Errorf("error unmarshaling JSON: %w", err) return AnthropicChatCompletionResponse{}, fmt.Errorf("error unmarshaling JSON: %w", err)
} }
var usedModelInfo ModelInfo
for mi := range ModelsInfos {
if ModelsInfos[mi].ID == model {
usedModelInfo = ModelsInfos[mi]
}
}
var inputCost float32 = float32(chatCompletionResponse.Usage.InputTokens) * usedModelInfo.InputPrice
var outputCost float32 = float32(chatCompletionResponse.Usage.OutputTokens) * usedModelInfo.OutputPrice
addUsage(inputCost, outputCost, chatCompletionResponse.Usage.InputTokens, chatCompletionResponse.Usage.OutputTokens, model)
return chatCompletionResponse, nil return chatCompletionResponse, nil
} }

View File

@ -31,9 +31,9 @@ type OpenaiChatCompletionResponse struct {
} }
type OpenaiUsage struct { type OpenaiUsage struct {
PromptTokens int `json:"prompt_tokens"` PromptTokens int32 `json:"prompt_tokens"`
CompletionTokens int `json:"completion_tokens"` CompletionTokens int32 `json:"completion_tokens"`
TotalTokens int `json:"total_tokens"` TotalTokens int32 `json:"total_tokens"`
} }
type OpenaiChoice struct { type OpenaiChoice struct {
@ -159,5 +159,15 @@ func RequestOpenai(model string, messages []Message, temperature float64) (Opena
return OpenaiChatCompletionResponse{}, fmt.Errorf("error unmarshaling JSON: %w", err) return OpenaiChatCompletionResponse{}, fmt.Errorf("error unmarshaling JSON: %w", err)
} }
var usedModelInfo ModelInfo
for mi := range ModelsInfos {
if ModelsInfos[mi].ID == model {
usedModelInfo = ModelsInfos[mi]
}
}
var inputCost float32 = float32(chatCompletionResponse.Usage.PromptTokens) * usedModelInfo.InputPrice
var outputCost float32 = float32(chatCompletionResponse.Usage.CompletionTokens) * usedModelInfo.OutputPrice
addUsage(inputCost, outputCost, chatCompletionResponse.Usage.PromptTokens, chatCompletionResponse.Usage.CompletionTokens, model)
return chatCompletionResponse, nil return chatCompletionResponse, nil
} }

View File

@ -88,10 +88,10 @@ func getLastArea() edgedb.UUID {
var inserted struct{ id edgedb.UUID } var inserted struct{ id edgedb.UUID }
err := edgeClient.QuerySingle(edgeCtx, ` err := edgeClient.QuerySingle(edgeCtx, `
select Area select Area
filter .conversation.name = 'Default' AND .conversation.user.id = <uuid>$0 filter .conversation.name = 'Default' AND .conversation.user = global currentUser
order by .position desc order by .position desc
limit 1 limit 1
`, &inserted, getCurrentUserID()) `, &inserted)
if err != nil { if err != nil {
fmt.Println("Error in edgedb.QuerySingle: in getLastArea") fmt.Println("Error in edgedb.QuerySingle: in getLastArea")
log.Fatal(err) log.Fatal(err)
@ -99,17 +99,6 @@ func getLastArea() edgedb.UUID {
return inserted.id return inserted.id
} }
func getCurrentUserID() edgedb.UUID {
var result User
err := edgeClient.QuerySingle(edgeCtx, "SELECT global currentUser LIMIT 1;", &result)
if err != nil {
fmt.Println("Error in edgedb.QuerySingle: in getCurrentUserID")
fmt.Println(err)
}
return result.ID
}
func checkIfLogin() bool { func checkIfLogin() bool {
var result User var result User
err := edgeClient.QuerySingle(edgeCtx, "SELECT global currentUser LIMIT 1;", &result) err := edgeClient.QuerySingle(edgeCtx, "SELECT global currentUser LIMIT 1;", &result)