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
import (
"fmt"
"log"
"sync"
"github.com/edgedb/edgedb-go"
@ -13,8 +15,8 @@ type ModelInfo struct {
Name string
Icon string
MaxToken int
InputPrice float64
OutputPrice float64
InputPrice float32
OutputPrice float32
}
type CompanyInfo struct {
@ -40,6 +42,24 @@ type BotContentMessage struct {
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 {
// Create a wait group to synchronize the goroutines
var wg sync.WaitGroup

View File

@ -36,8 +36,8 @@ type AnthropicContentItem struct {
}
type AnthropicUsage struct {
InputTokens int `json:"input_tokens"`
OutputTokens int `json:"output_tokens"`
InputTokens int32 `json:"input_tokens"`
OutputTokens int32 `json:"output_tokens"`
}
func init() {
@ -169,5 +169,15 @@ func RequestAnthropic(model string, messages []Message, maxTokens int, temperatu
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
}

View File

@ -31,9 +31,9 @@ type OpenaiChatCompletionResponse struct {
}
type OpenaiUsage struct {
PromptTokens int `json:"prompt_tokens"`
CompletionTokens int `json:"completion_tokens"`
TotalTokens int `json:"total_tokens"`
PromptTokens int32 `json:"prompt_tokens"`
CompletionTokens int32 `json:"completion_tokens"`
TotalTokens int32 `json:"total_tokens"`
}
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)
}
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
}

View File

@ -88,10 +88,10 @@ 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
filter .conversation.name = 'Default' AND .conversation.user = global currentUser
order by .position desc
limit 1
`, &inserted, getCurrentUserID())
`, &inserted)
if err != nil {
fmt.Println("Error in edgedb.QuerySingle: in getLastArea")
log.Fatal(err)
@ -99,17 +99,6 @@ func getLastArea() edgedb.UUID {
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 {
var result User
err := edgeClient.QuerySingle(edgeCtx, "SELECT global currentUser LIMIT 1;", &result)