mirrors2/opencatd-open
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix usage

Browse Source
This commit is contained in:
Sakurasan
2025-04-22 01:06:03 +08:00
parent 6662ea5e04
commit 000162b1b1
2 changed files with 29 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files

54
internal/controller/proxy/proxy.go
View File

@@ -26,7 +26,6 @@ import (
"github.com/lib/pq" "github.com/lib/pq"
"github.com/tidwall/gjson" "github.com/tidwall/gjson"
"gorm.io/gorm" "gorm.io/gorm"
"gorm.io/gorm/clause"
) )
type Proxy struct { type Proxy struct {
@@ -150,8 +149,8 @@ func (p *Proxy) Do(llmusage *llm.TokenUsage) error {
if err != nil { if err != nil {
return err return err
} }
// 1. 记录使用记录
if err := tx.WithContext(p.ctx).Create(&model.Usage{ usage := &model.Usage{
UserID: llmusage.User.ID, UserID: llmusage.User.ID,
TokenID: llmusage.TokenID, TokenID: llmusage.TokenID,
Date: now, Date: now,
@@ -161,35 +160,38 @@ func (p *Proxy) Do(llmusage *llm.TokenUsage) error {
CompletionTokens: llmusage.CompletionTokens, CompletionTokens: llmusage.CompletionTokens,
TotalTokens: llmusage.TotalTokens, TotalTokens: llmusage.TotalTokens,
Cost: fmt.Sprintf("%.8f", cost), Cost: fmt.Sprintf("%.8f", cost),
}).Error; err != nil { }
// 1. 记录使用记录
if err := tx.WithContext(p.ctx).Create(usage).Error; err != nil {
return fmt.Errorf("create usage error: %w", err) return fmt.Errorf("create usage error: %w", err)
} }
// 2. 更新每日统计（upsert 操作） // 2. 更新每日统计（upsert 操作）
dailyUsage := model.DailyUsage{ // dailyUsage := model.DailyUsage{
UserID: llmusage.User.ID, // UserID: llmusage.User.ID,
TokenID: llmusage.TokenID, // TokenID: llmusage.TokenID,
Date: time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, now.Location()), // Date: time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, now.Location()),
Model: llmusage.Model, // Model: llmusage.Model,
Stream: llmusage.Stream, // Stream: llmusage.Stream,
PromptTokens: llmusage.PromptTokens, // PromptTokens: llmusage.PromptTokens,
CompletionTokens: llmusage.CompletionTokens, // CompletionTokens: llmusage.CompletionTokens,
TotalTokens: llmusage.TotalTokens, // TotalTokens: llmusage.TotalTokens,
Cost: fmt.Sprintf("%.8f", cost), // Cost: fmt.Sprintf("%.8f", cost),
} // }
p.dailyUsageDAO.UpsertDailyUsage(p.ctx, usage)
// 使用 OnConflict 实现 upsert // 使用 OnConflict 实现 upsert
if err := tx.WithContext(p.ctx).Clauses(clause.OnConflict{ // if err := tx.WithContext(p.ctx).Clauses(clause.OnConflict{
Columns: []clause.Column{{Name: "user_id"}, {Name: "date"}}, // 唯一键 // Columns: []clause.Column{{Name: "user_id"}, {Name: "date"}}, // 唯一键
DoUpdates: clause.Assignments(map[string]interface{}{ // DoUpdates: clause.Assignments(map[string]interface{}{
"prompt_tokens": gorm.Expr("prompt_tokens + ?", llmusage.PromptTokens), // "prompt_tokens": gorm.Expr("prompt_tokens + ?", llmusage.PromptTokens),
"completion_tokens": gorm.Expr("completion_tokens + ?", llmusage.CompletionTokens), // "completion_tokens": gorm.Expr("completion_tokens + ?", llmusage.CompletionTokens),
"total_tokens": gorm.Expr("total_tokens + ?", llmusage.TotalTokens), // "total_tokens": gorm.Expr("total_tokens + ?", llmusage.TotalTokens),
"cost": gorm.Expr("cost + ?", fmt.Sprintf("%.8f", cost)), // "cost": gorm.Expr("cost + ?", fmt.Sprintf("%.8f", cost)),
}), // }),
}).Create(&dailyUsage).Error; err != nil { // }).Create(&dailyUsage).Error; err != nil {
return fmt.Errorf("upsert daily usage error: %w", err) // return fmt.Errorf("upsert daily usage error: %w", err)
} // }
// 3. 更新用户额度 // 3. 更新用户额度
if *llmusage.User.UnlimitedQuota { if *llmusage.User.UnlimitedQuota {

2
llm/openai_compatible/chat.go
View File

@@ -206,7 +206,7 @@ func (o *OpenAICompatible) StreamChat(ctx context.Context, chatReq llm.ChatReque
if err := json.Unmarshal(line, &streamResp); err != nil { if err := json.Unmarshal(line, &streamResp); err != nil {
continue continue
} }
fmt.Printf("%#v\n", streamResp.Usage)
if streamResp.Usage != nil { if streamResp.Usage != nil {
o.tokenUsage.PromptTokens += streamResp.Usage.PromptTokens o.tokenUsage.PromptTokens += streamResp.Usage.PromptTokens
o.tokenUsage.CompletionTokens += streamResp.Usage.CompletionTokens o.tokenUsage.CompletionTokens += streamResp.Usage.CompletionTokens