Merge branch 'main' into feat/baidu_channel_support_apikey

README.md

@@ -276,8 +276,10 @@ graph LR
    + 例子：`REDIS_CONN_STRING=redis://default:redispw@localhost:49153`
 2. `SESSION_SECRET`：设置之后将使用固定的会话密钥，这样系统重新启动后已登录用户的 cookie 将依旧有效。
    + 例子：`SESSION_SECRET=random_string`
-3. `SQL_DSN`：设置之后将使用指定数据库而非 SQLite，请使用 MySQL 8.0 版本。
+3. `SQL_DSN`：设置之后将使用指定数据库而非 SQLite，请使用 MySQL 或 PostgreSQL。
-   + 例子：`SQL_DSN=root:123456@tcp(localhost:3306)/oneapi`
+   + 例子：
+     + MySQL：`SQL_DSN=root:123456@tcp(localhost:3306)/oneapi`
+     + PostgreSQL：`SQL_DSN=postgres://postgres:123456@localhost:5432/oneapi`
    + 注意需要提前建立数据库 `oneapi`，无需手动建表，程序将自动建表。
    + 如果使用本地数据库：部署命令可添加 `--network="host"` 以使得容器内的程序可以访问到宿主机上的 MySQL。
    + 如果使用云数据库：如果云服务器需要验证身份，需要在连接参数中添加 `?tls=skip-verify`。
@@ -338,7 +340,8 @@ https://openai.justsong.cn
    + 上游通道 429 了。
 
 ## 相关项目
