collect usage

internal/controller/proxy/chat_proxy.go

@@ -4,10 +4,13 @@ import (
 	"fmt"
 	"net/http"
 	"opencatd-open/internal/dto"
+	"opencatd-open/internal/model"
 	"opencatd-open/llm"
 	"opencatd-open/llm/claude/v2"
 	"opencatd-open/llm/google/v2"
 	"opencatd-open/llm/openai_compatible"
+	"opencatd-open/pkg/tokenizer"
+	"strconv"
 
 	"github.com/gin-gonic/gin"
 )
@@ -57,4 +60,19 @@ func (h *Proxy) ChatHandler(c *gin.Context) {
 			c.SSEvent("", data)
 		}
 	}
+	llmusage := llm.GetTokenUsage()
+
+	cost := tokenizer.Cost(llmusage.Model, llmusage.PromptTokens+llmusage.ToolsTokens, llmusage.CompletionTokens)
+	userid, _ := strconv.ParseInt(c.GetString("user_id"), 10, 64)
+	usage := model.Usage{
+		UserID:           userid,
+		Model:            llmusage.Model,
+		Stream:           chatreq.Stream,
+		PromptTokens:     llmusage.PromptTokens + llmusage.ToolsTokens,
+		CompletionTokens: llmusage.CompletionTokens,
+		TotalTokens:      llmusage.TotalTokens,
+		Cost:             fmt.Sprintf("%f", cost),
+	}
+	h.SendUsage(&usage)
+	defer fmt.Println("cost:", cost, "prompt_tokens:", llmusage.PromptTokens, "completion_tokens:", llmusage.CompletionTokens, "total_tokens:", llmusage.TotalTokens)
 }

internal/model/usage.go

@@ -16,9 +16,9 @@ type Usage struct {
 	Date             time.Time `gorm:"column:date;autoCreateTime;index:idx_date"`
 	Model            string    `gorm:"column:model"`
 	Stream           bool      `gorm:"column:stream"`
-	PromptTokens     float64   `gorm:"column:prompt_tokens"`
-	CompletionTokens float64   `gorm:"column:completion_tokens"`
-	TotalTokens      float64   `gorm:"column:total_tokens"`
+	PromptTokens     int       `gorm:"column:prompt_tokens"`
+	CompletionTokens int       `gorm:"column:completion_tokens"`
+	TotalTokens      int       `gorm:"column:total_tokens"`
 	Cost             string    `gorm:"column:cost"`
 }
 
@@ -34,9 +34,9 @@ type DailyUsage struct {
 	Date             time.Time `gorm:"column:date;autoCreateTime;uniqueIndex:idx_daily_unique,priority:3"`
 	Model            string    `gorm:"column:model"`
 	Stream           bool      `gorm:"column:stream"`
-	PromptTokens     float64   `gorm:"column:prompt_tokens"`
-	CompletionTokens float64   `gorm:"column:completion_tokens"`
-	TotalTokens      float64   `gorm:"column:total_tokens"`
+	PromptTokens     int       `gorm:"column:prompt_tokens"`
+	CompletionTokens int       `gorm:"column:completion_tokens"`
+	TotalTokens      int       `gorm:"column:total_tokens"`
 	Cost             string    `gorm:"column:cost"`
 }