add claude proxy

deploy/docker/docker-compose.adminer.yml

@@ -0,0 +1,8 @@
+version: '3.9'
+
+services:
+  adminer:
+    image: adminer
+    restart: always
+    ports:
+      - 8080:8080

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

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

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

frontend/src/components/QRCodeCard.vue

@@ -1,5 +1,5 @@
 <template>
-    <div class="flex flex-col items-center space-y-4 p-6 bg-base-100 rounded-xl shadow-lg max-w-sm mx-auto backdrop-blur-xl glass">
+    <div class="flex flex-col items-center space-y-4 p-6 bg-base-100 rounded-xl shadow-lg mt-2 max-w-sm mx-auto backdrop-blur-xl glass">
 
         <div class="p-3 bg-white rounded-lg shadow-inner cursor-pointer" @click="toggleQRCode">
             <qrcode-vue :value="currentValue" :size="size" level="H" />

frontend/src/views/dashboard/KeyNew.vue

@@ -207,7 +207,7 @@ const newApiKey = ref({
   model_prefix: '',
   model_alias: '',
   parameters: '{}',
-  support_models: '',
+  support_models: '[]',
   support_models_array: [],
 })
 
@@ -224,7 +224,7 @@ const resetNewApiKey = () => {
     model_prefix: '',
     model_alias: '',
     parameters: '{}',
-    support_models: '',
+    support_models: '[]',
     support_models_array: [],
   }
 }

internal/controller/proxy/chat_proxy.go

