支持统计用户的使用量

2023-04-27 18:08:05 +08:00
parent ecffae0a19
commit 33034a9db0
7 changed files with 442 additions and 25 deletions
--- a/router/router.go
+++ b/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
-	baseUrl   = "https://api.openai.com"
+	rootToken     string
+	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) {
-	var localuser bool
-	auth := c.Request.Header.Get("Authorization")
-	if len(auth) > 7 && auth[:7] == "Bearer " {
-		localuser = store.IsExistAuthCache(auth[7:])
-	}
-	client := http.DefaultClient
+func getHttpClient() *http.Client {
 	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)
-	if err != nil {
-		log.Println(err)
-		c.JSON(http.StatusUnauthorized, gin.H{"error": err.Error()})
+	c.Writer.WriteHeader(resp.StatusCode)
+	writer := bufio.NewWriter(c.Writer)
+	defer writer.Flush()
+
+	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(c.Writer, resbody); err != nil {
+	if _, err := io.Copy(writer, reader); 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
+}