支持统计用户的使用量

go.mod

@@ -4,16 +4,20 @@ go 1.19
 
 require (
 	github.com/Sakurasan/to v0.0.0-20180919163141-e72657dd7c7d
+	github.com/duke-git/lancet/v2 v2.1.19
 	github.com/gin-gonic/gin v1.9.0
 	github.com/glebarez/sqlite v1.7.0
 	github.com/google/uuid v1.3.0
 	github.com/patrickmn/go-cache v2.1.0+incompatible
+	github.com/pkoukk/tiktoken-go v0.1.1-0.20230418101013-cae809389480
+	github.com/sashabaranov/go-openai v1.9.0
 	gorm.io/gorm v1.24.6
 )
 
 require (
 	github.com/bytedance/sonic v1.8.0 // indirect
 	github.com/chenzhuoyu/base64x v0.0.0-20221115062448-fe3a3abad311 // indirect
+	github.com/dlclark/regexp2 v1.8.1 // indirect
 	github.com/dustin/go-humanize v1.0.1 // indirect
 	github.com/gin-contrib/sse v0.1.0 // indirect
 	github.com/glebarez/go-sqlite v1.20.3 // indirect

go.sum

@@ -9,6 +9,10 @@ github.com/chenzhuoyu/base64x v0.0.0-20221115062448-fe3a3abad311/go.mod h1:b583j
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/dlclark/regexp2 v1.8.1 h1:6Lcdwya6GjPUNsBct8Lg/yRPwMhABj269AAzdGSiR+0=
+github.com/dlclark/regexp2 v1.8.1/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8=
+github.com/duke-git/lancet/v2 v2.1.19 h1:dbRB1m6wOMV1I0ax/3S6ngop8SYM6I7sr+7D9IXjS2E=
+github.com/duke-git/lancet/v2 v2.1.19/go.mod h1:hNcc06mV7qr+crH/0nP+rlC3TB0Q9g5OrVnO8/TGD4c=
 github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY=
 github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto=
 github.com/gin-contrib/sse v0.1.0 h1:Y/yl/+YNO8GZSjAhjMsSuLt29uWRFHdHYUb5lYOV9qE=
@@ -57,12 +61,16 @@ github.com/patrickmn/go-cache v2.1.0+incompatible h1:HRMgzkcYKYpi3C8ajMPV8OFXaaR
 github.com/patrickmn/go-cache v2.1.0+incompatible/go.mod h1:3Qf8kWWT7OJRJbdiICTKqZju1ZixQ/KpMGzzAfe6+WQ=
 github.com/pelletier/go-toml/v2 v2.0.6 h1:nrzqCb7j9cDFj2coyLNLaZuJTLjWjlaz6nvTvIwycIU=
 github.com/pelletier/go-toml/v2 v2.0.6/go.mod h1:eumQOmlWiOPt5WriQQqoM5y18pDHwha2N+QD+EUNTek=
+github.com/pkoukk/tiktoken-go v0.1.1-0.20230418101013-cae809389480 h1:IFhPCcB0/HtnEN+ZoUGDT55YgFCymbFJ15kXqs3nv5w=
+github.com/pkoukk/tiktoken-go v0.1.1-0.20230418101013-cae809389480/go.mod h1:BijIqAP84FMYC4XbdJgjyMpiSjusU8x0Y0W9K2t0QtU=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/remyoudompheng/bigfft v0.0.0-20200410134404-eec4a21b6bb0/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo=
 github.com/remyoudompheng/bigfft v0.0.0-20230126093431-47fa9a501578 h1:VstopitMQi3hZP0fzvnsLmzXZdQGc4bEcgu24cp+d4M=
 github.com/remyoudompheng/bigfft v0.0.0-20230126093431-47fa9a501578/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo=
 github.com/rogpeppe/go-internal v1.8.0 h1:FCbCCtXNOY3UtUuHUYaghJg4y7Fd14rXifAYUAtL9R8=
+github.com/sashabaranov/go-openai v1.9.0 h1:NoiO++IISxxJ1pRc0n7uZvMGMake0G+FJ1XPwXtprsA=
+github.com/sashabaranov/go-openai v1.9.0/go.mod h1:lj5b/K+zjTSFxVLijLSTDZuP7adOgerWeFyZLUhAKRg=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=
 github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo=
@@ -71,8 +79,8 @@ github.com/stretchr/testify v1.6.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/
 github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
-github.com/stretchr/testify v1.8.1 h1:w7B6lhMri9wdJUVmEZPGGhZzrYTPvgJArz7wNPgYKsk=
 github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
+github.com/stretchr/testify v1.8.2 h1:+h33VjcLVPDHtOdpUCuF+7gSuG3yGIftsP1YvFihtJ8=
 github.com/twitchyliquid64/golang-asm v0.15.1 h1:SU5vSMR7hnwNxj24w34ZyCi/FmDZTkS4MhqMhdFk5YI=
 github.com/twitchyliquid64/golang-asm v0.15.1/go.mod h1:a1lVb/DtPvCB8fslRZhAngC2+aY1QWCk3Cedj/Gdt08=
 github.com/ugorji/go/codec v1.2.9 h1:rmenucSohSTiyL09Y+l2OCk+FrMxGMzho2+tjr5ticU=

opencat.go

@@ -38,6 +38,8 @@ func main() {
 		// 获取所有用户信息
 		group.GET("/users", router.HandleUsers)
 
+		group.GET("/usages", router.HandleUsage)
+
 		// 添加Key
 		group.POST("/keys", router.HandleAddKey)
 
@@ -57,9 +59,12 @@ func main() {
 	// 初始化用户
 	r.POST("/1/users/init", router.Handleinit)
 
-	r.POST("/v1/chat/completions", router.HandleProy)
+	r.Any("/v1/*proxypath", router.HandleProy)
-	r.GET("/v1/models", router.HandleProy)
+
-	r.GET("/v1/dashboard/billing/subscription", router.HandleProy)
+	// r.POST("/v1/chat/completions", router.HandleProy)
+	// r.GET("/v1/models", router.HandleProy)
+	// r.GET("/v1/dashboard/billing/subscription", router.HandleProy)
+
 	r.GET("/", func(c *gin.Context) {
 		c.Writer.WriteHeader(http.StatusOK)
 		c.Writer.WriteString(`<h1><a href="https://github.com/mirrors2/opencatd-open" >opencatd-open</a> available</h1>Api-Keys:<a href=https://platform.openai.com/account/api-keys >https://platform.openai.com/account/api-keys</a>`)

router/router.go

@@ -1,8 +1,10 @@
 package router
 
 import (
+	"bufio"
 	"bytes"
 	"crypto/tls"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"io"
@@ -11,17 +13,25 @@ import (
 	"net/http"
 	"net/http/httputil"
 	"opencatd-open/store"
+	"strings"
+	"sync"
 	"time"
 
 	"github.com/Sakurasan/to"
+	"github.com/duke-git/lancet/v2/cryptor"
 	"github.com/gin-gonic/gin"
 	"github.com/google/uuid"
+	"github.com/pkoukk/tiktoken-go"
+	"github.com/sashabaranov/go-openai"
 	"gorm.io/gorm"
 )
 
 var (
-	rootToken string
+	rootToken     string
-	baseUrl   = "https://api.openai.com"
+	baseUrl       = "https://api.openai.com"
+	GPT3Dot5Turbo = "gpt-3.5-turbo"
+	GPT4          = "gpt-4"
+	client        = getHttpClient()
 )
 
 type User struct {
@@ -41,6 +51,46 @@ type Key struct {
 	CreatedAt string `json:"createdAt,omitempty"`
 }
 
+type ChatCompletionMessage struct {
+	Role    string `json:"role"`
+	Content string `json:"content"`
+	Name    string `json:"name,omitempty"`
+}
+
+type ChatCompletionRequest struct {
+	Model            string                  `json:"model"`
+	Messages         []ChatCompletionMessage `json:"messages"`
+	MaxTokens        int                     `json:"max_tokens,omitempty"`
+	Temperature      float32                 `json:"temperature,omitempty"`
+	TopP             float32                 `json:"top_p,omitempty"`
+	N                int                     `json:"n,omitempty"`
+	Stream           bool                    `json:"stream,omitempty"`
+	Stop             []string                `json:"stop,omitempty"`
+	PresencePenalty  float32                 `json:"presence_penalty,omitempty"`
+	FrequencyPenalty float32                 `json:"frequency_penalty,omitempty"`
+	LogitBias        map[string]int          `json:"logit_bias,omitempty"`
+	User             string                  `json:"user,omitempty"`
+}
+
+type ChatCompletionChoice struct {
+	Index        int                   `json:"index"`
+	Message      ChatCompletionMessage `json:"message"`
+	FinishReason string                `json:"finish_reason"`
+}
+
+type ChatCompletionResponse struct {
+	ID      string                 `json:"id"`
+	Object  string                 `json:"object"`
+	Created int64                  `json:"created"`
+	Model   string                 `json:"model"`
+	Choices []ChatCompletionChoice `json:"choices"`
+	Usage   struct {
+		PromptTokens     int `json:"prompt_tokens"`
+		CompletionTokens int `json:"completion_tokens"`
+		TotalTokens      int `json:"total_tokens"`
+	} `json:"usage"`
+}
+
 func AuthMiddleware() gin.HandlerFunc {
 	return func(c *gin.Context) {
 		if rootToken == "" {
@@ -249,13 +299,7 @@ func GenerateToken() string {
 	return token.String()
 }
 
-func HandleProy(c *gin.Context) {
+func getHttpClient() *http.Client {
-	var localuser bool
-	auth := c.Request.Header.Get("Authorization")
-	if len(auth) > 7 && auth[:7] == "Bearer " {
-		localuser = store.IsExistAuthCache(auth[7:])
-	}
-	client := http.DefaultClient
 	tr := &http.Transport{
 		Proxy: http.ProxyFromEnvironment,
 		DialContext: (&net.Dialer{
@@ -269,10 +313,43 @@ func HandleProy(c *gin.Context) {
 		ExpectContinueTimeout: 1 * time.Second,
 		TLSClientConfig:       &tls.Config{InsecureSkipVerify: true},
 	}
-	client.Transport = tr
+	return &http.Client{Transport: tr}
+}
 
+func HandleProy(c *gin.Context) {
+	var (
+		localuser  bool
+		isStream   bool
+		chatreq    = openai.ChatCompletionRequest{}
+		chatres    = openai.ChatCompletionResponse{}
+		chatlog    store.Tokens
+		pre_prompt string
+		wg         sync.WaitGroup
+	)
+	auth := c.Request.Header.Get("Authorization")
+	if len(auth) > 7 && auth[:7] == "Bearer " {
+		localuser = store.IsExistAuthCache(auth[7:])
+	}
+
+	if c.Request.URL.Path == "/v1/chat/completions" && localuser {
+
+		if err := c.BindJSON(&chatreq); err != nil {
+			c.AbortWithError(http.StatusBadRequest, err)
+			return
+		}
+		chatlog.Model = chatreq.Model
+		for _, m := range chatreq.Messages {
+			pre_prompt += m.Content + "\n"
+		}
+		chatlog.PromptHash = cryptor.Md5String(pre_prompt)
+		chatlog.PromptCount = NumTokensFromMessages(chatreq.Messages, chatreq.Model)
+		isStream = chatreq.Stream
+		chatlog.UserID, _ = store.GetUserID(auth[7:])
+	}
+	var body bytes.Buffer
+	json.NewEncoder(&body).Encode(chatreq)
 	// 创建 API 请求
-	req, err := http.NewRequest(c.Request.Method, baseUrl+c.Request.URL.Path, c.Request.Body)
+	req, err := http.NewRequest(c.Request.Method, baseUrl+c.Request.RequestURI, &body)
 	if err != nil {
 		log.Println(err)
 		c.JSON(http.StatusOK, gin.H{"error": err.Error()})
@@ -315,19 +392,55 @@ func HandleProy(c *gin.Context) {
 	resp.Header.Del("content-security-policy-report-only")
 	resp.Header.Del("clear-site-data")
 
-	bodyRes, err := io.ReadAll(resp.Body)
+	c.Writer.WriteHeader(resp.StatusCode)
-	if err != nil {
+	writer := bufio.NewWriter(c.Writer)
-		log.Println(err)
+	defer writer.Flush()
-		c.JSON(http.StatusUnauthorized, gin.H{"error": err.Error()})
+
+	reader := bufio.NewReader(resp.Body)
+	var resbuf = bytes.NewBuffer(nil)
+
+	if resp.StatusCode == 200 && localuser {
+		wg.Add(1)
+		if isStream {
+			contentCh := fetchResponseContent(resbuf, reader)
+			var buffer bytes.Buffer
+			for content := range contentCh {
+				buffer.WriteString(content)
+			}
+			chatlog.CompletionCount = NumTokensFromStr(buffer.String(), chatreq.Model)
+			chatlog.TotalTokens = chatlog.PromptCount + chatlog.CompletionCount
+		} else {
+			reader.WriteTo(resbuf)
+			json.NewDecoder(resbuf).Decode(&chatres)
+			chatlog.PromptCount = chatres.Usage.PromptTokens
+			chatlog.CompletionCount = chatres.Usage.CompletionTokens
+			chatlog.TotalTokens = chatres.Usage.TotalTokens
+		}
+		chatlog.Cost = fmt.Sprintf("%.6f", Cost(chatlog.Model, chatlog.PromptCount, chatlog.CompletionCount))
+		if err := store.Record(&chatlog); err != nil {
+			log.Println(err)
+		}
+		if err := store.SumDaily(chatlog.UserID); err != nil {
+			log.Println(err)
+		}
+		// 返回 API 响应主体
+		if _, err := io.Copy(writer, resbuf); err != nil {
+			c.JSON(http.StatusUnauthorized, gin.H{"error": err.Error()})
+			return
+		}
+		go func() {
+			defer wg.Done()
+			_, err = io.Copy(c.Writer, resp.Body)
+			if err != nil {
+				c.JSON(http.StatusUnauthorized, gin.H{"error": err.Error()})
+				return
+			}
+		}()
+		wg.Wait()
 		return
 	}
-	if resp.StatusCode == 200 {
-		// todo
-	}
-	resbody := io.NopCloser(bytes.NewReader(bodyRes))
 	// 返回 API 响应主体
-	c.Writer.WriteHeader(resp.StatusCode)
+	if _, err := io.Copy(writer, reader); err != nil {
-	if _, err := io.Copy(c.Writer, resbody); err != nil {
 		log.Println(err)
 		c.JSON(http.StatusUnauthorized, gin.H{"error": err.Error()})
 		return
@@ -381,3 +494,121 @@ func HandleReverseProxy(c *gin.Context) {
 	proxy.ServeHTTP(c.Writer, req)
 
 }
+func Cost(model string, promptCount, completionCount int) float64 {
+	var cost, prompt, completion float64
+	prompt = float64(promptCount)
+	completion = float64(completionCount)
+
+	switch model {
+	case "gpt-3.5-turbo", "gpt-3.5-turbo-0301":
+		cost = 0.002 * float64((prompt+completion)/1000)
+	case "gpt-4", "gpt-4-0314":
+		cost = 0.03*float64(prompt/1000) + 0.06*float64(completion/1000)
+	case "gpt-4-32k", "gpt-4-32k-0314":
+		cost = 0.06*float64(prompt/1000) + 0.12*float64(completion/1000)
+	}
+	return cost
+}
+
+func HandleUsage(c *gin.Context) {
+	fromStr := c.Query("from")
+	toStr := c.Query("to")
+
+	usage, err := store.QueryUsage(fromStr, toStr)
+	if err != nil {
+		c.JSON(http.StatusForbidden, gin.H{"error": err.Error()})
+		return
+	}
+
+	c.JSON(200, usage)
+}
+
+func fetchResponseContent(buf *bytes.Buffer, responseBody *bufio.Reader) <-chan string {
+	contentCh := make(chan string)
+	go func() {
+		defer close(contentCh)
+		for {
+			line, err := responseBody.ReadString('\n')
+			if err == nil {
+				buf.WriteString(line)
+				if line == "\n" {
+					continue
+				}
+				if strings.HasPrefix(line, "data:") {
+					line = strings.TrimSpace(strings.TrimPrefix(line, "data:"))
+					if strings.HasSuffix(line, "[DONE]") {
+						break
+					}
+					line = strings.TrimSpace(line)
+				}
+
+				dec := json.NewDecoder(strings.NewReader(line))
+				var data map[string]interface{}
+				if err := dec.Decode(&data); err == io.EOF {
+					log.Println("EOF:", err)
+					break
+				} else if err != nil {
+					fmt.Println("Error decoding response:", err)
+					return
+				}
+				if choices, ok := data["choices"].([]interface{}); ok {
+					for _, choice := range choices {
+						choiceMap := choice.(map[string]interface{})
+						if content, ok := choiceMap["delta"].(map[string]interface{})["content"]; ok {
+							contentCh <- content.(string)
+						}
+					}
+				}
+			} else {
+				break
+			}
+		}
+	}()
+	return contentCh
+}
+
+func NumTokensFromMessages(messages []openai.ChatCompletionMessage, model string) (num_tokens int) {
+	tkm, err := tiktoken.EncodingForModel(model)
+	if err != nil {
+		err = fmt.Errorf("EncodingForModel: %v", err)
+		fmt.Println(err)
+		return
+	}
+
+	var tokens_per_message int
+	var tokens_per_name int
+	if model == "gpt-3.5-turbo-0301" || model == "gpt-3.5-turbo" {
+		tokens_per_message = 4
+		tokens_per_name = -1
+	} else if model == "gpt-4-0314" || model == "gpt-4" {
+		tokens_per_message = 3
+		tokens_per_name = 1
+	} else {
+		fmt.Println("Warning: model not found. Using cl100k_base encoding.")
+		tokens_per_message = 3
+		tokens_per_name = 1
+	}
+
+	for _, message := range messages {
+		num_tokens += tokens_per_message
+		num_tokens += len(tkm.Encode(message.Content, nil, nil))
+		// num_tokens += len(tkm.Encode(message.Role, nil, nil))
+		if message.Name != "" {
+			num_tokens += tokens_per_name
+		}
+	}
+	num_tokens += 3
+	return num_tokens
+}
+
+func NumTokensFromStr(messages string, model string) (num_tokens int) {
+	tkm, err := tiktoken.EncodingForModel(model)
+	if err != nil {
+		err = fmt.Errorf("EncodingForModel: %v", err)
+		fmt.Println(err)
+		return
+	}
+
+	num_tokens += len(tkm.Encode(messages, nil, nil))
+	return num_tokens
+}

store/db.go

@@ -11,6 +11,8 @@ import (
 
 var db *gorm.DB
 
+var usage *gorm.DB
+
 func init() {
 	if _, err := os.Stat("db"); os.IsNotExist(err) {
 		errDir := os.MkdirAll("db", 0755)
@@ -31,4 +33,13 @@ func init() {
 	}
 	LoadKeysCache()
 	LoadAuthCache()
+
+	usage, err = gorm.Open(sqlite.Open("./db/usage.db"), &gorm.Config{})
+	if err != nil {
+		panic(err)
+	}
+	err = usage.AutoMigrate(&DailyUsage{}, &Usage{})
+	if err != nil {
+		panic(err)
+	}
 }

store/usage.go

@@ -0,0 +1,149 @@
+package store
+
+import (
+	"errors"
+	"time"
+
+	"github.com/Sakurasan/to"
+	"gorm.io/gorm"
+)
+
+type DailyUsage struct {
+	ID              int       `gorm:"column:id"`
+	UserID          int       `gorm:"column:user_id";primaryKey`
+	Date            time.Time `gorm:"column:date"`
+	SKU             string    `gorm:"column:sku"`
+	PromptUnits     int       `gorm:"column:prompt_units"`
+	CompletionUnits int       `gorm:"column:completion_units"`
+	TotalUnit       int       `gorm:"column:total_unit"`
+	Cost            string    `gorm:"column:cost"`
+}
+
+func (DailyUsage) TableName() string {
+	return "daily_usages"
+}
+
+type Usage struct {
+	ID              int       `gorm:"column:id"`
+	PromptHash      string    `gorm:"column:prompt_hash"`
+	UserID          int       `gorm:"column:user_id"`
+	SKU             string    `gorm:"column:sku"`
+	PromptUnits     int       `gorm:"column:prompt_units"`
+	CompletionUnits int       `gorm:"column:completion_units"`
+	TotalUnit       int       `gorm:"column:total_unit"`
+	Cost            string    `gorm:"column:cost"`
+	Date            time.Time `gorm:"column:date"`
+}
+
+func (Usage) TableName() string {
+	return "usages"
+}
+
+type Summary struct {
+	UserId             int     `gorm:"column:user_id"`
+	SumPromptUnits     int     `gorm:"column:sum_prompt_units"`
+	SumCompletionUnits int     `gorm:"column:sum_completion_units"`
+	SumTotalUnit       int     `gorm:"column:sum_total_unit"`
+	SumCost            float64 `gorm:"column:sum_cost"`
+}
+type CalcUsage struct {
+	UserID    int    `json:"userId,omitempty"`
+	TotalUnit int    `json:"totalUnit,omitempty"`
+	Cost      string `json:"cost,omitempty"`
+}
+
+func QueryUsage(from, to string) ([]CalcUsage, error) {
+	var results = []CalcUsage{}
+	err := usage.Model(&DailyUsage{}).Select(`user_id, 
+	--SUM(prompt_units) AS prompt_units,
+	-- SUM(completion_units) AS completion_units,
+	SUM(total_unit) AS total_unit,
+	printf('%.6f', SUM(cost)) AS cost`).
+		Group("user_id").
+		Where("date >= ? AND date < ?", from, to).
+		Find(&results).Error
+	if err != nil {
+		return nil, err
+	}
+	return results, nil
+}
+
+type Tokens struct {
+	UserID          int
+	PromptCount     int
+	CompletionCount int
+	TotalTokens     int
+	Cost            string
+	Model           string
+	PromptHash      string
+}
+
+func Record(chatlog *Tokens) (err error) {
+	u := &Usage{
+		UserID:          chatlog.UserID,
+		SKU:             chatlog.Model,
+		PromptHash:      chatlog.PromptHash,
+		PromptUnits:     chatlog.PromptCount,
+		CompletionUnits: chatlog.CompletionCount,
+		TotalUnit:       chatlog.TotalTokens,
+		Cost:            to.String(chatlog.Cost),
+		Date:            time.Now(),
+	}
+	err = usage.Create(u).Error
+	return
+
+}
+
+func SumDaily(userid int) error {
+	var count int64
+	err := usage.Model(&DailyUsage{}).Where("user_id = ? and date = ?", userid, time.Date(time.Now().Year(), time.Now().Month(), time.Now().Day(), 0, 0, 0, 0, time.UTC)).Count(&count).Error
+	if err != nil && !errors.Is(err, gorm.ErrRecordNotFound) {
+		return err
+	}
+	if count == 0 {
+		if err := insertSumDaily(userid); err != nil {
+			return err
+		}
+	} else {
+		if err := updateSumDaily(userid, time.Date(time.Now().Year(), time.Now().Month(), time.Now().Day(), 0, 0, 0, 0, time.UTC)); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func insertSumDaily(uid int) error {
+	nowstr := time.Date(time.Now().Year(), time.Now().Month(), time.Now().Day(), 0, 0, 0, 0, time.UTC)
+	err := usage.Exec(`INSERT INTO daily_usages
+		(user_id, date, sku, prompt_units, completion_units, total_unit, cost)
+		SELECT 
+		user_id,
+		?,
+		sku,
+		SUM(prompt_units) AS sum_prompt_units,
+		SUM(completion_units) AS sum_completion_units,
+		SUM(total_unit) AS sum_total_unit,
+		SUM(cost) AS sum_cost
+		FROM usages 
+		WHERE date >= ?
+		AND user_id = ?`, nowstr, nowstr, uid).Error
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+func updateSumDaily(uid int, date time.Time) error {
+	// var u = Summary{}
+	err := usage.Model(&Usage{}).Exec(`UPDATE daily_usages
+	SET 
+	prompt_units = (SELECT SUM(prompt_units) FROM usages WHERE user_id = daily_usages.user_id AND date >= daily_usages.date),
+	completion_units = (SELECT SUM(completion_units) FROM usages WHERE user_id = daily_usages.user_id AND date >= daily_usages.date),
+	total_unit = (SELECT SUM(total_unit) FROM usages WHERE user_id = daily_usages.user_id AND date >= daily_usages.date),
+	cost = (SELECT SUM(cost) FROM usages WHERE user_id = daily_usages.user_id AND date >= daily_usages.date)
+	WHERE user_id = ? AND date >= ?`, uid, date).Error
+	if err != nil {
+		return err
+	}
+	return nil
+}

store/userdb.go

@@ -73,6 +73,15 @@ func GetUserByName(name string) (*User, error) {
 	return &user, nil
 }
 
+func GetUserID(authkey string) (int, error) {
+	var user User
+	result := db.Where(&User{Token: authkey}).First(&user)
+	if result.Error != nil {
+		return 0, result.Error
+	}
+	return int(user.ID), nil
+}
+
 func GetAllUsers() ([]*User, error) {
 	var users []*User
 	result := db.Find(&users)