@@ -4,15 +4,23 @@ 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"
 
 	"github.com/gin-gonic/gin"
 )
 
 func (h *Proxy) ChatHandler(c *gin.Context) {
+	user := c.MustGet("user").(*model.User)
+	if user == nil {
+		dto.WrapErrorAsOpenAI(c, 401, "Unauthorized")
+		return
+	}
+
 	var chatreq llm.ChatRequest
 	if err := c.ShouldBindJSON(&chatreq); err != nil {
 		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
@@ -35,10 +43,10 @@ func (h *Proxy) ChatHandler(c *gin.Context) {
 		fallthrough
 	default:
 		llm, err = openai_compatible.NewOpenAICompatible(h.apikey)
-		if err != nil {
+	}
-			dto.WrapErrorAsOpenAI(c, 500, fmt.Errorf("create llm client error: %w", err).Error())
+	if err != nil {
-			return
+		dto.WrapErrorAsOpenAI(c, 500, fmt.Errorf("create llm client error: %w", err).Error())
-		}
+		return
 	}
 
 	if !chatreq.Stream {
@@ -57,4 +65,13 @@ func (h *Proxy) ChatHandler(c *gin.Context) {
 			c.SSEvent("", data)
 		}
 	}
+
+	llmusage := llm.GetTokenUsage()
+	llmusage.User = user
+	llmusage.TokenID = c.GetInt64("token_id")
+	cost := tokenizer.Cost(llmusage.Model, llmusage.PromptTokens+llmusage.ToolsTokens, llmusage.CompletionTokens)
+
+	h.SendUsage(llmusage)
+	defer fmt.Println("cost:", cost, "prompt_tokens:", llmusage.PromptTokens, "completion_tokens:", llmusage.CompletionTokens, "total_tokens:", llmusage.TotalTokens)
+
 }

internal/controller/proxy/proxy.go

@@ -13,7 +13,9 @@ import (
 	"opencatd-open/internal/dao"
 	"opencatd-open/internal/model"
 	"opencatd-open/internal/utils"
+	"opencatd-open/llm"
 	"opencatd-open/pkg/config"
+	"opencatd-open/pkg/tokenizer"
 	"os"
 	"strings"
 	"sync"
@@ -24,7 +26,6 @@ import (
 	"github.com/lib/pq"
 	"github.com/tidwall/gjson"
 	"gorm.io/gorm"
-	"gorm.io/gorm/clause"
 )
 
 type Proxy struct {
@@ -32,7 +33,7 @@ type Proxy struct {
 	cfg        *config.Config
 	db         *gorm.DB
 	wg         *sync.WaitGroup
-	usageChan  chan *model.Usage // 用于异步处理的channel
+	usageChan  chan *llm.TokenUsage // 用于异步处理的channel
 	apikey     *model.ApiKey
 	httpClient *http.Client
 	cache      gcache.Cache
@@ -63,7 +64,7 @@ func NewProxy(ctx context.Context, cfg *config.Config, db *gorm.DB, wg *sync.Wai
 		wg:            wg,
 		httpClient:    client,
 		cache:         gcache.New(1).Build(),
-		usageChan:     make(chan *model.Usage, cfg.UsageChanSize),
+		usageChan:     make(chan *llm.TokenUsage, cfg.UsageChanSize),
 		userDAO:       userDAO,
 		apiKeyDao:     apiKeyDAO,
 		tokenDAO:      tokenDAO,
@@ -82,9 +83,13 @@ func (p *Proxy) HandleProxy(c *gin.Context) {
 		p.ChatHandler(c)
 		return
 	}
+	if strings.HasPrefix(c.Request.URL.Path, "/v1/messages") {
+		p.ProxyClaude(c)
+		return
+	}
 }
 
-func (p *Proxy) SendUsage(usage *model.Usage) {
+func (p *Proxy) SendUsage(usage *llm.TokenUsage) {
 	select {
 	case p.usageChan <- usage:
 	default:
@@ -140,46 +145,90 @@ func (p *Proxy) ProcessUsage() {
 	}
 }
 
-func (p *Proxy) Do(usage *model.Usage) error {
+func (p *Proxy) Do(llmusage *llm.TokenUsage) error {
 	err := p.db.Transaction(func(tx *gorm.DB) error {
+		now := time.Now()
+		today, _ := time.Parse("2006-01-02", now.Format("2006-01-02"))
+
+		cost := tokenizer.Cost(llmusage.Model, llmusage.PromptTokens, llmusage.CompletionTokens)
+		token, err := p.tokenDAO.GetByID(p.ctx, llmusage.TokenID)
+		if err != nil {
+			return err
+		}
+
+		usage := &model.Usage{
+			UserID:           llmusage.User.ID,
+			TokenID:          llmusage.TokenID,
+			Date:             now,
+			Model:            llmusage.Model,
+			Stream:           llmusage.Stream,
+			PromptTokens:     llmusage.PromptTokens,
+			CompletionTokens: llmusage.CompletionTokens,
+			TotalTokens:      llmusage.TotalTokens,
+			Cost:             fmt.Sprintf("%.8f", cost),
+		}
 		// 1. 记录使用记录
 		if err := tx.WithContext(p.ctx).Create(usage).Error; err != nil {
 			return fmt.Errorf("create usage error: %w", err)
 		}
 
-		// 2. 更新每日统计（upsert 操作）
+		// 2. 更新每日统计
-		dailyUsage := model.DailyUsage{
+		var dailyUsage model.DailyUsage
-			UserID:           usage.UserID,
+		result := tx.WithContext(p.ctx).Where("user_id = ? and date = ?", llmusage.User.ID, today).First(&dailyUsage)
-			TokenID:          usage.TokenID,
+		if result.RowsAffected == 0 {
-			Capability:       usage.Capability,
+			dailyUsage.UserID = llmusage.User.ID
-			Date:             time.Date(usage.Date.Year(), usage.Date.Month(), usage.Date.Day(), 0, 0, 0, 0, usage.Date.Location()),
+			dailyUsage.TokenID = llmusage.TokenID
-			Model:            usage.Model,
+			dailyUsage.Date = today
-			Stream:           usage.Stream,
+			dailyUsage.Model = llmusage.Model
-			PromptTokens:     usage.PromptTokens,
+			dailyUsage.Stream = llmusage.Stream
-			CompletionTokens: usage.CompletionTokens,
+			dailyUsage.PromptTokens = llmusage.PromptTokens
-			TotalTokens:      usage.TotalTokens,
+			dailyUsage.CompletionTokens = llmusage.CompletionTokens
-			Cost:             usage.Cost,
+			dailyUsage.TotalTokens = llmusage.TotalTokens
-		}
+			dailyUsage.Cost = fmt.Sprintf("%.8f", cost)
-
+			if err := tx.WithContext(p.ctx).Create(&dailyUsage).Error; err != nil {
-		// 使用 OnConflict 实现 upsert
+				return fmt.Errorf("create daily usage error: %w", err)
-		if err := tx.WithContext(p.ctx).Clauses(clause.OnConflict{
+			}
-			Columns: []clause.Column{{Name: "user_id"}, {Name: "token_id"}, {Name: "capability"}, {Name: "date"}}, // 唯一键
+		} else {
-			DoUpdates: clause.Assignments(map[string]interface{}{
+			if err := tx.WithContext(p.ctx).Model(&model.DailyUsage{}).Where("user_id = ? and date = ?", llmusage.User.ID, today).
-				"prompt_tokens":     gorm.Expr("prompt_tokens + ?", usage.PromptTokens),
+				Updates(map[string]interface{}{
-				"completion_tokens": gorm.Expr("completion_tokens + ?", usage.CompletionTokens),
+					"prompt_tokens":     gorm.Expr("prompt_tokens + ?", llmusage.PromptTokens),
-				"total_tokens":      gorm.Expr("total_tokens + ?", usage.TotalTokens),
+					"completion_tokens": gorm.Expr("completion_tokens + ?", llmusage.CompletionTokens),
-				"cost":              gorm.Expr("cost + ?", usage.Cost),
+					"total_tokens":      gorm.Expr("total_tokens + ?", llmusage.TotalTokens),
-			}),
+				}).Error; err != nil {
-		}).Create(&dailyUsage).Error; err != nil {
+				return fmt.Errorf("update daily usage error: %w", err)
-			return fmt.Errorf("upsert daily usage error: %w", err)
+			}
 		}
 
 		// 3. 更新用户额度
-		if err := tx.WithContext(p.ctx).Model(&model.User{}).Where("id = ?", usage.UserID).Updates(map[string]interface{}{
+		if *llmusage.User.UnlimitedQuota {
-			"quota":      gorm.Expr("quota - ?", usage.Cost),
+			if err := tx.WithContext(p.ctx).Model(&model.User{}).Where("id = ?", llmusage.User.ID).Updates(map[string]interface{}{
-			"used_quota": gorm.Expr("used_quota + ?", usage.Cost),
+				"used_quota": gorm.Expr("used_quota + ?", fmt.Sprintf("%.8f", cost)),
-		}).Error; err != nil {
+			}).Error; err != nil {
-			return fmt.Errorf("update user quota and used_quota error: %w", err)
+				return fmt.Errorf("update user quota and used_quota error: %w", err)
+			}
+		} else {
+			if err := tx.WithContext(p.ctx).Model(&model.User{}).Where("id = ?", llmusage.User.ID).Updates(map[string]interface{}{
+				"quota":      gorm.Expr("quota - ?", fmt.Sprintf("%.8f", cost)),
+				"used_quota": gorm.Expr("used_quota + ?", fmt.Sprintf("%.8f", cost)),
+			}).Error; err != nil {
+				return fmt.Errorf("update user quota and used_quota error: %w", err)
+			}
+		}
+
+		//4 . 更新token额度
+		if *token.UnlimitedQuota {
+			if err := tx.WithContext(p.ctx).Model(&model.Token{}).Where("id = ?", llmusage.TokenID).Updates(map[string]interface{}{
+				"used_quota": gorm.Expr("used_quota + ?", fmt.Sprintf("%.8f", cost)),
+			}).Error; err != nil {
+				return fmt.Errorf("update token quota and used_quota error: %w", err)
+			}
+		} else {
+			if err := tx.WithContext(p.ctx).Model(&model.Token{}).Where("id = ?", llmusage.TokenID).Updates(map[string]interface{}{
+				"quota":      gorm.Expr("quota - ?", fmt.Sprintf("%.8f", cost)),
+				"used_quota": gorm.Expr("used_quota + ?", fmt.Sprintf("%.8f", cost)),
+			}).Error; err != nil {
+				return fmt.Errorf("update token quota and used_quota error: %w", err)
+			}
 		}
 
 		return nil
@@ -189,7 +238,6 @@ func (p *Proxy) Do(usage *model.Usage) error {
 
 func (p *Proxy) SelectApiKey(model string) error {
 	akpikeys, err := p.apiKeyDao.FindApiKeysBySupportModel(p.db, model)
-	fmt.Println(len(akpikeys), err)
 	if err != nil || len(akpikeys) == 0 {
 		if strings.HasPrefix(model, "gpt") || strings.HasPrefix(model, "o1") || strings.HasPrefix(model, "o3") || strings.HasPrefix(model, "o4") {
 			keys, err := p.apiKeyDao.FindKeys(map[string]any{"active = ?": true, "apitype = ?": "openai"})

internal/controller/proxy/proxy_claude.go

@@ -0,0 +1,14 @@
+package controller
+
+import (
+	"fmt"
+	"io"
+
+	"github.com/gin-gonic/gin"
+)
+
+func (p *Proxy) ProxyClaude(c *gin.Context) {
+	fmt.Println(c.Request.URL.String())
+	data, _ := io.ReadAll(c.Request.Body)
+	fmt.Println(string(data))
+}

internal/controller/user.go

@@ -1,7 +1,6 @@
 package controller
 
 import (
-	"fmt"
 	"net/http"
 	"opencatd-open/internal/dto"
 	"opencatd-open/internal/model"
@@ -142,7 +141,7 @@ func (a Api) CreateUser(c *gin.Context) {
 		dto.Fail(c, 400, err.Error())
 		return
 	}
-	fmt.Printf("user:%+v\n", user)
+
 	err = a.userService.Create(c, &user)
 	if err != nil {
 		dto.Fail(c, http.StatusInternalServerError, err.Error())

internal/controller/user_token.go

@@ -90,7 +90,7 @@ func (a Api) ResetToken(c *gin.Context) {
 		dto.Fail(c, http.StatusNotFound, "token not found")
 		return
 	}
-	token.UsedQuota = utils.ToPtr(int64(0))
+	token.UsedQuota = utils.ToPtr(float64(0))
 
 	err = a.tokenService.UpdateToken(c, token)
 	if err != nil {

internal/dao/apikey.go

@@ -3,6 +3,7 @@ package dao
 import (
 	"errors"
 	"opencatd-open/internal/model"
+	"opencatd-open/internal/utils"
 	"opencatd-open/pkg/config"
 
 	"gorm.io/gorm"
@@ -38,6 +39,9 @@ func (dao *ApiKeyDAO) Create(apiKey *model.ApiKey) error {
 	if apiKey == nil {
 		return errors.New("apiKey is nil")
 	}
+	if len(*apiKey.SupportModels) < 2 {
+		apiKey.SupportModels = utils.ToPtr("[]")
+	}
 	return dao.db.Create(apiKey).Error
 }
 

internal/dao/usage.go

@@ -212,11 +212,7 @@ func (d *DailyUsageDAO) UpsertDailyUsage(ctx context.Context, usage *model.Usage
 		return db.Clauses(clause.OnConflict{
 			Columns: []clause.Column{
 				{Name: "user_id"},
-				{Name: "token_id"},
-				{Name: "capability"},
 				{Name: "date"},
-				{Name: "model"},
-				{Name: "stream"},
 			},
 			DoUpdates: clause.Assignments(updateColumns),
 		}).Create(dailyUsage).Error
@@ -231,11 +227,7 @@ func (d *DailyUsageDAO) UpsertDailyUsage(ctx context.Context, usage *model.Usage
 		return db.Clauses(clause.OnConflict{
 			Columns: []clause.Column{
 				{Name: "user_id"},
-				{Name: "token_id"},
-				{Name: "capability"},
 				{Name: "date"},
-				{Name: "model"},
-				{Name: "stream"},
 			},
 			DoUpdates: clause.Assignments(updateColumns),
 		}).Create(dailyUsage).Error
@@ -244,8 +236,8 @@ func (d *DailyUsageDAO) UpsertDailyUsage(ctx context.Context, usage *model.Usage
 	default:
 		return db.Transaction(func(tx *gorm.DB) error {
 			var existing model.DailyUsage
-			err := tx.Where("user_id = ? AND token_id = ? AND capability = ? AND date = ? AND model = ? AND stream = ?",
+			err := tx.Where("user_id = ? AND date = ?",
-				usage.UserID, usage.TokenID, usage.Capability, date, usage.Model, usage.Stream).
+				usage.UserID, date).
 				First(&existing).Error
 
 			if err == gorm.ErrRecordNotFound {

internal/model/token.go

@@ -2,19 +2,19 @@ package model
 
 // 用户的token
 type Token struct {
-	ID             int64  `gorm:"column:id;primaryKey;autoIncrement" json:"id,omitempty"`
+	ID             int64    `gorm:"column:id;primaryKey;autoIncrement" json:"id,omitempty"`
-	UserID         int64  `gorm:"column:user_id;not null;index:idx_token_user_id" json:"userid,omitempty"`
+	UserID         int64    `gorm:"column:user_id;not null;index:idx_token_user_id" json:"userid,omitempty"`
-	Name           string `gorm:"column:name;not null;index:idx_token_name" json:"name,omitempty" binding:"required,min=1,max=20"`
+	Name           string   `gorm:"column:name;not null;index:idx_token_name" json:"name,omitempty" binding:"required,min=1,max=20"`
-	Key            string `gorm:"column:key;not null;uniqueIndex:idx_token_key;comment:token key" json:"key,omitempty"`
+	Key            string   `gorm:"column:key;not null;uniqueIndex:idx_token_key;comment:token key" json:"key,omitempty"`
-	Active         *bool  `gorm:"column:active;default:true" json:"active,omitempty"`                   //
+	Active         *bool    `gorm:"column:active;default:true" json:"active,omitempty"`                   //
-	Quota          *int64 `gorm:"column:quota;type:bigint;default:0" json:"quota,omitempty"`            // default 0
+	Quota          *float64 `gorm:"column:quota;type:bigint;default:0" json:"quota,omitempty"`            // default 0
-	UnlimitedQuota *bool  `gorm:"column:unlimited_quota;default:true" json:"unlimited_quota,omitempty"` // set Quota 1  unlimited
+	UnlimitedQuota *bool    `gorm:"column:unlimited_quota;default:true" json:"unlimited_quota,omitempty"` // set Quota 1  unlimited
-	UsedQuota      *int64 `gorm:"column:used_quota;type:bigint;default:0" json:"used_quota,omitempty"`
+	UsedQuota      *float64 `gorm:"column:used_quota;type:bigint;default:0" json:"used_quota,omitempty"`
-	ExpiredAt      *int64 `gorm:"column:expired_at;type:bigint;default:0" json:"expired_at,omitempty"`
+	ExpiredAt      *int64   `gorm:"column:expired_at;type:bigint;default:0" json:"expired_at,omitempty"`
-	NeverExpired   *bool  `gorm:"column:never_expires;type:bigint;" json:"never_expires,omitempty"`
+	NeverExpired   *bool    `gorm:"column:never_expires;type:bigint;" json:"never_expires,omitempty"`
-	CreatedAt      int64  `gorm:"column:created_at;type:bigint;autoCreateTime" json:"created_at,omitempty"`
+	CreatedAt      int64    `gorm:"column:created_at;type:bigint;autoCreateTime" json:"created_at,omitempty"`
-	LastUsedAt     int64  `gorm:"column:lastused_at;type:bigint;autoUpdateTime" json:"lastused_at,omitempty"`
+	LastUsedAt     int64    `gorm:"column:lastused_at;type:bigint;autoUpdateTime" json:"lastused_at,omitempty"`
-	User           *User  `gorm:"foreignKey:UserID;constraint:OnUpdate:CASCADE,OnDelete:CASCADE" json:"-"`
+	User           *User    `gorm:"foreignKey:UserID;constraint:OnUpdate:CASCADE,OnDelete:CASCADE" json:"-"`
 }
 
 func (Token) TableName() string {

internal/model/usage.go

@@ -1,11 +1,7 @@
 package model
 
 import (
-	"net/http"
-	"opencatd-open/store"
 	"time"
-
-	"github.com/gin-gonic/gin"
 )
 
 type Usage struct {
@@ -16,9 +12,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"`
+	PromptTokens     int       `gorm:"column:prompt_tokens"`
-	CompletionTokens float64   `gorm:"column:completion_tokens"`
+	CompletionTokens int       `gorm:"column:completion_tokens"`
-	TotalTokens      float64   `gorm:"column:total_tokens"`
+	TotalTokens      int       `gorm:"column:total_tokens"`
 	Cost             string    `gorm:"column:cost"`
 }
 
@@ -28,47 +24,18 @@ func (Usage) TableName() string {
 
 type DailyUsage struct {
 	ID               int64     `gorm:"column:id;primaryKey;autoIncrement"`
-	UserID           int64     `gorm:"column:user_id;uniqueIndex:idx_daily_unique,priority:1"`
+	UserID           int64     `gorm:"column:user_id;uniqueIndex:idx_daily_unique,priority:1"` // uniqueIndex:idx_daily_unique,priority:1
-	TokenID          int64     `gorm:"column:token_id;index:idx_daily_token_id"`
+	TokenID          int64     `gorm:"column:token_id;uniqueIndex:idx_daily_unique,priority:2"`
-	Capability       string    `gorm:"column:capability;uniqueIndex:idx_daily_unique,priority:2;comment:模型能力"`
+	Capability       string    `gorm:"column:capability;index:idx_daily_usage_capability;comment:模型能力"`
 	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"`
+	PromptTokens     int       `gorm:"column:prompt_tokens"`
-	CompletionTokens float64   `gorm:"column:completion_tokens"`
+	CompletionTokens int       `gorm:"column:completion_tokens"`
-	TotalTokens      float64   `gorm:"column:total_tokens"`
+	TotalTokens      int       `gorm:"column:total_tokens"`
 	Cost             string    `gorm:"column:cost"`
 }
 
 func (DailyUsage) TableName() string {
 	return "daily_usages"
 }
-
-func HandleUsage(c *gin.Context) {
-	fromStr := c.Query("from")
-	toStr := c.Query("to")
-	getMonthStartAndEnd := func() (start, end string) {
-		loc, _ := time.LoadLocation("Local")
-		now := time.Now().In(loc)
-
-		year, month, _ := now.Date()
-
-		startOfMonth := time.Date(year, month, 1, 0, 0, 0, 0, loc)
-		endOfMonth := startOfMonth.AddDate(0, 1, 0)
-
-		start = startOfMonth.Format("2006-01-02")
-		end = endOfMonth.Format("2006-01-02")
-		return
-	}
-	if fromStr == "" || toStr == "" {
-		fromStr, toStr = getMonthStartAndEnd()
-	}
-
-	usage, err := store.QueryUsage(fromStr, toStr)
-	if err != nil {
-		c.JSON(http.StatusForbidden, gin.H{"error": err.Error()})
-		return
-	}
-
-	c.JSON(200, usage)
-}

llm/claude/v2/chat.go

@@ -94,7 +94,9 @@ func (c *Claude) Chat(ctx context.Context, chatReq llm.ChatRequest) (*llm.ChatRe
 	if chatReq.MaxTokens > 0 {
 		maxTokens = chatReq.MaxTokens
 	} else {
-		if strings.Contains(chatReq.Model, "sonnet") || strings.Contains(chatReq.Model, "haiku") {
+		if strings.Contains(chatReq.Model, "3-7") {
+			maxTokens = 64000
+		} else if strings.Contains(chatReq.Model, "3-5") {
 			maxTokens = 8192
 		} else {
 			maxTokens = 4096
@@ -111,6 +113,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

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)

llm/llm.go

@@ -13,7 +13,7 @@ type LLM interface {
 
 type llm struct {
 	ApiKey   *model.ApiKey
-	Usage    *model.Usage
+	Usage    *TokenUsage
 	tools    any   // TODO
 	Messages []any // TODO
 	llm      LLM

llm/openai_compatible/chat.go

@@ -14,6 +14,8 @@ import (
 	"opencatd-open/llm"
 	"os"
 	"strings"
+
+	"github.com/sashabaranov/go-openai"
 )
 
 // https://learn.microsoft.com/en-us/azure/ai-services/openai/api-version-deprecation#latest-preview-api-releases
@@ -119,6 +121,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
@@ -127,6 +132,7 @@ func (o *OpenAICompatible) Chat(ctx context.Context, chatReq llm.ChatRequest) (*
 
 func (o *OpenAICompatible) StreamChat(ctx context.Context, chatReq llm.ChatRequest) (chan *llm.StreamChatResponse, error) {
 	chatReq.Stream = true
+	chatReq.StreamOptions = &openai.StreamOptions{IncludeUsage: true}
 	dst, err := utils.StructToMap(chatReq)
 	if err != nil {
 		return nil, err
@@ -200,6 +206,7 @@ func (o *OpenAICompatible) StreamChat(ctx context.Context, chatReq llm.ChatReque
 					if err := json.Unmarshal(line, &streamResp); err != nil {
 						continue
 					}
+
 					if streamResp.Usage != nil {
 						o.tokenUsage.PromptTokens += streamResp.Usage.PromptTokens
 						o.tokenUsage.CompletionTokens += streamResp.Usage.CompletionTokens

llm/types.go

@@ -2,6 +2,7 @@ package llm
 
 import (
 	"fmt"
+	"opencatd-open/internal/model"
 
 	"github.com/sashabaranov/go-openai"
 )
@@ -15,9 +16,13 @@ type StreamChatResponse openai.ChatCompletionStreamResponse
 type ChatMessage openai.ChatCompletionMessage
 
 type TokenUsage struct {
+	User             *model.User
+	TokenID          int64
+	Model            string `json:"model"`
+	Stream           bool
 	PromptTokens     int `json:"prompt_tokens"`
 	CompletionTokens int `json:"completion_tokens"`
-	ToolsTokens      int `json:"total_tokens"`
+	ToolsTokens      int `json:"tools_tokens"`
 	TotalTokens      int `json:"total_tokens"`
 }
 

middleware/auth_team.go

@@ -94,6 +94,8 @@ func AuthLLM(db *gorm.DB) gin.HandlerFunc {
 		}
 
 		c.Set("user", token.User)
+		c.Set("user_id", token.User.ID)
+		c.Set("token_id", token.ID)
 		c.Set("authed", true)
 		// 可以在这里对 token 进行验证并检查权限