From e112f3af12ed6e1903794abc07f4e329ad3268a2 Mon Sep 17 00:00:00 2001 From: Sakurasan <26715255+Sakurasan@users.noreply.github.com> Date: Mon, 21 Apr 2025 19:10:27 +0800 Subject: [PATCH] collect usage --- deploy/docker/docker-compose.adminer.yml | 8 ++++++++ deploy/docker/docker-compose.mariadb.yml | 8 ++------ deploy/docker/docker-compose.pg.yml | 5 ----- deploy/docker/docker-compose.sqlite.yml | 10 ++++++++++ internal/controller/proxy/chat_proxy.go | 18 ++++++++++++++++++ internal/model/usage.go | 12 ++++++------ llm/claude/v2/chat.go | 3 +++ llm/google/v2/chat.go | 3 +++ llm/openai_compatible/chat.go | 3 +++ llm/types.go | 9 +++++---- middleware/auth_team.go | 1 + 11 files changed, 59 insertions(+), 21 deletions(-) create mode 100644 deploy/docker/docker-compose.adminer.yml create mode 100644 deploy/docker/docker-compose.sqlite.yml diff --git a/deploy/docker/docker-compose.adminer.yml b/deploy/docker/docker-compose.adminer.yml new file mode 100644 index 0000000..1c39f42 --- /dev/null +++ b/deploy/docker/docker-compose.adminer.yml @@ -0,0 +1,8 @@ +version: '3.9' + +services: + adminer: + image: adminer + restart: always + ports: + - 8080:8080 diff --git a/deploy/docker/docker-compose.mariadb.yml b/deploy/docker/docker-compose.mariadb.yml index 928b1ba..8f8b074 100644 --- a/deploy/docker/docker-compose.mariadb.yml +++ b/deploy/docker/docker-compose.mariadb.yml @@ -4,6 +4,7 @@ services: mariadb: image: mariadb container_name: mysql + restart: unless-stopped ports: - "3306:3306" volumes: @@ -17,9 +18,4 @@ services: MYSQL_DATABASE: openteam MYSQL_USER: openteam MYSQL_PASSWORD: openteam - - # adminer: - # image: adminer - # restart: always - # ports: - # - 8080:8080 + diff --git a/deploy/docker/docker-compose.pg.yml b/deploy/docker/docker-compose.pg.yml index 8f23d13..47ae268 100644 --- a/deploy/docker/docker-compose.pg.yml +++ b/deploy/docker/docker-compose.pg.yml @@ -20,8 +20,3 @@ services: volumes: - $PWD/pgdata:/var/lib/postgresql/data - # adminer: - # image: adminer - # restart: always - # ports: - # - 8080:8080 diff --git a/deploy/docker/docker-compose.sqlite.yml b/deploy/docker/docker-compose.sqlite.yml new file mode 100644 index 0000000..c6ad4f8 --- /dev/null +++ b/deploy/docker/docker-compose.sqlite.yml @@ -0,0 +1,10 @@ +version: '3.7' +services: + sqlite-web: + image: vaalacat/sqlite-web + ports: + - 8800:8080 + volumes: + - $PWD/db:/data + environment: + - SQLITE_DATABASE=openteam.db \ No newline at end of file diff --git a/internal/controller/proxy/chat_proxy.go b/internal/controller/proxy/chat_proxy.go index 98d0c65..96851c3 100644 --- a/internal/controller/proxy/chat_proxy.go +++ b/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) } diff --git a/internal/model/usage.go b/internal/model/usage.go index f7204ca..d6e2be4 100644 --- a/internal/model/usage.go +++ b/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"` } diff --git a/llm/claude/v2/chat.go b/llm/claude/v2/chat.go index c8b2748..a2dc70b 100644 --- a/llm/claude/v2/chat.go +++ b/llm/claude/v2/chat.go @@ -111,6 +111,9 @@ func (c *Claude) Chat(ctx context.Context, chatReq llm.ChatRequest) (*llm.ChatRe return nil, err } + if c.tokenUsage.Model == "" && resp.Model != "" { + c.tokenUsage.Model = string(resp.Model) + } c.tokenUsage.PromptTokens += resp.Usage.InputTokens c.tokenUsage.CompletionTokens += resp.Usage.OutputTokens c.tokenUsage.TotalTokens += resp.Usage.InputTokens + resp.Usage.OutputTokens diff --git a/llm/google/v2/chat.go b/llm/google/v2/chat.go index 1bd883b..7b53ad4 100644 --- a/llm/google/v2/chat.go +++ b/llm/google/v2/chat.go @@ -110,6 +110,9 @@ func (g *Gemini) Chat(ctx context.Context, chatReq llm.ChatRequest) (*llm.ChatRe return nil, err } + if g.tokenUsage.Model == "" && response.ModelVersion != "" { + g.tokenUsage.Model = response.ModelVersion + } if response.UsageMetadata != nil { g.tokenUsage.PromptTokens += int(response.UsageMetadata.PromptTokenCount) g.tokenUsage.CompletionTokens += int(response.UsageMetadata.CandidatesTokenCount) diff --git a/llm/openai_compatible/chat.go b/llm/openai_compatible/chat.go index 1933d23..5a760ec 100644 --- a/llm/openai_compatible/chat.go +++ b/llm/openai_compatible/chat.go @@ -119,6 +119,9 @@ func (o *OpenAICompatible) Chat(ctx context.Context, chatReq llm.ChatRequest) (* return nil, err } + if o.tokenUsage.Model == "" && chatResp.Model != "" { + o.tokenUsage.Model = chatResp.Model + } o.tokenUsage.PromptTokens = chatResp.Usage.PromptTokens o.tokenUsage.CompletionTokens = chatResp.Usage.CompletionTokens o.tokenUsage.TotalTokens = chatResp.Usage.TotalTokens diff --git a/llm/types.go b/llm/types.go index 5c6d99f..0b8bbe9 100644 --- a/llm/types.go +++ b/llm/types.go @@ -15,10 +15,11 @@ type StreamChatResponse openai.ChatCompletionStreamResponse type ChatMessage openai.ChatCompletionMessage type TokenUsage struct { - PromptTokens int `json:"prompt_tokens"` - CompletionTokens int `json:"completion_tokens"` - ToolsTokens int `json:"total_tokens"` - TotalTokens int `json:"total_tokens"` + Model string `json:"model"` + PromptTokens int `json:"prompt_tokens"` + CompletionTokens int `json:"completion_tokens"` + ToolsTokens int `json:"tools_tokens"` + TotalTokens int `json:"total_tokens"` } type ErrorResponse struct { diff --git a/middleware/auth_team.go b/middleware/auth_team.go index a01757a..c87253c 100644 --- a/middleware/auth_team.go +++ b/middleware/auth_team.go @@ -94,6 +94,7 @@ func AuthLLM(db *gorm.DB) gin.HandlerFunc { } c.Set("user", token.User) + c.Set("user_id", token.User.ID) c.Set("authed", true) // 可以在这里对 token 进行验证并检查权限