bug fix

.drone.yml

@@ -4,7 +4,7 @@ name: api2gpt-mid
 
 trigger:
   branch:
-    - main
+    - dev
 
 steps:
   - name: build
@@ -25,7 +25,7 @@ steps:
       # 来源目录
       source: deploy/
       # 目标服务器目录
-      target: /data/wwwroot/api2gpt/mid
+      target: /data/wwwroot/api2gpt/mid2
       script:
-        - cd /data/wwwroot/api2gpt/mid
+        - cd /data/wwwroot/api2gpt/mid2
         - docker-compose up -d --build

.gitignore

@@ -1,2 +1,4 @@
 gin.log
 /logs
+.idea
+.vscode

README.md

@@ -7,9 +7,6 @@ docker build -t api2gpt-mid .
 docker run -d -p 8081:8080 --name=api2gpt-mid --env REDIS_ADDRESS=172.17.0.1:6379 api2gpt-mid
 docker run -p 8081:8080 --name=api2gpt-mid --env REDIS_ADDRESS=172.17.0.1:6379 api2gpt-mid
 
-
-
-
 ## nginx 配置
 proxy_set_header  Host   $host;
 proxy_set_header  X-Real-IP  $remote_addr;
@@ -22,4 +19,5 @@ proxy_buffering off;
 proxy_pass http://localhost:8081/;
 
 
-## ok
+## local run
+go run api2gpt-mid --port 8080 --log-dir ./logs

api.go

@@ -1,102 +0,0 @@
-package main
-
-import (
-	"bytes"
-	"encoding/json"
-	"io/ioutil"
-	"net/http"
-
-	"github.com/gin-gonic/gin"
-)
-
-type BalanceInfo struct {
-	ServerAddress string  `json:"server_address"`
-	AvailableKey  string  `json:"available_key"`
-	UserBalance   float64 `json:"user_balance"`
-	TokenRatio    float64 `json:"token_ratio"`
-}
-
-type Consumption struct {
-	SecretKey        string `json:"secretKey"`
-	Model            string `json:"model"`
-	MsgId            string `json:"msgId"`
-	PromptTokens     int    `json:"promptTokens"`
-	CompletionTokens int    `json:"completionTokens"`
-	TotalTokens      int    `json:"totalTokens"`
-}
-
-func mockBalanceInquiry(c *gin.Context) {
-	c.JSON(http.StatusOK, gin.H{
-		"server_address": "https://gptp.any-door.cn",
-		"available_key":  "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",
-		"user_balance":   10000,
-		"token_ratio":    1000,
-	})
-}
-
-func mockBalanceConsumption(c *gin.Context) {
-	c.JSON(http.StatusOK, gin.H{
-		"success": "true",
-	})
-}
-
-// 余额查询api调用
-func balanceInquiry(key string, model string) (*BalanceInfo, error) {
-	url := "http://localhost:8080/mock1?key=" + key + "&model=" + model
-	req, err := http.NewRequest("POST", url, nil)
-	req.Header.Set("Content-Type", "application/json")
-
-	client := &http.Client{}
-	resp, err := client.Do(req)
-	if err != nil {
-		panic(err)
-	}
-	defer resp.Body.Close()
-
-	body, err := ioutil.ReadAll(resp.Body)
-	if err != nil {
-		return nil, err
-	}
-
-	var balanceInfo BalanceInfo
-	if err := json.Unmarshal(body, &balanceInfo); err != nil {
-		return nil, err
-	}
-
-	return &balanceInfo, nil
-}
-
-// 余额消费
-func balanceConsumption(key string, model string, prompt_tokens int, completion_tokens int, total_tokens int, msg_id string) (string, error) {
-	var data = Consumption{
-		SecretKey:        key,
-		Model:            model,
-		MsgId:            msg_id,
-		PromptTokens:     prompt_tokens,
-		CompletionTokens: completion_tokens,
-		TotalTokens:      total_tokens,
-	}
-
-	jsonData, err := json.Marshal(data)
-	// 构造post请求的body
-	reqBody := bytes.NewBuffer(jsonData)
-
-	url := "http://172.17.0.1:8080/other/usageRecord"
-	req2, err := http.NewRequest("POST", url, reqBody)
-
-	// 设置http请求的header
-	req2.Header.Set("Content-Type", "application/json")
-
-	client := &http.Client{}
-	resp, err := client.Do(req2)
-	if err != nil {
-		return "", err
-	}
-	defer resp.Body.Close()
-
-	body, err := ioutil.ReadAll(resp.Body)
-	if err != nil {
-		return "", err
-	}
-	return string(body), nil
-}

api/server.go

@@ -0,0 +1,53 @@
+package api
+
+import (
+	"api2gpt-mid/model"
+	"bytes"
+	"encoding/json"
+	"io"
+	"net/http"
+	"os"
+)
+
+// BalanceConsumption 余额消费
+func BalanceConsumption(key string, modelStr string, prompt_tokens int, completion_tokens int, total_tokens int, msg_id string) (string, error) {
+	var data = model.Consumption{
+		SecretKey:        key,
+		Model:            modelStr,
+		MsgId:            msg_id,
+		PromptTokens:     prompt_tokens,
+		CompletionTokens: completion_tokens,
+		TotalTokens:      total_tokens,
+	}
+
+	jsonData, err := json.Marshal(data)
+	// 构造post请求的body
+	reqBody := bytes.NewBuffer(jsonData)
+
+	url := os.Getenv("SERVER_API_USAGE_RECORD_STRING")
+	if url == "" {
+		url = "http://172.17.0.1:8080/other/usageRecord"
+	}
+	req2, err := http.NewRequest("POST", url, reqBody)
+
+	// 设置http请求的header
+	req2.Header.Set("Content-Type", "application/json")
+
+	client := &http.Client{}
+	resp, err := client.Do(req2)
+	if err != nil {
+		return "", err
+	}
+	defer func(Body io.ReadCloser) {
+		err := Body.Close()
+		if err != nil {
+
+		}
+	}(resp.Body)
+
+	body, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return "", err
+	}
+	return string(body), nil
+}

common/constants.go

@@ -0,0 +1,12 @@
+package common
+
+import (
+	"os"
+	"time"
+)
+
+var StartTime = time.Now().Unix()
+var Version = "v0.0.0"
+var SystemName = "Api2gpt Mid"
+
+var DebugEnabled = os.Getenv("DEBUG") == "true"

common/init.go

@@ -0,0 +1,51 @@
+package common
+
+import (
+	"flag"
+	"fmt"
+	"log"
+	"os"
+	"path/filepath"
+)
+
+var (
+	Port         = flag.Int("port", 8080, "the listening port")
+	PrintVersion = flag.Bool("version", false, "print version and exit")
+	PrintHelp    = flag.Bool("help", false, "print help and exit")
+	LogDir       = flag.String("log-dir", "", "specify the log directory")
+)
+
+func printHelp() {
+	fmt.Println("Api2gpt " + Version + " - Mid Service.")
+	fmt.Println("Copyright (C) 2023 api2gpt. All rights reserved.")
+	fmt.Println("WebSite: https://www.api2gpt.com")
+	fmt.Println("Usage: one-api [--port <port>] [--log-dir <log directory>] [--version] [--help]")
+}
+
+func init() {
+	flag.Parse()
+
+	if *PrintVersion {
+		fmt.Println(Version)
+		os.Exit(0)
+	}
+
+	if *PrintHelp {
+		printHelp()
+		os.Exit(0)
+	}
+
+	if *LogDir != "" {
+		var err error
+		*LogDir, err = filepath.Abs(*LogDir)
+		if err != nil {
+			log.Fatal(err)
+		}
+		if _, err := os.Stat(*LogDir); os.IsNotExist(err) {
+			err = os.Mkdir(*LogDir, 0777)
+			if err != nil {
+				log.Fatal(err)
+			}
+		}
+	}
+}

common/logger.go