-[FastGPT](https://github.com/labring/FastGPT): 基于 LLM 大语言模型的知识库问答系统
+* [FastGPT](https://github.com/labring/FastGPT): 基于 LLM 大语言模型的知识库问答系统
+* [ChatGPT Next Web](https://github.com/Yidadaa/ChatGPT-Next-Web):  一键拥有你自己的跨平台 ChatGPT 应用
 
 ## 注意
 

common/constants.go

@@ -55,6 +55,8 @@ var EmailDomainWhitelist = []string{
 	"foxmail.com",
 }
 
+var DebugEnabled = os.Getenv("DEBUG") == "true"
+
 var LogConsumeEnabled = true
 
 var SMTPServer = ""

controller/relay-ali.go

@@ -166,11 +166,7 @@ func aliStreamHandler(c *gin.Context, resp *http.Response) (*OpenAIErrorWithStat
 		}
 		stopChan <- true
 	}()
-	c.Writer.Header().Set("Content-Type", "text/event-stream")
+	setEventStreamHeaders(c)
-	c.Writer.Header().Set("Cache-Control", "no-cache")
-	c.Writer.Header().Set("Connection", "keep-alive")
-	c.Writer.Header().Set("Transfer-Encoding", "chunked")
-	c.Writer.Header().Set("X-Accel-Buffering", "no")
 	lastResponseText := ""
 	c.Stream(func(w io.Writer) bool {
 		select {

controller/relay-baidu.go

@@ -201,11 +201,7 @@ func baiduStreamHandler(c *gin.Context, resp *http.Response) (*OpenAIErrorWithSt
 		}
 		stopChan <- true
 	}()
-	c.Writer.Header().Set("Content-Type", "text/event-stream")
+	setEventStreamHeaders(c)
-	c.Writer.Header().Set("Cache-Control", "no-cache")
-	c.Writer.Header().Set("Connection", "keep-alive")
-	c.Writer.Header().Set("Transfer-Encoding", "chunked")
-	c.Writer.Header().Set("X-Accel-Buffering", "no")
 	c.Stream(func(w io.Writer) bool {
 		select {
 		case data := <-dataChan:

controller/relay-claude.go

@@ -141,11 +141,7 @@ func claudeStreamHandler(c *gin.Context, resp *http.Response) (*OpenAIErrorWithS
 		}
 		stopChan <- true
 	}()
-	c.Writer.Header().Set("Content-Type", "text/event-stream")
+	setEventStreamHeaders(c)
-	c.Writer.Header().Set("Cache-Control", "no-cache")
-	c.Writer.Header().Set("Connection", "keep-alive")
-	c.Writer.Header().Set("Transfer-Encoding", "chunked")
-	c.Writer.Header().Set("X-Accel-Buffering", "no")
 	c.Stream(func(w io.Writer) bool {
 		select {
 		case data := <-dataChan:

controller/relay-openai.go

@@ -66,11 +66,7 @@ func openaiStreamHandler(c *gin.Context, resp *http.Response, relayMode int) (*O
 		}
 		stopChan <- true
 	}()
-	c.Writer.Header().Set("Content-Type", "text/event-stream")
+	setEventStreamHeaders(c)
-	c.Writer.Header().Set("Cache-Control", "no-cache")
-	c.Writer.Header().Set("Connection", "keep-alive")
-	c.Writer.Header().Set("Transfer-Encoding", "chunked")
-	c.Writer.Header().Set("X-Accel-Buffering", "no")
 	c.Stream(func(w io.Writer) bool {
 		select {
 		case data := <-dataChan:

controller/relay-palm.go

@@ -143,11 +143,7 @@ func palmStreamHandler(c *gin.Context, resp *http.Response) (*OpenAIErrorWithSta
 		dataChan <- string(jsonResponse)
 		stopChan <- true
 	}()
-	c.Writer.Header().Set("Content-Type", "text/event-stream")
+	setEventStreamHeaders(c)
-	c.Writer.Header().Set("Cache-Control", "no-cache")
-	c.Writer.Header().Set("Connection", "keep-alive")
-	c.Writer.Header().Set("Transfer-Encoding", "chunked")
-	c.Writer.Header().Set("X-Accel-Buffering", "no")
 	c.Stream(func(w io.Writer) bool {
 		select {
 		case data := <-dataChan:

controller/relay-text.go

@@ -314,50 +314,54 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
 	}
 
 	var textResponse TextResponse
+	tokenName := c.GetString("token_name")
+	channelId := c.GetInt("channel_id")
 
 	defer func() {
-		if consumeQuota {
+		// c.Writer.Flush()
-			quota := 0
+		go func() {
-			completionRatio := 1.0
+			if consumeQuota {
-			if strings.HasPrefix(textRequest.Model, "gpt-3.5") {
+				quota := 0
-				completionRatio = 1.333333
+				completionRatio := 1.0
-			}
+				if strings.HasPrefix(textRequest.Model, "gpt-3.5") {
-			if strings.HasPrefix(textRequest.Model, "gpt-4") {
+					completionRatio = 1.333333
-				completionRatio = 2
+				}
-			}
+				if strings.HasPrefix(textRequest.Model, "gpt-4") {
+					completionRatio = 2
+				}
 
-			promptTokens = textResponse.Usage.PromptTokens
+				promptTokens = textResponse.Usage.PromptTokens
-			completionTokens = textResponse.Usage.CompletionTokens
+				completionTokens = textResponse.Usage.CompletionTokens
 
-			quota = promptTokens + int(float64(completionTokens)*completionRatio)
+				quota = promptTokens + int(float64(completionTokens)*completionRatio)
-			quota = int(float64(quota) * ratio)
+				quota = int(float64(quota) * ratio)
-			if ratio != 0 && quota <= 0 {
+				if ratio != 0 && quota <= 0 {
-				quota = 1
+					quota = 1
+				}
+				totalTokens := promptTokens + completionTokens
+				if totalTokens == 0 {
+					// in this case, must be some error happened
+					// we cannot just return, because we may have to return the pre-consumed quota
+					quota = 0
+				}
+				quotaDelta := quota - preConsumedQuota
+				err := model.PostConsumeTokenQuota(tokenId, quotaDelta)
+				if err != nil {
+					common.SysError("error consuming token remain quota: " + err.Error())
+				}
+				err = model.CacheUpdateUserQuota(userId)
+				if err != nil {
+					common.SysError("error update user quota cache: " + err.Error())
+				}
+				if quota != 0 {
+					logContent := fmt.Sprintf("模型倍率 %.2f，分组倍率 %.2f", modelRatio, groupRatio)
+					model.RecordConsumeLog(userId, promptTokens, completionTokens, textRequest.Model, tokenName, quota, logContent)
+					model.UpdateUserUsedQuotaAndRequestCount(userId, quota)
+
+					model.UpdateChannelUsedQuota(channelId, quota)
+				}
 			}
-			totalTokens := promptTokens + completionTokens
+		}()
-			if totalTokens == 0 {
-				// in this case, must be some error happened
-				// we cannot just return, because we may have to return the pre-consumed quota
-				quota = 0
-			}
-			quotaDelta := quota - preConsumedQuota
-			err := model.PostConsumeTokenQuota(tokenId, quotaDelta)
-			if err != nil {
-				common.SysError("error consuming token remain quota: " + err.Error())
-			}
-			err = model.CacheUpdateUserQuota(userId)
-			if err != nil {
-				common.SysError("error update user quota cache: " + err.Error())
-			}
-			if quota != 0 {
-				tokenName := c.GetString("token_name")
-				logContent := fmt.Sprintf("模型倍率 %.2f，分组倍率 %.2f", modelRatio, groupRatio)
-				model.RecordConsumeLog(userId, promptTokens, completionTokens, textRequest.Model, tokenName, quota, logContent)
-				model.UpdateUserUsedQuotaAndRequestCount(userId, quota)
-				channelId := c.GetInt("channel_id")
-				model.UpdateChannelUsedQuota(channelId, quota)
-			}
-		}
 	}()
 	switch apiType {
 	case APITypeOpenAI:

controller/relay-utils.go

@@ -2,6 +2,7 @@ package controller
 
 import (
 	"fmt"
+	"github.com/gin-gonic/gin"
 	"github.com/pkoukk/tiktoken-go"
 	"one-api/common"
 )
@@ -106,3 +107,11 @@ func shouldDisableChannel(err *OpenAIError) bool {
 	}
 	return false
 }
+
+func setEventStreamHeaders(c *gin.Context) {
+	c.Writer.Header().Set("Content-Type", "text/event-stream")
+	c.Writer.Header().Set("Cache-Control", "no-cache")
+	c.Writer.Header().Set("Connection", "keep-alive")
+	c.Writer.Header().Set("Transfer-Encoding", "chunked")
+	c.Writer.Header().Set("X-Accel-Buffering", "no")
+}

controller/relay-xunfei.go

@@ -217,11 +217,7 @@ func xunfeiStreamHandler(c *gin.Context, textRequest GeneralOpenAIRequest, appId
 		}
 		stopChan <- true
 	}()
-	c.Writer.Header().Set("Content-Type", "text/event-stream")
+	setEventStreamHeaders(c)
-	c.Writer.Header().Set("Cache-Control", "no-cache")
-	c.Writer.Header().Set("Connection", "keep-alive")
-	c.Writer.Header().Set("Transfer-Encoding", "chunked")
-	c.Writer.Header().Set("X-Accel-Buffering", "no")
 	c.Stream(func(w io.Writer) bool {
 		select {
 		case xunfeiResponse := <-dataChan:

controller/relay-zhipu.go

@@ -224,11 +224,7 @@ func zhipuStreamHandler(c *gin.Context, resp *http.Response) (*OpenAIErrorWithSt
 		}
 		stopChan <- true
 	}()
-	c.Writer.Header().Set("Content-Type", "text/event-stream")
+	setEventStreamHeaders(c)
-	c.Writer.Header().Set("Cache-Control", "no-cache")
-	c.Writer.Header().Set("Connection", "keep-alive")
-	c.Writer.Header().Set("Transfer-Encoding", "chunked")
-	c.Writer.Header().Set("X-Accel-Buffering", "no")
 	c.Stream(func(w io.Writer) bool {
 		select {
 		case data := <-dataChan:

controller/token.go

@@ -109,7 +109,7 @@ func AddToken(c *gin.Context) {
 		})
 		return
 	}
-	if len(token.Name) == 0 || len(token.Name) > 30 {
+	if len(token.Name) > 30 {
 		c.JSON(http.StatusOK, gin.H{
 			"success": false,
 			"message": "令牌名称过长",
@@ -171,7 +171,7 @@ func UpdateToken(c *gin.Context) {
 		})
 		return
 	}
-	if len(token.Name) == 0 || len(token.Name) > 30 {
+	if len(token.Name) > 30 {
 		c.JSON(http.StatusOK, gin.H{
 			"success": false,
 			"message": "令牌名称过长",

go.mod

@@ -18,7 +18,7 @@ require (
 	golang.org/x/crypto v0.9.0
 	gorm.io/driver/mysql v1.4.3
 	gorm.io/driver/sqlite v1.4.3
-	gorm.io/gorm v1.24.0
+	gorm.io/gorm v1.25.0
 )
 
 require (
@@ -36,6 +36,9 @@ require (
 	github.com/gorilla/context v1.1.1 // indirect
 	github.com/gorilla/securecookie v1.1.1 // indirect
 	github.com/gorilla/sessions v1.2.1 // indirect
+	github.com/jackc/pgpassfile v1.0.0 // indirect
+	github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a // indirect
+	github.com/jackc/pgx/v5 v5.3.1 // indirect
 	github.com/jinzhu/inflection v1.0.0 // indirect
 	github.com/jinzhu/now v1.1.5 // indirect
 	github.com/json-iterator/go v1.1.12 // indirect
@@ -54,4 +57,5 @@ require (
 	golang.org/x/text v0.9.0 // indirect
 	google.golang.org/protobuf v1.30.0 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
+	gorm.io/driver/postgres v1.5.2 // indirect
 )

go.sum

@@ -69,6 +69,12 @@ github.com/gorilla/sessions v1.2.1 h1:DHd3rPN5lE3Ts3D8rKkQ8x/0kqfeNmBAaiSi+o7Fsg
 github.com/gorilla/sessions v1.2.1/go.mod h1:dk2InVEVJ0sfLlnXv9EAgkf6ecYs/i80K/zI+bUmuGM=
 github.com/gorilla/websocket v1.5.0 h1:PPwGk2jz7EePpoHN/+ClbZu8SPxiqlu12wZP/3sWmnc=
 github.com/gorilla/websocket v1.5.0/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
+github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM=
+github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg=
+github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a h1:bbPeKD0xmW/Y25WS6cokEszi5g+S0QxI/d45PkRi7Nk=
+github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a/go.mod h1:5TJZWKEWniPve33vlWYSoGYefn3gLQRzjfDlhSJ9ZKM=
+github.com/jackc/pgx/v5 v5.3.1 h1:Fcr8QJ1ZeLi5zsPZqQeUZhNhxfkkKBOgJuYkJHoBOtU=
+github.com/jackc/pgx/v5 v5.3.1/go.mod h1:t3JDKnCBlYIc0ewLF0Q7B8MXmoIaBOZj/ic7iHozM/8=
 github.com/jinzhu/inflection v1.0.0 h1:K317FqzuhWc8YvSVlFMCCUb36O/S9MCKRDI7QkRKD/E=
 github.com/jinzhu/inflection v1.0.0/go.mod h1:h+uFLlag+Qp1Va5pdKtLDYj+kHp5pxUVkryuEj+Srlc=
 github.com/jinzhu/now v1.1.4/go.mod h1:d3SSVoowX0Lcu0IBviAWJpolVfI5UJVZZ7cO71lE/z8=
@@ -187,9 +193,13 @@ gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
 gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
 gorm.io/driver/mysql v1.4.3 h1:/JhWJhO2v17d8hjApTltKNADm7K7YI2ogkR7avJUL3k=
 gorm.io/driver/mysql v1.4.3/go.mod h1:sSIebwZAVPiT+27jK9HIwvsqOGKx3YMPmrA3mBJR10c=
+gorm.io/driver/postgres v1.5.2 h1:ytTDxxEv+MplXOfFe3Lzm7SjG09fcdb3Z/c056DTBx0=
+gorm.io/driver/postgres v1.5.2/go.mod h1:fmpX0m2I1PKuR7mKZiEluwrP3hbs+ps7JIGMUBpCgl8=
 gorm.io/driver/sqlite v1.4.3 h1:HBBcZSDnWi5BW3B3rwvVTc510KGkBkexlOg0QrmLUuU=
 gorm.io/driver/sqlite v1.4.3/go.mod h1:0Aq3iPO+v9ZKbcdiz8gLWRw5VOPcBOPUQJFLq5e2ecI=
 gorm.io/gorm v1.23.8/go.mod h1:l2lP/RyAtc1ynaTjFksBde/O8v9oOGIApu2/xRitmZk=
 gorm.io/gorm v1.24.0 h1:j/CoiSm6xpRpmzbFJsQHYj+I8bGYWLXVHeYEyyKlF74=
 gorm.io/gorm v1.24.0/go.mod h1:DVrVomtaYTbqs7gB/x2uVvqnXzv0nqjB396B8cG4dBA=
+gorm.io/gorm v1.25.0 h1:+KtYtb2roDz14EQe4bla8CbQlmb9dN3VejSai3lprfU=
+gorm.io/gorm v1.25.0/go.mod h1:L4uxeKpfBml98NYqVqwAdmV1a2nBtAec/cf3fpucW/k=
 rsc.io/pdf v0.1.1/go.mod h1:n8OzWcQ6Sp37PL01nO98y4iUCRdTGarVfzxY20ICaU4=

main.go

@@ -26,6 +26,9 @@ func main() {
 	if os.Getenv("GIN_MODE") != "debug" {
 		gin.SetMode(gin.ReleaseMode)
 	}
+	if common.DebugEnabled {
+		common.SysLog("running in debug mode")
+	}
 	// Initialize SQL Database
 	err := model.InitDB()
 	if err != nil {

model/channel.go

@@ -141,7 +141,7 @@ func UpdateChannelStatusById(id int, status int) {
 }
 
 func UpdateChannelUsedQuota(id int, quota int) {
-	err := DB.Model(&Channel{}).Where("id = ?", id).Update("used_quota", gorm.Expr("used_quota + ?", quota)).Error
+	err := DB.Set("gorm:query_option", "FOR UPDATE").Model(&Channel{}).Where("id = ?", id).Update("used_quota", gorm.Expr("used_quota + ?", quota)).Error
 	if err != nil {
 		common.SysError("failed to update channel used quota: " + err.Error())
 	}

model/main.go

@@ -2,10 +2,12 @@ package model
 
 import (
 	"gorm.io/driver/mysql"
+	"gorm.io/driver/postgres"
 	"gorm.io/driver/sqlite"
 	"gorm.io/gorm"
 	"one-api/common"
 	"os"
+	"strings"
 	"time"
 )
 
@@ -34,29 +36,39 @@ func createRootAccountIfNeed() error {
 	return nil
 }
 
-func CountTable(tableName string) (num int64) {
+func chooseDB() (*gorm.DB, error) {
-	DB.Table(tableName).Count(&num)
-	return
-}
-
-func InitDB() (err error) {
-	var db *gorm.DB
 	if os.Getenv("SQL_DSN") != "" {
+		dsn := os.Getenv("SQL_DSN")
+		if strings.HasPrefix(dsn, "postgres://") {
+			// Use PostgreSQL
+			common.SysLog("using PostgreSQL as database")
+			return gorm.Open(postgres.New(postgres.Config{
+				DSN:                  dsn,
+				PreferSimpleProtocol: true, // disables implicit prepared statement usage
+			}), &gorm.Config{
+				PrepareStmt: true, // precompile SQL
+			})
+		}
 		// Use MySQL
 		common.SysLog("using MySQL as database")
-		db, err = gorm.Open(mysql.Open(os.Getenv("SQL_DSN")), &gorm.Config{
+		return gorm.Open(mysql.Open(dsn), &gorm.Config{
-			PrepareStmt: true, // precompile SQL
-		})
-	} else {
-		// Use SQLite
-		common.SysLog("SQL_DSN not set, using SQLite as database")
-		common.UsingSQLite = true
-		db, err = gorm.Open(sqlite.Open(common.SQLitePath), &gorm.Config{
 			PrepareStmt: true, // precompile SQL
 		})
 	}
-	common.SysLog("database connected")
+	// Use SQLite
+	common.SysLog("SQL_DSN not set, using SQLite as database")
+	common.UsingSQLite = true
+	return gorm.Open(sqlite.Open(common.SQLitePath), &gorm.Config{
+		PrepareStmt: true, // precompile SQL
+	})
+}
+
+func InitDB() (err error) {
+	db, err := chooseDB()
 	if err == nil {
+		if common.DebugEnabled {
+			db = db.Debug()
+		}
 		DB = db
 		sqlDB, err := DB.DB()
 		if err != nil {

model/token.go

@@ -131,7 +131,7 @@ func IncreaseTokenQuota(id int, quota int) (err error) {
 	if quota < 0 {
 		return errors.New("quota 不能为负数！")
 	}
-	err = DB.Model(&Token{}).Where("id = ?", id).Updates(
+	err = DB.Set("gorm:query_option", "FOR UPDATE").Model(&Token{}).Where("id = ?", id).Updates(
 		map[string]interface{}{
 			"remain_quota": gorm.Expr("remain_quota + ?", quota),
 			"used_quota":   gorm.Expr("used_quota - ?", quota),
@@ -144,7 +144,7 @@ func DecreaseTokenQuota(id int, quota int) (err error) {
 	if quota < 0 {
 		return errors.New("quota 不能为负数！")
 	}
-	err = DB.Model(&Token{}).Where("id = ?", id).Updates(
+	err = DB.Set("gorm:query_option", "FOR UPDATE").Model(&Token{}).Where("id = ?", id).Updates(
 		map[string]interface{}{
 			"remain_quota": gorm.Expr("remain_quota - ?", quota),
 			"used_quota":   gorm.Expr("used_quota + ?", quota),

model/user.go

@@ -275,7 +275,7 @@ func IncreaseUserQuota(id int, quota int) (err error) {
 	if quota < 0 {
 		return errors.New("quota 不能为负数！")
 	}
-	err = DB.Model(&User{}).Where("id = ?", id).Update("quota", gorm.Expr("quota + ?", quota)).Error
+	err = DB.Set("gorm:query_option", "FOR UPDATE").Model(&User{}).Where("id = ?", id).Update("quota", gorm.Expr("quota + ?", quota)).Error
 	return err
 }
 
@@ -283,7 +283,7 @@ func DecreaseUserQuota(id int, quota int) (err error) {
 	if quota < 0 {
 		return errors.New("quota 不能为负数！")
 	}
-	err = DB.Model(&User{}).Where("id = ?", id).Update("quota", gorm.Expr("quota - ?", quota)).Error
+	err = DB.Set("gorm:query_option", "FOR UPDATE").Model(&User{}).Where("id = ?", id).Update("quota", gorm.Expr("quota - ?", quota)).Error
 	return err
 }
 
@@ -293,7 +293,7 @@ func GetRootUserEmail() (email string) {
 }
 
 func UpdateUserUsedQuotaAndRequestCount(id int, quota int) {
-	err := DB.Model(&User{}).Where("id = ?", id).Updates(
+	err := DB.Set("gorm:query_option", "FOR UPDATE").Model(&User{}).Where("id = ?", id).Updates(
 		map[string]interface{}{
 			"used_quota":    gorm.Expr("used_quota + ?", quota),
 			"request_count": gorm.Expr("request_count + ?", 1),

web/src/pages/Home/index.js

@@ -65,7 +65,7 @@ const Home = () => {
                     <Card.Meta>系统信息总览</Card.Meta>
                     <Card.Description>
                       <p>名称：{statusState?.status?.system_name}</p>
-                      <p>版本：{statusState?.status?.version}</p>
+                      <p>版本：{statusState?.status?.version ? statusState?.status?.version : "unknown"}</p>
                       <p>
                         源码：
                         <a