From c363f126640908f8c1f950960a35dee4cff05c2d Mon Sep 17 00:00:00 2001
From: Adrien <adrien.bouvais.pro@gmail.com>
Date: Sun, 5 May 2024 21:33:46 +0200
Subject: [PATCH] Adding usage to db

---
 Request.go          | 24 ++++++++++++++++++++++--
 RequestAnthropic.go | 14 ++++++++++++--
 RequestOpenai.go    | 16 +++++++++++++---
 database.go         | 15 ++-------------
 4 files changed, 49 insertions(+), 20 deletions(-)

diff --git a/Request.go b/Request.go
index 2bd3c2a..bc22034 100644
--- a/Request.go
+++ b/Request.go
@@ -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
diff --git a/RequestAnthropic.go b/RequestAnthropic.go
index 983b3ae..99c20d1 100644
--- a/RequestAnthropic.go
+++ b/RequestAnthropic.go
@@ -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
 }
diff --git a/RequestOpenai.go b/RequestOpenai.go
index 18540f4..197a255 100644
--- a/RequestOpenai.go
+++ b/RequestOpenai.go
@@ -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
 }
diff --git a/database.go b/database.go
index 486e4df..2c6ecdc 100644
--- a/database.go
+++ b/database.go
@@ -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)