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)
 }