collect usage

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:
@@ -18,8 +19,3 @@ services:
       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

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"`
 }
 

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

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/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

llm/types.go

@@ -15,9 +15,10 @@ type StreamChatResponse openai.ChatCompletionStreamResponse
 type ChatMessage openai.ChatCompletionMessage
 
 type TokenUsage struct {
+	Model            string `json:"model"`
 	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,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 进行验证并检查权限