@@ -0,0 +1,44 @@
+package common
+
+import (
+	"fmt"
+	"github.com/gin-gonic/gin"
+	"io"
+	"log"
+	"os"
+	"path/filepath"
+	"time"
+)
+
+func SetupGinLog() {
+	if *LogDir != "" {
+		commonLogPath := filepath.Join(*LogDir, "common.log")
+		errorLogPath := filepath.Join(*LogDir, "error.log")
+		commonFd, err := os.OpenFile(commonLogPath, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
+		if err != nil {
+			log.Fatal("failed to open log file")
+		}
+		errorFd, err := os.OpenFile(errorLogPath, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
+		if err != nil {
+			log.Fatal("failed to open log file")
+		}
+		gin.DefaultWriter = io.MultiWriter(os.Stdout, commonFd)
+		gin.DefaultErrorWriter = io.MultiWriter(os.Stderr, errorFd)
+	}
+}
+
+func SysLog(s string) {
+	t := time.Now()
+	_, _ = fmt.Fprintf(gin.DefaultWriter, "[SYS] %v | %s \n", t.Format("2006/01/02 - 15:04:05"), s)
+}
+
+func SysError(s string) {
+	t := time.Now()
+	_, _ = fmt.Fprintf(gin.DefaultErrorWriter, "[SYS] %v | %s \n", t.Format("2006/01/02 - 15:04:05"), s)
+}
+
+func FatalLog(v ...any) {
+	t := time.Now()
+	_, _ = fmt.Fprintf(gin.DefaultErrorWriter, "[FATAL] %v | %v \n", t.Format("2006/01/02 - 15:04:05"), v)
+	os.Exit(1)
+}

common/redis.go

@@ -0,0 +1,77 @@
+package common
+
+import (
+	"context"
+	"github.com/redis/go-redis/v9"
+	"os"
+	"time"
+)
+
+var RDB *redis.Client
+
+// InitRedisClient This function is called after init()
+func InitRedisClient() (err error) {
+	SysLog("Redis start connection")
+	redisConnStr := os.Getenv("REDIS_CONN_STRING")
+	if redisConnStr == "" {
+		redisConnStr = "redis://@localhost:6379/0?dial_timeout=5s"
+	}
+	opt, err := redis.ParseURL(redisConnStr)
+	//opt, err := redis.ParseURL(os.Getenv("REDIS_CONN_STRING"))
+	if err != nil {
+		FatalLog("failed to parse Redis connection string: " + err.Error())
+	}
+	RDB = redis.NewClient(opt)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+
+	_, err = RDB.Ping(ctx).Result()
+	if err != nil {
+		FatalLog("Redis ping test failed: " + err.Error())
+	}
+	return err
+}
+
+func ParseRedisOption() *redis.Options {
+	opt, err := redis.ParseURL(os.Getenv("REDIS_CONN_STRING"))
+	if err != nil {
+		FatalLog("failed to parse Redis connection string: " + err.Error())
+	}
+	return opt
+}
+
+func RedisSet(key string, value string, expiration time.Duration) error {
+	ctx := context.Background()
+	return RDB.Set(ctx, key, value, expiration).Err()
+}
+
+func RedisGet(key string) (string, error) {
+	ctx := context.Background()
+	return RDB.Get(ctx, key).Result()
+}
+
+func RedisDel(key string) error {
+	ctx := context.Background()
+	return RDB.Del(ctx, key).Err()
+}
+
+func RedisDecrease(key string, value int64) error {
+	ctx := context.Background()
+	return RDB.DecrBy(ctx, key, value).Err()
+}
+
+func RedisIncrByFloat(key string, value float64) (float64, error) {
+	ctx := context.Background()
+	return RDB.IncrByFloat(ctx, key, value).Result()
+}
+
+func RedisIncr(key string) (int64, error) {
+	ctx := context.Background()
+	return RDB.Incr(ctx, key).Result()
+}
+
+func RedisExpire(key string, timeOut time.Duration) (bool, error) {
+	ctx := context.Background()
+	return RDB.Expire(ctx, key, timeOut).Result()
+}

common/utils.go

@@ -0,0 +1,50 @@
+package common
+
+import (
+	"api2gpt-mid/model"
+	"fmt"
+	"github.com/pkoukk/tiktoken-go"
+	"strings"
+)
+
+// NumTokensFromMessages 计算Messages中的token数量
+func NumTokensFromMessages(messages []model.Message, model string) int {
+	if strings.Contains(model, "gpt-3.5") {
+		model = "gpt-3.5-turbo"
+	}
+	if strings.Contains(model, "gpt-4") {
+		model = "gpt-4"
+	}
+	tkm, err := tiktoken.EncodingForModel(model)
+	if err != nil {
+		err = fmt.Errorf("getEncoding: %v", err)
+		return 0
+	}
+	numTokens := 0
+	for _, message := range messages {
+		numTokens += len(tkm.Encode(message.Content, nil, nil))
+		numTokens += 6
+	}
+	numTokens += 3
+	return numTokens
+}
+
+// NumTokensFromString 计算String中的token数量
+func NumTokensFromString(msg string, model string) int {
+	if strings.Contains(model, "gpt-3.5") {
+		model = "gpt-3.5-turbo"
+	}
+	if strings.Contains(model, "gpt-4") {
+		model = "gpt-4"
+	}
+	tkm, err := tiktoken.EncodingForModel(model)
+	if err != nil {
+		err = fmt.Errorf("getEncoding: %v", err)
+		return 0
+	}
+	if model == "text-davinci-003" {
+		return len(tkm.Encode(msg, nil, nil)) + 1
+	} else {
+		return len(tkm.Encode(msg, nil, nil)) + 9
+	}
+}

controller/billing.go

@@ -0,0 +1,30 @@
+package controller
+
+import (
+	"api2gpt-mid/service"
+	"github.com/gin-gonic/gin"
+	"strings"
+)
+
+type CreditSummary struct {
+	Object         string  `json:"object"`
+	TotalGranted   float64 `json:"total_granted"`
+	TotalUsed      float64 `json:"total_used"`
+	TotalRemaining float64 `json:"total_remaining"`
+}
+
+// Balance 余额查询
+func Balance(c *gin.Context) {
+	auth := c.Request.Header.Get("Authorization")
+	key := strings.Trim(auth, "Bearer ")
+	balance, err := service.QueryBlance(key)
+	if err != nil {
+		c.JSON(400, gin.H{"error": err.Error()})
+	}
+	var creditSummary CreditSummary
+	creditSummary.Object = "credit_grant"
+	creditSummary.TotalGranted = 999999
+	creditSummary.TotalUsed = 999999 - balance
+	creditSummary.TotalRemaining = balance
+	c.JSON(200, creditSummary)
+}

controller/model.go

@@ -0,0 +1,458 @@
+package controller
+
+import (
+	"fmt"
+
+	"github.com/gin-gonic/gin"
+)
+
+// https://platform.openai.com/docs/api-reference/models/list
+
+type OpenAIModelPermission struct {
+	Id                 string  `json:"id"`
+	Object             string  `json:"object"`
+	Created            int     `json:"created"`
+	AllowCreateEngine  bool    `json:"allow_create_engine"`
+	AllowSampling      bool    `json:"allow_sampling"`
+	AllowLogprobs      bool    `json:"allow_logprobs"`
+	AllowSearchIndices bool    `json:"allow_search_indices"`
+	AllowView          bool    `json:"allow_view"`
+	AllowFineTuning    bool    `json:"allow_fine_tuning"`
+	Organization       string  `json:"organization"`
+	Group              *string `json:"group"`
+	IsBlocking         bool    `json:"is_blocking"`
+}
+
+type OpenAIModels struct {
+	Id         string                  `json:"id"`
+	Object     string                  `json:"object"`
+	Created    int                     `json:"created"`
+	OwnedBy    string                  `json:"owned_by"`
+	Permission []OpenAIModelPermission `json:"permission"`
+	Root       string                  `json:"root"`
+	Parent     *string                 `json:"parent"`
+}
+
+type OpenAIError struct {
+	Message string `json:"message"`
+	Type    string `json:"type"`
+	Param   string `json:"param"`
+	Code    any    `json:"code"`
+}
+
+type OpenAIErrorWithStatusCode struct {
+	OpenAIError
+	StatusCode int `json:"status_code"`
+}
+
+var openAIModels []OpenAIModels
+var openAIModelsMap map[string]OpenAIModels
+
+func init() {
+	var permission []OpenAIModelPermission
+	permission = append(permission, OpenAIModelPermission{
+		Id:                 "modelperm-LwHkVFn8AcMItP432fKKDIKJ",
+		Object:             "model_permission",
+		Created:            1626777600,
+		AllowCreateEngine:  true,
+		AllowSampling:      true,
+		AllowLogprobs:      true,
+		AllowSearchIndices: false,
+		AllowView:          true,
+		AllowFineTuning:    false,
+		Organization:       "*",
+		Group:              nil,
+		IsBlocking:         false,
+	})
+	// https://platform.openai.com/docs/models/model-endpoint-compatibility
+	openAIModels = []OpenAIModels{
+		{
+			Id:         "dall-e",
+			Object:     "model",
+			Created:    1677649963,
+			OwnedBy:    "openai",
+			Permission: permission,
+			Root:       "dall-e",
+			Parent:     nil,
+		},
+		{
+			Id:         "whisper-1",
+			Object:     "model",
+			Created:    1677649963,
+			OwnedBy:    "openai",
+			Permission: permission,
+			Root:       "whisper-1",
+			Parent:     nil,
+		},
+		{
+			Id:         "gpt-3.5-turbo",
+			Object:     "model",
+			Created:    1677649963,
+			OwnedBy:    "openai",
+			Permission: permission,
+			Root:       "gpt-3.5-turbo",
+			Parent:     nil,
+		},
+		{
+			Id:         "gpt-3.5-turbo-0301",
+			Object:     "model",
+			Created:    1677649963,
+			OwnedBy:    "openai",
+			Permission: permission,
+			Root:       "gpt-3.5-turbo-0301",
+			Parent:     nil,
+		},
+		{
+			Id:         "gpt-3.5-turbo-0613",
+			Object:     "model",
+			Created:    1677649963,
+			OwnedBy:    "openai",
+			Permission: permission,
+			Root:       "gpt-3.5-turbo-0613",
+			Parent:     nil,
+		},
+		{
+			Id:         "gpt-3.5-turbo-16k",
+			Object:     "model",
+			Created:    1677649963,
+			OwnedBy:    "openai",
+			Permission: permission,
+			Root:       "gpt-3.5-turbo-16k",
+			Parent:     nil,
+		},
+		{
+			Id:         "gpt-3.5-turbo-16k-0613",
+			Object:     "model",
+			Created:    1677649963,
+			OwnedBy:    "openai",
+			Permission: permission,
+			Root:       "gpt-3.5-turbo-16k-0613",
+			Parent:     nil,
+		},
+		{
+			Id:         "gpt-4",
+			Object:     "model",
+			Created:    1677649963,
+			OwnedBy:    "openai",
+			Permission: permission,
+			Root:       "gpt-4",
+			Parent:     nil,
+		},
+		{
+			Id:         "gpt-4-0314",
+			Object:     "model",
+			Created:    1677649963,
+			OwnedBy:    "openai",
+			Permission: permission,
+			Root:       "gpt-4-0314",
+			Parent:     nil,
+		},
+		{
+			Id:         "gpt-4-0613",
+			Object:     "model",
+			Created:    1677649963,
+			OwnedBy:    "openai",
+			Permission: permission,
+			Root:       "gpt-4-0613",
+			Parent:     nil,
+		},
+		{
+			Id:         "gpt-4-32k",
+			Object:     "model",
+			Created:    1677649963,
+			OwnedBy:    "openai",
+			Permission: permission,
+			Root:       "gpt-4-32k",
+			Parent:     nil,
+		},
+		{
+			Id:         "gpt-4-32k-0314",
+			Object:     "model",
+			Created:    1677649963,
+			OwnedBy:    "openai",
+			Permission: permission,
+			Root:       "gpt-4-32k-0314",
+			Parent:     nil,
+		},
+		{
+			Id:         "gpt-4-32k-0613",
+			Object:     "model",
+			Created:    1677649963,
+			OwnedBy:    "openai",
+			Permission: permission,
+			Root:       "gpt-4-32k-0613",
+			Parent:     nil,
+		},
+		{
+			Id:         "text-embedding-ada-002",
+			Object:     "model",
+			Created:    1677649963,
+			OwnedBy:    "openai",
+			Permission: permission,
+			Root:       "text-embedding-ada-002",
+			Parent:     nil,
+		},
+		{
+			Id:         "text-davinci-003",
+			Object:     "model",
+			Created:    1677649963,
+			OwnedBy:    "openai",
+			Permission: permission,
+			Root:       "text-davinci-003",
+			Parent:     nil,
+		},
+		{
+			Id:         "text-davinci-002",
+			Object:     "model",
+			Created:    1677649963,
+			OwnedBy:    "openai",
+			Permission: permission,
+			Root:       "text-davinci-002",
+			Parent:     nil,
+		},
+		{
+			Id:         "text-curie-001",
+			Object:     "model",
+			Created:    1677649963,
+			OwnedBy:    "openai",
+			Permission: permission,
+			Root:       "text-curie-001",
+			Parent:     nil,
+		},
+		{
+			Id:         "text-babbage-001",
+			Object:     "model",
+			Created:    1677649963,
+			OwnedBy:    "openai",
+			Permission: permission,
+			Root:       "text-babbage-001",
+			Parent:     nil,
+		},
+		{
+			Id:         "text-ada-001",
+			Object:     "model",
+			Created:    1677649963,
+			OwnedBy:    "openai",
+			Permission: permission,
+			Root:       "text-ada-001",
+			Parent:     nil,
+		},
+		{
+			Id:         "text-moderation-latest",
+			Object:     "model",
+			Created:    1677649963,
+			OwnedBy:    "openai",
+			Permission: permission,
+			Root:       "text-moderation-latest",
+			Parent:     nil,
+		},
+		{
+			Id:         "text-moderation-stable",
+			Object:     "model",
+			Created:    1677649963,
+			OwnedBy:    "openai",
+			Permission: permission,
+			Root:       "text-moderation-stable",
+			Parent:     nil,
+		},
+		{
+			Id:         "text-davinci-edit-001",
+			Object:     "model",
+			Created:    1677649963,
+			OwnedBy:    "openai",
+			Permission: permission,
+			Root:       "text-davinci-edit-001",
+			Parent:     nil,
+		},
+		{
+			Id:         "code-davinci-edit-001",
+			Object:     "model",
+			Created:    1677649963,
+			OwnedBy:    "openai",
+			Permission: permission,
+			Root:       "code-davinci-edit-001",
+			Parent:     nil,
+		},
+		{
+			Id:         "claude-instant-1",
+			Object:     "model",
+			Created:    1677649963,
+			OwnedBy:    "anturopic",
+			Permission: permission,
+			Root:       "claude-instant-1",
+			Parent:     nil,
+		},
+		{
+			Id:         "claude-2",
+			Object:     "model",
+			Created:    1677649963,
+			OwnedBy:    "anturopic",
+			Permission: permission,
+			Root:       "claude-2",
+			Parent:     nil,
+		},
+		{
+			Id:         "ERNIE-Bot",
+			Object:     "model",
+			Created:    1677649963,
+			OwnedBy:    "baidu",
+			Permission: permission,
+			Root:       "ERNIE-Bot",
+			Parent:     nil,
+		},
+		{
+			Id:         "ERNIE-Bot-turbo",
+			Object:     "model",
+			Created:    1677649963,
+			OwnedBy:    "baidu",
+			Permission: permission,
+			Root:       "ERNIE-Bot-turbo",
+			Parent:     nil,
+		},
+		{
+			Id:         "Embedding-V1",
+			Object:     "model",
+			Created:    1677649963,
+			OwnedBy:    "baidu",
+			Permission: permission,
+			Root:       "Embedding-V1",
+			Parent:     nil,
+		},
+		{
+			Id:         "PaLM-2",
+			Object:     "model",
+			Created:    1677649963,
+			OwnedBy:    "google",
+			Permission: permission,
+			Root:       "PaLM-2",
+			Parent:     nil,
+		},
+		{
+			Id:         "chatglm_pro",
+			Object:     "model",
+			Created:    1677649963,
+			OwnedBy:    "zhipu",
+			Permission: permission,
+			Root:       "chatglm_pro",
+			Parent:     nil,
+		},
+		{
+			Id:         "chatglm_std",
+			Object:     "model",
+			Created:    1677649963,
+			OwnedBy:    "zhipu",
+			Permission: permission,
+			Root:       "chatglm_std",
+			Parent:     nil,
+		},
+		{
+			Id:         "chatglm_lite",
+			Object:     "model",
+			Created:    1677649963,
+			OwnedBy:    "zhipu",
+			Permission: permission,
+			Root:       "chatglm_lite",
+			Parent:     nil,
+		},
+		{
+			Id:         "qwen-v1",
+			Object:     "model",
+			Created:    1677649963,
+			OwnedBy:    "ali",
+			Permission: permission,
+			Root:       "qwen-v1",
+			Parent:     nil,
+		},
+		{
+			Id:         "qwen-plus-v1",
+			Object:     "model",
+			Created:    1677649963,
+			OwnedBy:    "ali",
+			Permission: permission,
+			Root:       "qwen-plus-v1",
+			Parent:     nil,
+		},
+		{
+			Id:         "SparkDesk",
+			Object:     "model",
+			Created:    1677649963,
+			OwnedBy:    "xunfei",
+			Permission: permission,
+			Root:       "SparkDesk",
+			Parent:     nil,
+		},
+		{
+			Id:         "360GPT_S2_V9",
+			Object:     "model",
+			Created:    1677649963,
+			OwnedBy:    "360",
+			Permission: permission,
+			Root:       "360GPT_S2_V9",
+			Parent:     nil,
+		},
+		{
+			Id:         "embedding-bert-512-v1",
+			Object:     "model",
+			Created:    1677649963,
+			OwnedBy:    "360",
+			Permission: permission,
+			Root:       "embedding-bert-512-v1",
+			Parent:     nil,
+		},
+		{
+			Id:         "embedding_s1_v1",
+			Object:     "model",
+			Created:    1677649963,
+			OwnedBy:    "360",
+			Permission: permission,
+			Root:       "embedding_s1_v1",
+			Parent:     nil,
+		},
+		{
+			Id:         "semantic_similarity_s1_v1",
+			Object:     "model",
+			Created:    1677649963,
+			OwnedBy:    "360",
+			Permission: permission,
+			Root:       "semantic_similarity_s1_v1",
+			Parent:     nil,
+		},
+		{
+			Id:         "360GPT_S2_V9.4",
+			Object:     "model",
+			Created:    1677649963,
+			OwnedBy:    "360",
+			Permission: permission,
+			Root:       "360GPT_S2_V9.4",
+			Parent:     nil,
+		},
+	}
+	openAIModelsMap = make(map[string]OpenAIModels)
+	for _, model := range openAIModels {
+		openAIModelsMap[model.Id] = model
+	}
+}
+
+func ListModels(c *gin.Context) {
+	c.JSON(200, gin.H{
+		"object": "list",
+		"data":   openAIModels,
+	})
+}
+
+func RetrieveModel(c *gin.Context) {
+	modelId := c.Param("model")
+	if model, ok := openAIModelsMap[modelId]; ok {
+		c.JSON(200, model)
+	} else {
+		openAIError := OpenAIError{
+			Message: fmt.Sprintf("The model '%s' does not exist", modelId),
+			Type:    "invalid_request_error",
+			Param:   "model",
+			Code:    "model_not_found",
+		}
+		c.JSON(200, gin.H{
+			"error": openAIError,
+		})
+	}
+}

controller/relay.go

@@ -0,0 +1,465 @@
+package controller
+
+import (
+	"api2gpt-mid/common"
+	"api2gpt-mid/model"
+	"api2gpt-mid/service"
+	"bufio"
+	"bytes"
+	"encoding/json"
+	"fmt"
+	"github.com/gin-gonic/gin"
+	"io"
+	"log"
+	"math/rand"
+	"net/http"
+	"net/http/httputil"
+	"net/url"
+	"strings"
+	"time"
+)
+
+func Images(c *gin.Context) {
+	var imagesRequest model.ImagesRequest
+	var modelStr = "images-generations"
+	if err := c.ShouldBindJSON(&imagesRequest); err != nil {
+		c.AbortWithStatusJSON(400, gin.H{"error": err.Error()})
+		return
+	}
+	auth := c.Request.Header.Get("Authorization")
+	key := auth[7:]
+
+	serverInfo, err := service.CheckBlanceForImages(key, modelStr, imagesRequest.N)
+	if err != nil {
+		c.AbortWithStatusJSON(403, gin.H{"error": err.Error()})
+		log.Printf("请求出错 KEY: %v Model: %v ERROR: %v", key, modelStr, err)
+		return
+	}
+
+	log.Printf("请求的KEY: %v Model: %v", key, modelStr)
+
+	remote, err := url.Parse(serverInfo.ServerAddress)
+	if err != nil {
+		c.AbortWithStatusJSON(400, gin.H{"error": err.Error()})
+		return
+	}
+
+	proxy := httputil.NewSingleHostReverseProxy(remote)
+	newReqBody, err := json.Marshal(imagesRequest)
+	if err != nil {
+		log.Printf("http request err: %v", err)
+		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": err.Error()})
+		return
+	}
+
+	proxy.Director = func(req *http.Request) {
+		req.Header = c.Request.Header
+		req.Host = remote.Host
+		req.URL.Scheme = remote.Scheme
+		req.URL.Host = remote.Host
+		req.URL.Path = c.Request.URL.Path
+		req.ContentLength = int64(len(newReqBody))
+		req.Header.Set("Content-Type", "application/json")
+		req.Header.Set("Accept-Encoding", "")
+		serverKey := serverInfo.AvailableKey
+		keyList := strings.Split(serverInfo.AvailableKey, ",")
+		if len(keyList) > 1 {
+			// 随机数种子
+			source := rand.NewSource(time.Now().UnixNano())
+			random := rand.New(source)
+			// 从数组中随机选择一个元素
+			serverKey = keyList[random.Intn(len(keyList))]
+		}
+		req.Header.Set("Authorization", "Bearer "+serverKey)
+		req.Body = io.NopCloser(bytes.NewReader(newReqBody))
+	}
+	sss, err := json.Marshal(imagesRequest)
+	log.Printf("开始处理返回逻辑 %d", string(sss))
+
+	proxy.ModifyResponse = func(resp *http.Response) error {
+		if resp.StatusCode != http.StatusOK {
+			//退回预扣除的余额
+			err := service.CheckBlanceReturnForImages(key, modelStr, imagesRequest.N)
+			if err != nil {
+				return err
+			}
+			return nil
+		}
+		resp.Header.Set("Openai-Organization", "api2gpt")
+		var imagesResponse model.ImagesResponse
+		body, err := io.ReadAll(resp.Body)
+		if err != nil {
+			log.Printf("读取返回数据出错: %v", err)
+			return err
+		}
+		err = json.Unmarshal(body, &imagesResponse)
+		if err != nil {
+			log.Printf("json解析数据出错: %v", err)
+			return err
+		}
+		resp.Body = io.NopCloser(bytes.NewReader(body))
+		log.Printf("image size:  %v", len(imagesResponse.Data))
+		timestamp := time.Now().Unix()
+		timestampID := "img-" + fmt.Sprintf("%d", timestamp)
+		//消费余额
+		_, err = service.ConsumptionForImages(key, modelStr, imagesRequest.N, len(imagesResponse.Data), timestampID)
+		if err != nil {
+			return err
+		}
+		return nil
+	}
+	proxy.ServeHTTP(c.Writer, c.Request)
+}
+
+func Edit(c *gin.Context) {
+	var promptTokens int
+	var complateTokens int
+	var totalTokens int
+	var chatRequest model.ChatRequest
+	if err := c.ShouldBindJSON(&chatRequest); err != nil {
+		c.AbortWithStatusJSON(400, gin.H{"error": err.Error()})
+		return
+	}
+
+	auth := c.Request.Header.Get("Authorization")
+	key := auth[7:]
+	//根据KEY调用用户余额接口，判断是否有足够的余额， 后期可考虑判断max_tokens参数来调整
+	serverInfo, err := service.CheckBlance(key, chatRequest.Model, chatRequest.MaxTokens)
+	if err != nil {
+		c.AbortWithStatusJSON(403, gin.H{"error": err.Error()})
+		log.Printf("请求出错 KEY: %v Model: %v ERROR: %v", key, chatRequest.Model, err)
+		return
+	}
+
+	log.Printf("请求的KEY: %v Model: %v", key, chatRequest.Model)
+
+	remote, err := url.Parse(serverInfo.ServerAddress)
+	if err != nil {
+		c.AbortWithStatusJSON(400, gin.H{"error": err.Error()})
+		return
+	}
+
+	proxy := httputil.NewSingleHostReverseProxy(remote)
+	newReqBody, err := json.Marshal(chatRequest)
+	if err != nil {
+		log.Printf("http request err: %v", err)
+		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": err.Error()})
+		return
+	}
+
+	proxy.Director = func(req *http.Request) {
+		req.Header = c.Request.Header
+		req.Host = remote.Host
+		req.URL.Scheme = remote.Scheme
+		req.URL.Host = remote.Host
+		req.URL.Path = c.Request.URL.Path
+		req.ContentLength = int64(len(newReqBody))
+		req.Header.Set("Content-Type", "application/json")
+		req.Header.Set("Accept-Encoding", "")
+		serverKey := serverInfo.AvailableKey
+		keyList := strings.Split(serverInfo.AvailableKey, ",")
+		if len(keyList) > 1 {
+			// 随机数种子
+			source := rand.NewSource(time.Now().UnixNano())
+			random := rand.New(source)
+			// 从数组中随机选择一个元素
+			serverKey = keyList[random.Intn(len(keyList))]
+		}
+		req.Header.Set("Authorization", "Bearer "+serverKey)
+		req.Body = io.NopCloser(bytes.NewReader(newReqBody))
+	}
+	sss, err := json.Marshal(chatRequest)
+	log.Printf("开始处理返回逻辑 %d", string(sss))
+
+	proxy.ModifyResponse = func(resp *http.Response) error {
+		resp.Header.Set("Openai-Organization", "api2gpt")
+		var chatResponse model.ChatResponse
+		body, err := io.ReadAll(resp.Body)
+		if err != nil {
+			log.Printf("读取返回数据出错: %v", err)
+			return err
+		}
+		err = json.Unmarshal(body, &chatResponse)
+		if err != nil {
+			log.Printf("json解析数据出错: %v", err)
+			return err
+		}
+		promptTokens = chatResponse.Usage.PromptTokens
+		complateTokens = chatResponse.Usage.CompletionTokens
+		totalTokens = chatResponse.Usage.TotalTokens
+		resp.Body = io.NopCloser(bytes.NewReader(body))
+		log.Printf("prompt_tokens:  %v complate_tokens: %v total_tokens: %v", promptTokens, complateTokens, totalTokens)
+		timestamp := time.Now().Unix()
+		timestampID := "edit-" + fmt.Sprintf("%d", timestamp)
+		//消费余额
+		_, err = service.Consumption(key, chatRequest.Model, promptTokens, 0, totalTokens, timestampID)
+		if err != nil {
+			return err
+		}
+		return nil
+	}
+
+	proxy.ServeHTTP(c.Writer, c.Request)
+
+}
+
+func Embeddings(c *gin.Context) {
+	var promptTokens int
+	var totalTokens int
+	var chatRequest model.ChatRequest
+	if err := c.ShouldBindJSON(&chatRequest); err != nil {
+		c.AbortWithStatusJSON(400, gin.H{"error": err.Error()})
+		return
+	}
+	auth := c.Request.Header.Get("Authorization")
+	key := auth[7:]
+	serverInfo, err := service.CheckBlance(key, chatRequest.Model, chatRequest.MaxTokens)
+	if err != nil {
+		c.AbortWithStatusJSON(403, gin.H{"error": err.Error()})
+		log.Printf("请求出错 KEY: %v Model: %v ERROR: %v", key, chatRequest.Model, err)
+		return
+	}
+
+	log.Printf("请求的KEY: %v Model: %v", key, chatRequest.Model)
+
+	remote, err := url.Parse(serverInfo.ServerAddress)
+	if err != nil {
+		c.AbortWithStatusJSON(400, gin.H{"error": err.Error()})
+		return
+	}
+
+	proxy := httputil.NewSingleHostReverseProxy(remote)
+	newReqBody, err := json.Marshal(chatRequest)
+	if err != nil {
+		log.Printf("http request err: %v", err)
+		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": err.Error()})
+		return
+	}
+
+	proxy.Director = func(req *http.Request) {
+		req.Header = c.Request.Header
+		req.Host = remote.Host
+		req.URL.Scheme = remote.Scheme
+		req.URL.Host = remote.Host
+		req.URL.Path = c.Request.URL.Path
+		req.ContentLength = int64(len(newReqBody))
+		req.Header.Set("Content-Type", "application/json")
+		req.Header.Set("Accept-Encoding", "")
+		serverKey := serverInfo.AvailableKey
+		keyList := strings.Split(serverInfo.AvailableKey, ",")
+		if len(keyList) > 1 {
+			// 随机数种子
+			source := rand.NewSource(time.Now().UnixNano())
+			random := rand.New(source)
+			// 从数组中随机选择一个元素
+			serverKey = keyList[random.Intn(len(keyList))]
+		}
+		req.Header.Set("Authorization", "Bearer "+serverKey)
+		req.Body = io.NopCloser(bytes.NewReader(newReqBody))
+	}
+	sss, err := json.Marshal(chatRequest)
+	log.Printf("开始处理返回逻辑 %d", string(sss))
+
+	proxy.ModifyResponse = func(resp *http.Response) error {
+		resp.Header.Set("Openai-Organization", "api2gpt")
+		var chatResponse model.ChatResponse
+		body, err := io.ReadAll(resp.Body)
+		if err != nil {
+			log.Printf("读取返回数据出错: %v", err)
+			return err
+		}
+		err = json.Unmarshal(body, &chatResponse)
+		if err != nil {
+			log.Printf("json解析数据出错: %v", err)
+			return err
+		}
+		promptTokens = chatResponse.Usage.PromptTokens
+		totalTokens = chatResponse.Usage.TotalTokens
+		resp.Body = io.NopCloser(bytes.NewReader(body))
+		log.Printf("prompt_tokens:  %v total_tokens: %v", promptTokens, totalTokens)
+		timestamp := time.Now().Unix()
+		timestampID := "emb-" + fmt.Sprintf("%d", timestamp)
+		//消费余额
+		_, err = service.Consumption(key, chatRequest.Model, promptTokens, 0, totalTokens, timestampID)
+		if err != nil {
+			return err
+		}
+		return nil
+	}
+
+	proxy.ServeHTTP(c.Writer, c.Request)
+
+}
+
+func Completions(c *gin.Context) {
+	var promptTokens int
+	var completionTokens int
+	var totalTokens int
+	var chatRequest model.ChatRequest
+	if err := c.ShouldBindJSON(&chatRequest); err != nil {
+		c.AbortWithStatusJSON(400, gin.H{"error": err.Error()})
+		return
+	}
+
+	auth := c.Request.Header.Get("Authorization")
+	key := auth[7:]
+	serverInfo, err := service.CheckBlance(key, chatRequest.Model, chatRequest.MaxTokens)
+	if err != nil {
+		c.AbortWithStatusJSON(403, gin.H{"error": err.Error()})
+		log.Printf("请求出错 KEY: %v Model: %v ERROR: %v", key, chatRequest.Model, err)
+		return
+	}
+
+	log.Printf("请求的KEY: %v Model: %v", key, chatRequest.Model)
+
+	remote, err := url.Parse(serverInfo.ServerAddress)
+	if err != nil {
+		c.AbortWithStatusJSON(400, gin.H{"error": err.Error()})
+		return
+	}
+
+	proxy := httputil.NewSingleHostReverseProxy(remote)
+	newReqBody, err := json.Marshal(chatRequest)
+	if err != nil {
+		log.Printf("http request err: %v", err)
+		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": err.Error()})
+		return
+	}
+	proxy.Director = func(req *http.Request) {
+		req.Header = c.Request.Header
+		req.Host = remote.Host
+		req.URL.Scheme = remote.Scheme
+		req.URL.Host = remote.Host
+		req.URL.Path = c.Request.URL.Path
+		req.ContentLength = int64(len(newReqBody))
+		req.Header.Set("Content-Type", "application/json")
+		req.Header.Set("Accept-Encoding", "")
+		serverKey := serverInfo.AvailableKey
+		keyList := strings.Split(serverInfo.AvailableKey, ",")
+		if len(keyList) > 1 {
+			// 随机数种子
+			source := rand.NewSource(time.Now().UnixNano())
+			random := rand.New(source)
+			// 从数组中随机选择一个元素
+			serverKey = keyList[random.Intn(len(keyList))]
+		}
+		req.Header.Set("Authorization", "Bearer "+serverKey)
+		req.Body = io.NopCloser(bytes.NewReader(newReqBody))
+	}
+	sss, err := json.Marshal(chatRequest)
+	if err != nil {
+		log.Printf("chatRequest 转化出错 %v", err)
+	}
+	log.Printf("开始处理返回逻辑: %v", string(sss))
+	if chatRequest.Stream {
+		// 流式回应，处理
+		proxy.ModifyResponse = func(resp *http.Response) error {
+			log.Printf("流式回应 http status code: %v", resp.StatusCode)
+			if resp.StatusCode != http.StatusOK {
+				//退回预扣除的余额
+				err = service.CheckBlanceReturn(key, chatRequest.Model, chatRequest.MaxTokens)
+				if err != nil {
+					return err
+				}
+				return nil
+			}
+			chatRequestId := ""
+			reqContent := ""
+			reader := bufio.NewReader(resp.Body)
+			headers := resp.Header
+			for k, v := range headers {
+				c.Writer.Header().Set(k, v[0])
+			}
+			c.Writer.Header().Set("Openai-Organization", "api2gpt")
+			for {
+				chunk, err := reader.ReadBytes('\n')
+				if err != nil {
+					if err == io.EOF {
+						break
+					}
+					log.Printf("流式回应，处理 err %v:", err.Error())
+					break
+					//return err
+				}
+				var chatResponse model.ChatResponse
+				//去除回应中的data:前缀
+				var trimStr = strings.Trim(string(chunk), "data: ")
+				if trimStr != "\n" {
+					err := json.Unmarshal([]byte(trimStr), &chatResponse)
+					if err != nil {
+						log.Printf(err.Error())
+					}
+					if chatResponse.Choices != nil {
+						if chatResponse.Choices[0].Text != "" {
+							reqContent += chatResponse.Choices[0].Text
+						} else {
+							reqContent += chatResponse.Choices[0].Delta.Content
+						}
+
+						chatRequestId = chatResponse.Id
+					}
+					// 写回数据
+					_, err = c.Writer.Write([]byte(string(chunk) + "\n"))
+					if err != nil {
+						log.Printf("写回数据 err: %v", err.Error())
+						return err
+					}
+					c.Writer.(http.Flusher).Flush()
+				}
+			}
+			if chatRequest.Model == "text-davinci-003" {
+				promptTokens = common.NumTokensFromString(chatRequest.Prompt, chatRequest.Model)
+			} else {
+				promptTokens = common.NumTokensFromMessages(chatRequest.Messages, chatRequest.Model)
+			}
+			completionTokens = common.NumTokensFromString(reqContent, chatRequest.Model)
+			log.Printf("返回内容：%v", reqContent)
+			totalTokens = promptTokens + completionTokens
+			log.Printf("prompt_tokens:  %v completion_tokens:  %v total_tokens: %v", promptTokens, completionTokens, totalTokens)
+			//消费余额
+			_, err := service.Consumption(key, chatRequest.Model, promptTokens, completionTokens, totalTokens, chatRequestId)
+			if err != nil {
+				return err
+			}
+			return nil
+		}
+	} else {
+		// 非流式回应，处理
+		proxy.ModifyResponse = func(resp *http.Response) error {
+			resp.Header.Set("Openai-Organization", "api2gpt")
+			var chatResponse model.ChatResponse
+			body, err := io.ReadAll(resp.Body)
+			if err != nil {
+				log.Printf("非流式回应，处理 err: %v", err)
+				return err
+			}
+			err = json.Unmarshal(body, &chatResponse)
+			if err != nil {
+				log.Printf(err.Error())
+			}
+			promptTokens = chatResponse.Usage.PromptTokens
+			completionTokens = chatResponse.Usage.CompletionTokens
+			totalTokens = chatResponse.Usage.TotalTokens
+			resp.Body = io.NopCloser(bytes.NewReader(body))
+			log.Printf("prompt_tokens:  %v completion_tokens:  %v total_tokens: %v", promptTokens, completionTokens, totalTokens)
+			//消费余额
+			_, err = service.Consumption(key, chatRequest.Model, promptTokens, completionTokens, totalTokens, chatResponse.Id)
+			if err != nil {
+				return err
+			}
+			return nil
+		}
+	}
+
+	proxy.ServeHTTP(c.Writer, c.Request)
+}
+
+// HandleOptions 针对接口预检OPTIONS的处理
+func HandleOptions(c *gin.Context) {
+	// BUGFIX: fix options request, see https://github.com/diemus/azure-openai-proxy/issues/1
+	c.Header("Access-Control-Allow-Origin", "*")
+	c.Header("Access-Control-Allow-Methods", "POST, GET, OPTIONS, PUT, DELETE")
+	c.Header("Access-Control-Allow-Headers", "Content-Type, Authorization")
+	c.Status(200)
+	return
+}

deploy/Dockerfile

@@ -9,4 +9,4 @@ RUN adduser -S -D -H -h /app -s /sbin/nologin -u 1000 app
 
 ADD api2gpt-mid /app/
 
-CMD ["/app/api2gpt-mid"]
+CMD ["/app/api2gpt-mid", "--log-dir", "/app/logs"]

deploy/docker-compose.yml

@@ -1,17 +1,19 @@
 version: "3"
 services:
-  server:
+  api2gpt-mid:
     build:
       context: .
       dockerfile: Dockerfile
-    container_name: api2gpt-mid
+    container_name: api2gpt-mid-dev
     environment:
       # 时区上海
       TZ: Asia/Shanghai
       REDIS_ADDRESS: 172.17.0.1:6379
+      REDIS_CONN_STRING: redis://@172.17.0.1:6379/0?dial_timeout=5s
+      SERVER_API_USAGE_RECORD_STRING: http://172.17.0.1:8080/other/usageRecord
     privileged: true
     restart: always
     ports:
-      - 8082:8080
+      - 8083:8080
     volumes:
       - ./logs:/app/logs

go.mod

@@ -1,37 +1,40 @@
-module main
+module api2gpt-mid
 
 go 1.20
 
 require (
-	github.com/bytedance/sonic v1.8.0 // indirect
+	github.com/gin-gonic/gin v1.9.1
+	github.com/pkoukk/tiktoken-go v0.1.0
+	github.com/redis/go-redis/v9 v9.0.3
+)
+
+require (
+	github.com/bytedance/sonic v1.9.1 // indirect
 	github.com/cespare/xxhash/v2 v2.2.0 // indirect
 	github.com/chenzhuoyu/base64x v0.0.0-20221115062448-fe3a3abad311 // indirect
 	github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect
 	github.com/dlclark/regexp2 v1.8.1 // indirect
+	github.com/gabriel-vasile/mimetype v1.4.2 // indirect
 	github.com/gin-contrib/sse v0.1.0 // indirect
-	github.com/gin-gonic/gin v1.9.0 // indirect
 	github.com/go-playground/locales v0.14.1 // indirect
 	github.com/go-playground/universal-translator v0.18.1 // indirect
-	github.com/go-playground/validator/v10 v10.11.2 // indirect
+	github.com/go-playground/validator/v10 v10.14.0 // indirect
-	github.com/go-redis/redis v6.15.9+incompatible // indirect
+	github.com/goccy/go-json v0.10.2 // indirect
-	github.com/goccy/go-json v0.10.0 // indirect
 	github.com/google/uuid v1.3.0 // indirect
 	github.com/json-iterator/go v1.1.12 // indirect
-	github.com/klauspost/cpuid/v2 v2.0.9 // indirect
+	github.com/klauspost/cpuid/v2 v2.2.4 // indirect
-	github.com/leodido/go-urn v1.2.1 // indirect
+	github.com/leodido/go-urn v1.2.4 // indirect
-	github.com/mattn/go-isatty v0.0.17 // indirect
+	github.com/mattn/go-isatty v0.0.19 // indirect
-	github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421 // indirect
+	github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
 	github.com/modern-go/reflect2 v1.0.2 // indirect
-	github.com/pelletier/go-toml/v2 v2.0.6 // indirect
+	github.com/pelletier/go-toml/v2 v2.0.8 // indirect
-	github.com/pkoukk/tiktoken-go v0.1.0 // indirect
-	github.com/redis/go-redis/v9 v9.0.3 // indirect
 	github.com/twitchyliquid64/golang-asm v0.15.1 // indirect
-	github.com/ugorji/go/codec v1.2.9 // indirect
+	github.com/ugorji/go/codec v1.2.11 // indirect
-	golang.org/x/arch v0.0.0-20210923205945-b76863e36670 // indirect
+	golang.org/x/arch v0.3.0 // indirect
-	golang.org/x/crypto v0.5.0 // indirect
+	golang.org/x/crypto v0.9.0 // indirect
-	golang.org/x/net v0.7.0 // indirect
+	golang.org/x/net v0.10.0 // indirect
-	golang.org/x/sys v0.5.0 // indirect
+	golang.org/x/sys v0.8.0 // indirect
-	golang.org/x/text v0.7.0 // indirect
+	golang.org/x/text v0.9.0 // indirect
-	google.golang.org/protobuf v1.28.1 // indirect
+	google.golang.org/protobuf v1.30.0 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )

go.sum

@@ -1,85 +1,97 @@
+github.com/bsm/ginkgo/v2 v2.7.0 h1:ItPMPH90RbmZJt5GtkcNvIRuGEdwlBItdNVoyzaNQao=
+github.com/bsm/gomega v1.26.0 h1:LhQm+AFcgV2M0WyKroMASzAzCAJVpAxQXv4SaI9a69Y=
 github.com/bytedance/sonic v1.5.0/go.mod h1:ED5hyg4y6t3/9Ku1R6dU/4KyJ48DZ4jPhfY1O2AihPM=
-github.com/bytedance/sonic v1.8.0 h1:ea0Xadu+sHlu7x5O3gKhRpQ1IKiMrSiHttPF0ybECuA=
+github.com/bytedance/sonic v1.9.1 h1:6iJ6NqdoxCDr6mbY8h18oSO+cShGSMRGCEo7F2h0x8s=
-github.com/bytedance/sonic v1.8.0/go.mod h1:i736AoUSYt75HyZLoJW9ERYxcy6eaN6h4BZXU064P/U=
+github.com/bytedance/sonic v1.9.1/go.mod h1:i736AoUSYt75HyZLoJW9ERYxcy6eaN6h4BZXU064P/U=
 github.com/cespare/xxhash/v2 v2.2.0 h1:DC2CZ1Ep5Y4k3ZQ899DldepgrayRUGE6BBZ/cd9Cj44=
 github.com/cespare/xxhash/v2 v2.2.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
 github.com/chenzhuoyu/base64x v0.0.0-20211019084208-fb5309c8db06/go.mod h1:DH46F32mSOjUmXrMHnKwZdA8wcEefY7UVqBKYGjpdQY=
 github.com/chenzhuoyu/base64x v0.0.0-20221115062448-fe3a3abad311 h1:qSGYFH7+jGhDF8vLC+iwCD4WpbV1EBDSzWkJODFLams=
 github.com/chenzhuoyu/base64x v0.0.0-20221115062448-fe3a3abad311/go.mod h1:b583jCggY9gE99b6G5LEC39OIiVsWj+R97kbl5odCEk=
 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/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/rVNCu3HqELle0jiPLLBs70cWOduZpkS1E78=
 github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc=
 github.com/dlclark/regexp2 v1.8.1 h1:6Lcdwya6GjPUNsBct8Lg/yRPwMhABj269AAzdGSiR+0=
 github.com/dlclark/regexp2 v1.8.1/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8=
+github.com/gabriel-vasile/mimetype v1.4.2 h1:w5qFW6JKBz9Y393Y4q372O9A7cUSequkh1Q7OhCmWKU=
+github.com/gabriel-vasile/mimetype v1.4.2/go.mod h1:zApsH/mKG4w07erKIaJPFiX0Tsq9BFQgN3qGY5GnNgA=
 github.com/gin-contrib/sse v0.1.0 h1:Y/yl/+YNO8GZSjAhjMsSuLt29uWRFHdHYUb5lYOV9qE=
 github.com/gin-contrib/sse v0.1.0/go.mod h1:RHrZQHXnP2xjPF+u1gW/2HnVO7nvIa9PG3Gm+fLHvGI=
-github.com/gin-gonic/gin v1.9.0 h1:OjyFBKICoexlu99ctXNR2gg+c5pKrKMuyjgARg9qeY8=
+github.com/gin-gonic/gin v1.9.1 h1:4idEAncQnU5cB7BeOkPtxjfCSye0AAm1R0RVIqJ+Jmg=
-github.com/gin-gonic/gin v1.9.0/go.mod h1:W1Me9+hsUSyj3CePGrd1/QrKJMSJ1Tu/0hFEH89961k=
+github.com/gin-gonic/gin v1.9.1/go.mod h1:hPrL7YrpYKXt5YId3A/Tnip5kqbEAP+KLuI3SUcPTeU=
+github.com/go-playground/assert/v2 v2.2.0 h1:JvknZsQTYeFEAhQwI4qEt9cyV5ONwRHC+lYKSsYSR8s=
 github.com/go-playground/locales v0.14.1 h1:EWaQ/wswjilfKLTECiXz7Rh+3BjFhfDFKv/oXslEjJA=
 github.com/go-playground/locales v0.14.1/go.mod h1:hxrqLVvrK65+Rwrd5Fc6F2O76J/NuW9t0sjnWqG1slY=
 github.com/go-playground/universal-translator v0.18.1 h1:Bcnm0ZwsGyWbCzImXv+pAJnYK9S473LQFuzCbDbfSFY=
 github.com/go-playground/universal-translator v0.18.1/go.mod h1:xekY+UJKNuX9WP91TpwSH2VMlDf28Uj24BCp08ZFTUY=
-github.com/go-playground/validator/v10 v10.11.2 h1:q3SHpufmypg+erIExEKUmsgmhDTyhcJ38oeKGACXohU=
+github.com/go-playground/validator/v10 v10.14.0 h1:vgvQWe3XCz3gIeFDm/HnTIbj6UGmg/+t63MyGU2n5js=
-github.com/go-playground/validator/v10 v10.11.2/go.mod h1:NieE624vt4SCTJtD87arVLvdmjPAeV8BQlHtMnw9D7s=
+github.com/go-playground/validator/v10 v10.14.0/go.mod h1:9iXMNT7sEkjXb0I+enO7QXmzG6QCsPWY4zveKFVRSyU=
-github.com/go-redis/redis v6.15.9+incompatible h1:K0pv1D7EQUjfyoMql+r/jZqCLizCGKFlFgcHWWmHQjg=
+github.com/goccy/go-json v0.10.2 h1:CrxCmQqYDkv1z7lO7Wbh2HN93uovUHgrECaO5ZrCXAU=
-github.com/go-redis/redis v6.15.9+incompatible/go.mod h1:NAIEuMOZ/fxfXJIrKDQDz8wamY7mA7PouImQ2Jvg6kA=
+github.com/goccy/go-json v0.10.2/go.mod h1:6MelG93GURQebXPDq3khkgXZkazVtN9CRI+MGFi0w8I=
-github.com/goccy/go-json v0.10.0 h1:mXKd9Qw4NuzShiRlOXKews24ufknHO7gx30lsDyokKA=
-github.com/goccy/go-json v0.10.0/go.mod h1:6MelG93GURQebXPDq3khkgXZkazVtN9CRI+MGFi0w8I=
 github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaSAoJOfIk=
+github.com/google/go-cmp v0.5.5 h1:Khx7svrCpmxxtHBq5j2mp/xVjsi8hQMfNLvJFAlrGgU=
 github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
 github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=
 github.com/google/uuid v1.3.0 h1:t6JiXgmwXMjEs8VusXIJk2BXHsn+wx8BZdTaoZ5fu7I=
 github.com/google/uuid v1.3.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM=
 github.com/json-iterator/go v1.1.12/go.mod h1:e30LSqwooZae/UwlEbR2852Gd8hjQvJoHmT4TnhNGBo=
-github.com/klauspost/cpuid/v2 v2.0.9 h1:lgaqFMSdTdQYdZ04uHyN2d/eKdOMyi2YLSvlQIBFYa4=
 github.com/klauspost/cpuid/v2 v2.0.9/go.mod h1:FInQzS24/EEf25PyTYn52gqo7WaD8xa0213Md/qVLRg=
-github.com/leodido/go-urn v1.2.1 h1:BqpAaACuzVSgi/VLzGZIobT2z4v53pjosyNd9Yv6n/w=
+github.com/klauspost/cpuid/v2 v2.2.4 h1:acbojRNwl3o09bUq+yDCtZFc1aiwaAAxtcn8YkZXnvk=
-github.com/leodido/go-urn v1.2.1/go.mod h1:zt4jvISO2HfUBqxjfIshjdMTYS56ZS/qv49ictyFfxY=
+github.com/klauspost/cpuid/v2 v2.2.4/go.mod h1:RVVoqg1df56z8g3pUjL/3lE5UfnlrJX8tyFgg4nqhuY=
-github.com/mattn/go-isatty v0.0.17 h1:BTarxUcIeDqL27Mc+vyvdWYSL28zpIhv3RoTdsLMPng=
+github.com/leodido/go-urn v1.2.4 h1:XlAE/cm/ms7TE/VMVoduSpNBoyc2dOxHs5MZSwAN63Q=
-github.com/mattn/go-isatty v0.0.17/go.mod h1:kYGgaQfpe5nmfYZH+SKPsOc2e4SrIfOl2e/yFXSvRLM=
+github.com/leodido/go-urn v1.2.4/go.mod h1:7ZrI8mTSeBSHl/UaRyKQW1qZeMgak41ANeCNaVckg+4=
-github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421 h1:ZqeYNhU3OHLH3mGKHDcjJRFFRrJa6eAM5H+CtDdOsPc=
+github.com/mattn/go-isatty v0.0.19 h1:JITubQf0MOLdlGRuRq+jtsDlekdYPia9ZFsB8h/APPA=
+github.com/mattn/go-isatty v0.0.19/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
 github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q=
+github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd h1:TRLaZ9cD/w8PVh93nsPXa1VrQ6jlwL5oN8l14QlcNfg=
+github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q=
 github.com/modern-go/reflect2 v1.0.2 h1:xBagoLtFs94CBntxluKeaWgTMpvLxC4ur3nMaC9Gz0M=
 github.com/modern-go/reflect2 v1.0.2/go.mod h1:yWuevngMOJpCy52FWWMvUC8ws7m/LJsjYzDa0/r8luk=
-github.com/pelletier/go-toml/v2 v2.0.6 h1:nrzqCb7j9cDFj2coyLNLaZuJTLjWjlaz6nvTvIwycIU=
+github.com/pelletier/go-toml/v2 v2.0.8 h1:0ctb6s9mE31h0/lhu+J6OPmVeDxJn+kYnJc2jZR9tGQ=
-github.com/pelletier/go-toml/v2 v2.0.6/go.mod h1:eumQOmlWiOPt5WriQQqoM5y18pDHwha2N+QD+EUNTek=
+github.com/pelletier/go-toml/v2 v2.0.8/go.mod h1:vuYfssBdrU2XDZ9bYydBu6t+6a6PYNcZljzZR9VXg+4=
 github.com/pkoukk/tiktoken-go v0.1.0 h1:X1uP3+Nd8C3xe6AIGRWjchrylyaye0FDDTG22cxNQZs=
 github.com/pkoukk/tiktoken-go v0.1.0/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/redis/go-redis v6.15.9+incompatible h1:F+tnlesQSl3h9V8DdmtcYFdvkHLhbb7AgcLW6UJxnC4=
-github.com/redis/go-redis v6.15.9+incompatible/go.mod h1:ic6dLmR0d9rkHSzaa0Ab3QVRZcjopJ9hSSPCrecj/+s=
 github.com/redis/go-redis/v9 v9.0.3 h1:+7mmR26M0IvyLxGZUHxu4GiBkJkVDid0Un+j4ScYu4k=
 github.com/redis/go-redis/v9 v9.0.3/go.mod h1:WqMKv5vnQbRuZstUwxQI195wHy+t4PuXDOjzMvcuQHk=
 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=
 github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
-github.com/stretchr/testify v1.6.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 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/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
+github.com/stretchr/testify v1.8.2/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
+github.com/stretchr/testify v1.8.3 h1:RP3t2pwF7cMEbC1dqtB6poj3niw/9gnV4Cjg5oW5gtY=
+github.com/stretchr/testify v1.8.3/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo=
 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=
+github.com/ugorji/go/codec v1.2.11 h1:BMaWp1Bb6fHwEtbplGBGJ498wD+LKlNSl25MjdZY4dU=
-github.com/ugorji/go/codec v1.2.9/go.mod h1:UNopzCgEMSXjBc6AOMqYvWC1ktqTAfzJZUZgYf6w6lg=
+github.com/ugorji/go/codec v1.2.11/go.mod h1:UNopzCgEMSXjBc6AOMqYvWC1ktqTAfzJZUZgYf6w6lg=
-golang.org/x/arch v0.0.0-20210923205945-b76863e36670 h1:18EFjUmQOcUvxNYSkA6jO9VAiXCnxFY6NyDX0bHDmkU=
 golang.org/x/arch v0.0.0-20210923205945-b76863e36670/go.mod h1:5om86z9Hs0C8fWVUuoMHwpExlXzs5Tkyp9hOrfG7pp8=
-golang.org/x/crypto v0.5.0 h1:U/0M97KRkSFvyD/3FSmdP5W5swImpNgle/EHFhOsQPE=
+golang.org/x/arch v0.3.0 h1:02VY4/ZcO/gBOH6PUaoiptASxtXU10jazRCP865E97k=
-golang.org/x/crypto v0.5.0/go.mod h1:NK/OQwhpMQP3MwtdjgLlYHnH9ebylxKWv3e0fK+mkQU=
+golang.org/x/arch v0.3.0/go.mod h1:5om86z9Hs0C8fWVUuoMHwpExlXzs5Tkyp9hOrfG7pp8=
-golang.org/x/net v0.7.0 h1:rJrUqqhjsgNp7KqAIc25s9pZnjU7TUcSY7HcVZjdn1g=
+golang.org/x/crypto v0.9.0 h1:LF6fAI+IutBocDJ2OT0Q1g8plpYljMZ4+lty+dsqw3g=
-golang.org/x/net v0.7.0/go.mod h1:2Tu9+aMcznHK/AK1HMvgo6xiTLG5rD5rZLDS+rp2Bjs=
+golang.org/x/crypto v0.9.0/go.mod h1:yrmDGqONDYtNj3tH8X9dzUun2m2lzPa9ngI6/RUPGR0=
-golang.org/x/sys v0.0.0-20220811171246-fbc7d0a398ab/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/net v0.10.0 h1:X2//UzNDwYmtCLn7To6G58Wr6f5ahEAQgKNzv9Y951M=
-golang.org/x/sys v0.5.0 h1:MUK/U/4lj1t1oPg0HfuXDN/Z1wv31ZJ/YcPiGccS4DU=
+golang.org/x/net v0.10.0/go.mod h1:0qNGK6F8kojg2nk9dLZ2mShWaEBan6FAoqfSigmmuDg=
-golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.0.0-20220704084225-05e143d24a9e/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
-golang.org/x/text v0.7.0 h1:4BRB4x83lYWy72KwLD/qYDuTu7q9PjSagHvijDw7cLo=
+golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
-golang.org/x/text v0.7.0/go.mod h1:mrYo+phRRbMaCq/xk9113O4dZlRixOauAjOtrjsXDZ8=
+golang.org/x/sys v0.8.0 h1:EBmGv8NaZBZTWvrbjNoL6HVt+IVy3QDQpJs7VRIw3tU=
+golang.org/x/sys v0.8.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/text v0.9.0 h1:2sjJmO8cDvYveuX97RDLsxlyUxLl+GHoLxBiRdHllBE=
+golang.org/x/text v0.9.0/go.mod h1:e1OnstbJyHTd6l/uOt8jFFHp6TRDWZR/bV3emEE/zU8=
+golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543 h1:E7g+9GITq07hpfrRu66IVDexMakfv52eLZ2CXBWiKr4=
 golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw=
-google.golang.org/protobuf v1.28.1 h1:d0NfwRgPtno5B1Wa6L2DAG+KivqkdutMf1UhdNx175w=
+google.golang.org/protobuf v1.30.0 h1:kPPoIgf3TsEvrm0PFe15JQ+570QVxYzEvvHqChK+cng=
-google.golang.org/protobuf v1.28.1/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I=
+google.golang.org/protobuf v1.30.0/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=

main.go

@@ -1,739 +1,48 @@
 package main
 
 import (
-	"bufio"
+	"api2gpt-mid/common"
-	"bytes"
+	"api2gpt-mid/middleware"
-	"context"
+	"api2gpt-mid/model"
+	"api2gpt-mid/router"
 	"encoding/json"
-	"fmt"
-	"io"
-	"io/ioutil"
-	"log"
-	"math/rand"
-	"net/http"
-	"net/http/httputil"
-	"net/url"
-	"os"
-	"strings"
-	"time"
-
 	"github.com/gin-gonic/gin"
-	"github.com/pkoukk/tiktoken-go"
+	"os"
-	"github.com/redis/go-redis/v9"
+	"strconv"
-)
-
-type Message struct {
-	Role    string `json:"role,omitempty"`
-	Name    string `json:"name,omitempty"`
-	Content string `json:"content,omitempty"`
-}
-
-type ChatRequest struct {
-	Stream          bool        `json:"stream,omitempty"`
-	Model           string      `json:"model,omitempty"`
-	PresencePenalty float64     `json:"presence_penalty,omitempty"`
-	Temperature     float64     `json:"temperature,omitempty"`
-	Messages        []Message   `json:"messages,omitempty"`
-	Prompt          string      `json:"prompt,omitempty"`
-	Input           interface{} `json:"input,omitempty"`
-	Instruction     string      `json:"instruction,omitempty"`
-	N               int         `json:"n,omitempty"`
-	Functions       interface{} `json:"functions,omitempty"`
-}
-
-type ImagesRequest struct {
-	Prompt string `json:"prompt,omitempty"`
-	N      int    `json:"n,omitempty"`
-	Size   string `json:"size,omitempty"`
-}
-
-type ChatResponse struct {
-	Id      string   `json:"id,omitempty"`
-	Object  string   `json:"object,omitempty"`
-	Created int64    `json:"created,omitempty"`
-	Model   string   `json:"model,omitempty"`
-	Usage   Usage    `json:"usage,omitempty"`
-	Choices []Choice `json:"choices,omitempty"`
-}
-type ImagesResponse struct {
-	Created int64        `json:"created,omitempty"`
-	Data    []ImagesData `json:"data,omitempty"`
-}
-
-type ImagesData struct {
-	Url string `json:"url,omitempty"`
-}
-
-type Usage struct {
-	PromptTokens     int `json:"prompt_tokens,omitempty"`
-	CompletionTokens int `json:"completion_tokens,omitempty"`
-	TotalTokens      int `json:"total_tokens,omitempty"`
-}
-
-type Choice struct {
-	Message      Message `json:"message,omitempty"`
-	Delta        Delta   `json:"delta,omitempty"`
-	FinishReason string  `json:"finish_reason,omitempty"`
-	Index        int     `json:"index,omitempty"`
-	Text         string  `json:"text,omitempty"`
-}
-
-type Delta struct {
-	Content string `json:"content,omitempty"`
-}
-
-type CreditSummary struct {
-	Object         string  `json:"object"`
-	TotalGranted   float64 `json:"total_granted"`
-	TotalUsed      float64 `json:"total_used"`
-	TotalRemaining float64 `json:"total_remaining"`
-}
-
-type ListModelResponse struct {
-	Object string  `json:"object"`
-	Data   []Model `json:"data"`
-}
-
-type Model struct {
-	ID         string            `json:"id"`
-	Object     string            `json:"object"`
-	Created    int               `json:"created"`
-	OwnedBy    string            `json:"owned_by"`
-	Permission []ModelPermission `json:"permission"`
-	Root       string            `json:"root"`
-	Parent     any               `json:"parent"`
-}
-
-type ModelPermission struct {
-	ID                 string `json:"id"`
-	Object             string `json:"object"`
-	Created            int    `json:"created"`
-	AllowCreateEngine  bool   `json:"allow_create_engine"`
-	AllowSampling      bool   `json:"allow_sampling"`
-	AllowLogprobs      bool   `json:"allow_logprobs"`
-	AllowSearchIndices bool   `json:"allow_search_indices"`
-	AllowView          bool   `json:"allow_view"`
-	AllowFineTuning    bool   `json:"allow_fine_tuning"`
-	Organization       string `json:"organization"`
-	Group              any    `json:"group"`
-	IsBlocking         bool   `json:"is_blocking"`
-}
-
-type UserInfo struct {
-	UID string `json:"uid"`
-	SID string `json:"sid"`
-}
-
-type ServerInfo struct {
-	ServerAddress string `json:"server_address"`
-	AvailableKey  string `json:"available_key"`
-}
-
-type ModelInfo struct {
-	ModelName       string  `json:"model_name"`
-	ModelPrice      float64 `json:"model_price"`
-	ModelPrice2     float64 `json:"model_price2"`
-	ModelPrepayment int     `json:"model_prepayment"`
-	ServerId        int     `json:"server_id"`
-}
-
-var (
-	Redis        *redis.Client
-	RedisAddress = "localhost:6379"
 )
 
-func init() {
+func test_redis() {
-	//gin.SetMode(gin.ReleaseMode)
+	//添加reids测试数据
-	if v := os.Getenv("REDIS_ADDRESS"); v != "" {
+	var serverInfo model.ServerInfo = model.ServerInfo{
-		RedisAddress = v
+		ServerAddress: "https://gptp.any-door.cn",
-	}
+		AvailableKey:  "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",
-	log.Printf("loading redis address: %s", RedisAddress)
-}
-
-// redis 初始化
-func InitRedis() *redis.Client {
-	rdb := redis.NewClient(&redis.Options{
-		Addr:     RedisAddress,
-		Password: "", // no password set
-		DB:       0,  // use default DB
-		PoolSize: 10,
-	})
-	result := rdb.Ping(context.Background())
-	fmt.Println("redis ping:", result.Val())
-	if result.Val() != "PONG" {
-		// 连接有问题
-		return nil
-	}
-	return rdb
-}
-
-// 计算Messages中的token数量
-func numTokensFromMessages(messages []Message, model string) int {
-	if strings.Contains(model, "gpt-3.5") {
-		model = "gpt-3.5-turbo"
-	}
-	if strings.Contains(model, "gpt-4") {
-		model = "gpt-4"
-	}
-	tkm, err := tiktoken.EncodingForModel(model)
-	if err != nil {
-		err = fmt.Errorf("getEncoding: %v", err)
-		return 0
-	}
-	numTokens := 0
-	for _, message := range messages {
-		numTokens += len(tkm.Encode(message.Content, nil, nil))
-		numTokens += 6
-	}
-	numTokens += 3
-	return numTokens
-}
-
-// 计算String中的token数量
-func numTokensFromString(msg string, model string) int {
-	if strings.Contains(model, "gpt-3.5") {
-		model = "gpt-3.5-turbo"
-	}
-	if strings.Contains(model, "gpt-4") {
-		model = "gpt-4"
-	}
-	tkm, err := tiktoken.EncodingForModel(model)
-	if err != nil {
-		err = fmt.Errorf("getEncoding: %v", err)
-		return 0
-	}
-	if model == "text-davinci-003" {
-		return len(tkm.Encode(msg, nil, nil)) + 1
-	} else {
-		return len(tkm.Encode(msg, nil, nil)) + 9
-	}
-}
-
-func images(c *gin.Context) {
-	var imagesRequest ImagesRequest
-	var model = "images-generations"
-	if err := c.ShouldBindJSON(&imagesRequest); err != nil {
-		c.AbortWithStatusJSON(400, gin.H{"error": err.Error()})
-		return
-	}
-	auth := c.Request.Header.Get("Authorization")
-	key := auth[7:]
-
-	serverInfo, err := checkBlanceForImages(key, model, imagesRequest.N)
-	if err != nil {
-		c.AbortWithStatusJSON(403, gin.H{"error": err.Error()})
-		log.Printf("请求出错 KEY: %v Model: %v ERROR: %v", key, model, err)
-		return
-	}
-
-	log.Printf("请求的KEY: %v Model: %v", key, model)
-
-	remote, err := url.Parse(serverInfo.ServerAddress)
-	if err != nil {
-		c.AbortWithStatusJSON(400, gin.H{"error": err.Error()})
-		return
-	}
-
-	proxy := httputil.NewSingleHostReverseProxy(remote)
-	newReqBody, err := json.Marshal(imagesRequest)
-	if err != nil {
-		log.Printf("http request err: %v", err)
-		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": err.Error()})
-		return
-	}
-
-	proxy.Director = func(req *http.Request) {
-		req.Header = c.Request.Header
-		req.Host = remote.Host
-		req.URL.Scheme = remote.Scheme
-		req.URL.Host = remote.Host
-		req.URL.Path = c.Request.URL.Path
-		req.ContentLength = int64(len(newReqBody))
-		req.Header.Set("Content-Type", "application/json")
-		req.Header.Set("Accept-Encoding", "")
-		serverKey := serverInfo.AvailableKey
-		keyList := strings.Split(serverInfo.AvailableKey, ",")
-		if len(keyList) > 1 {
-			// 随机数种子
-			rand.Seed(time.Now().UnixNano())
-			// 从数组中随机选择一个元素
-			serverKey = keyList[rand.Intn(len(keyList))]
-		}
-		req.Header.Set("Authorization", "Bearer "+serverKey)
-		req.Body = ioutil.NopCloser(bytes.NewReader(newReqBody))
-	}
-	sss, err := json.Marshal(imagesRequest)
-	log.Printf("开始处理返回逻辑 %d", string(sss))
-
-	proxy.ModifyResponse = func(resp *http.Response) error {
-		if resp.StatusCode != http.StatusOK {
-			//退回预扣除的余额
-			checkBlanceReturnForImages(key, model, imagesRequest.N)
-			return nil
-		}
-		resp.Header.Set("Openai-Organization", "api2gpt")
-		var imagesResponse ImagesResponse
-		body, err := ioutil.ReadAll(resp.Body)
-		if err != nil {
-			log.Printf("读取返回数据出错: %v", err)
-			return err
-		}
-		json.Unmarshal(body, &imagesResponse)
-		resp.Body = ioutil.NopCloser(bytes.NewReader(body))
-		log.Printf("image size:  %v", len(imagesResponse.Data))
-		timestamp := time.Now().Unix()
-		timestampID := "img-" + fmt.Sprintf("%d", timestamp)
-		//消费余额
-		consumptionForImages(key, model, imagesRequest.N, len(imagesResponse.Data), timestampID)
-		return nil
-	}
-
-	proxy.ServeHTTP(c.Writer, c.Request)
-
-}
-
-func edit(c *gin.Context) {
-	var prompt_tokens int
-	var complate_tokens int
-	var total_tokens int
-	var chatRequest ChatRequest
-	if err := c.ShouldBindJSON(&chatRequest); err != nil {
-		c.AbortWithStatusJSON(400, gin.H{"error": err.Error()})
-		return
-	}
-
-	auth := c.Request.Header.Get("Authorization")
-	key := auth[7:]
-	//根据KEY调用用户余额接口，判断是否有足够的余额， 后期可考虑判断max_tokens参数来调整
-	serverInfo, err := checkBlance(key, chatRequest.Model)
-	if err != nil {
-		c.AbortWithStatusJSON(403, gin.H{"error": err.Error()})
-		log.Printf("请求出错 KEY: %v Model: %v ERROR: %v", key, chatRequest.Model, err)
-		return
-	}
-
-	log.Printf("请求的KEY: %v Model: %v", key, chatRequest.Model)
-
-	remote, err := url.Parse(serverInfo.ServerAddress)
-	if err != nil {
-		c.AbortWithStatusJSON(400, gin.H{"error": err.Error()})
-		return
-	}
-
-	proxy := httputil.NewSingleHostReverseProxy(remote)
-	newReqBody, err := json.Marshal(chatRequest)
-	if err != nil {
-		log.Printf("http request err: %v", err)
-		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": err.Error()})
-		return
-	}
-
-	proxy.Director = func(req *http.Request) {
-		req.Header = c.Request.Header
-		req.Host = remote.Host
-		req.URL.Scheme = remote.Scheme
-		req.URL.Host = remote.Host
-		req.URL.Path = c.Request.URL.Path
-		req.ContentLength = int64(len(newReqBody))
-		req.Header.Set("Content-Type", "application/json")
-		req.Header.Set("Accept-Encoding", "")
-		serverKey := serverInfo.AvailableKey
-		keyList := strings.Split(serverInfo.AvailableKey, ",")
-		if len(keyList) > 1 {
-			// 随机数种子
-			rand.Seed(time.Now().UnixNano())
-			// 从数组中随机选择一个元素
-			serverKey = keyList[rand.Intn(len(keyList))]
-		}
-		req.Header.Set("Authorization", "Bearer "+serverKey)
-		req.Body = ioutil.NopCloser(bytes.NewReader(newReqBody))
-	}
-	sss, err := json.Marshal(chatRequest)
-	log.Printf("开始处理返回逻辑 %d", string(sss))
-
-	proxy.ModifyResponse = func(resp *http.Response) error {
-		resp.Header.Set("Openai-Organization", "api2gpt")
-		var chatResponse ChatResponse
-		body, err := ioutil.ReadAll(resp.Body)
-		if err != nil {
-			log.Printf("读取返回数据出错: %v", err)
-			return err
-		}
-		json.Unmarshal(body, &chatResponse)
-		prompt_tokens = chatResponse.Usage.PromptTokens
-		complate_tokens = chatResponse.Usage.CompletionTokens
-		total_tokens = chatResponse.Usage.TotalTokens
-		resp.Body = ioutil.NopCloser(bytes.NewReader(body))
-		log.Printf("prompt_tokens:  %v complate_tokens: %v total_tokens: %v", prompt_tokens, complate_tokens, total_tokens)
-		timestamp := time.Now().Unix()
-		timestampID := "edit-" + fmt.Sprintf("%d", timestamp)
-		//消费余额
-		consumption(key, chatRequest.Model, prompt_tokens, 0, total_tokens, timestampID)
-		return nil
-	}
-
-	proxy.ServeHTTP(c.Writer, c.Request)
-
-}
-
-func embeddings(c *gin.Context) {
-	var prompt_tokens int
-	var total_tokens int
-	var chatRequest ChatRequest
-	if err := c.ShouldBindJSON(&chatRequest); err != nil {
-		c.AbortWithStatusJSON(400, gin.H{"error": err.Error()})
-		return
-	}
-
-	auth := c.Request.Header.Get("Authorization")
-	//key := strings.Trim(auth, "Bearer ")
-	key := auth[7:]
-	//根据KEY调用用户余额接口，判断是否有足够的余额， 后期可考虑判断max_tokens参数来调整
-	serverInfo, err := checkBlance(key, chatRequest.Model)
-	if err != nil {
-		c.AbortWithStatusJSON(403, gin.H{"error": err.Error()})
-		log.Printf("请求出错 KEY: %v Model: %v ERROR: %v", key, chatRequest.Model, err)
-		return
 	}
-
+	// var serverInfo2 ServerInfo = ServerInfo{
-	log.Printf("请求的KEY: %v Model: %v", key, chatRequest.Model)
+	// 	ServerAddress: "https://azure.any-door.cn",
-
+	// 	AvailableKey:  "6c4d2c65970b40e482e7cd27adb0d119",
-	remote, err := url.Parse(serverInfo.ServerAddress)
+	// }
+	serverInfoStr, _ := json.Marshal(&serverInfo)
+	// serverInfoStr2, _ := json.Marshal(&serverInfo2)
+	err := common.RedisSet("server:0", string(serverInfoStr), 0)
 	if err != nil {
-		c.AbortWithStatusJSON(400, gin.H{"error": err.Error()})
 		return
 	}
-
+	err = common.RedisSet("server:1", string(serverInfoStr), 0)
-	proxy := httputil.NewSingleHostReverseProxy(remote)
-	newReqBody, err := json.Marshal(chatRequest)
 	if err != nil {
-		log.Printf("http request err: %v", err)
-		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": err.Error()})
 		return
 	}
+	// Redis.Set(context.Background(), "server:2", serverInfoStr2, 0)
 
-	proxy.Director = func(req *http.Request) {
+	var modelInfo model.ModelInfo = model.ModelInfo{
-		req.Header = c.Request.Header
+		ModelName:       "images-generations",
-		req.Host = remote.Host
+		ModelPrice:      0.0001,
-		req.URL.Scheme = remote.Scheme
+		ModelPrepayment: 4000,
-		req.URL.Host = remote.Host
-		req.URL.Path = c.Request.URL.Path
-		req.ContentLength = int64(len(newReqBody))
-		req.Header.Set("Content-Type", "application/json")
-		req.Header.Set("Accept-Encoding", "")
-		serverKey := serverInfo.AvailableKey
-		keyList := strings.Split(serverInfo.AvailableKey, ",")
-		if len(keyList) > 1 {
-			// 随机数种子
-			rand.Seed(time.Now().UnixNano())
-			// 从数组中随机选择一个元素
-			serverKey = keyList[rand.Intn(len(keyList))]
-		}
-		req.Header.Set("Authorization", "Bearer "+serverKey)
-		req.Body = ioutil.NopCloser(bytes.NewReader(newReqBody))
-	}
-	sss, err := json.Marshal(chatRequest)
-	log.Printf("开始处理返回逻辑 %d", string(sss))
-
-	proxy.ModifyResponse = func(resp *http.Response) error {
-		resp.Header.Set("Openai-Organization", "api2gpt")
-		var chatResponse ChatResponse
-		body, err := ioutil.ReadAll(resp.Body)
-		if err != nil {
-			log.Printf("读取返回数据出错: %v", err)
-			return err
-		}
-		json.Unmarshal(body, &chatResponse)
-		prompt_tokens = chatResponse.Usage.PromptTokens
-		total_tokens = chatResponse.Usage.TotalTokens
-		resp.Body = ioutil.NopCloser(bytes.NewReader(body))
-		log.Printf("prompt_tokens:  %v total_tokens: %v", prompt_tokens, total_tokens)
-		timestamp := time.Now().Unix()
-		timestampID := "emb-" + fmt.Sprintf("%d", timestamp)
-		//消费余额
-		consumption(key, chatRequest.Model, prompt_tokens, 0, total_tokens, timestampID)
-		return nil
-	}
-
-	proxy.ServeHTTP(c.Writer, c.Request)
-
-}
-
-func completions(c *gin.Context) {
-	var prompt_tokens int
-	var completion_tokens int
-	var total_tokens int
-	var chatRequest ChatRequest
-	if err := c.ShouldBindJSON(&chatRequest); err != nil {
-		c.AbortWithStatusJSON(400, gin.H{"error": err.Error()})
-		return
-	}
-
-	auth := c.Request.Header.Get("Authorization")
-	key := auth[7:]
-	//根据KEY调用用户余额接口，判断是否有足够的余额， 后期可考虑判断max_tokens参数来调整
-	serverInfo, err := checkBlance(key, chatRequest.Model)
-	if err != nil {
-		c.AbortWithStatusJSON(403, gin.H{"error": err.Error()})
-		log.Printf("请求出错 KEY: %v Model: %v ERROR: %v", key, chatRequest.Model, err)
-		return
 	}
-
+	modelInfoStr, _ := json.Marshal(&modelInfo)
-	log.Printf("请求的KEY: %v Model: %v", key, chatRequest.Model)
+	err = common.RedisSet("model:images-generations", string(modelInfoStr), 0)
-
-	remote, err := url.Parse(serverInfo.ServerAddress)
 	if err != nil {
-		c.AbortWithStatusJSON(400, gin.H{"error": err.Error()})
 		return
 	}
-
-	proxy := httputil.NewSingleHostReverseProxy(remote)
-	newReqBody, err := json.Marshal(chatRequest)
-	if err != nil {
-		log.Printf("http request err: %v", err)
-		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": err.Error()})
-		return
-	}
-	proxy.Director = func(req *http.Request) {
-		req.Header = c.Request.Header
-		req.Host = remote.Host
-		req.URL.Scheme = remote.Scheme
-		req.URL.Host = remote.Host
-		req.URL.Path = c.Request.URL.Path
-		req.ContentLength = int64(len(newReqBody))
-		req.Header.Set("Content-Type", "application/json")
-		req.Header.Set("Accept-Encoding", "")
-		serverKey := serverInfo.AvailableKey
-		keyList := strings.Split(serverInfo.AvailableKey, ",")
-		if len(keyList) > 1 {
-			// 随机数种子
-			rand.Seed(time.Now().UnixNano())
-			// 从数组中随机选择一个元素
-			serverKey = keyList[rand.Intn(len(keyList))]
-		}
-		req.Header.Set("Authorization", "Bearer "+serverKey)
-		req.Body = ioutil.NopCloser(bytes.NewReader(newReqBody))
-	}
-	sss, err := json.Marshal(chatRequest)
-	if err != nil {
-		log.Printf("chatRequest 转化出错 %v", err)
-	}
-	log.Printf("开始处理返回逻辑: %v", string(sss))
-	if chatRequest.Stream {
-		// 流式回应，处理
-		proxy.ModifyResponse = func(resp *http.Response) error {
-			log.Printf("流式回应 http status code: %v", resp.StatusCode)
-			if resp.StatusCode != http.StatusOK {
-				//退回预扣除的余额
-				checkBlanceReturn(key, chatRequest.Model)
-				return nil
-			}
-			chatRequestId := ""
-			reqContent := ""
-			reader := bufio.NewReader(resp.Body)
-			headers := resp.Header
-			for k, v := range headers {
-				c.Writer.Header().Set(k, v[0])
-			}
-			c.Writer.Header().Set("Openai-Organization", "api2gpt")
-			for {
-				chunk, err := reader.ReadBytes('\n')
-				if err != nil {
-					if err == io.EOF {
-						break
-					}
-					log.Printf("流式回应，处理 err %v:", err.Error())
-					break
-					//return err
-				}
-				var chatResponse ChatResponse
-				//去除回应中的data:前缀
-				var trimStr = strings.Trim(string(chunk), "data: ")
-				if trimStr != "\n" {
-					json.Unmarshal([]byte(trimStr), &chatResponse)
-					if chatResponse.Choices != nil {
-						if chatResponse.Choices[0].Text != "" {
-							reqContent += chatResponse.Choices[0].Text
-						} else {
-							reqContent += chatResponse.Choices[0].Delta.Content
-						}
-
-						chatRequestId = chatResponse.Id
-					}
-
-					// 写回数据
-					_, err = c.Writer.Write([]byte(string(chunk) + "\n"))
-					if err != nil {
-						log.Printf("写回数据 err: %v", err.Error())
-						return err
-					}
-					c.Writer.(http.Flusher).Flush()
-				}
-			}
-			if chatRequest.Model == "text-davinci-003" {
-				prompt_tokens = numTokensFromString(chatRequest.Prompt, chatRequest.Model)
-			} else {
-				prompt_tokens = numTokensFromMessages(chatRequest.Messages, chatRequest.Model)
-			}
-			completion_tokens = numTokensFromString(reqContent, chatRequest.Model)
-			log.Printf("返回内容：%v", reqContent)
-			total_tokens = prompt_tokens + completion_tokens
-			log.Printf("prompt_tokens:  %v completion_tokens:  %v total_tokens: %v", prompt_tokens, completion_tokens, total_tokens)
-			//消费余额
-			consumption(key, chatRequest.Model, prompt_tokens, completion_tokens, total_tokens, chatRequestId)
-			return nil
-		}
-	} else {
-		// 非流式回应，处理
-		proxy.ModifyResponse = func(resp *http.Response) error {
-			resp.Header.Set("Openai-Organization", "api2gpt")
-			var chatResponse ChatResponse
-			body, err := ioutil.ReadAll(resp.Body)
-			if err != nil {
-				log.Printf("非流式回应，处理 err: %v", err)
-				return err
-			}
-			json.Unmarshal(body, &chatResponse)
-			prompt_tokens = chatResponse.Usage.PromptTokens
-			completion_tokens = chatResponse.Usage.CompletionTokens
-			total_tokens = chatResponse.Usage.TotalTokens
-			resp.Body = ioutil.NopCloser(bytes.NewReader(body))
-			log.Printf("prompt_tokens:  %v completion_tokens:  %v total_tokens: %v", prompt_tokens, completion_tokens, total_tokens)
-			//消费余额
-			consumption(key, chatRequest.Model, prompt_tokens, completion_tokens, total_tokens, chatResponse.Id)
-			return nil
-		}
-	}
-
-	proxy.ServeHTTP(c.Writer, c.Request)
-}
-
-// model查询列表
-func handleGetModels(c *gin.Context) {
-	// BUGFIX: fix options request, see https://github.com/diemus/azure-openai-proxy/issues/3
-	//models := []string{"gpt-4", "gpt-4-0314", "gpt-4-32k", "gpt-4-32k-0314", "gpt-3.5-turbo", "gpt-3.5-turbo-0301", "text-davinci-003", "text-embedding-ada-002"}
-	models := []string{"gpt-4", "gpt-4-0314", "gpt-4-0613", "gpt-3.5-turbo", "gpt-3.5-turbo-0301", "gpt-3.5-turbo-0613", "gpt-3.5-turbo-16k", "gpt-3.5-turbo-16k-0613", "text-davinci-003", "text-embedding-ada-002", "text-davinci-edit-001", "code-davinci-edit-001", "images-generations"}
-	result := ListModelResponse{
-		Object: "list",
-	}
-	for _, model := range models {
-		result.Data = append(result.Data, Model{
-			ID:      model,
-			Object:  "model",
-			Created: 1677649963,
-			OwnedBy: "openai",
-			Permission: []ModelPermission{
-				{
-					ID:                 "",
-					Object:             "model",
-					Created:            1679602087,
-					AllowCreateEngine:  true,
-					AllowSampling:      true,
-					AllowLogprobs:      true,
-					AllowSearchIndices: true,
-					AllowView:          true,
-					AllowFineTuning:    true,
-					Organization:       "*",
-					Group:              nil,
-					IsBlocking:         false,
-				},
-			},
-			Root:   model,
-			Parent: nil,
-		})
-	}
-	c.JSON(200, result)
-}
-
-// 针对接口预检OPTIONS的处理
-func handleOptions(c *gin.Context) {
-	// BUGFIX: fix options request, see https://github.com/diemus/azure-openai-proxy/issues/1
-	c.Header("Access-Control-Allow-Origin", "*")
-	c.Header("Access-Control-Allow-Methods", "POST, GET, OPTIONS, PUT, DELETE")
-	c.Header("Access-Control-Allow-Headers", "Content-Type, Authorization")
-	c.Status(200)
-	return
-}
-
-// 余额查询
-func balance(c *gin.Context) {
-	auth := c.Request.Header.Get("Authorization")
-	key := strings.Trim(auth, "Bearer ")
-	balance, err := queryBlance(key)
-	if err != nil {
-		c.JSON(400, gin.H{"error": err.Error()})
-	}
-	var creditSummary CreditSummary
-	creditSummary.Object = "credit_grant"
-	creditSummary.TotalGranted = 999999
-	creditSummary.TotalUsed = 999999 - balance
-	creditSummary.TotalRemaining = balance
-	c.JSON(200, creditSummary)
-}
-
-// key 检测中间件
-func checkKeyMid() gin.HandlerFunc {
-	return func(c *gin.Context) {
-		auth := c.Request.Header.Get("Authorization")
-		log.Printf("auth: %v", auth)
-		if auth == "" {
-			c.AbortWithStatusJSON(401, gin.H{"code": 40001})
-		} else {
-			//key := strings.Trim(auth, "Bearer ")
-			key := auth[7:]
-			log.Printf("key: %v", key)
-			msg, err := checkKeyAndTimeCount(key)
-			if err != nil {
-				log.Printf("checkKeyMid err: %v", err)
-				c.AbortWithStatusJSON(msg, gin.H{"code": err.Error()})
-			}
-		}
-		log.Printf("auth check end")
-		// 执行函数
-		c.Next()
-	}
-}
-
-func Cors() gin.HandlerFunc {
-	return func(c *gin.Context) {
-		method := c.Request.Method
-		origin := c.Request.Header.Get("Origin")
-		if origin != "" {
-			c.Header("Access-Control-Allow-Origin", "*") // 可将将 * 替换为指定的域名
-			c.Header("Access-Control-Allow-Methods", "POST, GET, OPTIONS, PUT, DELETE, UPDATE")
-			c.Header("Access-Control-Allow-Headers", "Origin, X-Requested-With, Content-Type, Accept, Authorization")
-			c.Header("Access-Control-Expose-Headers", "Content-Length, Access-Control-Allow-Origin, Access-Control-Allow-Headers, Cache-Control, Content-Language, Content-Type")
-			c.Header("Access-Control-Allow-Credentials", "true")
-		}
-		if method == "OPTIONS" {
-			c.AbortWithStatus(http.StatusNoContent)
-		}
-		c.Next()
-	}
-}
-
-func test_redis() {
-	//添加reids测试数据
-	//var serverInfo ServerInfo = ServerInfo{
-	//	ServerAddress: "https://gptp.any-door.cn",
-	//	AvailableKey:  "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",
-	//}
-	// var serverInfo2 ServerInfo = ServerInfo{
-	// 	ServerAddress: "https://azure.any-door.cn",
-	// 	AvailableKey:  "6c4d2c65970b40e482e7cd27adb0d119",
-	// }
-	//serverInfoStr, _ := json.Marshal(&serverInfo)
-	// serverInfoStr2, _ := json.Marshal(&serverInfo2)
-	//Redis.Set(context.Background(), "server:1", serverInfoStr, 0)
-	// Redis.Set(context.Background(), "server:2", serverInfoStr2, 0)
-
-	// var modelInfo ModelInfo = ModelInfo{
-	// 	ModelName:       "gpt-3.5-turbo-0613",
-	// 	ModelPrice:      0.0001,
-	// 	ModelPrepayment: 4000,
-	// }
-	// modelInfoStr, _ := json.Marshal(&modelInfo)
-	// Redis.Set(context.Background(), "model:gpt-3.5-turbo-0613", modelInfoStr, 0)
 	// var modelInfo1 ModelInfo = ModelInfo{
 	// 	ModelName:       "gpt-3.5-turbo-16k",
 	// 	ModelPrice:      0.0001,
@@ -772,7 +81,7 @@ func test_redis() {
 	// modelInfoStr2, _ := json.Marshal(&modelInfo2)
 	// Redis.Set(context.Background(), "model:images-generations", modelInfoStr2, 0)
 
-	var userInfo UserInfo = UserInfo{
+	var userInfo model.UserInfo = model.UserInfo{
 		UID: "0",
 		SID: "1",
 	}
@@ -782,55 +91,51 @@ func test_redis() {
 	// }
 	userInfoStr, _ := json.Marshal(&userInfo)
 	// userInfoStr2, _ := json.Marshal(&userInfo2)
-	Redis.Set(context.Background(), "user:key0", userInfoStr, 0)
+	err = common.RedisSet("user:key0", string(userInfoStr), 0)
+	if err != nil {
+		return
+	}
+	//Redis.Set(context.Background(), "user:key0", userInfoStr, 0)
 	// Redis.Set(context.Background(), "user:AK-7d8ab782-a152-4cc1-9972-568713465c96", userInfoStr2, 0)
-
+	_, err = common.RedisIncrByFloat("user:0:balance", 1000)
-	Redis.IncrByFloat(context.Background(), "user:0:balance", 1000).Result()
+	if err != nil {
+		return
+	}
+	//Redis.IncrByFloat(context.Background(), "user:0:balance", 1000).Result()
 	// Redis.IncrByFloat(context.Background(), "user:2:balance", 1000).Result()
 }
 
 func main() {
+	common.SetupGinLog()
+	common.SysLog("Api2gpt Mid Service " + common.Version + " started")
+	if os.Getenv("GIN_MODE") != "debug" {
+		gin.SetMode(gin.ReleaseMode)
+	}
+	if common.DebugEnabled {
+		common.SysLog("running in debug mode")
+	}
 
-	// 禁用控制台颜色，将日志写入文件时不需要控制台颜色。
+	// Initialize Redis
-	gin.DisableConsoleColor()
+	err := common.InitRedisClient()
-
+	if err != nil {
-	// 记录到文件。
+		common.FatalLog("failed to initialize Redis: " + err.Error())
-	filename := time.Now().Format("20060102150405") + ".log"
+	}
-	f, _ := os.Create("logs/gin" + filename)
-	//gin.DefaultWriter = io.MultiWriter(f)
-	// 如果需要同时将日志写入文件和控制台，请使用以下代码。
-	gin.DefaultWriter = io.MultiWriter(f, os.Stdout)
-
-	log.SetOutput(gin.DefaultWriter)
-
-	r := gin.Default()
-
-	//添加跨域支持
-	r.Use(Cors())
-
-	r.GET("/dashboard/billing/credit_grants", checkKeyMid(), balance)
-
-	r.GET("/v1/models", handleGetModels)
-	r.OPTIONS("/v1/*path", handleOptions)
-	r.POST("/v1/chat/completions", checkKeyMid(), completions)
-	r.POST("/v1/completions", checkKeyMid(), completions)
-	r.POST("/v1/embeddings", checkKeyMid(), embeddings)
-	r.POST("/v1/edits", checkKeyMid(), edit)
-	r.POST("/v1/images/generations", checkKeyMid(), images)
 
-	//r.POST("/mock1", mockBalanceInquiry)
+	// Initialize HTTP server
-	//r.POST("/mock2", mockBalanceConsumption)
+	server := gin.Default()
+	server.Use(middleware.CORS())
 
-	// 定义一个GET请求测试接口
+	router.SetRouter(server)
-	r.GET("/ping", func(c *gin.Context) {
-		c.JSON(200, gin.H{
-			"message": "pong from api2gpt",
-		})
-	})
 
-	Redis = InitRedis()
 	//添加测试数据
 	//test_redis()
 
-	r.Run("0.0.0.0:8080")
+	var port = os.Getenv("PORT")
+	if port == "" {
+		port = strconv.Itoa(*common.Port)
+	}
+	err = server.Run(":" + port)
+	if err != nil {
+		common.FatalLog("failed to start HTTP server: " + err.Error())
+	}
 }

middleware/auth.go

@@ -0,0 +1,29 @@
+package middleware
+
+import (
+	"api2gpt-mid/service"
+	"github.com/gin-gonic/gin"
+	"log"
+)
+
+func TokenAuth() func(c *gin.Context) {
+	return func(c *gin.Context) {
+		auth := c.Request.Header.Get("Authorization")
+		log.Printf("auth: %v", auth)
+		if auth == "" {
+			c.AbortWithStatusJSON(401, gin.H{"code": 40001})
+		} else {
+			//key := strings.Trim(auth, "Bearer ")
+			key := auth[7:]
+			log.Printf("key: %v", key)
+			msg, err := service.CheckKeyAndTimeCount(key)
+			if err != nil {
+				log.Printf("checkKeyMid err: %v", err)
+				c.AbortWithStatusJSON(msg, gin.H{"code": err.Error()})
+			}
+		}
+		log.Printf("auth check end")
+		// 执行函数
+		c.Next()
+	}
+}

middleware/cors.go

@@ -0,0 +1,24 @@
+package middleware
+
+import (
+	"github.com/gin-gonic/gin"
+	"net/http"
+)
+
+func CORS() gin.HandlerFunc {
+	return func(c *gin.Context) {
+		method := c.Request.Method
+		origin := c.Request.Header.Get("Origin")
+		if origin != "" {
+			c.Header("Access-Control-Allow-Origin", "*") // 可将将 * 替换为指定的域名
+			c.Header("Access-Control-Allow-Methods", "POST, GET, OPTIONS, PUT, DELETE, UPDATE")
+			c.Header("Access-Control-Allow-Headers", "Origin, X-Requested-With, Content-Type, Accept, Authorization")
+			c.Header("Access-Control-Expose-Headers", "Content-Length, Access-Control-Allow-Origin, Access-Control-Allow-Headers, Cache-Control, Content-Language, Content-Type")
+			c.Header("Access-Control-Allow-Credentials", "true")
+		}
+		if method == "OPTIONS" {
+			c.AbortWithStatus(http.StatusNoContent)
+		}
+		c.Next()
+	}
+}

model/consumption.go

@@ -0,0 +1,10 @@
+package model
+
+type Consumption struct {
+	SecretKey        string `json:"secretKey"`
+	Model            string `json:"model"`
+	MsgId            string `json:"msgId"`
+	PromptTokens     int    `json:"promptTokens"`
+	CompletionTokens int    `json:"completionTokens"`
+	TotalTokens      int    `json:"totalTokens"`
+}

model/model.go

@@ -0,0 +1,9 @@
+package model
+
+type ModelInfo struct {
+	ModelName       string  `json:"model_name"`
+	ModelPrice      float64 `json:"model_price"`
+	ModelPrice2     float64 `json:"model_price2"`
+	ModelPrepayment int     `json:"model_prepayment"`
+	ServerId        int     `json:"server_id"`
+}

model/openai.go

@@ -0,0 +1,99 @@
+package model
+
+type Message struct {
+	Role    string `json:"role,omitempty"`
+	Name    string `json:"name,omitempty"`
+	Content string `json:"content,omitempty"`
+}
+
+type ChatRequest struct {
+	Stream          bool        `json:"stream,omitempty"`
+	Model           string      `json:"model,omitempty"`
+	PresencePenalty float64     `json:"presence_penalty,omitempty"`
+	Temperature     float64     `json:"temperature,omitempty"`
+	Messages        []Message   `json:"messages,omitempty"`
+	Prompt          string      `json:"prompt,omitempty"`
+	Input           interface{} `json:"input,omitempty"`
+	Instruction     string      `json:"instruction,omitempty"`
+	N               int         `json:"n,omitempty"`
+	Functions       interface{} `json:"functions,omitempty"`
+	MaxTokens       int         `json:"max_tokens,omitempty"`
+}
+
+type ImagesRequest struct {
+	Prompt string `json:"prompt,omitempty"`
+	N      int    `json:"n,omitempty"`
+	Size   string `json:"size,omitempty"`
+}
+
+type ChatResponse struct {
+	Id      string   `json:"id,omitempty"`
+	Object  string   `json:"object,omitempty"`
+	Created int64    `json:"created,omitempty"`
+	Model   string   `json:"model,omitempty"`
+	Usage   Usage    `json:"usage,omitempty"`
+	Choices []Choice `json:"choices,omitempty"`
+}
+type ImagesResponse struct {
+	Created int64        `json:"created,omitempty"`
+	Data    []ImagesData `json:"data,omitempty"`
+}
+
+type ImagesData struct {
+	Url string `json:"url,omitempty"`
+}
+
+type Usage struct {
+	PromptTokens     int `json:"prompt_tokens,omitempty"`
+	CompletionTokens int `json:"completion_tokens,omitempty"`
+	TotalTokens      int `json:"total_tokens,omitempty"`
+}
+
+type Choice struct {
+	Message      Message `json:"message,omitempty"`
+	Delta        Delta   `json:"delta,omitempty"`
+	FinishReason string  `json:"finish_reason,omitempty"`
+	Index        int     `json:"index,omitempty"`
+	Text         string  `json:"text,omitempty"`
+}
+
+type Delta struct {
+	Content string `json:"content,omitempty"`
+}
+
+type CreditSummary struct {
+	Object         string  `json:"object"`
+	TotalGranted   float64 `json:"total_granted"`
+	TotalUsed      float64 `json:"total_used"`
+	TotalRemaining float64 `json:"total_remaining"`
+}
+
+type ListModelResponse struct {
+	Object string  `json:"object"`
+	Data   []Model `json:"data"`
+}
+
+type Model struct {
+	ID         string            `json:"id"`
+	Object     string            `json:"object"`
+	Created    int               `json:"created"`
+	OwnedBy    string            `json:"owned_by"`
+	Permission []ModelPermission `json:"permission"`
+	Root       string            `json:"root"`
+	Parent     any               `json:"parent"`
+}
+
+type ModelPermission struct {
+	ID                 string `json:"id"`
+	Object             string `json:"object"`
+	Created            int    `json:"created"`
+	AllowCreateEngine  bool   `json:"allow_create_engine"`
+	AllowSampling      bool   `json:"allow_sampling"`
+	AllowLogprobs      bool   `json:"allow_logprobs"`
+	AllowSearchIndices bool   `json:"allow_search_indices"`
+	AllowView          bool   `json:"allow_view"`
+	AllowFineTuning    bool   `json:"allow_fine_tuning"`
+	Organization       string `json:"organization"`
+	Group              any    `json:"group"`
+	IsBlocking         bool   `json:"is_blocking"`
+}

model/server.go

@@ -0,0 +1,6 @@
+package model
+
+type ServerInfo struct {
+	ServerAddress string `json:"server_address"`
+	AvailableKey  string `json:"available_key"`
+}

model/user.go

@@ -0,0 +1,8 @@
+package model
+
+type UserInfo struct {
+	UID string `json:"uid"`
+	SID string `json:"sid"`
+	RPM int    `json:"rpm"`
+	TPM int    `json:"tpm"`
+}

router/api-router.go

@@ -0,0 +1,36 @@
+package router
+
+import (
+	"api2gpt-mid/controller"
+	"api2gpt-mid/middleware"
+	"github.com/gin-gonic/gin"
+)
+
+func SetApiRouter(router *gin.Engine) {
+	modelsRouter := router.Group("/v1/models")
+	modelsRouter.Use(middleware.TokenAuth())
+	{
+		modelsRouter.GET("", controller.ListModels)
+		modelsRouter.GET("/:model", controller.RetrieveModel)
+	}
+	relayV1Router := router.Group("/v1")
+	relayV1Router.Use(middleware.TokenAuth())
+	{
+		relayV1Router.POST("/completions", controller.Completions)
+		relayV1Router.POST("/chat/completions", controller.Completions)
+		relayV1Router.POST("/embeddings", controller.Embeddings)
+		relayV1Router.POST("/edits", controller.Edit)
+		relayV1Router.POST("/images/generations", controller.Images)
+	}
+	dashboardRouter := router.Group("/dashboard")
+	dashboardRouter.Use(middleware.TokenAuth())
+	{
+		dashboardRouter.GET("/dashboard/billing/credit_grants", controller.Balance)
+	}
+	router.OPTIONS("/v1/*path", controller.HandleOptions)
+	router.GET("/ping", func(c *gin.Context) {
+		c.JSON(200, gin.H{
+			"message": "pong from api2gpt",
+		})
+	})
+}

router/main.go

@@ -0,0 +1,7 @@
+package router
+
+import "github.com/gin-gonic/gin"
+
+func SetRouter(router *gin.Engine) {
+	SetApiRouter(router)
+}

service.go

@@ -1,255 +0,0 @@
-package main
-
-import (
-	"context"
-	"encoding/json"
-	"errors"
-	"log"
-	"strconv"
-	"time"
-)
-
-// 检测key是否存在，是否超出每分钟请求次数
-func checkKeyAndTimeCount(key string) (int, error) {
-	var timeOut = 60 * time.Second
-	var timeCount = 30
-	userInfoStr, err := Redis.Get(context.Background(), "user:"+key).Result()
-	log.Printf("用户信息 %s", userInfoStr)
-	if err != nil {
-		//用户不存在
-		return 401, errors.New("40003")
-	}
-	var userInfo UserInfo
-	err = json.Unmarshal([]byte(userInfoStr), &userInfo)
-	if err != nil {
-		//用户状态异常
-		return 401, errors.New("40004")
-	}
-	count, err := Redis.Incr(context.Background(), "user:count:"+key).Result()
-	log.Printf("用户请求次数 %d", count)
-	if err != nil {
-		log.Printf("系统计数器设置异常 %s", err.Error())
-		return 500, errors.New("系统计数器设置异常")
-	}
-	if count == 1 {
-		_, err := Redis.Expire(context.Background(), "user:count:"+key, timeOut).Result()
-		if err != nil {
-			log.Printf("系统计数器异常 %s", err.Error())
-			return 500, errors.New("系统计数器异常")
-		}
-	}
-	// 如果请求次数超出限制，则中断请求
-	if count > int64(timeCount) {
-		//您的账户请求次数过多，超过分钟配额
-		return 429, errors.New("42901")
-	}
-	return 200, nil
-}
-
-func queryBlance(key string) (float64, error) {
-	userInfoStr, err := Redis.Get(context.Background(), "user:"+key).Result()
-	var userInfo UserInfo
-	err = json.Unmarshal([]byte(userInfoStr), &userInfo)
-	balance, err := Redis.IncrByFloat(context.Background(), "user:"+userInfo.UID+":balance", 0).Result()
-	if err != nil {
-		return 0, errors.New("余额计算失败")
-	}
-	return balance, nil
-}
-
-// 余额查询
-func checkBlance(key string, model string) (ServerInfo, error) {
-	var serverInfo ServerInfo
-
-	//获取模型价格
-	modelPriceStr, err := Redis.Get(context.Background(), "model:"+model).Result()
-	if err != nil {
-		return serverInfo, errors.New("模型信息不存在")
-	}
-	var modelInfo ModelInfo
-	err = json.Unmarshal([]byte(modelPriceStr), &modelInfo)
-	if err != nil {
-		return serverInfo, errors.New("模型信息解析失败")
-	}
-
-	//获取用户信息
-	userInfoStr, err := Redis.Get(context.Background(), "user:"+key).Result()
-	var userInfo UserInfo
-	err = json.Unmarshal([]byte(userInfoStr), &userInfo)
-	//获取服务器信息
-	serverInfoStr, err := Redis.Get(context.Background(), "server:"+strconv.Itoa(modelInfo.ServerId)).Result()
-	if err != nil {
-		return serverInfo, errors.New("服务器信息不存在")
-	}
-	err = json.Unmarshal([]byte(serverInfoStr), &serverInfo)
-	if err != nil {
-		return serverInfo, errors.New("服务器信息解析失败")
-	}
-
-	//计算余额-先扣除指定金额
-	balance, err := Redis.IncrByFloat(context.Background(), "user:"+userInfo.UID+":balance", -(float64(modelInfo.ModelPrepayment) * modelInfo.ModelPrice)).Result()
-	if err != nil {
-		return serverInfo, errors.New("余额计算失败")
-	}
-	log.Printf("用户余额 %f key: %v 预扣了：%f", balance, key, (float64(modelInfo.ModelPrepayment) * modelInfo.ModelPrice))
-	if balance < 0 {
-		Redis.IncrByFloat(context.Background(), "user:"+userInfo.UID+":balance", float64(modelInfo.ModelPrepayment)*modelInfo.ModelPrice).Result()
-		return serverInfo, errors.New("用户余额不足")
-	}
-
-	return serverInfo, nil
-}
-
-// 余额查询 for images
-func checkBlanceForImages(key string, model string, n int) (ServerInfo, error) {
-	var serverInfo ServerInfo
-
-	//获取模型价格
-	modelPriceStr, err := Redis.Get(context.Background(), "model:"+model).Result()
-	if err != nil {
-		return serverInfo, errors.New("模型信息不存在")
-	}
-	var modelInfo ModelInfo
-	err = json.Unmarshal([]byte(modelPriceStr), &modelInfo)
-	if err != nil {
-		return serverInfo, errors.New("模型信息解析失败")
-	}
-
-	//获取用户信息
-	userInfoStr, err := Redis.Get(context.Background(), "user:"+key).Result()
-	var userInfo UserInfo
-	err = json.Unmarshal([]byte(userInfoStr), &userInfo)
-	//获取服务器信息
-	serverInfoStr, err := Redis.Get(context.Background(), "server:"+strconv.Itoa(modelInfo.ServerId)).Result()
-	if err != nil {
-		return serverInfo, errors.New("服务器信息不存在")
-	}
-	err = json.Unmarshal([]byte(serverInfoStr), &serverInfo)
-	if err != nil {
-		return serverInfo, errors.New("服务器信息解析失败")
-	}
-
-	//计算余额-先扣除指定金额
-	balance, err := Redis.IncrByFloat(context.Background(), "user:"+userInfo.UID+":balance", -(float64(modelInfo.ModelPrepayment*n) * modelInfo.ModelPrice)).Result()
-	if err != nil {
-		return serverInfo, errors.New("余额计算失败")
-	}
-	log.Printf("用户余额 %f key: %v 预扣了：%f", balance, key, (float64(modelInfo.ModelPrepayment*n) * modelInfo.ModelPrice))
-	if balance < 0 {
-		Redis.IncrByFloat(context.Background(), "user:"+userInfo.UID+":balance", float64(modelInfo.ModelPrepayment*n)*modelInfo.ModelPrice).Result()
-		return serverInfo, errors.New("用户余额不足")
-	}
-
-	return serverInfo, nil
-}
-
-// 预扣返还
-func checkBlanceReturn(key string, model string) error {
-	//获取用户信息
-	userInfoStr, err := Redis.Get(context.Background(), "user:"+key).Result()
-	var userInfo UserInfo
-	err = json.Unmarshal([]byte(userInfoStr), &userInfo)
-
-	//获取模型价格
-	modelPriceStr, err := Redis.Get(context.Background(), "model:"+model).Result()
-	if err != nil {
-		return errors.New("模型信息不存在")
-	}
-	var modelInfo ModelInfo
-	err = json.Unmarshal([]byte(modelPriceStr), &modelInfo)
-	if err != nil {
-		return errors.New("模型信息解析失败")
-	}
-	balance, err := Redis.IncrByFloat(context.Background(), "user:"+userInfo.UID+":balance", (float64(modelInfo.ModelPrepayment) * modelInfo.ModelPrice)).Result()
-	log.Printf("用户余额 %f key: %v 返还预扣：%f", balance, key, (float64(modelInfo.ModelPrepayment) * modelInfo.ModelPrice))
-	return nil
-}
-
-// 预扣返还 for images
-func checkBlanceReturnForImages(key string, model string, n int) error {
-	//获取用户信息
-	userInfoStr, err := Redis.Get(context.Background(), "user:"+key).Result()
-	var userInfo UserInfo
-	err = json.Unmarshal([]byte(userInfoStr), &userInfo)
-	//获取模型价格
-	modelPriceStr, err := Redis.Get(context.Background(), "model:"+model).Result()
-	if err != nil {
-		return errors.New("模型信息不存在")
-	}
-	var modelInfo ModelInfo
-	err = json.Unmarshal([]byte(modelPriceStr), &modelInfo)
-	if err != nil {
-		return errors.New("模型信息解析失败")
-	}
-	balance, err := Redis.IncrByFloat(context.Background(), "user:"+userInfo.UID+":balance", (float64(modelInfo.ModelPrepayment*n) * modelInfo.ModelPrice)).Result()
-	log.Printf("用户余额 %f key: %v 返还预扣：%f", balance, key, (float64(modelInfo.ModelPrepayment*n) * modelInfo.ModelPrice))
-	return nil
-}
-
-// 余额消费
-func consumption(key string, model string, prompt_tokens int, completion_tokens int, total_tokens int, msg_id string) (string, error) {
-	//获取用户信息
-	userInfoStr, err := Redis.Get(context.Background(), "user:"+key).Result()
-	if err != nil {
-		return "", errors.New("用户信息不存在")
-	}
-	var userInfo UserInfo
-	err = json.Unmarshal([]byte(userInfoStr), &userInfo)
-	if err != nil {
-		return "", errors.New("用户信息解析失败")
-	}
-	//获取模型价格
-	modelPriceStr, err := Redis.Get(context.Background(), "model:"+model).Result()
-	if err != nil {
-		return "", errors.New("模型信息不存在")
-	}
-	var modelInfo ModelInfo
-	err = json.Unmarshal([]byte(modelPriceStr), &modelInfo)
-	if err != nil {
-		return "", errors.New("模型信息解析失败")
-	}
-	balance, err := Redis.IncrByFloat(context.Background(), "user:"+userInfo.UID+":balance", float64(modelInfo.ModelPrepayment)*modelInfo.ModelPrice-(float64(prompt_tokens)*modelInfo.ModelPrice+float64(completion_tokens)*modelInfo.ModelPrice2)).Result()
-
-	// 余额消费日志请求
-	result, err := balanceConsumption(key, model, prompt_tokens, completion_tokens, total_tokens, msg_id)
-	log.Printf("用户余额：%f 扣费KEY: %s 扣费token数: %d 扣费：%f 扣费日志发送结果 %s", balance, key, total_tokens, float64(modelInfo.ModelPrepayment)*modelInfo.ModelPrice-(float64(prompt_tokens)*modelInfo.ModelPrice+float64(completion_tokens)*modelInfo.ModelPrice2), result)
-	if err != nil {
-		log.Printf("%s 余额消费日志请求失败 %v", key, err)
-		return "", err
-	}
-	return result, nil
-}
-
-// 余额消费 for images
-func consumptionForImages(key string, model string, n int, dataNum int, msg_id string) (string, error) {
-	//获取用户信息
-	userInfoStr, err := Redis.Get(context.Background(), "user:"+key).Result()
-	if err != nil {
-		return "", errors.New("用户信息不存在")
-	}
-	var userInfo UserInfo
-	err = json.Unmarshal([]byte(userInfoStr), &userInfo)
-	if err != nil {
-		return "", errors.New("用户信息解析失败")
-	}
-	//获取模型价格
-	modelPriceStr, err := Redis.Get(context.Background(), "model:"+model).Result()
-	if err != nil {
-		return "", errors.New("模型信息不存在")
-	}
-	var modelInfo ModelInfo
-	err = json.Unmarshal([]byte(modelPriceStr), &modelInfo)
-	if err != nil {
-		return "", errors.New("模型信息解析失败")
-	}
-	balance, err := Redis.IncrByFloat(context.Background(), "user:"+userInfo.UID+":balance", float64(modelInfo.ModelPrepayment*n)*modelInfo.ModelPrice-(float64(1000*dataNum)*modelInfo.ModelPrice)).Result()
-
-	// 余额消费日志请求
-	result, err := balanceConsumption(key, model, 0, 1000*dataNum, 1000*dataNum, msg_id)
-	log.Printf("用户余额：%f 扣费KEY: %s 扣费token数: %d 扣费：%f 扣费日志发送结果 %s", balance, key, 1000*dataNum, float64(1000*dataNum)*modelInfo.ModelPrice, result)
-	if err != nil {
-		log.Printf("%s 余额消费日志请求失败 %v", key, err)
-		return "", err
-	}
-	return result, nil
-}

service/business.go

@@ -0,0 +1,272 @@
+package service
+
+import (
+	"api2gpt-mid/api"
+	"api2gpt-mid/common"
+	"api2gpt-mid/model"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"strconv"
+	"time"
+)
+
+// CheckKeyAndTimeCount 检测key是否存在，是否超出每分钟请求次数
+func CheckKeyAndTimeCount(key string) (int, error) {
+	var timeOut = 60 * time.Second
+	var timeCount = 60
+	userInfoStr, err := common.RedisGet("user:" + key)
+	common.SysLog("用户信息" + userInfoStr)
+	if err != nil {
+		//用户不存在
+		return 401, errors.New("40003")
+	}
+	var userInfo model.UserInfo
+	err = json.Unmarshal([]byte(userInfoStr), &userInfo)
+	if err != nil {
+		//用户状态异常
+		return 401, errors.New("40004")
+	}
+	if userInfo.RPM > 0 {
+		timeCount = userInfo.RPM
+	}
+	count, err := common.RedisIncr("user:count:" + key)
+	common.SysLog("用户请求次数" + strconv.FormatInt(count, 10))
+	if err != nil {
+		fmt.Printf("系统计数器设置异常 %s", err.Error())
+		return 500, errors.New("系统计数器设置异常")
+	}
+	if count == 1 {
+		_, err := common.RedisExpire("user:count:"+key, timeOut)
+		if err != nil {
+			fmt.Printf("系统计数器异常 %s", err.Error())
+			return 500, errors.New("系统计数器异常")
+		}
+	}
+	// 如果请求次数超出限制，则中断请求
+	if count > int64(timeCount) {
+		//您的账户请求次数过多，超过分钟配额
+		return 429, errors.New("42901")
+	}
+	return 200, nil
+}
+
+func QueryBlance(key string) (float64, error) {
+	userInfoStr, err := common.RedisGet("user:" + key)
+	var userInfo model.UserInfo
+	err = json.Unmarshal([]byte(userInfoStr), &userInfo)
+	balance, err := common.RedisIncrByFloat("user:"+userInfo.UID+":balance", 0)
+	if err != nil {
+		return 0, errors.New("余额计算失败")
+	}
+	return balance, nil
+}
+
+// 余额查询
+func CheckBlance(key string, modelStr string, maxTokens int) (model.ServerInfo, error) {
+	var serverInfo model.ServerInfo
+
+	//获取模型价格
+	modelPriceStr, err := common.RedisGet("model:" + modelStr)
+	if err != nil {
+		return serverInfo, errors.New("模型信息不存在")
+	}
+	var modelInfo model.ModelInfo
+	err = json.Unmarshal([]byte(modelPriceStr), &modelInfo)
+	if err != nil {
+		return serverInfo, errors.New("模型信息解析失败")
+	}
+
+	//获取用户信息
+	userInfoStr, err := common.RedisGet("user:" + key)
+	var userInfo model.UserInfo
+	err = json.Unmarshal([]byte(userInfoStr), &userInfo)
+	//获取服务器信息
+	serverInfoStr, err := common.RedisGet("server:" + strconv.Itoa(modelInfo.ServerId))
+	if err != nil {
+		return serverInfo, errors.New("服务器信息不存在")
+	}
+	err = json.Unmarshal([]byte(serverInfoStr), &serverInfo)
+	if err != nil {
+		return serverInfo, errors.New("服务器信息解析失败")
+	}
+
+	//计算余额-先扣除指定金额
+	if maxTokens == 0 {
+		maxTokens = modelInfo.ModelPrepayment
+	}
+	balance, err := common.RedisIncrByFloat("user:"+userInfo.UID+":balance", -(float64(maxTokens) * modelInfo.ModelPrice))
+	if err != nil {
+		return serverInfo, errors.New("余额计算失败")
+	}
+	fmt.Printf("用户余额 %f key: %v 预扣了：%f", balance, key, (float64(maxTokens) * modelInfo.ModelPrice))
+	if balance < 0 {
+		_, err := common.RedisIncrByFloat("user:"+userInfo.UID+":balance", float64(maxTokens)*modelInfo.ModelPrice)
+		if err != nil {
+			return serverInfo, errors.New("用户缓存出错")
+		}
+		return serverInfo, errors.New("用户余额不足")
+	}
+
+	return serverInfo, nil
+}
+
+// 余额查询 for images
+func CheckBlanceForImages(key string, modelStr string, n int) (model.ServerInfo, error) {
+	var serverInfo model.ServerInfo
+
+	//获取模型价格
+	modelPriceStr, err := common.RedisGet("model:" + modelStr)
+	if err != nil {
+		return serverInfo, errors.New("模型信息不存在")
+	}
+	var modelInfo model.ModelInfo
+	err = json.Unmarshal([]byte(modelPriceStr), &modelInfo)
+	if err != nil {
+		return serverInfo, errors.New("模型信息解析失败")
+	}
+
+	//获取用户信息
+	userInfoStr, err := common.RedisGet("user:" + key)
+	var userInfo model.UserInfo
+	err = json.Unmarshal([]byte(userInfoStr), &userInfo)
+	//获取服务器信息
+	serverInfoStr, err := common.RedisGet("server:" + strconv.Itoa(modelInfo.ServerId))
+	if err != nil {
+		return serverInfo, errors.New("服务器信息不存在")
+	}
+	err = json.Unmarshal([]byte(serverInfoStr), &serverInfo)
+	if err != nil {
+		return serverInfo, errors.New("服务器信息解析失败")
+	}
+
+	//计算余额-先扣除指定金额
+	balance, err := common.RedisIncrByFloat("user:"+userInfo.UID+":balance", -(float64(modelInfo.ModelPrepayment*n) * modelInfo.ModelPrice))
+	if err != nil {
+		return serverInfo, errors.New("余额计算失败")
+	}
+	fmt.Printf("用户余额 %f key: %v 预扣了：%f", balance, key, (float64(modelInfo.ModelPrepayment*n) * modelInfo.ModelPrice))
+	if balance < 0 {
+		_, err := common.RedisIncrByFloat("user:"+userInfo.UID+":balance", float64(modelInfo.ModelPrepayment*n)*modelInfo.ModelPrice)
+		if err != nil {
+			return serverInfo, errors.New("用户缓存出错")
+		}
+		return serverInfo, errors.New("用户余额不足")
+	}
+
+	return serverInfo, nil
+}
+
+// 预扣返还
+func CheckBlanceReturn(key string, modelStr string, maxTokens int) error {
+	//获取用户信息
+	userInfoStr, err := common.RedisGet("user:" + key)
+	var userInfo model.UserInfo
+	err = json.Unmarshal([]byte(userInfoStr), &userInfo)
+
+	//获取模型价格
+	modelPriceStr, err := common.RedisGet("model:" + modelStr)
+	if err != nil {
+		return errors.New("模型信息不存在")
+	}
+	var modelInfo model.ModelInfo
+	err = json.Unmarshal([]byte(modelPriceStr), &modelInfo)
+	if err != nil {
+		return errors.New("模型信息解析失败")
+	}
+	if maxTokens == 0 {
+		maxTokens = modelInfo.ModelPrepayment
+	}
+	balance, err := common.RedisIncrByFloat("user:"+userInfo.UID+":balance", (float64(maxTokens) * modelInfo.ModelPrice))
+	fmt.Printf("用户余额 %f key: %v 返还预扣：%f", balance, key, (float64(maxTokens) * modelInfo.ModelPrice))
+	return nil
+}
+
+// 预扣返还 for images
+func CheckBlanceReturnForImages(key string, modelStr string, n int) error {
+	//获取用户信息
+	userInfoStr, err := common.RedisGet("user:" + key)
+	var userInfo model.UserInfo
+	err = json.Unmarshal([]byte(userInfoStr), &userInfo)
+	//获取模型价格
+	modelPriceStr, err := common.RedisGet("model:" + modelStr)
+	if err != nil {
+		return errors.New("模型信息不存在")
+	}
+	var modelInfo model.ModelInfo
+	err = json.Unmarshal([]byte(modelPriceStr), &modelInfo)
+	if err != nil {
+		return errors.New("模型信息解析失败")
+	}
+	balance, err := common.RedisIncrByFloat("user:"+userInfo.UID+":balance", (float64(modelInfo.ModelPrepayment*n) * modelInfo.ModelPrice))
+	fmt.Printf("用户余额 %f key: %v 返还预扣：%f", balance, key, (float64(modelInfo.ModelPrepayment*n) * modelInfo.ModelPrice))
+	return nil
+}
+
+// 余额消费
+func Consumption(key string, modelStr string, prompt_tokens int, completion_tokens int, total_tokens int, msg_id string) (string, error) {
+	//获取用户信息
+	userInfoStr, err := common.RedisGet("user:" + key)
+	if err != nil {
+		return "", errors.New("用户信息不存在")
+	}
+	var userInfo model.UserInfo
+	err = json.Unmarshal([]byte(userInfoStr), &userInfo)
+	if err != nil {
+		return "", errors.New("用户信息解析失败")
+	}
+	//获取模型价格
+	modelPriceStr, err := common.RedisGet("model:" + modelStr)
+	if err != nil {
+		return "", errors.New("模型信息不存在")
+	}
+	var modelInfo model.ModelInfo
+	err = json.Unmarshal([]byte(modelPriceStr), &modelInfo)
+	if err != nil {
+		return "", errors.New("模型信息解析失败")
+	}
+	balance, err := common.RedisIncrByFloat("user:"+userInfo.UID+":balance", float64(modelInfo.ModelPrepayment)*modelInfo.ModelPrice-(float64(prompt_tokens)*modelInfo.ModelPrice+float64(completion_tokens)*modelInfo.ModelPrice2))
+
+	// 余额消费日志请求
+	result, err := api.BalanceConsumption(key, modelStr, prompt_tokens, completion_tokens, total_tokens, msg_id)
+	fmt.Printf("用户余额：%f 扣费KEY: %s 扣费token数: %d 扣费：%f 扣费日志发送结果 %s", balance, key, total_tokens, float64(modelInfo.ModelPrepayment)*modelInfo.ModelPrice-(float64(prompt_tokens)*modelInfo.ModelPrice+float64(completion_tokens)*modelInfo.ModelPrice2), result)
+	if err != nil {
+		fmt.Printf("%s 余额消费日志请求失败 %v", key, err)
+		return "", err
+	}
+	return result, nil
+}
+
+// 余额消费 for images
+func ConsumptionForImages(key string, modelStr string, n int, dataNum int, msg_id string) (string, error) {
+	//获取用户信息
+	userInfoStr, err := common.RedisGet("user:" + key)
+	if err != nil {
+		return "", errors.New("用户信息不存在")
+	}
+	var userInfo model.UserInfo
+	err = json.Unmarshal([]byte(userInfoStr), &userInfo)
+	if err != nil {
+		return "", errors.New("用户信息解析失败")
+	}
+	//获取模型价格
+	modelPriceStr, err := common.RedisGet("model:" + modelStr)
+	if err != nil {
+		return "", errors.New("模型信息不存在")
+	}
+	var modelInfo model.ModelInfo
+	err = json.Unmarshal([]byte(modelPriceStr), &modelInfo)
+	if err != nil {
+		return "", errors.New("模型信息解析失败")
+	}
+	balance, err := common.RedisIncrByFloat("user:"+userInfo.UID+":balance", float64(modelInfo.ModelPrepayment*n)*modelInfo.ModelPrice-(float64(1000*dataNum)*modelInfo.ModelPrice))
+
+	// 余额消费日志请求
+	result, err := api.BalanceConsumption(key, modelStr, 0, 1000*dataNum, 1000*dataNum, msg_id)
+	fmt.Printf("用户余额：%f 扣费KEY: %s 扣费token数: %d 扣费：%f 扣费日志发送结果 %s", balance, key, 1000*dataNum, float64(1000*dataNum)*modelInfo.ModelPrice, result)
+	if err != nil {
+		fmt.Printf("%s 余额消费日志请求失败 %v", key, err)
+		return "", err
+	}
+	return result, nil
+}