diff --git a/common/config/config.go b/common/config/config.go new file mode 100644 index 00000000..34d871fe --- /dev/null +++ b/common/config/config.go @@ -0,0 +1,127 @@ +package config + +import ( + "one-api/common/helper" + "os" + "strconv" + "sync" + "time" + + "github.com/google/uuid" +) + +var SystemName = "One API" +var ServerAddress = "http://localhost:3000" +var Footer = "" +var Logo = "" +var TopUpLink = "" +var ChatLink = "" +var QuotaPerUnit = 500 * 1000.0 // $0.002 / 1K tokens +var DisplayInCurrencyEnabled = true +var DisplayTokenStatEnabled = true + +// Any options with "Secret", "Token" in its key won't be return by GetOptions + +var SessionSecret = uuid.New().String() + +var OptionMap map[string]string +var OptionMapRWMutex sync.RWMutex + +var ItemsPerPage = 10 +var MaxRecentItems = 100 + +var PasswordLoginEnabled = true +var PasswordRegisterEnabled = true +var EmailVerificationEnabled = false +var GitHubOAuthEnabled = false +var WeChatAuthEnabled = false +var TurnstileCheckEnabled = false +var RegisterEnabled = true + +var EmailDomainRestrictionEnabled = false +var EmailDomainWhitelist = []string{ + "gmail.com", + "163.com", + "126.com", + "qq.com", + "outlook.com", + "hotmail.com", + "icloud.com", + "yahoo.com", + "foxmail.com", +} + +var DebugEnabled = os.Getenv("DEBUG") == "true" +var MemoryCacheEnabled = os.Getenv("MEMORY_CACHE_ENABLED") == "true" + +var LogConsumeEnabled = true + +var SMTPServer = "" +var SMTPPort = 587 +var SMTPAccount = "" +var SMTPFrom = "" +var SMTPToken = "" + +var GitHubClientId = "" +var GitHubClientSecret = "" + +var WeChatServerAddress = "" +var WeChatServerToken = "" +var WeChatAccountQRCodeImageURL = "" + +var TurnstileSiteKey = "" +var TurnstileSecretKey = "" + +var QuotaForNewUser = 0 +var QuotaForInviter = 0 +var QuotaForInvitee = 0 +var ChannelDisableThreshold = 5.0 +var AutomaticDisableChannelEnabled = false +var AutomaticEnableChannelEnabled = false +var QuotaRemindThreshold = 1000 +var PreConsumedQuota = 500 +var ApproximateTokenEnabled = false +var RetryTimes = 0 + +var RootUserEmail = "" + +var IsMasterNode = os.Getenv("NODE_TYPE") != "slave" + +var requestInterval, _ = strconv.Atoi(os.Getenv("POLLING_INTERVAL")) +var RequestInterval = time.Duration(requestInterval) * time.Second + +var SyncFrequency = helper.GetOrDefaultEnvInt("SYNC_FREQUENCY", 10*60) // unit is second + +var BatchUpdateEnabled = false +var BatchUpdateInterval = helper.GetOrDefaultEnvInt("BATCH_UPDATE_INTERVAL", 5) + +var RelayTimeout = helper.GetOrDefaultEnvInt("RELAY_TIMEOUT", 0) // unit is second + +var GeminiSafetySetting = helper.GetOrDefaultEnvString("GEMINI_SAFETY_SETTING", "BLOCK_NONE") + +var Theme = helper.GetOrDefaultEnvString("THEME", "default") +var ValidThemes = map[string]bool{ + "default": true, + "berry": true, +} + +// All duration's unit is seconds +// Shouldn't larger then RateLimitKeyExpirationDuration +var ( + GlobalApiRateLimitNum = helper.GetOrDefaultEnvInt("GLOBAL_API_RATE_LIMIT", 180) + GlobalApiRateLimitDuration int64 = 3 * 60 + + GlobalWebRateLimitNum = helper.GetOrDefaultEnvInt("GLOBAL_WEB_RATE_LIMIT", 60) + GlobalWebRateLimitDuration int64 = 3 * 60 + + UploadRateLimitNum = 10 + UploadRateLimitDuration int64 = 60 + + DownloadRateLimitNum = 10 + DownloadRateLimitDuration int64 = 60 + + CriticalRateLimitNum = 20 + CriticalRateLimitDuration int64 = 20 * 60 +) + +var RateLimitKeyExpirationDuration = 20 * time.Minute diff --git a/common/constants.go b/common/constants.go index 9ee791df..325454d4 100644 --- a/common/constants.go +++ b/common/constants.go @@ -1,114 +1,9 @@ package common -import ( - "os" - "strconv" - "sync" - "time" - - "github.com/google/uuid" -) +import "time" var StartTime = time.Now().Unix() // unit: second var Version = "v0.0.0" // this hard coding will be replaced automatically when building, no need to manually change -var SystemName = "One API" -var ServerAddress = "http://localhost:3000" -var Footer = "" -var Logo = "" -var TopUpLink = "" -var ChatLink = "" -var QuotaPerUnit = 500 * 1000.0 // $0.002 / 1K tokens -var DisplayInCurrencyEnabled = true -var DisplayTokenStatEnabled = true - -// Any options with "Secret", "Token" in its key won't be return by GetOptions - -var SessionSecret = uuid.New().String() - -var OptionMap map[string]string -var OptionMapRWMutex sync.RWMutex - -var ItemsPerPage = 10 -var MaxRecentItems = 100 - -var PasswordLoginEnabled = true -var PasswordRegisterEnabled = true -var EmailVerificationEnabled = false -var GitHubOAuthEnabled = false -var WeChatAuthEnabled = false -var TurnstileCheckEnabled = false -var RegisterEnabled = true - -var EmailDomainRestrictionEnabled = false -var EmailDomainWhitelist = []string{ - "gmail.com", - "163.com", - "126.com", - "qq.com", - "outlook.com", - "hotmail.com", - "icloud.com", - "yahoo.com", - "foxmail.com", -} - -var DebugEnabled = os.Getenv("DEBUG") == "true" -var MemoryCacheEnabled = os.Getenv("MEMORY_CACHE_ENABLED") == "true" - -var LogConsumeEnabled = true - -var SMTPServer = "" -var SMTPPort = 587 -var SMTPAccount = "" -var SMTPFrom = "" -var SMTPToken = "" - -var GitHubClientId = "" -var GitHubClientSecret = "" - -var WeChatServerAddress = "" -var WeChatServerToken = "" -var WeChatAccountQRCodeImageURL = "" - -var TurnstileSiteKey = "" -var TurnstileSecretKey = "" - -var QuotaForNewUser = 0 -var QuotaForInviter = 0 -var QuotaForInvitee = 0 -var ChannelDisableThreshold = 5.0 -var AutomaticDisableChannelEnabled = false -var AutomaticEnableChannelEnabled = false -var QuotaRemindThreshold = 1000 -var PreConsumedQuota = 500 -var ApproximateTokenEnabled = false -var RetryTimes = 0 - -var RootUserEmail = "" - -var IsMasterNode = os.Getenv("NODE_TYPE") != "slave" - -var requestInterval, _ = strconv.Atoi(os.Getenv("POLLING_INTERVAL")) -var RequestInterval = time.Duration(requestInterval) * time.Second - -var SyncFrequency = GetOrDefault("SYNC_FREQUENCY", 10*60) // unit is second - -var BatchUpdateEnabled = false -var BatchUpdateInterval = GetOrDefault("BATCH_UPDATE_INTERVAL", 5) - -var RelayTimeout = GetOrDefault("RELAY_TIMEOUT", 0) // unit is second - -var GeminiSafetySetting = GetOrDefaultString("GEMINI_SAFETY_SETTING", "BLOCK_NONE") - -var Theme = GetOrDefaultString("THEME", "default") -var ValidThemes = map[string]bool{ - "default": true, - "berry": true, -} - -const ( - RequestIdKey = "X-Oneapi-Request-Id" -) const ( RoleGuestUser = 0 @@ -117,34 +12,6 @@ const ( RoleRootUser = 100 ) -var ( - FileUploadPermission = RoleGuestUser - FileDownloadPermission = RoleGuestUser - ImageUploadPermission = RoleGuestUser - ImageDownloadPermission = RoleGuestUser -) - -// All duration's unit is seconds -// Shouldn't larger then RateLimitKeyExpirationDuration -var ( - GlobalApiRateLimitNum = GetOrDefault("GLOBAL_API_RATE_LIMIT", 180) - GlobalApiRateLimitDuration int64 = 3 * 60 - - GlobalWebRateLimitNum = GetOrDefault("GLOBAL_WEB_RATE_LIMIT", 60) - GlobalWebRateLimitDuration int64 = 3 * 60 - - UploadRateLimitNum = 10 - UploadRateLimitDuration int64 = 60 - - DownloadRateLimitNum = 10 - DownloadRateLimitDuration int64 = 60 - - CriticalRateLimitNum = 20 - CriticalRateLimitDuration int64 = 20 * 60 -) - -var RateLimitKeyExpirationDuration = 20 * time.Minute - const ( UserStatusEnabled = 1 // don't use 0, 0 is the default value! UserStatusDisabled = 2 // also don't use 0 @@ -199,29 +66,29 @@ const ( ) var ChannelBaseURLs = []string{ - "", // 0 - "https://api.openai.com", // 1 - "https://oa.api2d.net", // 2 - "", // 3 - "https://api.closeai-proxy.xyz", // 4 - "https://api.openai-sb.com", // 5 - "https://api.openaimax.com", // 6 - "https://api.ohmygpt.com", // 7 - "", // 8 - "https://api.caipacity.com", // 9 - "https://api.aiproxy.io", // 10 - "", // 11 - "https://api.api2gpt.com", // 12 - "https://api.aigc2d.com", // 13 - "https://api.anthropic.com", // 14 - "https://aip.baidubce.com", // 15 - "https://open.bigmodel.cn", // 16 - "https://dashscope.aliyuncs.com", // 17 - "", // 18 - "https://ai.360.cn", // 19 - "https://openrouter.ai/api", // 20 - "https://api.aiproxy.io", // 21 - "https://fastgpt.run/api/openapi", // 22 - "https://hunyuan.cloud.tencent.com", //23 - "", //24 + "", // 0 + "https://api.openai.com", // 1 + "https://oa.api2d.net", // 2 + "", // 3 + "https://api.closeai-proxy.xyz", // 4 + "https://api.openai-sb.com", // 5 + "https://api.openaimax.com", // 6 + "https://api.ohmygpt.com", // 7 + "", // 8 + "https://api.caipacity.com", // 9 + "https://api.aiproxy.io", // 10 + "https://generativelanguage.googleapis.com", // 11 + "https://api.api2gpt.com", // 12 + "https://api.aigc2d.com", // 13 + "https://api.anthropic.com", // 14 + "https://aip.baidubce.com", // 15 + "https://open.bigmodel.cn", // 16 + "https://dashscope.aliyuncs.com", // 17 + "", // 18 + "https://ai.360.cn", // 19 + "https://openrouter.ai/api", // 20 + "https://api.aiproxy.io", // 21 + "https://fastgpt.run/api/openapi", // 22 + "https://hunyuan.cloud.tencent.com", // 23 + "https://generativelanguage.googleapis.com", // 24 } diff --git a/common/database.go b/common/database.go index 76f2cd55..4c2c2717 100644 --- a/common/database.go +++ b/common/database.go @@ -1,7 +1,9 @@ package common +import "one-api/common/helper" + var UsingSQLite = false var UsingPostgreSQL = false var SQLitePath = "one-api.db" -var SQLiteBusyTimeout = GetOrDefault("SQLITE_BUSY_TIMEOUT", 3000) +var SQLiteBusyTimeout = helper.GetOrDefaultEnvInt("SQLITE_BUSY_TIMEOUT", 3000) diff --git a/common/email.go b/common/email.go index b915f0f9..28719742 100644 --- a/common/email.go +++ b/common/email.go @@ -6,18 +6,19 @@ import ( "encoding/base64" "fmt" "net/smtp" + "one-api/common/config" "strings" "time" ) func SendEmail(subject string, receiver string, content string) error { - if SMTPFrom == "" { // for compatibility - SMTPFrom = SMTPAccount + if config.SMTPFrom == "" { // for compatibility + config.SMTPFrom = config.SMTPAccount } encodedSubject := fmt.Sprintf("=?UTF-8?B?%s?=", base64.StdEncoding.EncodeToString([]byte(subject))) // Extract domain from SMTPFrom - parts := strings.Split(SMTPFrom, "@") + parts := strings.Split(config.SMTPFrom, "@") var domain string if len(parts) > 1 { domain = parts[1] @@ -36,21 +37,21 @@ func SendEmail(subject string, receiver string, content string) error { "Message-ID: %s\r\n"+ // add Message-ID header to avoid being treated as spam, RFC 5322 "Date: %s\r\n"+ "Content-Type: text/html; charset=UTF-8\r\n\r\n%s\r\n", - receiver, SystemName, SMTPFrom, encodedSubject, messageId, time.Now().Format(time.RFC1123Z), content)) - auth := smtp.PlainAuth("", SMTPAccount, SMTPToken, SMTPServer) - addr := fmt.Sprintf("%s:%d", SMTPServer, SMTPPort) + receiver, config.SystemName, config.SMTPFrom, encodedSubject, messageId, time.Now().Format(time.RFC1123Z), content)) + auth := smtp.PlainAuth("", config.SMTPAccount, config.SMTPToken, config.SMTPServer) + addr := fmt.Sprintf("%s:%d", config.SMTPServer, config.SMTPPort) to := strings.Split(receiver, ";") - if SMTPPort == 465 { + if config.SMTPPort == 465 { tlsConfig := &tls.Config{ InsecureSkipVerify: true, - ServerName: SMTPServer, + ServerName: config.SMTPServer, } - conn, err := tls.Dial("tcp", fmt.Sprintf("%s:%d", SMTPServer, SMTPPort), tlsConfig) + conn, err := tls.Dial("tcp", fmt.Sprintf("%s:%d", config.SMTPServer, config.SMTPPort), tlsConfig) if err != nil { return err } - client, err := smtp.NewClient(conn, SMTPServer) + client, err := smtp.NewClient(conn, config.SMTPServer) if err != nil { return err } @@ -58,7 +59,7 @@ func SendEmail(subject string, receiver string, content string) error { if err = client.Auth(auth); err != nil { return err } - if err = client.Mail(SMTPFrom); err != nil { + if err = client.Mail(config.SMTPFrom); err != nil { return err } receiverEmails := strings.Split(receiver, ";") @@ -80,7 +81,7 @@ func SendEmail(subject string, receiver string, content string) error { return err } } else { - err = smtp.SendMail(addr, auth, SMTPAccount, to, mail) + err = smtp.SendMail(addr, auth, config.SMTPAccount, to, mail) } return err } diff --git a/common/group-ratio.go b/common/group-ratio.go index 1ec73c78..86e9e1f6 100644 --- a/common/group-ratio.go +++ b/common/group-ratio.go @@ -1,6 +1,9 @@ package common -import "encoding/json" +import ( + "encoding/json" + "one-api/common/logger" +) var GroupRatio = map[string]float64{ "default": 1, @@ -11,7 +14,7 @@ var GroupRatio = map[string]float64{ func GroupRatio2JSONString() string { jsonBytes, err := json.Marshal(GroupRatio) if err != nil { - SysError("error marshalling model ratio: " + err.Error()) + logger.SysError("error marshalling model ratio: " + err.Error()) } return string(jsonBytes) } @@ -24,7 +27,7 @@ func UpdateGroupRatioByJSONString(jsonStr string) error { func GetGroupRatio(name string) float64 { ratio, ok := GroupRatio[name] if !ok { - SysError("group ratio not found: " + name) + logger.SysError("group ratio not found: " + name) return 1 } return ratio diff --git a/common/helper/helper.go b/common/helper/helper.go new file mode 100644 index 00000000..09f1df29 --- /dev/null +++ b/common/helper/helper.go @@ -0,0 +1,224 @@ +package helper + +import ( + "fmt" + "github.com/google/uuid" + "html/template" + "log" + "math/rand" + "net" + "one-api/common/logger" + "os" + "os/exec" + "runtime" + "strconv" + "strings" + "time" +) + +func OpenBrowser(url string) { + var err error + + switch runtime.GOOS { + case "linux": + err = exec.Command("xdg-open", url).Start() + case "windows": + err = exec.Command("rundll32", "url.dll,FileProtocolHandler", url).Start() + case "darwin": + err = exec.Command("open", url).Start() + } + if err != nil { + log.Println(err) + } +} + +func GetIp() (ip string) { + ips, err := net.InterfaceAddrs() + if err != nil { + log.Println(err) + return ip + } + + for _, a := range ips { + if ipNet, ok := a.(*net.IPNet); ok && !ipNet.IP.IsLoopback() { + if ipNet.IP.To4() != nil { + ip = ipNet.IP.String() + if strings.HasPrefix(ip, "10") { + return + } + if strings.HasPrefix(ip, "172") { + return + } + if strings.HasPrefix(ip, "192.168") { + return + } + ip = "" + } + } + } + return +} + +var sizeKB = 1024 +var sizeMB = sizeKB * 1024 +var sizeGB = sizeMB * 1024 + +func Bytes2Size(num int64) string { + numStr := "" + unit := "B" + if num/int64(sizeGB) > 1 { + numStr = fmt.Sprintf("%.2f", float64(num)/float64(sizeGB)) + unit = "GB" + } else if num/int64(sizeMB) > 1 { + numStr = fmt.Sprintf("%d", int(float64(num)/float64(sizeMB))) + unit = "MB" + } else if num/int64(sizeKB) > 1 { + numStr = fmt.Sprintf("%d", int(float64(num)/float64(sizeKB))) + unit = "KB" + } else { + numStr = fmt.Sprintf("%d", num) + } + return numStr + " " + unit +} + +func Seconds2Time(num int) (time string) { + if num/31104000 > 0 { + time += strconv.Itoa(num/31104000) + " 年 " + num %= 31104000 + } + if num/2592000 > 0 { + time += strconv.Itoa(num/2592000) + " 个月 " + num %= 2592000 + } + if num/86400 > 0 { + time += strconv.Itoa(num/86400) + " 天 " + num %= 86400 + } + if num/3600 > 0 { + time += strconv.Itoa(num/3600) + " 小时 " + num %= 3600 + } + if num/60 > 0 { + time += strconv.Itoa(num/60) + " 分钟 " + num %= 60 + } + time += strconv.Itoa(num) + " 秒" + return +} + +func Interface2String(inter interface{}) string { + switch inter.(type) { + case string: + return inter.(string) + case int: + return fmt.Sprintf("%d", inter.(int)) + case float64: + return fmt.Sprintf("%f", inter.(float64)) + } + return "Not Implemented" +} + +func UnescapeHTML(x string) interface{} { + return template.HTML(x) +} + +func IntMax(a int, b int) int { + if a >= b { + return a + } else { + return b + } +} + +func GetUUID() string { + code := uuid.New().String() + code = strings.Replace(code, "-", "", -1) + return code +} + +const keyChars = "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ" + +func init() { + rand.Seed(time.Now().UnixNano()) +} + +func GenerateKey() string { + rand.Seed(time.Now().UnixNano()) + key := make([]byte, 48) + for i := 0; i < 16; i++ { + key[i] = keyChars[rand.Intn(len(keyChars))] + } + uuid_ := GetUUID() + for i := 0; i < 32; i++ { + c := uuid_[i] + if i%2 == 0 && c >= 'a' && c <= 'z' { + c = c - 'a' + 'A' + } + key[i+16] = c + } + return string(key) +} + +func GetRandomString(length int) string { + rand.Seed(time.Now().UnixNano()) + key := make([]byte, length) + for i := 0; i < length; i++ { + key[i] = keyChars[rand.Intn(len(keyChars))] + } + return string(key) +} + +func GetTimestamp() int64 { + return time.Now().Unix() +} + +func GetTimeString() string { + now := time.Now() + return fmt.Sprintf("%s%d", now.Format("20060102150405"), now.UnixNano()%1e9) +} + +func Max(a int, b int) int { + if a >= b { + return a + } else { + return b + } +} + +func GetOrDefaultEnvInt(env string, defaultValue int) int { + if env == "" || os.Getenv(env) == "" { + return defaultValue + } + num, err := strconv.Atoi(os.Getenv(env)) + if err != nil { + logger.SysError(fmt.Sprintf("failed to parse %s: %s, using default value: %d", env, err.Error(), defaultValue)) + return defaultValue + } + return num +} + +func GetOrDefaultEnvString(env string, defaultValue string) string { + if env == "" || os.Getenv(env) == "" { + return defaultValue + } + return os.Getenv(env) +} + +func AssignOrDefault(value string, defaultValue string) string { + if len(value) != 0 { + return value + } + return defaultValue +} + +func MessageWithRequestId(message string, id string) string { + return fmt.Sprintf("%s (request id: %s)", message, id) +} + +func String2Int(str string) int { + num, err := strconv.Atoi(str) + if err != nil { + return 0 + } + return num +} diff --git a/common/init.go b/common/init.go index 12df5f51..abc71108 100644 --- a/common/init.go +++ b/common/init.go @@ -4,6 +4,8 @@ import ( "flag" "fmt" "log" + "one-api/common/config" + "one-api/common/logger" "os" "path/filepath" ) @@ -37,9 +39,9 @@ func init() { if os.Getenv("SESSION_SECRET") != "" { if os.Getenv("SESSION_SECRET") == "random_string" { - SysError("SESSION_SECRET is set to an example value, please change it to a random string.") + logger.SysError("SESSION_SECRET is set to an example value, please change it to a random string.") } else { - SessionSecret = os.Getenv("SESSION_SECRET") + config.SessionSecret = os.Getenv("SESSION_SECRET") } } if os.Getenv("SQLITE_PATH") != "" { @@ -57,5 +59,6 @@ func init() { log.Fatal(err) } } + logger.LogDir = *LogDir } } diff --git a/common/logger/constants.go b/common/logger/constants.go new file mode 100644 index 00000000..78d32062 --- /dev/null +++ b/common/logger/constants.go @@ -0,0 +1,7 @@ +package logger + +const ( + RequestIdKey = "X-Oneapi-Request-Id" +) + +var LogDir string diff --git a/common/logger.go b/common/logger/logger.go similarity index 76% rename from common/logger.go rename to common/logger/logger.go index 61627217..b89dbdb7 100644 --- a/common/logger.go +++ b/common/logger/logger.go @@ -1,4 +1,4 @@ -package common +package logger import ( "context" @@ -25,7 +25,7 @@ var setupLogLock sync.Mutex var setupLogWorking bool func SetupLogger() { - if *LogDir != "" { + if LogDir != "" { ok := setupLogLock.TryLock() if !ok { log.Println("setup log is already working") @@ -35,7 +35,7 @@ func SetupLogger() { setupLogLock.Unlock() setupLogWorking = false }() - logPath := filepath.Join(*LogDir, fmt.Sprintf("oneapi-%s.log", time.Now().Format("20060102"))) + logPath := filepath.Join(LogDir, fmt.Sprintf("oneapi-%s.log", time.Now().Format("20060102"))) fd, err := os.OpenFile(logPath, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644) if err != nil { log.Fatal("failed to open log file") @@ -55,18 +55,30 @@ func SysError(s string) { _, _ = fmt.Fprintf(gin.DefaultErrorWriter, "[SYS] %v | %s \n", t.Format("2006/01/02 - 15:04:05"), s) } -func LogInfo(ctx context.Context, msg string) { +func Info(ctx context.Context, msg string) { logHelper(ctx, loggerINFO, msg) } -func LogWarn(ctx context.Context, msg string) { +func Warn(ctx context.Context, msg string) { logHelper(ctx, loggerWarn, msg) } -func LogError(ctx context.Context, msg string) { +func Error(ctx context.Context, msg string) { logHelper(ctx, loggerError, msg) } +func Infof(ctx context.Context, format string, a ...any) { + Info(ctx, fmt.Sprintf(format, a)) +} + +func Warnf(ctx context.Context, format string, a ...any) { + Warn(ctx, fmt.Sprintf(format, a)) +} + +func Errorf(ctx context.Context, format string, a ...any) { + Error(ctx, fmt.Sprintf(format, a)) +} + func logHelper(ctx context.Context, level string, msg string) { writer := gin.DefaultErrorWriter if level == loggerINFO { @@ -90,11 +102,3 @@ func FatalLog(v ...any) { _, _ = fmt.Fprintf(gin.DefaultErrorWriter, "[FATAL] %v | %v \n", t.Format("2006/01/02 - 15:04:05"), v) os.Exit(1) } - -func LogQuota(quota int) string { - if DisplayInCurrencyEnabled { - return fmt.Sprintf("$%.6f 额度", float64(quota)/QuotaPerUnit) - } else { - return fmt.Sprintf("%d 点额度", quota) - } -} diff --git a/common/model-ratio.go b/common/model-ratio.go index 97cb060d..9f31e0d7 100644 --- a/common/model-ratio.go +++ b/common/model-ratio.go @@ -2,6 +2,7 @@ package common import ( "encoding/json" + "one-api/common/logger" "strings" "time" ) @@ -107,7 +108,7 @@ var ModelRatio = map[string]float64{ func ModelRatio2JSONString() string { jsonBytes, err := json.Marshal(ModelRatio) if err != nil { - SysError("error marshalling model ratio: " + err.Error()) + logger.SysError("error marshalling model ratio: " + err.Error()) } return string(jsonBytes) } @@ -123,7 +124,7 @@ func GetModelRatio(name string) float64 { } ratio, ok := ModelRatio[name] if !ok { - SysError("model ratio not found: " + name) + logger.SysError("model ratio not found: " + name) return 30 } return ratio diff --git a/common/redis.go b/common/redis.go index 12c477b8..ed3fcd9d 100644 --- a/common/redis.go +++ b/common/redis.go @@ -3,6 +3,7 @@ package common import ( "context" "github.com/go-redis/redis/v8" + "one-api/common/logger" "os" "time" ) @@ -14,18 +15,18 @@ var RedisEnabled = true func InitRedisClient() (err error) { if os.Getenv("REDIS_CONN_STRING") == "" { RedisEnabled = false - SysLog("REDIS_CONN_STRING not set, Redis is not enabled") + logger.SysLog("REDIS_CONN_STRING not set, Redis is not enabled") return nil } if os.Getenv("SYNC_FREQUENCY") == "" { RedisEnabled = false - SysLog("SYNC_FREQUENCY not set, Redis is disabled") + logger.SysLog("SYNC_FREQUENCY not set, Redis is disabled") return nil } - SysLog("Redis is enabled") + logger.SysLog("Redis is enabled") opt, err := redis.ParseURL(os.Getenv("REDIS_CONN_STRING")) if err != nil { - FatalLog("failed to parse Redis connection string: " + err.Error()) + logger.FatalLog("failed to parse Redis connection string: " + err.Error()) } RDB = redis.NewClient(opt) @@ -34,7 +35,7 @@ func InitRedisClient() (err error) { _, err = RDB.Ping(ctx).Result() if err != nil { - FatalLog("Redis ping test failed: " + err.Error()) + logger.FatalLog("Redis ping test failed: " + err.Error()) } return err } @@ -42,7 +43,7 @@ func InitRedisClient() (err error) { 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()) + logger.FatalLog("failed to parse Redis connection string: " + err.Error()) } return opt } diff --git a/common/utils.go b/common/utils.go index 9a7038e2..41de9367 100644 --- a/common/utils.go +++ b/common/utils.go @@ -2,215 +2,13 @@ package common import ( "fmt" - "github.com/google/uuid" - "html/template" - "log" - "math/rand" - "net" - "os" - "os/exec" - "runtime" - "strconv" - "strings" - "time" + "one-api/common/config" ) -func OpenBrowser(url string) { - var err error - - switch runtime.GOOS { - case "linux": - err = exec.Command("xdg-open", url).Start() - case "windows": - err = exec.Command("rundll32", "url.dll,FileProtocolHandler", url).Start() - case "darwin": - err = exec.Command("open", url).Start() - } - if err != nil { - log.Println(err) - } -} - -func GetIp() (ip string) { - ips, err := net.InterfaceAddrs() - if err != nil { - log.Println(err) - return ip - } - - for _, a := range ips { - if ipNet, ok := a.(*net.IPNet); ok && !ipNet.IP.IsLoopback() { - if ipNet.IP.To4() != nil { - ip = ipNet.IP.String() - if strings.HasPrefix(ip, "10") { - return - } - if strings.HasPrefix(ip, "172") { - return - } - if strings.HasPrefix(ip, "192.168") { - return - } - ip = "" - } - } - } - return -} - -var sizeKB = 1024 -var sizeMB = sizeKB * 1024 -var sizeGB = sizeMB * 1024 - -func Bytes2Size(num int64) string { - numStr := "" - unit := "B" - if num/int64(sizeGB) > 1 { - numStr = fmt.Sprintf("%.2f", float64(num)/float64(sizeGB)) - unit = "GB" - } else if num/int64(sizeMB) > 1 { - numStr = fmt.Sprintf("%d", int(float64(num)/float64(sizeMB))) - unit = "MB" - } else if num/int64(sizeKB) > 1 { - numStr = fmt.Sprintf("%d", int(float64(num)/float64(sizeKB))) - unit = "KB" +func LogQuota(quota int) string { + if config.DisplayInCurrencyEnabled { + return fmt.Sprintf("$%.6f 额度", float64(quota)/config.QuotaPerUnit) } else { - numStr = fmt.Sprintf("%d", num) - } - return numStr + " " + unit -} - -func Seconds2Time(num int) (time string) { - if num/31104000 > 0 { - time += strconv.Itoa(num/31104000) + " 年 " - num %= 31104000 - } - if num/2592000 > 0 { - time += strconv.Itoa(num/2592000) + " 个月 " - num %= 2592000 - } - if num/86400 > 0 { - time += strconv.Itoa(num/86400) + " 天 " - num %= 86400 - } - if num/3600 > 0 { - time += strconv.Itoa(num/3600) + " 小时 " - num %= 3600 - } - if num/60 > 0 { - time += strconv.Itoa(num/60) + " 分钟 " - num %= 60 - } - time += strconv.Itoa(num) + " 秒" - return -} - -func Interface2String(inter interface{}) string { - switch inter.(type) { - case string: - return inter.(string) - case int: - return fmt.Sprintf("%d", inter.(int)) - case float64: - return fmt.Sprintf("%f", inter.(float64)) - } - return "Not Implemented" -} - -func UnescapeHTML(x string) interface{} { - return template.HTML(x) -} - -func IntMax(a int, b int) int { - if a >= b { - return a - } else { - return b + return fmt.Sprintf("%d 点额度", quota) } } - -func GetUUID() string { - code := uuid.New().String() - code = strings.Replace(code, "-", "", -1) - return code -} - -const keyChars = "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ" - -func init() { - rand.Seed(time.Now().UnixNano()) -} - -func GenerateKey() string { - rand.Seed(time.Now().UnixNano()) - key := make([]byte, 48) - for i := 0; i < 16; i++ { - key[i] = keyChars[rand.Intn(len(keyChars))] - } - uuid_ := GetUUID() - for i := 0; i < 32; i++ { - c := uuid_[i] - if i%2 == 0 && c >= 'a' && c <= 'z' { - c = c - 'a' + 'A' - } - key[i+16] = c - } - return string(key) -} - -func GetRandomString(length int) string { - rand.Seed(time.Now().UnixNano()) - key := make([]byte, length) - for i := 0; i < length; i++ { - key[i] = keyChars[rand.Intn(len(keyChars))] - } - return string(key) -} - -func GetTimestamp() int64 { - return time.Now().Unix() -} - -func GetTimeString() string { - now := time.Now() - return fmt.Sprintf("%s%d", now.Format("20060102150405"), now.UnixNano()%1e9) -} - -func Max(a int, b int) int { - if a >= b { - return a - } else { - return b - } -} - -func GetOrDefault(env string, defaultValue int) int { - if env == "" || os.Getenv(env) == "" { - return defaultValue - } - num, err := strconv.Atoi(os.Getenv(env)) - if err != nil { - SysError(fmt.Sprintf("failed to parse %s: %s, using default value: %d", env, err.Error(), defaultValue)) - return defaultValue - } - return num -} - -func GetOrDefaultString(env string, defaultValue string) string { - if env == "" || os.Getenv(env) == "" { - return defaultValue - } - return os.Getenv(env) -} - -func MessageWithRequestId(message string, id string) string { - return fmt.Sprintf("%s (request id: %s)", message, id) -} - -func String2Int(str string) int { - num, err := strconv.Atoi(str) - if err != nil { - return 0 - } - return num -} diff --git a/controller/billing.go b/controller/billing.go index e27fd614..1003bf10 100644 --- a/controller/billing.go +++ b/controller/billing.go @@ -2,7 +2,7 @@ package controller import ( "github.com/gin-gonic/gin" - "one-api/common" + "one-api/common/config" "one-api/model" "one-api/relay/channel/openai" ) @@ -13,7 +13,7 @@ func GetSubscription(c *gin.Context) { var err error var token *model.Token var expiredTime int64 - if common.DisplayTokenStatEnabled { + if config.DisplayTokenStatEnabled { tokenId := c.GetInt("token_id") token, err = model.GetTokenById(tokenId) expiredTime = token.ExpiredTime @@ -39,8 +39,8 @@ func GetSubscription(c *gin.Context) { } quota := remainQuota + usedQuota amount := float64(quota) - if common.DisplayInCurrencyEnabled { - amount /= common.QuotaPerUnit + if config.DisplayInCurrencyEnabled { + amount /= config.QuotaPerUnit } if token != nil && token.UnlimitedQuota { amount = 100000000 @@ -61,7 +61,7 @@ func GetUsage(c *gin.Context) { var quota int var err error var token *model.Token - if common.DisplayTokenStatEnabled { + if config.DisplayTokenStatEnabled { tokenId := c.GetInt("token_id") token, err = model.GetTokenById(tokenId) quota = token.UsedQuota @@ -80,8 +80,8 @@ func GetUsage(c *gin.Context) { return } amount := float64(quota) - if common.DisplayInCurrencyEnabled { - amount /= common.QuotaPerUnit + if config.DisplayInCurrencyEnabled { + amount /= config.QuotaPerUnit } usage := OpenAIUsageResponse{ Object: "list", diff --git a/controller/channel-billing.go b/controller/channel-billing.go index 29346cde..49c76760 100644 --- a/controller/channel-billing.go +++ b/controller/channel-billing.go @@ -7,6 +7,8 @@ import ( "io" "net/http" "one-api/common" + "one-api/common/config" + "one-api/common/logger" "one-api/model" "one-api/relay/util" "strconv" @@ -314,7 +316,7 @@ func updateAllChannelsBalance() error { disableChannel(channel.Id, channel.Name, "余额不足") } } - time.Sleep(common.RequestInterval) + time.Sleep(config.RequestInterval) } return nil } @@ -339,8 +341,8 @@ func UpdateAllChannelsBalance(c *gin.Context) { func AutomaticallyUpdateChannels(frequency int) { for { time.Sleep(time.Duration(frequency) * time.Minute) - common.SysLog("updating all channels") + logger.SysLog("updating all channels") _ = updateAllChannelsBalance() - common.SysLog("channels update done") + logger.SysLog("channels update done") } } diff --git a/controller/channel-test.go b/controller/channel-test.go index f64f0ee3..88d6e3f2 100644 --- a/controller/channel-test.go +++ b/controller/channel-test.go @@ -8,6 +8,8 @@ import ( "io" "net/http" "one-api/common" + "one-api/common/config" + "one-api/common/logger" "one-api/model" "one-api/relay/channel/openai" "one-api/relay/util" @@ -150,12 +152,12 @@ var testAllChannelsLock sync.Mutex var testAllChannelsRunning bool = false func notifyRootUser(subject string, content string) { - if common.RootUserEmail == "" { - common.RootUserEmail = model.GetRootUserEmail() + if config.RootUserEmail == "" { + config.RootUserEmail = model.GetRootUserEmail() } - err := common.SendEmail(subject, common.RootUserEmail, content) + err := common.SendEmail(subject, config.RootUserEmail, content) if err != nil { - common.SysError(fmt.Sprintf("failed to send email: %s", err.Error())) + logger.SysError(fmt.Sprintf("failed to send email: %s", err.Error())) } } @@ -176,8 +178,8 @@ func enableChannel(channelId int, channelName string) { } func testAllChannels(notify bool) error { - if common.RootUserEmail == "" { - common.RootUserEmail = model.GetRootUserEmail() + if config.RootUserEmail == "" { + config.RootUserEmail = model.GetRootUserEmail() } testAllChannelsLock.Lock() if testAllChannelsRunning { @@ -191,7 +193,7 @@ func testAllChannels(notify bool) error { return err } testRequest := buildTestRequest() - var disableThreshold = int64(common.ChannelDisableThreshold * 1000) + var disableThreshold = int64(config.ChannelDisableThreshold * 1000) if disableThreshold == 0 { disableThreshold = 10000000 // a impossible value } @@ -213,15 +215,15 @@ func testAllChannels(notify bool) error { enableChannel(channel.Id, channel.Name) } channel.UpdateResponseTime(milliseconds) - time.Sleep(common.RequestInterval) + time.Sleep(config.RequestInterval) } testAllChannelsLock.Lock() testAllChannelsRunning = false testAllChannelsLock.Unlock() if notify { - err := common.SendEmail("通道测试完成", common.RootUserEmail, "通道测试完成,如果没有收到禁用通知,说明所有通道都正常") + err := common.SendEmail("通道测试完成", config.RootUserEmail, "通道测试完成,如果没有收到禁用通知,说明所有通道都正常") if err != nil { - common.SysError(fmt.Sprintf("failed to send email: %s", err.Error())) + logger.SysError(fmt.Sprintf("failed to send email: %s", err.Error())) } } }() @@ -247,8 +249,8 @@ func TestAllChannels(c *gin.Context) { func AutomaticallyTestChannels(frequency int) { for { time.Sleep(time.Duration(frequency) * time.Minute) - common.SysLog("testing all channels") + logger.SysLog("testing all channels") _ = testAllChannels(false) - common.SysLog("channel test finished") + logger.SysLog("channel test finished") } } diff --git a/controller/channel.go b/controller/channel.go index 904abc23..6d368066 100644 --- a/controller/channel.go +++ b/controller/channel.go @@ -3,7 +3,8 @@ package controller import ( "github.com/gin-gonic/gin" "net/http" - "one-api/common" + "one-api/common/config" + "one-api/common/helper" "one-api/model" "strconv" "strings" @@ -14,7 +15,7 @@ func GetAllChannels(c *gin.Context) { if p < 0 { p = 0 } - channels, err := model.GetAllChannels(p*common.ItemsPerPage, common.ItemsPerPage, false) + channels, err := model.GetAllChannels(p*config.ItemsPerPage, config.ItemsPerPage, false) if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, @@ -83,7 +84,7 @@ func AddChannel(c *gin.Context) { }) return } - channel.CreatedTime = common.GetTimestamp() + channel.CreatedTime = helper.GetTimestamp() keys := strings.Split(channel.Key, "\n") channels := make([]model.Channel, 0, len(keys)) for _, key := range keys { diff --git a/controller/github.go b/controller/github.go index ee995379..f6049ddb 100644 --- a/controller/github.go +++ b/controller/github.go @@ -9,6 +9,9 @@ import ( "github.com/gin-gonic/gin" "net/http" "one-api/common" + "one-api/common/config" + "one-api/common/helper" + "one-api/common/logger" "one-api/model" "strconv" "time" @@ -30,7 +33,7 @@ func getGitHubUserInfoByCode(code string) (*GitHubUser, error) { if code == "" { return nil, errors.New("无效的参数") } - values := map[string]string{"client_id": common.GitHubClientId, "client_secret": common.GitHubClientSecret, "code": code} + values := map[string]string{"client_id": config.GitHubClientId, "client_secret": config.GitHubClientSecret, "code": code} jsonData, err := json.Marshal(values) if err != nil { return nil, err @@ -46,7 +49,7 @@ func getGitHubUserInfoByCode(code string) (*GitHubUser, error) { } res, err := client.Do(req) if err != nil { - common.SysLog(err.Error()) + logger.SysLog(err.Error()) return nil, errors.New("无法连接至 GitHub 服务器,请稍后重试!") } defer res.Body.Close() @@ -62,7 +65,7 @@ func getGitHubUserInfoByCode(code string) (*GitHubUser, error) { req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", oAuthResponse.AccessToken)) res2, err := client.Do(req) if err != nil { - common.SysLog(err.Error()) + logger.SysLog(err.Error()) return nil, errors.New("无法连接至 GitHub 服务器,请稍后重试!") } defer res2.Body.Close() @@ -93,7 +96,7 @@ func GitHubOAuth(c *gin.Context) { return } - if !common.GitHubOAuthEnabled { + if !config.GitHubOAuthEnabled { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "管理员未开启通过 GitHub 登录以及注册", @@ -122,7 +125,7 @@ func GitHubOAuth(c *gin.Context) { return } } else { - if common.RegisterEnabled { + if config.RegisterEnabled { user.Username = "github_" + strconv.Itoa(model.GetMaxUserId()+1) if githubUser.Name != "" { user.DisplayName = githubUser.Name @@ -160,7 +163,7 @@ func GitHubOAuth(c *gin.Context) { } func GitHubBind(c *gin.Context) { - if !common.GitHubOAuthEnabled { + if !config.GitHubOAuthEnabled { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "管理员未开启通过 GitHub 登录以及注册", @@ -216,7 +219,7 @@ func GitHubBind(c *gin.Context) { func GenerateOAuthCode(c *gin.Context) { session := sessions.Default(c) - state := common.GetRandomString(12) + state := helper.GetRandomString(12) session.Set("oauth_state", state) err := session.Save() if err != nil { diff --git a/controller/log.go b/controller/log.go index b65867fe..418c6859 100644 --- a/controller/log.go +++ b/controller/log.go @@ -3,7 +3,7 @@ package controller import ( "github.com/gin-gonic/gin" "net/http" - "one-api/common" + "one-api/common/config" "one-api/model" "strconv" ) @@ -20,7 +20,7 @@ func GetAllLogs(c *gin.Context) { tokenName := c.Query("token_name") modelName := c.Query("model_name") channel, _ := strconv.Atoi(c.Query("channel")) - logs, err := model.GetAllLogs(logType, startTimestamp, endTimestamp, modelName, username, tokenName, p*common.ItemsPerPage, common.ItemsPerPage, channel) + logs, err := model.GetAllLogs(logType, startTimestamp, endTimestamp, modelName, username, tokenName, p*config.ItemsPerPage, config.ItemsPerPage, channel) if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, @@ -47,7 +47,7 @@ func GetUserLogs(c *gin.Context) { endTimestamp, _ := strconv.ParseInt(c.Query("end_timestamp"), 10, 64) tokenName := c.Query("token_name") modelName := c.Query("model_name") - logs, err := model.GetUserLogs(userId, logType, startTimestamp, endTimestamp, modelName, tokenName, p*common.ItemsPerPage, common.ItemsPerPage) + logs, err := model.GetUserLogs(userId, logType, startTimestamp, endTimestamp, modelName, tokenName, p*config.ItemsPerPage, config.ItemsPerPage) if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, diff --git a/controller/misc.go b/controller/misc.go index 2bcbb41f..df7c0728 100644 --- a/controller/misc.go +++ b/controller/misc.go @@ -5,6 +5,7 @@ import ( "fmt" "net/http" "one-api/common" + "one-api/common/config" "one-api/model" "strings" @@ -18,55 +19,55 @@ func GetStatus(c *gin.Context) { "data": gin.H{ "version": common.Version, "start_time": common.StartTime, - "email_verification": common.EmailVerificationEnabled, - "github_oauth": common.GitHubOAuthEnabled, - "github_client_id": common.GitHubClientId, - "system_name": common.SystemName, - "logo": common.Logo, - "footer_html": common.Footer, - "wechat_qrcode": common.WeChatAccountQRCodeImageURL, - "wechat_login": common.WeChatAuthEnabled, - "server_address": common.ServerAddress, - "turnstile_check": common.TurnstileCheckEnabled, - "turnstile_site_key": common.TurnstileSiteKey, - "top_up_link": common.TopUpLink, - "chat_link": common.ChatLink, - "quota_per_unit": common.QuotaPerUnit, - "display_in_currency": common.DisplayInCurrencyEnabled, + "email_verification": config.EmailVerificationEnabled, + "github_oauth": config.GitHubOAuthEnabled, + "github_client_id": config.GitHubClientId, + "system_name": config.SystemName, + "logo": config.Logo, + "footer_html": config.Footer, + "wechat_qrcode": config.WeChatAccountQRCodeImageURL, + "wechat_login": config.WeChatAuthEnabled, + "server_address": config.ServerAddress, + "turnstile_check": config.TurnstileCheckEnabled, + "turnstile_site_key": config.TurnstileSiteKey, + "top_up_link": config.TopUpLink, + "chat_link": config.ChatLink, + "quota_per_unit": config.QuotaPerUnit, + "display_in_currency": config.DisplayInCurrencyEnabled, }, }) return } func GetNotice(c *gin.Context) { - common.OptionMapRWMutex.RLock() - defer common.OptionMapRWMutex.RUnlock() + config.OptionMapRWMutex.RLock() + defer config.OptionMapRWMutex.RUnlock() c.JSON(http.StatusOK, gin.H{ "success": true, "message": "", - "data": common.OptionMap["Notice"], + "data": config.OptionMap["Notice"], }) return } func GetAbout(c *gin.Context) { - common.OptionMapRWMutex.RLock() - defer common.OptionMapRWMutex.RUnlock() + config.OptionMapRWMutex.RLock() + defer config.OptionMapRWMutex.RUnlock() c.JSON(http.StatusOK, gin.H{ "success": true, "message": "", - "data": common.OptionMap["About"], + "data": config.OptionMap["About"], }) return } func GetHomePageContent(c *gin.Context) { - common.OptionMapRWMutex.RLock() - defer common.OptionMapRWMutex.RUnlock() + config.OptionMapRWMutex.RLock() + defer config.OptionMapRWMutex.RUnlock() c.JSON(http.StatusOK, gin.H{ "success": true, "message": "", - "data": common.OptionMap["HomePageContent"], + "data": config.OptionMap["HomePageContent"], }) return } @@ -80,9 +81,9 @@ func SendEmailVerification(c *gin.Context) { }) return } - if common.EmailDomainRestrictionEnabled { + if config.EmailDomainRestrictionEnabled { allowed := false - for _, domain := range common.EmailDomainWhitelist { + for _, domain := range config.EmailDomainWhitelist { if strings.HasSuffix(email, "@"+domain) { allowed = true break @@ -105,10 +106,10 @@ func SendEmailVerification(c *gin.Context) { } code := common.GenerateVerificationCode(6) common.RegisterVerificationCodeWithKey(email, code, common.EmailVerificationPurpose) - subject := fmt.Sprintf("%s邮箱验证邮件", common.SystemName) + subject := fmt.Sprintf("%s邮箱验证邮件", config.SystemName) content := fmt.Sprintf("

您好,你正在进行%s邮箱验证。

"+ "

您的验证码为: %s

"+ - "

验证码 %d 分钟内有效,如果不是本人操作,请忽略。

", common.SystemName, code, common.VerificationValidMinutes) + "

验证码 %d 分钟内有效,如果不是本人操作,请忽略。

", config.SystemName, code, common.VerificationValidMinutes) err := common.SendEmail(subject, email, content) if err != nil { c.JSON(http.StatusOK, gin.H{ @@ -142,12 +143,12 @@ func SendPasswordResetEmail(c *gin.Context) { } code := common.GenerateVerificationCode(0) common.RegisterVerificationCodeWithKey(email, code, common.PasswordResetPurpose) - link := fmt.Sprintf("%s/user/reset?email=%s&token=%s", common.ServerAddress, email, code) - subject := fmt.Sprintf("%s密码重置", common.SystemName) + link := fmt.Sprintf("%s/user/reset?email=%s&token=%s", config.ServerAddress, email, code) + subject := fmt.Sprintf("%s密码重置", config.SystemName) content := fmt.Sprintf("

您好,你正在进行%s密码重置。

"+ "

点击 此处 进行密码重置。

"+ "

如果链接无法点击,请尝试点击下面的链接或将其复制到浏览器中打开:
%s

"+ - "

重置链接 %d 分钟内有效,如果不是本人操作,请忽略。

", common.SystemName, link, link, common.VerificationValidMinutes) + "

重置链接 %d 分钟内有效,如果不是本人操作,请忽略。

", config.SystemName, link, link, common.VerificationValidMinutes) err := common.SendEmail(subject, email, content) if err != nil { c.JSON(http.StatusOK, gin.H{ diff --git a/controller/option.go b/controller/option.go index 3b1cbad2..593ac2ae 100644 --- a/controller/option.go +++ b/controller/option.go @@ -3,7 +3,8 @@ package controller import ( "encoding/json" "net/http" - "one-api/common" + "one-api/common/config" + "one-api/common/helper" "one-api/model" "strings" @@ -12,17 +13,17 @@ import ( func GetOptions(c *gin.Context) { var options []*model.Option - common.OptionMapRWMutex.Lock() - for k, v := range common.OptionMap { + config.OptionMapRWMutex.Lock() + for k, v := range config.OptionMap { if strings.HasSuffix(k, "Token") || strings.HasSuffix(k, "Secret") { continue } options = append(options, &model.Option{ Key: k, - Value: common.Interface2String(v), + Value: helper.Interface2String(v), }) } - common.OptionMapRWMutex.Unlock() + config.OptionMapRWMutex.Unlock() c.JSON(http.StatusOK, gin.H{ "success": true, "message": "", @@ -43,7 +44,7 @@ func UpdateOption(c *gin.Context) { } switch option.Key { case "Theme": - if !common.ValidThemes[option.Value] { + if !config.ValidThemes[option.Value] { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "无效的主题", @@ -51,7 +52,7 @@ func UpdateOption(c *gin.Context) { return } case "GitHubOAuthEnabled": - if option.Value == "true" && common.GitHubClientId == "" { + if option.Value == "true" && config.GitHubClientId == "" { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "无法启用 GitHub OAuth,请先填入 GitHub Client Id 以及 GitHub Client Secret!", @@ -59,7 +60,7 @@ func UpdateOption(c *gin.Context) { return } case "EmailDomainRestrictionEnabled": - if option.Value == "true" && len(common.EmailDomainWhitelist) == 0 { + if option.Value == "true" && len(config.EmailDomainWhitelist) == 0 { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "无法启用邮箱域名限制,请先填入限制的邮箱域名!", @@ -67,7 +68,7 @@ func UpdateOption(c *gin.Context) { return } case "WeChatAuthEnabled": - if option.Value == "true" && common.WeChatServerAddress == "" { + if option.Value == "true" && config.WeChatServerAddress == "" { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "无法启用微信登录,请先填入微信登录相关配置信息!", @@ -75,7 +76,7 @@ func UpdateOption(c *gin.Context) { return } case "TurnstileCheckEnabled": - if option.Value == "true" && common.TurnstileSiteKey == "" { + if option.Value == "true" && config.TurnstileSiteKey == "" { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "无法启用 Turnstile 校验,请先填入 Turnstile 校验相关配置信息!", diff --git a/controller/redemption.go b/controller/redemption.go index 0f656be0..9eeeb943 100644 --- a/controller/redemption.go +++ b/controller/redemption.go @@ -3,7 +3,8 @@ package controller import ( "github.com/gin-gonic/gin" "net/http" - "one-api/common" + "one-api/common/config" + "one-api/common/helper" "one-api/model" "strconv" ) @@ -13,7 +14,7 @@ func GetAllRedemptions(c *gin.Context) { if p < 0 { p = 0 } - redemptions, err := model.GetAllRedemptions(p*common.ItemsPerPage, common.ItemsPerPage) + redemptions, err := model.GetAllRedemptions(p*config.ItemsPerPage, config.ItemsPerPage) if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, @@ -105,12 +106,12 @@ func AddRedemption(c *gin.Context) { } var keys []string for i := 0; i < redemption.Count; i++ { - key := common.GetUUID() + key := helper.GetUUID() cleanRedemption := model.Redemption{ UserId: c.GetInt("id"), Name: redemption.Name, Key: key, - CreatedTime: common.GetTimestamp(), + CreatedTime: helper.GetTimestamp(), Quota: redemption.Quota, } err = cleanRedemption.Insert() diff --git a/controller/relay.go b/controller/relay.go index 198d6c9a..46fedc7e 100644 --- a/controller/relay.go +++ b/controller/relay.go @@ -2,43 +2,22 @@ package controller import ( "fmt" + "github.com/gin-gonic/gin" "net/http" - "one-api/common" + "one-api/common/config" + "one-api/common/helper" + "one-api/common/logger" "one-api/relay/channel/openai" "one-api/relay/constant" "one-api/relay/controller" "one-api/relay/util" "strconv" - "strings" - - "github.com/gin-gonic/gin" ) // https://platform.openai.com/docs/api-reference/chat func Relay(c *gin.Context) { - relayMode := constant.RelayModeUnknown - if strings.HasPrefix(c.Request.URL.Path, "/v1/chat/completions") { - relayMode = constant.RelayModeChatCompletions - } else if strings.HasPrefix(c.Request.URL.Path, "/v1/completions") { - relayMode = constant.RelayModeCompletions - } else if strings.HasPrefix(c.Request.URL.Path, "/v1/embeddings") { - relayMode = constant.RelayModeEmbeddings - } else if strings.HasSuffix(c.Request.URL.Path, "embeddings") { - relayMode = constant.RelayModeEmbeddings - } else if strings.HasPrefix(c.Request.URL.Path, "/v1/moderations") { - relayMode = constant.RelayModeModerations - } else if strings.HasPrefix(c.Request.URL.Path, "/v1/images/generations") { - relayMode = constant.RelayModeImagesGenerations - } else if strings.HasPrefix(c.Request.URL.Path, "/v1/edits") { - relayMode = constant.RelayModeEdits - } else if strings.HasPrefix(c.Request.URL.Path, "/v1/audio/speech") { - relayMode = constant.RelayModeAudioSpeech - } else if strings.HasPrefix(c.Request.URL.Path, "/v1/audio/transcriptions") { - relayMode = constant.RelayModeAudioTranscription - } else if strings.HasPrefix(c.Request.URL.Path, "/v1/audio/translations") { - relayMode = constant.RelayModeAudioTranslation - } + relayMode := constant.Path2RelayMode(c.Request.URL.Path) var err *openai.ErrorWithStatusCode switch relayMode { case constant.RelayModeImagesGenerations: @@ -53,11 +32,11 @@ func Relay(c *gin.Context) { err = controller.RelayTextHelper(c, relayMode) } if err != nil { - requestId := c.GetString(common.RequestIdKey) + requestId := c.GetString(logger.RequestIdKey) retryTimesStr := c.Query("retry") retryTimes, _ := strconv.Atoi(retryTimesStr) if retryTimesStr == "" { - retryTimes = common.RetryTimes + retryTimes = config.RetryTimes } if retryTimes > 0 { c.Redirect(http.StatusTemporaryRedirect, fmt.Sprintf("%s?retry=%d", c.Request.URL.Path, retryTimes-1)) @@ -65,13 +44,13 @@ func Relay(c *gin.Context) { if err.StatusCode == http.StatusTooManyRequests { err.Error.Message = "当前分组上游负载已饱和,请稍后再试" } - err.Error.Message = common.MessageWithRequestId(err.Error.Message, requestId) + err.Error.Message = helper.MessageWithRequestId(err.Error.Message, requestId) c.JSON(err.StatusCode, gin.H{ "error": err.Error, }) } channelId := c.GetInt("channel_id") - common.LogError(c.Request.Context(), fmt.Sprintf("relay error (channel #%d): %s", channelId, err.Message)) + logger.Error(c.Request.Context(), fmt.Sprintf("relay error (channel #%d): %s", channelId, err.Message)) // https://platform.openai.com/docs/guides/error-codes/api-errors if util.ShouldDisableChannel(&err.Error, err.StatusCode) { channelId := c.GetInt("channel_id") diff --git a/controller/token.go b/controller/token.go index 8642122c..d6554abe 100644 --- a/controller/token.go +++ b/controller/token.go @@ -4,6 +4,8 @@ import ( "github.com/gin-gonic/gin" "net/http" "one-api/common" + "one-api/common/config" + "one-api/common/helper" "one-api/model" "strconv" ) @@ -14,7 +16,7 @@ func GetAllTokens(c *gin.Context) { if p < 0 { p = 0 } - tokens, err := model.GetAllUserTokens(userId, p*common.ItemsPerPage, common.ItemsPerPage) + tokens, err := model.GetAllUserTokens(userId, p*config.ItemsPerPage, config.ItemsPerPage) if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, @@ -119,9 +121,9 @@ func AddToken(c *gin.Context) { cleanToken := model.Token{ UserId: c.GetInt("id"), Name: token.Name, - Key: common.GenerateKey(), - CreatedTime: common.GetTimestamp(), - AccessedTime: common.GetTimestamp(), + Key: helper.GenerateKey(), + CreatedTime: helper.GetTimestamp(), + AccessedTime: helper.GetTimestamp(), ExpiredTime: token.ExpiredTime, RemainQuota: token.RemainQuota, UnlimitedQuota: token.UnlimitedQuota, @@ -187,7 +189,7 @@ func UpdateToken(c *gin.Context) { return } if token.Status == common.TokenStatusEnabled { - if cleanToken.Status == common.TokenStatusExpired && cleanToken.ExpiredTime <= common.GetTimestamp() && cleanToken.ExpiredTime != -1 { + if cleanToken.Status == common.TokenStatusExpired && cleanToken.ExpiredTime <= helper.GetTimestamp() && cleanToken.ExpiredTime != -1 { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "令牌已过期,无法启用,请先修改令牌过期时间,或者设置为永不过期", diff --git a/controller/user.go b/controller/user.go index 174300ed..d13acddd 100644 --- a/controller/user.go +++ b/controller/user.go @@ -5,6 +5,8 @@ import ( "fmt" "net/http" "one-api/common" + "one-api/common/config" + "one-api/common/helper" "one-api/model" "strconv" "time" @@ -19,7 +21,7 @@ type LoginRequest struct { } func Login(c *gin.Context) { - if !common.PasswordLoginEnabled { + if !config.PasswordLoginEnabled { c.JSON(http.StatusOK, gin.H{ "message": "管理员关闭了密码登录", "success": false, @@ -106,14 +108,14 @@ func Logout(c *gin.Context) { } func Register(c *gin.Context) { - if !common.RegisterEnabled { + if !config.RegisterEnabled { c.JSON(http.StatusOK, gin.H{ "message": "管理员关闭了新用户注册", "success": false, }) return } - if !common.PasswordRegisterEnabled { + if !config.PasswordRegisterEnabled { c.JSON(http.StatusOK, gin.H{ "message": "管理员关闭了通过密码进行注册,请使用第三方账户验证的形式进行注册", "success": false, @@ -136,7 +138,7 @@ func Register(c *gin.Context) { }) return } - if common.EmailVerificationEnabled { + if config.EmailVerificationEnabled { if user.Email == "" || user.VerificationCode == "" { c.JSON(http.StatusOK, gin.H{ "success": false, @@ -160,7 +162,7 @@ func Register(c *gin.Context) { DisplayName: user.Username, InviterId: inviterId, } - if common.EmailVerificationEnabled { + if config.EmailVerificationEnabled { cleanUser.Email = user.Email } if err := cleanUser.Insert(inviterId); err != nil { @@ -182,7 +184,7 @@ func GetAllUsers(c *gin.Context) { if p < 0 { p = 0 } - users, err := model.GetAllUsers(p*common.ItemsPerPage, common.ItemsPerPage) + users, err := model.GetAllUsers(p*config.ItemsPerPage, config.ItemsPerPage) if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, @@ -282,7 +284,7 @@ func GenerateAccessToken(c *gin.Context) { }) return } - user.AccessToken = common.GetUUID() + user.AccessToken = helper.GetUUID() if model.DB.Where("access_token = ?", user.AccessToken).First(user).RowsAffected != 0 { c.JSON(http.StatusOK, gin.H{ @@ -319,7 +321,7 @@ func GetAffCode(c *gin.Context) { return } if user.AffCode == "" { - user.AffCode = common.GetRandomString(4) + user.AffCode = helper.GetRandomString(4) if err := user.Update(false); err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, @@ -726,7 +728,7 @@ func EmailBind(c *gin.Context) { return } if user.Role == common.RoleRootUser { - common.RootUserEmail = email + config.RootUserEmail = email } c.JSON(http.StatusOK, gin.H{ "success": true, diff --git a/controller/wechat.go b/controller/wechat.go index ff4c9fb6..fc231a24 100644 --- a/controller/wechat.go +++ b/controller/wechat.go @@ -7,6 +7,7 @@ import ( "github.com/gin-gonic/gin" "net/http" "one-api/common" + "one-api/common/config" "one-api/model" "strconv" "time" @@ -22,11 +23,11 @@ func getWeChatIdByCode(code string) (string, error) { if code == "" { return "", errors.New("无效的参数") } - req, err := http.NewRequest("GET", fmt.Sprintf("%s/api/wechat/user?code=%s", common.WeChatServerAddress, code), nil) + req, err := http.NewRequest("GET", fmt.Sprintf("%s/api/wechat/user?code=%s", config.WeChatServerAddress, code), nil) if err != nil { return "", err } - req.Header.Set("Authorization", common.WeChatServerToken) + req.Header.Set("Authorization", config.WeChatServerToken) client := http.Client{ Timeout: 5 * time.Second, } @@ -50,7 +51,7 @@ func getWeChatIdByCode(code string) (string, error) { } func WeChatAuth(c *gin.Context) { - if !common.WeChatAuthEnabled { + if !config.WeChatAuthEnabled { c.JSON(http.StatusOK, gin.H{ "message": "管理员未开启通过微信登录以及注册", "success": false, @@ -79,7 +80,7 @@ func WeChatAuth(c *gin.Context) { return } } else { - if common.RegisterEnabled { + if config.RegisterEnabled { user.Username = "wechat_" + strconv.Itoa(model.GetMaxUserId()+1) user.DisplayName = "WeChat User" user.Role = common.RoleCommonUser @@ -112,7 +113,7 @@ func WeChatAuth(c *gin.Context) { } func WeChatBind(c *gin.Context) { - if !common.WeChatAuthEnabled { + if !config.WeChatAuthEnabled { c.JSON(http.StatusOK, gin.H{ "message": "管理员未开启通过微信登录以及注册", "success": false, diff --git a/main.go b/main.go index 28a41287..b79c7bf7 100644 --- a/main.go +++ b/main.go @@ -7,6 +7,8 @@ import ( "github.com/gin-contrib/sessions/cookie" "github.com/gin-gonic/gin" "one-api/common" + "one-api/common/config" + "one-api/common/logger" "one-api/controller" "one-api/middleware" "one-api/model" @@ -20,65 +22,65 @@ import ( var buildFS embed.FS func main() { - common.SetupLogger() - common.SysLog(fmt.Sprintf("One API %s started", common.Version)) + logger.SetupLogger() + logger.SysLog(fmt.Sprintf("One API %s started", common.Version)) if os.Getenv("GIN_MODE") != "debug" { gin.SetMode(gin.ReleaseMode) } - if common.DebugEnabled { - common.SysLog("running in debug mode") + if config.DebugEnabled { + logger.SysLog("running in debug mode") } // Initialize SQL Database err := model.InitDB() if err != nil { - common.FatalLog("failed to initialize database: " + err.Error()) + logger.FatalLog("failed to initialize database: " + err.Error()) } defer func() { err := model.CloseDB() if err != nil { - common.FatalLog("failed to close database: " + err.Error()) + logger.FatalLog("failed to close database: " + err.Error()) } }() // Initialize Redis err = common.InitRedisClient() if err != nil { - common.FatalLog("failed to initialize Redis: " + err.Error()) + logger.FatalLog("failed to initialize Redis: " + err.Error()) } // Initialize options model.InitOptionMap() - common.SysLog(fmt.Sprintf("using theme %s", common.Theme)) + logger.SysLog(fmt.Sprintf("using theme %s", config.Theme)) if common.RedisEnabled { // for compatibility with old versions - common.MemoryCacheEnabled = true + config.MemoryCacheEnabled = true } - if common.MemoryCacheEnabled { - common.SysLog("memory cache enabled") - common.SysError(fmt.Sprintf("sync frequency: %d seconds", common.SyncFrequency)) + if config.MemoryCacheEnabled { + logger.SysLog("memory cache enabled") + logger.SysError(fmt.Sprintf("sync frequency: %d seconds", config.SyncFrequency)) model.InitChannelCache() } - if common.MemoryCacheEnabled { - go model.SyncOptions(common.SyncFrequency) - go model.SyncChannelCache(common.SyncFrequency) + if config.MemoryCacheEnabled { + go model.SyncOptions(config.SyncFrequency) + go model.SyncChannelCache(config.SyncFrequency) } if os.Getenv("CHANNEL_UPDATE_FREQUENCY") != "" { frequency, err := strconv.Atoi(os.Getenv("CHANNEL_UPDATE_FREQUENCY")) if err != nil { - common.FatalLog("failed to parse CHANNEL_UPDATE_FREQUENCY: " + err.Error()) + logger.FatalLog("failed to parse CHANNEL_UPDATE_FREQUENCY: " + err.Error()) } go controller.AutomaticallyUpdateChannels(frequency) } if os.Getenv("CHANNEL_TEST_FREQUENCY") != "" { frequency, err := strconv.Atoi(os.Getenv("CHANNEL_TEST_FREQUENCY")) if err != nil { - common.FatalLog("failed to parse CHANNEL_TEST_FREQUENCY: " + err.Error()) + logger.FatalLog("failed to parse CHANNEL_TEST_FREQUENCY: " + err.Error()) } go controller.AutomaticallyTestChannels(frequency) } if os.Getenv("BATCH_UPDATE_ENABLED") == "true" { - common.BatchUpdateEnabled = true - common.SysLog("batch update enabled with interval " + strconv.Itoa(common.BatchUpdateInterval) + "s") + config.BatchUpdateEnabled = true + logger.SysLog("batch update enabled with interval " + strconv.Itoa(config.BatchUpdateInterval) + "s") model.InitBatchUpdater() } openai.InitTokenEncoders() @@ -91,7 +93,7 @@ func main() { server.Use(middleware.RequestId()) middleware.SetUpLogger(server) // Initialize session store - store := cookie.NewStore([]byte(common.SessionSecret)) + store := cookie.NewStore([]byte(config.SessionSecret)) server.Use(sessions.Sessions("session", store)) router.SetRouter(server, buildFS) @@ -101,6 +103,6 @@ func main() { } err = server.Run(":" + port) if err != nil { - common.FatalLog("failed to start HTTP server: " + err.Error()) + logger.FatalLog("failed to start HTTP server: " + err.Error()) } } diff --git a/middleware/distributor.go b/middleware/distributor.go index 81338130..6b607d68 100644 --- a/middleware/distributor.go +++ b/middleware/distributor.go @@ -4,6 +4,7 @@ import ( "fmt" "net/http" "one-api/common" + "one-api/common/logger" "one-api/model" "strconv" "strings" @@ -69,7 +70,7 @@ func Distribute() func(c *gin.Context) { if err != nil { message := fmt.Sprintf("当前分组 %s 下对于模型 %s 无可用渠道", userGroup, modelRequest.Model) if channel != nil { - common.SysError(fmt.Sprintf("渠道不存在:%d", channel.Id)) + logger.SysError(fmt.Sprintf("渠道不存在:%d", channel.Id)) message = "数据库一致性已被破坏,请联系管理员" } abortWithMessage(c, http.StatusServiceUnavailable, message) diff --git a/middleware/logger.go b/middleware/logger.go index 02f2e0a9..ca372c52 100644 --- a/middleware/logger.go +++ b/middleware/logger.go @@ -3,14 +3,14 @@ package middleware import ( "fmt" "github.com/gin-gonic/gin" - "one-api/common" + "one-api/common/logger" ) func SetUpLogger(server *gin.Engine) { server.Use(gin.LoggerWithFormatter(func(param gin.LogFormatterParams) string { var requestID string if param.Keys != nil { - requestID = param.Keys[common.RequestIdKey].(string) + requestID = param.Keys[logger.RequestIdKey].(string) } return fmt.Sprintf("[GIN] %s | %s | %3d | %13v | %15s | %7s %s\n", param.TimeStamp.Format("2006/01/02 - 15:04:05"), diff --git a/middleware/rate-limit.go b/middleware/rate-limit.go index 8e5cff6c..e89ef8d6 100644 --- a/middleware/rate-limit.go +++ b/middleware/rate-limit.go @@ -6,6 +6,7 @@ import ( "github.com/gin-gonic/gin" "net/http" "one-api/common" + "one-api/common/config" "time" ) @@ -26,7 +27,7 @@ func redisRateLimiter(c *gin.Context, maxRequestNum int, duration int64, mark st } if listLength < int64(maxRequestNum) { rdb.LPush(ctx, key, time.Now().Format(timeFormat)) - rdb.Expire(ctx, key, common.RateLimitKeyExpirationDuration) + rdb.Expire(ctx, key, config.RateLimitKeyExpirationDuration) } else { oldTimeStr, _ := rdb.LIndex(ctx, key, -1).Result() oldTime, err := time.Parse(timeFormat, oldTimeStr) @@ -47,14 +48,14 @@ func redisRateLimiter(c *gin.Context, maxRequestNum int, duration int64, mark st // time.Since will return negative number! // See: https://stackoverflow.com/questions/50970900/why-is-time-since-returning-negative-durations-on-windows if int64(nowTime.Sub(oldTime).Seconds()) < duration { - rdb.Expire(ctx, key, common.RateLimitKeyExpirationDuration) + rdb.Expire(ctx, key, config.RateLimitKeyExpirationDuration) c.Status(http.StatusTooManyRequests) c.Abort() return } else { rdb.LPush(ctx, key, time.Now().Format(timeFormat)) rdb.LTrim(ctx, key, 0, int64(maxRequestNum-1)) - rdb.Expire(ctx, key, common.RateLimitKeyExpirationDuration) + rdb.Expire(ctx, key, config.RateLimitKeyExpirationDuration) } } } @@ -75,7 +76,7 @@ func rateLimitFactory(maxRequestNum int, duration int64, mark string) func(c *gi } } else { // It's safe to call multi times. - inMemoryRateLimiter.Init(common.RateLimitKeyExpirationDuration) + inMemoryRateLimiter.Init(config.RateLimitKeyExpirationDuration) return func(c *gin.Context) { memoryRateLimiter(c, maxRequestNum, duration, mark) } @@ -83,21 +84,21 @@ func rateLimitFactory(maxRequestNum int, duration int64, mark string) func(c *gi } func GlobalWebRateLimit() func(c *gin.Context) { - return rateLimitFactory(common.GlobalWebRateLimitNum, common.GlobalWebRateLimitDuration, "GW") + return rateLimitFactory(config.GlobalWebRateLimitNum, config.GlobalWebRateLimitDuration, "GW") } func GlobalAPIRateLimit() func(c *gin.Context) { - return rateLimitFactory(common.GlobalApiRateLimitNum, common.GlobalApiRateLimitDuration, "GA") + return rateLimitFactory(config.GlobalApiRateLimitNum, config.GlobalApiRateLimitDuration, "GA") } func CriticalRateLimit() func(c *gin.Context) { - return rateLimitFactory(common.CriticalRateLimitNum, common.CriticalRateLimitDuration, "CT") + return rateLimitFactory(config.CriticalRateLimitNum, config.CriticalRateLimitDuration, "CT") } func DownloadRateLimit() func(c *gin.Context) { - return rateLimitFactory(common.DownloadRateLimitNum, common.DownloadRateLimitDuration, "DW") + return rateLimitFactory(config.DownloadRateLimitNum, config.DownloadRateLimitDuration, "DW") } func UploadRateLimit() func(c *gin.Context) { - return rateLimitFactory(common.UploadRateLimitNum, common.UploadRateLimitDuration, "UP") + return rateLimitFactory(config.UploadRateLimitNum, config.UploadRateLimitDuration, "UP") } diff --git a/middleware/recover.go b/middleware/recover.go index 8338a514..9d3edc27 100644 --- a/middleware/recover.go +++ b/middleware/recover.go @@ -4,7 +4,7 @@ import ( "fmt" "github.com/gin-gonic/gin" "net/http" - "one-api/common" + "one-api/common/logger" "runtime/debug" ) @@ -12,8 +12,8 @@ func RelayPanicRecover() gin.HandlerFunc { return func(c *gin.Context) { defer func() { if err := recover(); err != nil { - common.SysError(fmt.Sprintf("panic detected: %v", err)) - common.SysError(fmt.Sprintf("stacktrace from panic: %s", string(debug.Stack()))) + logger.SysError(fmt.Sprintf("panic detected: %v", err)) + logger.SysError(fmt.Sprintf("stacktrace from panic: %s", string(debug.Stack()))) c.JSON(http.StatusInternalServerError, gin.H{ "error": gin.H{ "message": fmt.Sprintf("Panic detected, error: %v. Please submit a issue here: https://github.com/songquanpeng/one-api", err), diff --git a/middleware/request-id.go b/middleware/request-id.go index e623be7a..811a7ad8 100644 --- a/middleware/request-id.go +++ b/middleware/request-id.go @@ -3,16 +3,17 @@ package middleware import ( "context" "github.com/gin-gonic/gin" - "one-api/common" + "one-api/common/helper" + "one-api/common/logger" ) func RequestId() func(c *gin.Context) { return func(c *gin.Context) { - id := common.GetTimeString() + common.GetRandomString(8) - c.Set(common.RequestIdKey, id) - ctx := context.WithValue(c.Request.Context(), common.RequestIdKey, id) + id := helper.GetTimeString() + helper.GetRandomString(8) + c.Set(logger.RequestIdKey, id) + ctx := context.WithValue(c.Request.Context(), logger.RequestIdKey, id) c.Request = c.Request.WithContext(ctx) - c.Header(common.RequestIdKey, id) + c.Header(logger.RequestIdKey, id) c.Next() } } diff --git a/middleware/turnstile-check.go b/middleware/turnstile-check.go index 26688810..629395e7 100644 --- a/middleware/turnstile-check.go +++ b/middleware/turnstile-check.go @@ -6,7 +6,8 @@ import ( "github.com/gin-gonic/gin" "net/http" "net/url" - "one-api/common" + "one-api/common/config" + "one-api/common/logger" ) type turnstileCheckResponse struct { @@ -15,7 +16,7 @@ type turnstileCheckResponse struct { func TurnstileCheck() gin.HandlerFunc { return func(c *gin.Context) { - if common.TurnstileCheckEnabled { + if config.TurnstileCheckEnabled { session := sessions.Default(c) turnstileChecked := session.Get("turnstile") if turnstileChecked != nil { @@ -32,12 +33,12 @@ func TurnstileCheck() gin.HandlerFunc { return } rawRes, err := http.PostForm("https://challenges.cloudflare.com/turnstile/v0/siteverify", url.Values{ - "secret": {common.TurnstileSecretKey}, + "secret": {config.TurnstileSecretKey}, "response": {response}, "remoteip": {c.ClientIP()}, }) if err != nil { - common.SysError(err.Error()) + logger.SysError(err.Error()) c.JSON(http.StatusOK, gin.H{ "success": false, "message": err.Error(), @@ -49,7 +50,7 @@ func TurnstileCheck() gin.HandlerFunc { var res turnstileCheckResponse err = json.NewDecoder(rawRes.Body).Decode(&res) if err != nil { - common.SysError(err.Error()) + logger.SysError(err.Error()) c.JSON(http.StatusOK, gin.H{ "success": false, "message": err.Error(), diff --git a/middleware/utils.go b/middleware/utils.go index 536125cc..d866d75b 100644 --- a/middleware/utils.go +++ b/middleware/utils.go @@ -2,16 +2,17 @@ package middleware import ( "github.com/gin-gonic/gin" - "one-api/common" + "one-api/common/helper" + "one-api/common/logger" ) func abortWithMessage(c *gin.Context, statusCode int, message string) { c.JSON(statusCode, gin.H{ "error": gin.H{ - "message": common.MessageWithRequestId(message, c.GetString(common.RequestIdKey)), + "message": helper.MessageWithRequestId(message, c.GetString(logger.RequestIdKey)), "type": "one_api_error", }, }) c.Abort() - common.LogError(c.Request.Context(), message) + logger.Error(c.Request.Context(), message) } diff --git a/model/cache.go b/model/cache.go index c6d0c70a..a81bdddd 100644 --- a/model/cache.go +++ b/model/cache.go @@ -6,6 +6,8 @@ import ( "fmt" "math/rand" "one-api/common" + "one-api/common/config" + "one-api/common/logger" "sort" "strconv" "strings" @@ -14,10 +16,10 @@ import ( ) var ( - TokenCacheSeconds = common.SyncFrequency - UserId2GroupCacheSeconds = common.SyncFrequency - UserId2QuotaCacheSeconds = common.SyncFrequency - UserId2StatusCacheSeconds = common.SyncFrequency + TokenCacheSeconds = config.SyncFrequency + UserId2GroupCacheSeconds = config.SyncFrequency + UserId2QuotaCacheSeconds = config.SyncFrequency + UserId2StatusCacheSeconds = config.SyncFrequency ) func CacheGetTokenByKey(key string) (*Token, error) { @@ -42,7 +44,7 @@ func CacheGetTokenByKey(key string) (*Token, error) { } err = common.RedisSet(fmt.Sprintf("token:%s", key), string(jsonBytes), time.Duration(TokenCacheSeconds)*time.Second) if err != nil { - common.SysError("Redis set token error: " + err.Error()) + logger.SysError("Redis set token error: " + err.Error()) } return &token, nil } @@ -62,7 +64,7 @@ func CacheGetUserGroup(id int) (group string, err error) { } err = common.RedisSet(fmt.Sprintf("user_group:%d", id), group, time.Duration(UserId2GroupCacheSeconds)*time.Second) if err != nil { - common.SysError("Redis set user group error: " + err.Error()) + logger.SysError("Redis set user group error: " + err.Error()) } } return group, err @@ -80,7 +82,7 @@ func CacheGetUserQuota(id int) (quota int, err error) { } err = common.RedisSet(fmt.Sprintf("user_quota:%d", id), fmt.Sprintf("%d", quota), time.Duration(UserId2QuotaCacheSeconds)*time.Second) if err != nil { - common.SysError("Redis set user quota error: " + err.Error()) + logger.SysError("Redis set user quota error: " + err.Error()) } return quota, err } @@ -127,7 +129,7 @@ func CacheIsUserEnabled(userId int) (bool, error) { } err = common.RedisSet(fmt.Sprintf("user_enabled:%d", userId), enabled, time.Duration(UserId2StatusCacheSeconds)*time.Second) if err != nil { - common.SysError("Redis set user enabled error: " + err.Error()) + logger.SysError("Redis set user enabled error: " + err.Error()) } return userEnabled, err } @@ -178,19 +180,19 @@ func InitChannelCache() { channelSyncLock.Lock() group2model2channels = newGroup2model2channels channelSyncLock.Unlock() - common.SysLog("channels synced from database") + logger.SysLog("channels synced from database") } func SyncChannelCache(frequency int) { for { time.Sleep(time.Duration(frequency) * time.Second) - common.SysLog("syncing channels from database") + logger.SysLog("syncing channels from database") InitChannelCache() } } func CacheGetRandomSatisfiedChannel(group string, model string) (*Channel, error) { - if !common.MemoryCacheEnabled { + if !config.MemoryCacheEnabled { return GetRandomSatisfiedChannel(group, model) } channelSyncLock.RLock() diff --git a/model/channel.go b/model/channel.go index 7e7b42e6..966b7c39 100644 --- a/model/channel.go +++ b/model/channel.go @@ -1,8 +1,13 @@ package model import ( + "encoding/json" + "fmt" "gorm.io/gorm" "one-api/common" + "one-api/common/config" + "one-api/common/helper" + "one-api/common/logger" ) type Channel struct { @@ -42,7 +47,7 @@ func SearchChannels(keyword string) (channels []*Channel, err error) { if common.UsingPostgreSQL { keyCol = `"key"` } - err = DB.Omit("key").Where("id = ? or name LIKE ? or "+keyCol+" = ?", common.String2Int(keyword), keyword+"%", keyword).Find(&channels).Error + err = DB.Omit("key").Where("id = ? or name LIKE ? or "+keyCol+" = ?", helper.String2Int(keyword), keyword+"%", keyword).Find(&channels).Error return channels, err } @@ -86,11 +91,17 @@ func (channel *Channel) GetBaseURL() string { return *channel.BaseURL } -func (channel *Channel) GetModelMapping() string { - if channel.ModelMapping == nil { - return "" +func (channel *Channel) GetModelMapping() map[string]string { + if channel.ModelMapping == nil || *channel.ModelMapping == "" || *channel.ModelMapping == "{}" { + return nil } - return *channel.ModelMapping + modelMapping := make(map[string]string) + err := json.Unmarshal([]byte(*channel.ModelMapping), &modelMapping) + if err != nil { + logger.SysError(fmt.Sprintf("failed to unmarshal model mapping for channel %d, error: %s", channel.Id, err.Error())) + return nil + } + return modelMapping } func (channel *Channel) Insert() error { @@ -116,21 +127,21 @@ func (channel *Channel) Update() error { func (channel *Channel) UpdateResponseTime(responseTime int64) { err := DB.Model(channel).Select("response_time", "test_time").Updates(Channel{ - TestTime: common.GetTimestamp(), + TestTime: helper.GetTimestamp(), ResponseTime: int(responseTime), }).Error if err != nil { - common.SysError("failed to update response time: " + err.Error()) + logger.SysError("failed to update response time: " + err.Error()) } } func (channel *Channel) UpdateBalance(balance float64) { err := DB.Model(channel).Select("balance_updated_time", "balance").Updates(Channel{ - BalanceUpdatedTime: common.GetTimestamp(), + BalanceUpdatedTime: helper.GetTimestamp(), Balance: balance, }).Error if err != nil { - common.SysError("failed to update balance: " + err.Error()) + logger.SysError("failed to update balance: " + err.Error()) } } @@ -147,16 +158,16 @@ func (channel *Channel) Delete() error { func UpdateChannelStatusById(id int, status int) { err := UpdateAbilityStatus(id, status == common.ChannelStatusEnabled) if err != nil { - common.SysError("failed to update ability status: " + err.Error()) + logger.SysError("failed to update ability status: " + err.Error()) } err = DB.Model(&Channel{}).Where("id = ?", id).Update("status", status).Error if err != nil { - common.SysError("failed to update channel status: " + err.Error()) + logger.SysError("failed to update channel status: " + err.Error()) } } func UpdateChannelUsedQuota(id int, quota int) { - if common.BatchUpdateEnabled { + if config.BatchUpdateEnabled { addNewRecord(BatchUpdateTypeChannelUsedQuota, id, quota) return } @@ -166,7 +177,7 @@ func UpdateChannelUsedQuota(id int, quota int) { func updateChannelUsedQuota(id int, quota int) { err := DB.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()) + logger.SysError("failed to update channel used quota: " + err.Error()) } } diff --git a/model/log.go b/model/log.go index 06085acf..78b6d9b3 100644 --- a/model/log.go +++ b/model/log.go @@ -4,6 +4,9 @@ import ( "context" "fmt" "one-api/common" + "one-api/common/config" + "one-api/common/helper" + "one-api/common/logger" "gorm.io/gorm" ) @@ -32,31 +35,31 @@ const ( ) func RecordLog(userId int, logType int, content string) { - if logType == LogTypeConsume && !common.LogConsumeEnabled { + if logType == LogTypeConsume && !config.LogConsumeEnabled { return } log := &Log{ UserId: userId, Username: GetUsernameById(userId), - CreatedAt: common.GetTimestamp(), + CreatedAt: helper.GetTimestamp(), Type: logType, Content: content, } err := DB.Create(log).Error if err != nil { - common.SysError("failed to record log: " + err.Error()) + logger.SysError("failed to record log: " + err.Error()) } } func RecordConsumeLog(ctx context.Context, userId int, channelId int, promptTokens int, completionTokens int, modelName string, tokenName string, quota int, content string) { - common.LogInfo(ctx, fmt.Sprintf("record consume log: userId=%d, channelId=%d, promptTokens=%d, completionTokens=%d, modelName=%s, tokenName=%s, quota=%d, content=%s", userId, channelId, promptTokens, completionTokens, modelName, tokenName, quota, content)) - if !common.LogConsumeEnabled { + logger.Info(ctx, fmt.Sprintf("record consume log: userId=%d, channelId=%d, promptTokens=%d, completionTokens=%d, modelName=%s, tokenName=%s, quota=%d, content=%s", userId, channelId, promptTokens, completionTokens, modelName, tokenName, quota, content)) + if !config.LogConsumeEnabled { return } log := &Log{ UserId: userId, Username: GetUsernameById(userId), - CreatedAt: common.GetTimestamp(), + CreatedAt: helper.GetTimestamp(), Type: LogTypeConsume, Content: content, PromptTokens: promptTokens, @@ -68,7 +71,7 @@ func RecordConsumeLog(ctx context.Context, userId int, channelId int, promptToke } err := DB.Create(log).Error if err != nil { - common.LogError(ctx, "failed to record log: "+err.Error()) + logger.Error(ctx, "failed to record log: "+err.Error()) } } @@ -125,12 +128,12 @@ func GetUserLogs(userId int, logType int, startTimestamp int64, endTimestamp int } func SearchAllLogs(keyword string) (logs []*Log, err error) { - err = DB.Where("type = ? or content LIKE ?", keyword, keyword+"%").Order("id desc").Limit(common.MaxRecentItems).Find(&logs).Error + err = DB.Where("type = ? or content LIKE ?", keyword, keyword+"%").Order("id desc").Limit(config.MaxRecentItems).Find(&logs).Error return logs, err } func SearchUserLogs(userId int, keyword string) (logs []*Log, err error) { - err = DB.Where("user_id = ? and type = ?", userId, keyword).Order("id desc").Limit(common.MaxRecentItems).Omit("id").Find(&logs).Error + err = DB.Where("user_id = ? and type = ?", userId, keyword).Order("id desc").Limit(config.MaxRecentItems).Omit("id").Find(&logs).Error return logs, err } diff --git a/model/main.go b/model/main.go index 9723e638..2ed6f0e3 100644 --- a/model/main.go +++ b/model/main.go @@ -7,6 +7,9 @@ import ( "gorm.io/driver/sqlite" "gorm.io/gorm" "one-api/common" + "one-api/common/config" + "one-api/common/helper" + "one-api/common/logger" "os" "strings" "time" @@ -18,7 +21,7 @@ func createRootAccountIfNeed() error { var user User //if user.Status != util.UserStatusEnabled { if err := DB.First(&user).Error; err != nil { - common.SysLog("no user exists, create a root user for you: username is root, password is 123456") + logger.SysLog("no user exists, create a root user for you: username is root, password is 123456") hashedPassword, err := common.Password2Hash("123456") if err != nil { return err @@ -29,7 +32,7 @@ func createRootAccountIfNeed() error { Role: common.RoleRootUser, Status: common.UserStatusEnabled, DisplayName: "Root User", - AccessToken: common.GetUUID(), + AccessToken: helper.GetUUID(), Quota: 100000000, } DB.Create(&rootUser) @@ -42,7 +45,7 @@ func chooseDB() (*gorm.DB, error) { dsn := os.Getenv("SQL_DSN") if strings.HasPrefix(dsn, "postgres://") { // Use PostgreSQL - common.SysLog("using PostgreSQL as database") + logger.SysLog("using PostgreSQL as database") common.UsingPostgreSQL = true return gorm.Open(postgres.New(postgres.Config{ DSN: dsn, @@ -52,13 +55,13 @@ func chooseDB() (*gorm.DB, error) { }) } // Use MySQL - common.SysLog("using MySQL as database") + logger.SysLog("using MySQL as database") return gorm.Open(mysql.Open(dsn), &gorm.Config{ PrepareStmt: true, // precompile SQL }) } // Use SQLite - common.SysLog("SQL_DSN not set, using SQLite as database") + logger.SysLog("SQL_DSN not set, using SQLite as database") common.UsingSQLite = true config := fmt.Sprintf("?_busy_timeout=%d", common.SQLiteBusyTimeout) return gorm.Open(sqlite.Open(common.SQLitePath+config), &gorm.Config{ @@ -69,7 +72,7 @@ func chooseDB() (*gorm.DB, error) { func InitDB() (err error) { db, err := chooseDB() if err == nil { - if common.DebugEnabled { + if config.DebugEnabled { db = db.Debug() } DB = db @@ -77,14 +80,14 @@ func InitDB() (err error) { if err != nil { return err } - sqlDB.SetMaxIdleConns(common.GetOrDefault("SQL_MAX_IDLE_CONNS", 100)) - sqlDB.SetMaxOpenConns(common.GetOrDefault("SQL_MAX_OPEN_CONNS", 1000)) - sqlDB.SetConnMaxLifetime(time.Second * time.Duration(common.GetOrDefault("SQL_MAX_LIFETIME", 60))) + sqlDB.SetMaxIdleConns(helper.GetOrDefaultEnvInt("SQL_MAX_IDLE_CONNS", 100)) + sqlDB.SetMaxOpenConns(helper.GetOrDefaultEnvInt("SQL_MAX_OPEN_CONNS", 1000)) + sqlDB.SetConnMaxLifetime(time.Second * time.Duration(helper.GetOrDefaultEnvInt("SQL_MAX_LIFETIME", 60))) - if !common.IsMasterNode { + if !config.IsMasterNode { return nil } - common.SysLog("database migration started") + logger.SysLog("database migration started") err = db.AutoMigrate(&Channel{}) if err != nil { return err @@ -113,11 +116,11 @@ func InitDB() (err error) { if err != nil { return err } - common.SysLog("database migrated") + logger.SysLog("database migrated") err = createRootAccountIfNeed() return err } else { - common.FatalLog(err) + logger.FatalLog(err) } return err } diff --git a/model/option.go b/model/option.go index 20575c9a..e211264c 100644 --- a/model/option.go +++ b/model/option.go @@ -2,6 +2,8 @@ package model import ( "one-api/common" + "one-api/common/config" + "one-api/common/logger" "strconv" "strings" "time" @@ -20,60 +22,56 @@ func AllOption() ([]*Option, error) { } func InitOptionMap() { - common.OptionMapRWMutex.Lock() - common.OptionMap = make(map[string]string) - common.OptionMap["FileUploadPermission"] = strconv.Itoa(common.FileUploadPermission) - common.OptionMap["FileDownloadPermission"] = strconv.Itoa(common.FileDownloadPermission) - common.OptionMap["ImageUploadPermission"] = strconv.Itoa(common.ImageUploadPermission) - common.OptionMap["ImageDownloadPermission"] = strconv.Itoa(common.ImageDownloadPermission) - common.OptionMap["PasswordLoginEnabled"] = strconv.FormatBool(common.PasswordLoginEnabled) - common.OptionMap["PasswordRegisterEnabled"] = strconv.FormatBool(common.PasswordRegisterEnabled) - common.OptionMap["EmailVerificationEnabled"] = strconv.FormatBool(common.EmailVerificationEnabled) - common.OptionMap["GitHubOAuthEnabled"] = strconv.FormatBool(common.GitHubOAuthEnabled) - common.OptionMap["WeChatAuthEnabled"] = strconv.FormatBool(common.WeChatAuthEnabled) - common.OptionMap["TurnstileCheckEnabled"] = strconv.FormatBool(common.TurnstileCheckEnabled) - common.OptionMap["RegisterEnabled"] = strconv.FormatBool(common.RegisterEnabled) - common.OptionMap["AutomaticDisableChannelEnabled"] = strconv.FormatBool(common.AutomaticDisableChannelEnabled) - common.OptionMap["AutomaticEnableChannelEnabled"] = strconv.FormatBool(common.AutomaticEnableChannelEnabled) - common.OptionMap["ApproximateTokenEnabled"] = strconv.FormatBool(common.ApproximateTokenEnabled) - common.OptionMap["LogConsumeEnabled"] = strconv.FormatBool(common.LogConsumeEnabled) - common.OptionMap["DisplayInCurrencyEnabled"] = strconv.FormatBool(common.DisplayInCurrencyEnabled) - common.OptionMap["DisplayTokenStatEnabled"] = strconv.FormatBool(common.DisplayTokenStatEnabled) - common.OptionMap["ChannelDisableThreshold"] = strconv.FormatFloat(common.ChannelDisableThreshold, 'f', -1, 64) - common.OptionMap["EmailDomainRestrictionEnabled"] = strconv.FormatBool(common.EmailDomainRestrictionEnabled) - common.OptionMap["EmailDomainWhitelist"] = strings.Join(common.EmailDomainWhitelist, ",") - common.OptionMap["SMTPServer"] = "" - common.OptionMap["SMTPFrom"] = "" - common.OptionMap["SMTPPort"] = strconv.Itoa(common.SMTPPort) - common.OptionMap["SMTPAccount"] = "" - common.OptionMap["SMTPToken"] = "" - common.OptionMap["Notice"] = "" - common.OptionMap["About"] = "" - common.OptionMap["HomePageContent"] = "" - common.OptionMap["Footer"] = common.Footer - common.OptionMap["SystemName"] = common.SystemName - common.OptionMap["Logo"] = common.Logo - common.OptionMap["ServerAddress"] = "" - common.OptionMap["GitHubClientId"] = "" - common.OptionMap["GitHubClientSecret"] = "" - common.OptionMap["WeChatServerAddress"] = "" - common.OptionMap["WeChatServerToken"] = "" - common.OptionMap["WeChatAccountQRCodeImageURL"] = "" - common.OptionMap["TurnstileSiteKey"] = "" - common.OptionMap["TurnstileSecretKey"] = "" - common.OptionMap["QuotaForNewUser"] = strconv.Itoa(common.QuotaForNewUser) - common.OptionMap["QuotaForInviter"] = strconv.Itoa(common.QuotaForInviter) - common.OptionMap["QuotaForInvitee"] = strconv.Itoa(common.QuotaForInvitee) - common.OptionMap["QuotaRemindThreshold"] = strconv.Itoa(common.QuotaRemindThreshold) - common.OptionMap["PreConsumedQuota"] = strconv.Itoa(common.PreConsumedQuota) - common.OptionMap["ModelRatio"] = common.ModelRatio2JSONString() - common.OptionMap["GroupRatio"] = common.GroupRatio2JSONString() - common.OptionMap["TopUpLink"] = common.TopUpLink - common.OptionMap["ChatLink"] = common.ChatLink - common.OptionMap["QuotaPerUnit"] = strconv.FormatFloat(common.QuotaPerUnit, 'f', -1, 64) - common.OptionMap["RetryTimes"] = strconv.Itoa(common.RetryTimes) - common.OptionMap["Theme"] = common.Theme - common.OptionMapRWMutex.Unlock() + config.OptionMapRWMutex.Lock() + config.OptionMap = make(map[string]string) + config.OptionMap["PasswordLoginEnabled"] = strconv.FormatBool(config.PasswordLoginEnabled) + config.OptionMap["PasswordRegisterEnabled"] = strconv.FormatBool(config.PasswordRegisterEnabled) + config.OptionMap["EmailVerificationEnabled"] = strconv.FormatBool(config.EmailVerificationEnabled) + config.OptionMap["GitHubOAuthEnabled"] = strconv.FormatBool(config.GitHubOAuthEnabled) + config.OptionMap["WeChatAuthEnabled"] = strconv.FormatBool(config.WeChatAuthEnabled) + config.OptionMap["TurnstileCheckEnabled"] = strconv.FormatBool(config.TurnstileCheckEnabled) + config.OptionMap["RegisterEnabled"] = strconv.FormatBool(config.RegisterEnabled) + config.OptionMap["AutomaticDisableChannelEnabled"] = strconv.FormatBool(config.AutomaticDisableChannelEnabled) + config.OptionMap["AutomaticEnableChannelEnabled"] = strconv.FormatBool(config.AutomaticEnableChannelEnabled) + config.OptionMap["ApproximateTokenEnabled"] = strconv.FormatBool(config.ApproximateTokenEnabled) + config.OptionMap["LogConsumeEnabled"] = strconv.FormatBool(config.LogConsumeEnabled) + config.OptionMap["DisplayInCurrencyEnabled"] = strconv.FormatBool(config.DisplayInCurrencyEnabled) + config.OptionMap["DisplayTokenStatEnabled"] = strconv.FormatBool(config.DisplayTokenStatEnabled) + config.OptionMap["ChannelDisableThreshold"] = strconv.FormatFloat(config.ChannelDisableThreshold, 'f', -1, 64) + config.OptionMap["EmailDomainRestrictionEnabled"] = strconv.FormatBool(config.EmailDomainRestrictionEnabled) + config.OptionMap["EmailDomainWhitelist"] = strings.Join(config.EmailDomainWhitelist, ",") + config.OptionMap["SMTPServer"] = "" + config.OptionMap["SMTPFrom"] = "" + config.OptionMap["SMTPPort"] = strconv.Itoa(config.SMTPPort) + config.OptionMap["SMTPAccount"] = "" + config.OptionMap["SMTPToken"] = "" + config.OptionMap["Notice"] = "" + config.OptionMap["About"] = "" + config.OptionMap["HomePageContent"] = "" + config.OptionMap["Footer"] = config.Footer + config.OptionMap["SystemName"] = config.SystemName + config.OptionMap["Logo"] = config.Logo + config.OptionMap["ServerAddress"] = "" + config.OptionMap["GitHubClientId"] = "" + config.OptionMap["GitHubClientSecret"] = "" + config.OptionMap["WeChatServerAddress"] = "" + config.OptionMap["WeChatServerToken"] = "" + config.OptionMap["WeChatAccountQRCodeImageURL"] = "" + config.OptionMap["TurnstileSiteKey"] = "" + config.OptionMap["TurnstileSecretKey"] = "" + config.OptionMap["QuotaForNewUser"] = strconv.Itoa(config.QuotaForNewUser) + config.OptionMap["QuotaForInviter"] = strconv.Itoa(config.QuotaForInviter) + config.OptionMap["QuotaForInvitee"] = strconv.Itoa(config.QuotaForInvitee) + config.OptionMap["QuotaRemindThreshold"] = strconv.Itoa(config.QuotaRemindThreshold) + config.OptionMap["PreConsumedQuota"] = strconv.Itoa(config.PreConsumedQuota) + config.OptionMap["ModelRatio"] = common.ModelRatio2JSONString() + config.OptionMap["GroupRatio"] = common.GroupRatio2JSONString() + config.OptionMap["TopUpLink"] = config.TopUpLink + config.OptionMap["ChatLink"] = config.ChatLink + config.OptionMap["QuotaPerUnit"] = strconv.FormatFloat(config.QuotaPerUnit, 'f', -1, 64) + config.OptionMap["RetryTimes"] = strconv.Itoa(config.RetryTimes) + config.OptionMap["Theme"] = config.Theme + config.OptionMapRWMutex.Unlock() loadOptionsFromDatabase() } @@ -82,7 +80,7 @@ func loadOptionsFromDatabase() { for _, option := range options { err := updateOptionMap(option.Key, option.Value) if err != nil { - common.SysError("failed to update option map: " + err.Error()) + logger.SysError("failed to update option map: " + err.Error()) } } } @@ -90,7 +88,7 @@ func loadOptionsFromDatabase() { func SyncOptions(frequency int) { for { time.Sleep(time.Duration(frequency) * time.Second) - common.SysLog("syncing options from database") + logger.SysLog("syncing options from database") loadOptionsFromDatabase() } } @@ -112,117 +110,104 @@ func UpdateOption(key string, value string) error { } func updateOptionMap(key string, value string) (err error) { - common.OptionMapRWMutex.Lock() - defer common.OptionMapRWMutex.Unlock() - common.OptionMap[key] = value - if strings.HasSuffix(key, "Permission") { - intValue, _ := strconv.Atoi(value) - switch key { - case "FileUploadPermission": - common.FileUploadPermission = intValue - case "FileDownloadPermission": - common.FileDownloadPermission = intValue - case "ImageUploadPermission": - common.ImageUploadPermission = intValue - case "ImageDownloadPermission": - common.ImageDownloadPermission = intValue - } - } + config.OptionMapRWMutex.Lock() + defer config.OptionMapRWMutex.Unlock() + config.OptionMap[key] = value if strings.HasSuffix(key, "Enabled") { boolValue := value == "true" switch key { case "PasswordRegisterEnabled": - common.PasswordRegisterEnabled = boolValue + config.PasswordRegisterEnabled = boolValue case "PasswordLoginEnabled": - common.PasswordLoginEnabled = boolValue + config.PasswordLoginEnabled = boolValue case "EmailVerificationEnabled": - common.EmailVerificationEnabled = boolValue + config.EmailVerificationEnabled = boolValue case "GitHubOAuthEnabled": - common.GitHubOAuthEnabled = boolValue + config.GitHubOAuthEnabled = boolValue case "WeChatAuthEnabled": - common.WeChatAuthEnabled = boolValue + config.WeChatAuthEnabled = boolValue case "TurnstileCheckEnabled": - common.TurnstileCheckEnabled = boolValue + config.TurnstileCheckEnabled = boolValue case "RegisterEnabled": - common.RegisterEnabled = boolValue + config.RegisterEnabled = boolValue case "EmailDomainRestrictionEnabled": - common.EmailDomainRestrictionEnabled = boolValue + config.EmailDomainRestrictionEnabled = boolValue case "AutomaticDisableChannelEnabled": - common.AutomaticDisableChannelEnabled = boolValue + config.AutomaticDisableChannelEnabled = boolValue case "AutomaticEnableChannelEnabled": - common.AutomaticEnableChannelEnabled = boolValue + config.AutomaticEnableChannelEnabled = boolValue case "ApproximateTokenEnabled": - common.ApproximateTokenEnabled = boolValue + config.ApproximateTokenEnabled = boolValue case "LogConsumeEnabled": - common.LogConsumeEnabled = boolValue + config.LogConsumeEnabled = boolValue case "DisplayInCurrencyEnabled": - common.DisplayInCurrencyEnabled = boolValue + config.DisplayInCurrencyEnabled = boolValue case "DisplayTokenStatEnabled": - common.DisplayTokenStatEnabled = boolValue + config.DisplayTokenStatEnabled = boolValue } } switch key { case "EmailDomainWhitelist": - common.EmailDomainWhitelist = strings.Split(value, ",") + config.EmailDomainWhitelist = strings.Split(value, ",") case "SMTPServer": - common.SMTPServer = value + config.SMTPServer = value case "SMTPPort": intValue, _ := strconv.Atoi(value) - common.SMTPPort = intValue + config.SMTPPort = intValue case "SMTPAccount": - common.SMTPAccount = value + config.SMTPAccount = value case "SMTPFrom": - common.SMTPFrom = value + config.SMTPFrom = value case "SMTPToken": - common.SMTPToken = value + config.SMTPToken = value case "ServerAddress": - common.ServerAddress = value + config.ServerAddress = value case "GitHubClientId": - common.GitHubClientId = value + config.GitHubClientId = value case "GitHubClientSecret": - common.GitHubClientSecret = value + config.GitHubClientSecret = value case "Footer": - common.Footer = value + config.Footer = value case "SystemName": - common.SystemName = value + config.SystemName = value case "Logo": - common.Logo = value + config.Logo = value case "WeChatServerAddress": - common.WeChatServerAddress = value + config.WeChatServerAddress = value case "WeChatServerToken": - common.WeChatServerToken = value + config.WeChatServerToken = value case "WeChatAccountQRCodeImageURL": - common.WeChatAccountQRCodeImageURL = value + config.WeChatAccountQRCodeImageURL = value case "TurnstileSiteKey": - common.TurnstileSiteKey = value + config.TurnstileSiteKey = value case "TurnstileSecretKey": - common.TurnstileSecretKey = value + config.TurnstileSecretKey = value case "QuotaForNewUser": - common.QuotaForNewUser, _ = strconv.Atoi(value) + config.QuotaForNewUser, _ = strconv.Atoi(value) case "QuotaForInviter": - common.QuotaForInviter, _ = strconv.Atoi(value) + config.QuotaForInviter, _ = strconv.Atoi(value) case "QuotaForInvitee": - common.QuotaForInvitee, _ = strconv.Atoi(value) + config.QuotaForInvitee, _ = strconv.Atoi(value) case "QuotaRemindThreshold": - common.QuotaRemindThreshold, _ = strconv.Atoi(value) + config.QuotaRemindThreshold, _ = strconv.Atoi(value) case "PreConsumedQuota": - common.PreConsumedQuota, _ = strconv.Atoi(value) + config.PreConsumedQuota, _ = strconv.Atoi(value) case "RetryTimes": - common.RetryTimes, _ = strconv.Atoi(value) + config.RetryTimes, _ = strconv.Atoi(value) case "ModelRatio": err = common.UpdateModelRatioByJSONString(value) case "GroupRatio": err = common.UpdateGroupRatioByJSONString(value) case "TopUpLink": - common.TopUpLink = value + config.TopUpLink = value case "ChatLink": - common.ChatLink = value + config.ChatLink = value case "ChannelDisableThreshold": - common.ChannelDisableThreshold, _ = strconv.ParseFloat(value, 64) + config.ChannelDisableThreshold, _ = strconv.ParseFloat(value, 64) case "QuotaPerUnit": - common.QuotaPerUnit, _ = strconv.ParseFloat(value, 64) + config.QuotaPerUnit, _ = strconv.ParseFloat(value, 64) case "Theme": - common.Theme = value + config.Theme = value } return err } diff --git a/model/redemption.go b/model/redemption.go index f16412b5..026794e0 100644 --- a/model/redemption.go +++ b/model/redemption.go @@ -5,6 +5,7 @@ import ( "fmt" "gorm.io/gorm" "one-api/common" + "one-api/common/helper" ) type Redemption struct { @@ -67,7 +68,7 @@ func Redeem(key string, userId int) (quota int, err error) { if err != nil { return err } - redemption.RedeemedTime = common.GetTimestamp() + redemption.RedeemedTime = helper.GetTimestamp() redemption.Status = common.RedemptionCodeStatusUsed err = tx.Save(redemption).Error return err diff --git a/model/token.go b/model/token.go index 2e53ac0b..2087225b 100644 --- a/model/token.go +++ b/model/token.go @@ -5,6 +5,9 @@ import ( "fmt" "gorm.io/gorm" "one-api/common" + "one-api/common/config" + "one-api/common/helper" + "one-api/common/logger" ) type Token struct { @@ -39,7 +42,7 @@ func ValidateUserToken(key string) (token *Token, err error) { } token, err = CacheGetTokenByKey(key) if err != nil { - common.SysError("CacheGetTokenByKey failed: " + err.Error()) + logger.SysError("CacheGetTokenByKey failed: " + err.Error()) if errors.Is(err, gorm.ErrRecordNotFound) { return nil, errors.New("无效的令牌") } @@ -53,12 +56,12 @@ func ValidateUserToken(key string) (token *Token, err error) { if token.Status != common.TokenStatusEnabled { return nil, errors.New("该令牌状态不可用") } - if token.ExpiredTime != -1 && token.ExpiredTime < common.GetTimestamp() { + if token.ExpiredTime != -1 && token.ExpiredTime < helper.GetTimestamp() { if !common.RedisEnabled { token.Status = common.TokenStatusExpired err := token.SelectUpdate() if err != nil { - common.SysError("failed to update token status" + err.Error()) + logger.SysError("failed to update token status" + err.Error()) } } return nil, errors.New("该令牌已过期") @@ -69,7 +72,7 @@ func ValidateUserToken(key string) (token *Token, err error) { token.Status = common.TokenStatusExhausted err := token.SelectUpdate() if err != nil { - common.SysError("failed to update token status" + err.Error()) + logger.SysError("failed to update token status" + err.Error()) } } return nil, errors.New("该令牌额度已用尽") @@ -138,7 +141,7 @@ func IncreaseTokenQuota(id int, quota int) (err error) { if quota < 0 { return errors.New("quota 不能为负数!") } - if common.BatchUpdateEnabled { + if config.BatchUpdateEnabled { addNewRecord(BatchUpdateTypeTokenQuota, id, quota) return nil } @@ -150,7 +153,7 @@ func increaseTokenQuota(id int, quota int) (err error) { map[string]interface{}{ "remain_quota": gorm.Expr("remain_quota + ?", quota), "used_quota": gorm.Expr("used_quota - ?", quota), - "accessed_time": common.GetTimestamp(), + "accessed_time": helper.GetTimestamp(), }, ).Error return err @@ -160,7 +163,7 @@ func DecreaseTokenQuota(id int, quota int) (err error) { if quota < 0 { return errors.New("quota 不能为负数!") } - if common.BatchUpdateEnabled { + if config.BatchUpdateEnabled { addNewRecord(BatchUpdateTypeTokenQuota, id, -quota) return nil } @@ -172,7 +175,7 @@ func decreaseTokenQuota(id int, quota int) (err error) { map[string]interface{}{ "remain_quota": gorm.Expr("remain_quota - ?", quota), "used_quota": gorm.Expr("used_quota + ?", quota), - "accessed_time": common.GetTimestamp(), + "accessed_time": helper.GetTimestamp(), }, ).Error return err @@ -196,24 +199,24 @@ func PreConsumeTokenQuota(tokenId int, quota int) (err error) { if userQuota < quota { return errors.New("用户额度不足") } - quotaTooLow := userQuota >= common.QuotaRemindThreshold && userQuota-quota < common.QuotaRemindThreshold + quotaTooLow := userQuota >= config.QuotaRemindThreshold && userQuota-quota < config.QuotaRemindThreshold noMoreQuota := userQuota-quota <= 0 if quotaTooLow || noMoreQuota { go func() { email, err := GetUserEmail(token.UserId) if err != nil { - common.SysError("failed to fetch user email: " + err.Error()) + logger.SysError("failed to fetch user email: " + err.Error()) } prompt := "您的额度即将用尽" if noMoreQuota { prompt = "您的额度已用尽" } if email != "" { - topUpLink := fmt.Sprintf("%s/topup", common.ServerAddress) + topUpLink := fmt.Sprintf("%s/topup", config.ServerAddress) err = common.SendEmail(prompt, email, fmt.Sprintf("%s,当前剩余额度为 %d,为了不影响您的使用,请及时充值。
充值链接:%s", prompt, userQuota, topUpLink, topUpLink)) if err != nil { - common.SysError("failed to send email" + err.Error()) + logger.SysError("failed to send email" + err.Error()) } } }() diff --git a/model/user.go b/model/user.go index 1c2c0a75..82e9707b 100644 --- a/model/user.go +++ b/model/user.go @@ -5,6 +5,9 @@ import ( "fmt" "gorm.io/gorm" "one-api/common" + "one-api/common/config" + "one-api/common/helper" + "one-api/common/logger" "strings" ) @@ -89,24 +92,24 @@ func (user *User) Insert(inviterId int) error { return err } } - user.Quota = common.QuotaForNewUser - user.AccessToken = common.GetUUID() - user.AffCode = common.GetRandomString(4) + user.Quota = config.QuotaForNewUser + user.AccessToken = helper.GetUUID() + user.AffCode = helper.GetRandomString(4) result := DB.Create(user) if result.Error != nil { return result.Error } - if common.QuotaForNewUser > 0 { - RecordLog(user.Id, LogTypeSystem, fmt.Sprintf("新用户注册赠送 %s", common.LogQuota(common.QuotaForNewUser))) + if config.QuotaForNewUser > 0 { + RecordLog(user.Id, LogTypeSystem, fmt.Sprintf("新用户注册赠送 %s", common.LogQuota(config.QuotaForNewUser))) } if inviterId != 0 { - if common.QuotaForInvitee > 0 { - _ = IncreaseUserQuota(user.Id, common.QuotaForInvitee) - RecordLog(user.Id, LogTypeSystem, fmt.Sprintf("使用邀请码赠送 %s", common.LogQuota(common.QuotaForInvitee))) + if config.QuotaForInvitee > 0 { + _ = IncreaseUserQuota(user.Id, config.QuotaForInvitee) + RecordLog(user.Id, LogTypeSystem, fmt.Sprintf("使用邀请码赠送 %s", common.LogQuota(config.QuotaForInvitee))) } - if common.QuotaForInviter > 0 { - _ = IncreaseUserQuota(inviterId, common.QuotaForInviter) - RecordLog(inviterId, LogTypeSystem, fmt.Sprintf("邀请用户赠送 %s", common.LogQuota(common.QuotaForInviter))) + if config.QuotaForInviter > 0 { + _ = IncreaseUserQuota(inviterId, config.QuotaForInviter) + RecordLog(inviterId, LogTypeSystem, fmt.Sprintf("邀请用户赠送 %s", common.LogQuota(config.QuotaForInviter))) } } return nil @@ -232,7 +235,7 @@ func IsAdmin(userId int) bool { var user User err := DB.Where("id = ?", userId).Select("role").Find(&user).Error if err != nil { - common.SysError("no such user " + err.Error()) + logger.SysError("no such user " + err.Error()) return false } return user.Role >= common.RoleAdminUser @@ -291,7 +294,7 @@ func IncreaseUserQuota(id int, quota int) (err error) { if quota < 0 { return errors.New("quota 不能为负数!") } - if common.BatchUpdateEnabled { + if config.BatchUpdateEnabled { addNewRecord(BatchUpdateTypeUserQuota, id, quota) return nil } @@ -307,7 +310,7 @@ func DecreaseUserQuota(id int, quota int) (err error) { if quota < 0 { return errors.New("quota 不能为负数!") } - if common.BatchUpdateEnabled { + if config.BatchUpdateEnabled { addNewRecord(BatchUpdateTypeUserQuota, id, -quota) return nil } @@ -325,7 +328,7 @@ func GetRootUserEmail() (email string) { } func UpdateUserUsedQuotaAndRequestCount(id int, quota int) { - if common.BatchUpdateEnabled { + if config.BatchUpdateEnabled { addNewRecord(BatchUpdateTypeUsedQuota, id, quota) addNewRecord(BatchUpdateTypeRequestCount, id, 1) return @@ -341,7 +344,7 @@ func updateUserUsedQuotaAndRequestCount(id int, quota int, count int) { }, ).Error if err != nil { - common.SysError("failed to update user used quota and request count: " + err.Error()) + logger.SysError("failed to update user used quota and request count: " + err.Error()) } } @@ -352,14 +355,14 @@ func updateUserUsedQuota(id int, quota int) { }, ).Error if err != nil { - common.SysError("failed to update user used quota: " + err.Error()) + logger.SysError("failed to update user used quota: " + err.Error()) } } func updateUserRequestCount(id int, count int) { err := DB.Model(&User{}).Where("id = ?", id).Update("request_count", gorm.Expr("request_count + ?", count)).Error if err != nil { - common.SysError("failed to update user request count: " + err.Error()) + logger.SysError("failed to update user request count: " + err.Error()) } } diff --git a/model/utils.go b/model/utils.go index 1c28340b..e0826e0d 100644 --- a/model/utils.go +++ b/model/utils.go @@ -1,7 +1,8 @@ package model import ( - "one-api/common" + "one-api/common/config" + "one-api/common/logger" "sync" "time" ) @@ -28,7 +29,7 @@ func init() { func InitBatchUpdater() { go func() { for { - time.Sleep(time.Duration(common.BatchUpdateInterval) * time.Second) + time.Sleep(time.Duration(config.BatchUpdateInterval) * time.Second) batchUpdate() } }() @@ -45,7 +46,7 @@ func addNewRecord(type_ int, id int, value int) { } func batchUpdate() { - common.SysLog("batch update started") + logger.SysLog("batch update started") for i := 0; i < BatchUpdateTypeCount; i++ { batchUpdateLocks[i].Lock() store := batchUpdateStores[i] @@ -57,12 +58,12 @@ func batchUpdate() { case BatchUpdateTypeUserQuota: err := increaseUserQuota(key, value) if err != nil { - common.SysError("failed to batch update user quota: " + err.Error()) + logger.SysError("failed to batch update user quota: " + err.Error()) } case BatchUpdateTypeTokenQuota: err := increaseTokenQuota(key, value) if err != nil { - common.SysError("failed to batch update token quota: " + err.Error()) + logger.SysError("failed to batch update token quota: " + err.Error()) } case BatchUpdateTypeUsedQuota: updateUserUsedQuota(key, value) @@ -73,5 +74,5 @@ func batchUpdate() { } } } - common.SysLog("batch update finished") + logger.SysLog("batch update finished") } diff --git a/relay/channel/aiproxy/adaptor.go b/relay/channel/aiproxy/adaptor.go new file mode 100644 index 00000000..44b6f58d --- /dev/null +++ b/relay/channel/aiproxy/adaptor.go @@ -0,0 +1,22 @@ +package aiproxy + +import ( + "github.com/gin-gonic/gin" + "net/http" + "one-api/relay/channel/openai" +) + +type Adaptor struct { +} + +func (a *Adaptor) Auth(c *gin.Context) error { + return nil +} + +func (a *Adaptor) ConvertRequest(request *openai.GeneralOpenAIRequest) (any, error) { + return nil, nil +} + +func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response) (*openai.ErrorWithStatusCode, *openai.Usage, error) { + return nil, nil, nil +} diff --git a/relay/channel/aiproxy/main.go b/relay/channel/aiproxy/main.go index bee4d9d3..af9cd6f6 100644 --- a/relay/channel/aiproxy/main.go +++ b/relay/channel/aiproxy/main.go @@ -8,6 +8,8 @@ import ( "io" "net/http" "one-api/common" + "one-api/common/helper" + "one-api/common/logger" "one-api/relay/channel/openai" "one-api/relay/constant" "strconv" @@ -50,9 +52,9 @@ func responseAIProxyLibrary2OpenAI(response *LibraryResponse) *openai.TextRespon FinishReason: "stop", } fullTextResponse := openai.TextResponse{ - Id: common.GetUUID(), + Id: helper.GetUUID(), Object: "chat.completion", - Created: common.GetTimestamp(), + Created: helper.GetTimestamp(), Choices: []openai.TextResponseChoice{choice}, } return &fullTextResponse @@ -63,9 +65,9 @@ func documentsAIProxyLibrary(documents []LibraryDocument) *openai.ChatCompletion choice.Delta.Content = aiProxyDocuments2Markdown(documents) choice.FinishReason = &constant.StopFinishReason return &openai.ChatCompletionsStreamResponse{ - Id: common.GetUUID(), + Id: helper.GetUUID(), Object: "chat.completion.chunk", - Created: common.GetTimestamp(), + Created: helper.GetTimestamp(), Model: "", Choices: []openai.ChatCompletionsStreamResponseChoice{choice}, } @@ -75,9 +77,9 @@ func streamResponseAIProxyLibrary2OpenAI(response *LibraryStreamResponse) *opena var choice openai.ChatCompletionsStreamResponseChoice choice.Delta.Content = response.Content return &openai.ChatCompletionsStreamResponse{ - Id: common.GetUUID(), + Id: helper.GetUUID(), Object: "chat.completion.chunk", - Created: common.GetTimestamp(), + Created: helper.GetTimestamp(), Model: response.Model, Choices: []openai.ChatCompletionsStreamResponseChoice{choice}, } @@ -122,7 +124,7 @@ func StreamHandler(c *gin.Context, resp *http.Response) (*openai.ErrorWithStatus var AIProxyLibraryResponse LibraryStreamResponse err := json.Unmarshal([]byte(data), &AIProxyLibraryResponse) if err != nil { - common.SysError("error unmarshalling stream response: " + err.Error()) + logger.SysError("error unmarshalling stream response: " + err.Error()) return true } if len(AIProxyLibraryResponse.Documents) != 0 { @@ -131,7 +133,7 @@ func StreamHandler(c *gin.Context, resp *http.Response) (*openai.ErrorWithStatus response := streamResponseAIProxyLibrary2OpenAI(&AIProxyLibraryResponse) jsonResponse, err := json.Marshal(response) if err != nil { - common.SysError("error marshalling stream response: " + err.Error()) + logger.SysError("error marshalling stream response: " + err.Error()) return true } c.Render(-1, common.CustomEvent{Data: "data: " + string(jsonResponse)}) @@ -140,7 +142,7 @@ func StreamHandler(c *gin.Context, resp *http.Response) (*openai.ErrorWithStatus response := documentsAIProxyLibrary(documents) jsonResponse, err := json.Marshal(response) if err != nil { - common.SysError("error marshalling stream response: " + err.Error()) + logger.SysError("error marshalling stream response: " + err.Error()) return true } c.Render(-1, common.CustomEvent{Data: "data: " + string(jsonResponse)}) diff --git a/relay/channel/ali/adaptor.go b/relay/channel/ali/adaptor.go new file mode 100644 index 00000000..49022cfc --- /dev/null +++ b/relay/channel/ali/adaptor.go @@ -0,0 +1,22 @@ +package ali + +import ( + "github.com/gin-gonic/gin" + "net/http" + "one-api/relay/channel/openai" +) + +type Adaptor struct { +} + +func (a *Adaptor) Auth(c *gin.Context) error { + return nil +} + +func (a *Adaptor) ConvertRequest(request *openai.GeneralOpenAIRequest) (any, error) { + return nil, nil +} + +func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response) (*openai.ErrorWithStatusCode, *openai.Usage, error) { + return nil, nil, nil +} diff --git a/relay/channel/ali/main.go b/relay/channel/ali/main.go index f45a515a..81dc5370 100644 --- a/relay/channel/ali/main.go +++ b/relay/channel/ali/main.go @@ -7,6 +7,8 @@ import ( "io" "net/http" "one-api/common" + "one-api/common/helper" + "one-api/common/logger" "one-api/relay/channel/openai" "strings" ) @@ -118,7 +120,7 @@ func responseAli2OpenAI(response *ChatResponse) *openai.TextResponse { fullTextResponse := openai.TextResponse{ Id: response.RequestId, Object: "chat.completion", - Created: common.GetTimestamp(), + Created: helper.GetTimestamp(), Choices: []openai.TextResponseChoice{choice}, Usage: openai.Usage{ PromptTokens: response.Usage.InputTokens, @@ -139,7 +141,7 @@ func streamResponseAli2OpenAI(aliResponse *ChatResponse) *openai.ChatCompletions response := openai.ChatCompletionsStreamResponse{ Id: aliResponse.RequestId, Object: "chat.completion.chunk", - Created: common.GetTimestamp(), + Created: helper.GetTimestamp(), Model: "qwen", Choices: []openai.ChatCompletionsStreamResponseChoice{choice}, } @@ -185,7 +187,7 @@ func StreamHandler(c *gin.Context, resp *http.Response) (*openai.ErrorWithStatus var aliResponse ChatResponse err := json.Unmarshal([]byte(data), &aliResponse) if err != nil { - common.SysError("error unmarshalling stream response: " + err.Error()) + logger.SysError("error unmarshalling stream response: " + err.Error()) return true } if aliResponse.Usage.OutputTokens != 0 { @@ -198,7 +200,7 @@ func StreamHandler(c *gin.Context, resp *http.Response) (*openai.ErrorWithStatus //lastResponseText = aliResponse.Output.Text jsonResponse, err := json.Marshal(response) if err != nil { - common.SysError("error marshalling stream response: " + err.Error()) + logger.SysError("error marshalling stream response: " + err.Error()) return true } c.Render(-1, common.CustomEvent{Data: "data: " + string(jsonResponse)}) diff --git a/relay/channel/anthropic/adaptor.go b/relay/channel/anthropic/adaptor.go new file mode 100644 index 00000000..55577228 --- /dev/null +++ b/relay/channel/anthropic/adaptor.go @@ -0,0 +1,22 @@ +package anthropic + +import ( + "github.com/gin-gonic/gin" + "net/http" + "one-api/relay/channel/openai" +) + +type Adaptor struct { +} + +func (a *Adaptor) Auth(c *gin.Context) error { + return nil +} + +func (a *Adaptor) ConvertRequest(request *openai.GeneralOpenAIRequest) (any, error) { + return nil, nil +} + +func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response) (*openai.ErrorWithStatusCode, *openai.Usage, error) { + return nil, nil, nil +} diff --git a/relay/channel/anthropic/main.go b/relay/channel/anthropic/main.go index a4272d7b..060fcde8 100644 --- a/relay/channel/anthropic/main.go +++ b/relay/channel/anthropic/main.go @@ -8,6 +8,8 @@ import ( "io" "net/http" "one-api/common" + "one-api/common/helper" + "one-api/common/logger" "one-api/relay/channel/openai" "strings" ) @@ -78,9 +80,9 @@ func responseClaude2OpenAI(claudeResponse *Response) *openai.TextResponse { FinishReason: stopReasonClaude2OpenAI(claudeResponse.StopReason), } fullTextResponse := openai.TextResponse{ - Id: fmt.Sprintf("chatcmpl-%s", common.GetUUID()), + Id: fmt.Sprintf("chatcmpl-%s", helper.GetUUID()), Object: "chat.completion", - Created: common.GetTimestamp(), + Created: helper.GetTimestamp(), Choices: []openai.TextResponseChoice{choice}, } return &fullTextResponse @@ -88,8 +90,8 @@ func responseClaude2OpenAI(claudeResponse *Response) *openai.TextResponse { func StreamHandler(c *gin.Context, resp *http.Response) (*openai.ErrorWithStatusCode, string) { responseText := "" - responseId := fmt.Sprintf("chatcmpl-%s", common.GetUUID()) - createdTime := common.GetTimestamp() + responseId := fmt.Sprintf("chatcmpl-%s", helper.GetUUID()) + createdTime := helper.GetTimestamp() scanner := bufio.NewScanner(resp.Body) scanner.Split(func(data []byte, atEOF bool) (advance int, token []byte, err error) { if atEOF && len(data) == 0 { @@ -125,7 +127,7 @@ func StreamHandler(c *gin.Context, resp *http.Response) (*openai.ErrorWithStatus var claudeResponse Response err := json.Unmarshal([]byte(data), &claudeResponse) if err != nil { - common.SysError("error unmarshalling stream response: " + err.Error()) + logger.SysError("error unmarshalling stream response: " + err.Error()) return true } responseText += claudeResponse.Completion @@ -134,7 +136,7 @@ func StreamHandler(c *gin.Context, resp *http.Response) (*openai.ErrorWithStatus response.Created = createdTime jsonStr, err := json.Marshal(response) if err != nil { - common.SysError("error marshalling stream response: " + err.Error()) + logger.SysError("error marshalling stream response: " + err.Error()) return true } c.Render(-1, common.CustomEvent{Data: "data: " + string(jsonStr)}) diff --git a/relay/channel/baidu/adaptor.go b/relay/channel/baidu/adaptor.go new file mode 100644 index 00000000..498b664a --- /dev/null +++ b/relay/channel/baidu/adaptor.go @@ -0,0 +1,22 @@ +package baidu + +import ( + "github.com/gin-gonic/gin" + "net/http" + "one-api/relay/channel/openai" +) + +type Adaptor struct { +} + +func (a *Adaptor) Auth(c *gin.Context) error { + return nil +} + +func (a *Adaptor) ConvertRequest(request *openai.GeneralOpenAIRequest) (any, error) { + return nil, nil +} + +func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response) (*openai.ErrorWithStatusCode, *openai.Usage, error) { + return nil, nil, nil +} diff --git a/relay/channel/baidu/main.go b/relay/channel/baidu/main.go index 47969492..f5b98155 100644 --- a/relay/channel/baidu/main.go +++ b/relay/channel/baidu/main.go @@ -9,6 +9,7 @@ import ( "io" "net/http" "one-api/common" + "one-api/common/logger" "one-api/relay/channel/openai" "one-api/relay/constant" "one-api/relay/util" @@ -19,49 +20,49 @@ import ( // https://cloud.baidu.com/doc/WENXINWORKSHOP/s/flfmc9do2 -type BaiduTokenResponse struct { +type TokenResponse struct { ExpiresIn int `json:"expires_in"` AccessToken string `json:"access_token"` } -type BaiduMessage struct { +type Message struct { Role string `json:"role"` Content string `json:"content"` } -type BaiduChatRequest struct { - Messages []BaiduMessage `json:"messages"` - Stream bool `json:"stream"` - UserId string `json:"user_id,omitempty"` +type ChatRequest struct { + Messages []Message `json:"messages"` + Stream bool `json:"stream"` + UserId string `json:"user_id,omitempty"` } -type BaiduError struct { +type Error struct { ErrorCode int `json:"error_code"` ErrorMsg string `json:"error_msg"` } var baiduTokenStore sync.Map -func ConvertRequest(request openai.GeneralOpenAIRequest) *BaiduChatRequest { - messages := make([]BaiduMessage, 0, len(request.Messages)) +func ConvertRequest(request openai.GeneralOpenAIRequest) *ChatRequest { + messages := make([]Message, 0, len(request.Messages)) for _, message := range request.Messages { if message.Role == "system" { - messages = append(messages, BaiduMessage{ + messages = append(messages, Message{ Role: "user", Content: message.StringContent(), }) - messages = append(messages, BaiduMessage{ + messages = append(messages, Message{ Role: "assistant", Content: "Okay", }) } else { - messages = append(messages, BaiduMessage{ + messages = append(messages, Message{ Role: message.Role, Content: message.StringContent(), }) } } - return &BaiduChatRequest{ + return &ChatRequest{ Messages: messages, Stream: request.Stream, } @@ -160,7 +161,7 @@ func StreamHandler(c *gin.Context, resp *http.Response) (*openai.ErrorWithStatus var baiduResponse ChatStreamResponse err := json.Unmarshal([]byte(data), &baiduResponse) if err != nil { - common.SysError("error unmarshalling stream response: " + err.Error()) + logger.SysError("error unmarshalling stream response: " + err.Error()) return true } if baiduResponse.Usage.TotalTokens != 0 { @@ -171,7 +172,7 @@ func StreamHandler(c *gin.Context, resp *http.Response) (*openai.ErrorWithStatus response := streamResponseBaidu2OpenAI(&baiduResponse) jsonResponse, err := json.Marshal(response) if err != nil { - common.SysError("error marshalling stream response: " + err.Error()) + logger.SysError("error marshalling stream response: " + err.Error()) return true } c.Render(-1, common.CustomEvent{Data: "data: " + string(jsonResponse)}) diff --git a/relay/channel/baidu/model.go b/relay/channel/baidu/model.go index caaebafb..e182f5dd 100644 --- a/relay/channel/baidu/model.go +++ b/relay/channel/baidu/model.go @@ -13,7 +13,7 @@ type ChatResponse struct { IsTruncated bool `json:"is_truncated"` NeedClearHistory bool `json:"need_clear_history"` Usage openai.Usage `json:"usage"` - BaiduError + Error } type ChatStreamResponse struct { @@ -38,7 +38,7 @@ type EmbeddingResponse struct { Created int64 `json:"created"` Data []EmbeddingData `json:"data"` Usage openai.Usage `json:"usage"` - BaiduError + Error } type AccessToken struct { diff --git a/relay/channel/google/adaptor.go b/relay/channel/google/adaptor.go new file mode 100644 index 00000000..b328db32 --- /dev/null +++ b/relay/channel/google/adaptor.go @@ -0,0 +1,22 @@ +package google + +import ( + "github.com/gin-gonic/gin" + "net/http" + "one-api/relay/channel/openai" +) + +type Adaptor struct { +} + +func (a *Adaptor) Auth(c *gin.Context) error { + return nil +} + +func (a *Adaptor) ConvertRequest(request *openai.GeneralOpenAIRequest) (any, error) { + return nil, nil +} + +func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response) (*openai.ErrorWithStatusCode, *openai.Usage, error) { + return nil, nil, nil +} diff --git a/relay/channel/google/gemini.go b/relay/channel/google/gemini.go index f49caadf..3adc3fdd 100644 --- a/relay/channel/google/gemini.go +++ b/relay/channel/google/gemini.go @@ -7,7 +7,10 @@ import ( "io" "net/http" "one-api/common" + "one-api/common/config" + "one-api/common/helper" "one-api/common/image" + "one-api/common/logger" "one-api/relay/channel/openai" "one-api/relay/constant" "strings" @@ -28,19 +31,19 @@ func ConvertGeminiRequest(textRequest openai.GeneralOpenAIRequest) *GeminiChatRe SafetySettings: []GeminiChatSafetySettings{ { Category: "HARM_CATEGORY_HARASSMENT", - Threshold: common.GeminiSafetySetting, + Threshold: config.GeminiSafetySetting, }, { Category: "HARM_CATEGORY_HATE_SPEECH", - Threshold: common.GeminiSafetySetting, + Threshold: config.GeminiSafetySetting, }, { Category: "HARM_CATEGORY_SEXUALLY_EXPLICIT", - Threshold: common.GeminiSafetySetting, + Threshold: config.GeminiSafetySetting, }, { Category: "HARM_CATEGORY_DANGEROUS_CONTENT", - Threshold: common.GeminiSafetySetting, + Threshold: config.GeminiSafetySetting, }, }, GenerationConfig: GeminiChatGenerationConfig{ @@ -151,9 +154,9 @@ type GeminiChatPromptFeedback struct { func responseGeminiChat2OpenAI(response *GeminiChatResponse) *openai.TextResponse { fullTextResponse := openai.TextResponse{ - Id: fmt.Sprintf("chatcmpl-%s", common.GetUUID()), + Id: fmt.Sprintf("chatcmpl-%s", helper.GetUUID()), Object: "chat.completion", - Created: common.GetTimestamp(), + Created: helper.GetTimestamp(), Choices: make([]openai.TextResponseChoice, 0, len(response.Candidates)), } for i, candidate := range response.Candidates { @@ -229,15 +232,15 @@ func StreamHandler(c *gin.Context, resp *http.Response) (*openai.ErrorWithStatus var choice openai.ChatCompletionsStreamResponseChoice choice.Delta.Content = dummy.Content response := openai.ChatCompletionsStreamResponse{ - Id: fmt.Sprintf("chatcmpl-%s", common.GetUUID()), + Id: fmt.Sprintf("chatcmpl-%s", helper.GetUUID()), Object: "chat.completion.chunk", - Created: common.GetTimestamp(), + Created: helper.GetTimestamp(), Model: "gemini-pro", Choices: []openai.ChatCompletionsStreamResponseChoice{choice}, } jsonResponse, err := json.Marshal(response) if err != nil { - common.SysError("error marshalling stream response: " + err.Error()) + logger.SysError("error marshalling stream response: " + err.Error()) return true } c.Render(-1, common.CustomEvent{Data: "data: " + string(jsonResponse)}) diff --git a/relay/channel/google/palm.go b/relay/channel/google/palm.go index 77d8cbd6..3c86a432 100644 --- a/relay/channel/google/palm.go +++ b/relay/channel/google/palm.go @@ -7,6 +7,8 @@ import ( "io" "net/http" "one-api/common" + "one-api/common/helper" + "one-api/common/logger" "one-api/relay/channel/openai" "one-api/relay/constant" ) @@ -71,27 +73,27 @@ func streamResponsePaLM2OpenAI(palmResponse *PaLMChatResponse) *openai.ChatCompl func PaLMStreamHandler(c *gin.Context, resp *http.Response) (*openai.ErrorWithStatusCode, string) { responseText := "" - responseId := fmt.Sprintf("chatcmpl-%s", common.GetUUID()) - createdTime := common.GetTimestamp() + responseId := fmt.Sprintf("chatcmpl-%s", helper.GetUUID()) + createdTime := helper.GetTimestamp() dataChan := make(chan string) stopChan := make(chan bool) go func() { responseBody, err := io.ReadAll(resp.Body) if err != nil { - common.SysError("error reading stream response: " + err.Error()) + logger.SysError("error reading stream response: " + err.Error()) stopChan <- true return } err = resp.Body.Close() if err != nil { - common.SysError("error closing stream response: " + err.Error()) + logger.SysError("error closing stream response: " + err.Error()) stopChan <- true return } var palmResponse PaLMChatResponse err = json.Unmarshal(responseBody, &palmResponse) if err != nil { - common.SysError("error unmarshalling stream response: " + err.Error()) + logger.SysError("error unmarshalling stream response: " + err.Error()) stopChan <- true return } @@ -103,7 +105,7 @@ func PaLMStreamHandler(c *gin.Context, resp *http.Response) (*openai.ErrorWithSt } jsonResponse, err := json.Marshal(fullTextResponse) if err != nil { - common.SysError("error marshalling stream response: " + err.Error()) + logger.SysError("error marshalling stream response: " + err.Error()) stopChan <- true return } diff --git a/relay/channel/interface.go b/relay/channel/interface.go new file mode 100644 index 00000000..7a0fcbd3 --- /dev/null +++ b/relay/channel/interface.go @@ -0,0 +1,15 @@ +package channel + +import ( + "github.com/gin-gonic/gin" + "net/http" + "one-api/relay/channel/openai" +) + +type Adaptor interface { + GetRequestURL() string + Auth(c *gin.Context) error + ConvertRequest(request *openai.GeneralOpenAIRequest) (any, error) + DoRequest(request *openai.GeneralOpenAIRequest) error + DoResponse(c *gin.Context, resp *http.Response) (*openai.ErrorWithStatusCode, *openai.Usage, error) +} diff --git a/relay/channel/openai/adaptor.go b/relay/channel/openai/adaptor.go new file mode 100644 index 00000000..cc302611 --- /dev/null +++ b/relay/channel/openai/adaptor.go @@ -0,0 +1,21 @@ +package openai + +import ( + "github.com/gin-gonic/gin" + "net/http" +) + +type Adaptor struct { +} + +func (a *Adaptor) Auth(c *gin.Context) error { + return nil +} + +func (a *Adaptor) ConvertRequest(request *GeneralOpenAIRequest) (any, error) { + return nil, nil +} + +func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response) (*ErrorWithStatusCode, *Usage, error) { + return nil, nil, nil +} diff --git a/relay/channel/openai/main.go b/relay/channel/openai/main.go index 848a6fa4..5f464249 100644 --- a/relay/channel/openai/main.go +++ b/relay/channel/openai/main.go @@ -8,6 +8,7 @@ import ( "io" "net/http" "one-api/common" + "one-api/common/logger" "one-api/relay/constant" "strings" ) @@ -46,7 +47,7 @@ func StreamHandler(c *gin.Context, resp *http.Response, relayMode int) (*ErrorWi var streamResponse ChatCompletionsStreamResponse err := json.Unmarshal([]byte(data), &streamResponse) if err != nil { - common.SysError("error unmarshalling stream response: " + err.Error()) + logger.SysError("error unmarshalling stream response: " + err.Error()) continue // just ignore the error } for _, choice := range streamResponse.Choices { @@ -56,7 +57,7 @@ func StreamHandler(c *gin.Context, resp *http.Response, relayMode int) (*ErrorWi var streamResponse CompletionsStreamResponse err := json.Unmarshal([]byte(data), &streamResponse) if err != nil { - common.SysError("error unmarshalling stream response: " + err.Error()) + logger.SysError("error unmarshalling stream response: " + err.Error()) continue } for _, choice := range streamResponse.Choices { diff --git a/relay/channel/openai/model.go b/relay/channel/openai/model.go index c831ce19..937fb424 100644 --- a/relay/channel/openai/model.go +++ b/relay/channel/openai/model.go @@ -207,6 +207,11 @@ type Usage struct { TotalTokens int `json:"total_tokens"` } +type UsageOrResponseText struct { + *Usage + ResponseText string +} + type Error struct { Message string `json:"message"` Type string `json:"type"` diff --git a/relay/channel/openai/token.go b/relay/channel/openai/token.go index 4b40b228..6803770e 100644 --- a/relay/channel/openai/token.go +++ b/relay/channel/openai/token.go @@ -6,7 +6,9 @@ import ( "github.com/pkoukk/tiktoken-go" "math" "one-api/common" + "one-api/common/config" "one-api/common/image" + "one-api/common/logger" "strings" ) @@ -15,15 +17,15 @@ var tokenEncoderMap = map[string]*tiktoken.Tiktoken{} var defaultTokenEncoder *tiktoken.Tiktoken func InitTokenEncoders() { - common.SysLog("initializing token encoders") + logger.SysLog("initializing token encoders") gpt35TokenEncoder, err := tiktoken.EncodingForModel("gpt-3.5-turbo") if err != nil { - common.FatalLog(fmt.Sprintf("failed to get gpt-3.5-turbo token encoder: %s", err.Error())) + logger.FatalLog(fmt.Sprintf("failed to get gpt-3.5-turbo token encoder: %s", err.Error())) } defaultTokenEncoder = gpt35TokenEncoder gpt4TokenEncoder, err := tiktoken.EncodingForModel("gpt-4") if err != nil { - common.FatalLog(fmt.Sprintf("failed to get gpt-4 token encoder: %s", err.Error())) + logger.FatalLog(fmt.Sprintf("failed to get gpt-4 token encoder: %s", err.Error())) } for model, _ := range common.ModelRatio { if strings.HasPrefix(model, "gpt-3.5") { @@ -34,7 +36,7 @@ func InitTokenEncoders() { tokenEncoderMap[model] = nil } } - common.SysLog("token encoders initialized") + logger.SysLog("token encoders initialized") } func getTokenEncoder(model string) *tiktoken.Tiktoken { @@ -45,7 +47,7 @@ func getTokenEncoder(model string) *tiktoken.Tiktoken { if ok { tokenEncoder, err := tiktoken.EncodingForModel(model) if err != nil { - common.SysError(fmt.Sprintf("failed to get token encoder for model %s: %s, using encoder for gpt-3.5-turbo", model, err.Error())) + logger.SysError(fmt.Sprintf("failed to get token encoder for model %s: %s, using encoder for gpt-3.5-turbo", model, err.Error())) tokenEncoder = defaultTokenEncoder } tokenEncoderMap[model] = tokenEncoder @@ -55,7 +57,7 @@ func getTokenEncoder(model string) *tiktoken.Tiktoken { } func getTokenNum(tokenEncoder *tiktoken.Tiktoken, text string) int { - if common.ApproximateTokenEnabled { + if config.ApproximateTokenEnabled { return int(float64(len(text)) * 0.38) } return len(tokenEncoder.Encode(text, nil, nil)) @@ -99,7 +101,7 @@ func CountTokenMessages(messages []Message, model string) int { } imageTokens, err := countImageTokens(url, detail) if err != nil { - common.SysError("error counting image tokens: " + err.Error()) + logger.SysError("error counting image tokens: " + err.Error()) } else { tokenNum += imageTokens } diff --git a/relay/channel/tencent/adaptor.go b/relay/channel/tencent/adaptor.go new file mode 100644 index 00000000..e9f86aff --- /dev/null +++ b/relay/channel/tencent/adaptor.go @@ -0,0 +1,22 @@ +package tencent + +import ( + "github.com/gin-gonic/gin" + "net/http" + "one-api/relay/channel/openai" +) + +type Adaptor struct { +} + +func (a *Adaptor) Auth(c *gin.Context) error { + return nil +} + +func (a *Adaptor) ConvertRequest(request *openai.GeneralOpenAIRequest) (any, error) { + return nil, nil +} + +func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response) (*openai.ErrorWithStatusCode, *openai.Usage, error) { + return nil, nil, nil +} diff --git a/relay/channel/tencent/main.go b/relay/channel/tencent/main.go index 60e275a9..64091d62 100644 --- a/relay/channel/tencent/main.go +++ b/relay/channel/tencent/main.go @@ -12,6 +12,8 @@ import ( "io" "net/http" "one-api/common" + "one-api/common/helper" + "one-api/common/logger" "one-api/relay/channel/openai" "one-api/relay/constant" "sort" @@ -46,9 +48,9 @@ func ConvertRequest(request openai.GeneralOpenAIRequest) *ChatRequest { stream = 1 } return &ChatRequest{ - Timestamp: common.GetTimestamp(), - Expired: common.GetTimestamp() + 24*60*60, - QueryID: common.GetUUID(), + Timestamp: helper.GetTimestamp(), + Expired: helper.GetTimestamp() + 24*60*60, + QueryID: helper.GetUUID(), Temperature: request.Temperature, TopP: request.TopP, Stream: stream, @@ -59,7 +61,7 @@ func ConvertRequest(request openai.GeneralOpenAIRequest) *ChatRequest { func responseTencent2OpenAI(response *ChatResponse) *openai.TextResponse { fullTextResponse := openai.TextResponse{ Object: "chat.completion", - Created: common.GetTimestamp(), + Created: helper.GetTimestamp(), Usage: response.Usage, } if len(response.Choices) > 0 { @@ -79,7 +81,7 @@ func responseTencent2OpenAI(response *ChatResponse) *openai.TextResponse { func streamResponseTencent2OpenAI(TencentResponse *ChatResponse) *openai.ChatCompletionsStreamResponse { response := openai.ChatCompletionsStreamResponse{ Object: "chat.completion.chunk", - Created: common.GetTimestamp(), + Created: helper.GetTimestamp(), Model: "tencent-hunyuan", } if len(TencentResponse.Choices) > 0 { @@ -131,7 +133,7 @@ func StreamHandler(c *gin.Context, resp *http.Response) (*openai.ErrorWithStatus var TencentResponse ChatResponse err := json.Unmarshal([]byte(data), &TencentResponse) if err != nil { - common.SysError("error unmarshalling stream response: " + err.Error()) + logger.SysError("error unmarshalling stream response: " + err.Error()) return true } response := streamResponseTencent2OpenAI(&TencentResponse) @@ -140,7 +142,7 @@ func StreamHandler(c *gin.Context, resp *http.Response) (*openai.ErrorWithStatus } jsonResponse, err := json.Marshal(response) if err != nil { - common.SysError("error marshalling stream response: " + err.Error()) + logger.SysError("error marshalling stream response: " + err.Error()) return true } c.Render(-1, common.CustomEvent{Data: "data: " + string(jsonResponse)}) diff --git a/relay/channel/xunfei/adaptor.go b/relay/channel/xunfei/adaptor.go new file mode 100644 index 00000000..89e58485 --- /dev/null +++ b/relay/channel/xunfei/adaptor.go @@ -0,0 +1,22 @@ +package xunfei + +import ( + "github.com/gin-gonic/gin" + "net/http" + "one-api/relay/channel/openai" +) + +type Adaptor struct { +} + +func (a *Adaptor) Auth(c *gin.Context) error { + return nil +} + +func (a *Adaptor) ConvertRequest(request *openai.GeneralOpenAIRequest) (any, error) { + return nil, nil +} + +func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response) (*openai.ErrorWithStatusCode, *openai.Usage, error) { + return nil, nil, nil +} diff --git a/relay/channel/xunfei/main.go b/relay/channel/xunfei/main.go index 1cc0b664..906d2844 100644 --- a/relay/channel/xunfei/main.go +++ b/relay/channel/xunfei/main.go @@ -12,6 +12,8 @@ import ( "net/http" "net/url" "one-api/common" + "one-api/common/helper" + "one-api/common/logger" "one-api/relay/channel/openai" "one-api/relay/constant" "strings" @@ -68,7 +70,7 @@ func responseXunfei2OpenAI(response *ChatResponse) *openai.TextResponse { } fullTextResponse := openai.TextResponse{ Object: "chat.completion", - Created: common.GetTimestamp(), + Created: helper.GetTimestamp(), Choices: []openai.TextResponseChoice{choice}, Usage: response.Payload.Usage.Text, } @@ -90,7 +92,7 @@ func streamResponseXunfei2OpenAI(xunfeiResponse *ChatResponse) *openai.ChatCompl } response := openai.ChatCompletionsStreamResponse{ Object: "chat.completion.chunk", - Created: common.GetTimestamp(), + Created: helper.GetTimestamp(), Model: "SparkDesk", Choices: []openai.ChatCompletionsStreamResponseChoice{choice}, } @@ -140,7 +142,7 @@ func StreamHandler(c *gin.Context, textRequest openai.GeneralOpenAIRequest, appI response := streamResponseXunfei2OpenAI(&xunfeiResponse) jsonResponse, err := json.Marshal(response) if err != nil { - common.SysError("error marshalling stream response: " + err.Error()) + logger.SysError("error marshalling stream response: " + err.Error()) return true } c.Render(-1, common.CustomEvent{Data: "data: " + string(jsonResponse)}) @@ -215,20 +217,20 @@ func xunfeiMakeRequest(textRequest openai.GeneralOpenAIRequest, domain, authUrl, for { _, msg, err := conn.ReadMessage() if err != nil { - common.SysError("error reading stream response: " + err.Error()) + logger.SysError("error reading stream response: " + err.Error()) break } var response ChatResponse err = json.Unmarshal(msg, &response) if err != nil { - common.SysError("error unmarshalling stream response: " + err.Error()) + logger.SysError("error unmarshalling stream response: " + err.Error()) break } dataChan <- response if response.Payload.Choices.Status == 2 { err := conn.Close() if err != nil { - common.SysError("error closing websocket connection: " + err.Error()) + logger.SysError("error closing websocket connection: " + err.Error()) } break } @@ -247,7 +249,7 @@ func getXunfeiAuthUrl(c *gin.Context, apiKey string, apiSecret string) (string, } if apiVersion == "" { apiVersion = "v1.1" - common.SysLog("api_version not found, use default: " + apiVersion) + logger.SysLog("api_version not found, use default: " + apiVersion) } domain := "general" if apiVersion != "v1.1" { diff --git a/relay/channel/zhipu/adaptor.go b/relay/channel/zhipu/adaptor.go new file mode 100644 index 00000000..6d901bc3 --- /dev/null +++ b/relay/channel/zhipu/adaptor.go @@ -0,0 +1,22 @@ +package zhipu + +import ( + "github.com/gin-gonic/gin" + "net/http" + "one-api/relay/channel/openai" +) + +type Adaptor struct { +} + +func (a *Adaptor) Auth(c *gin.Context) error { + return nil +} + +func (a *Adaptor) ConvertRequest(request *openai.GeneralOpenAIRequest) (any, error) { + return nil, nil +} + +func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response) (*openai.ErrorWithStatusCode, *openai.Usage, error) { + return nil, nil, nil +} diff --git a/relay/channel/zhipu/main.go b/relay/channel/zhipu/main.go index 3dc613a4..d831f57a 100644 --- a/relay/channel/zhipu/main.go +++ b/relay/channel/zhipu/main.go @@ -8,6 +8,8 @@ import ( "io" "net/http" "one-api/common" + "one-api/common/helper" + "one-api/common/logger" "one-api/relay/channel/openai" "one-api/relay/constant" "strings" @@ -34,7 +36,7 @@ func GetToken(apikey string) string { split := strings.Split(apikey, ".") if len(split) != 2 { - common.SysError("invalid zhipu key: " + apikey) + logger.SysError("invalid zhipu key: " + apikey) return "" } @@ -101,7 +103,7 @@ func responseZhipu2OpenAI(response *Response) *openai.TextResponse { fullTextResponse := openai.TextResponse{ Id: response.Data.TaskId, Object: "chat.completion", - Created: common.GetTimestamp(), + Created: helper.GetTimestamp(), Choices: make([]openai.TextResponseChoice, 0, len(response.Data.Choices)), Usage: response.Data.Usage, } @@ -127,7 +129,7 @@ func streamResponseZhipu2OpenAI(zhipuResponse string) *openai.ChatCompletionsStr choice.Delta.Content = zhipuResponse response := openai.ChatCompletionsStreamResponse{ Object: "chat.completion.chunk", - Created: common.GetTimestamp(), + Created: helper.GetTimestamp(), Model: "chatglm", Choices: []openai.ChatCompletionsStreamResponseChoice{choice}, } @@ -141,7 +143,7 @@ func streamMetaResponseZhipu2OpenAI(zhipuResponse *StreamMetaResponse) (*openai. response := openai.ChatCompletionsStreamResponse{ Id: zhipuResponse.RequestId, Object: "chat.completion.chunk", - Created: common.GetTimestamp(), + Created: helper.GetTimestamp(), Model: "chatglm", Choices: []openai.ChatCompletionsStreamResponseChoice{choice}, } @@ -193,7 +195,7 @@ func StreamHandler(c *gin.Context, resp *http.Response) (*openai.ErrorWithStatus response := streamResponseZhipu2OpenAI(data) jsonResponse, err := json.Marshal(response) if err != nil { - common.SysError("error marshalling stream response: " + err.Error()) + logger.SysError("error marshalling stream response: " + err.Error()) return true } c.Render(-1, common.CustomEvent{Data: "data: " + string(jsonResponse)}) @@ -202,13 +204,13 @@ func StreamHandler(c *gin.Context, resp *http.Response) (*openai.ErrorWithStatus var zhipuResponse StreamMetaResponse err := json.Unmarshal([]byte(data), &zhipuResponse) if err != nil { - common.SysError("error unmarshalling stream response: " + err.Error()) + logger.SysError("error unmarshalling stream response: " + err.Error()) return true } response, zhipuUsage := streamMetaResponseZhipu2OpenAI(&zhipuResponse) jsonResponse, err := json.Marshal(response) if err != nil { - common.SysError("error marshalling stream response: " + err.Error()) + logger.SysError("error marshalling stream response: " + err.Error()) return true } usage = zhipuUsage diff --git a/relay/constant/api_type.go b/relay/constant/api_type.go new file mode 100644 index 00000000..658bfb90 --- /dev/null +++ b/relay/constant/api_type.go @@ -0,0 +1,69 @@ +package constant + +import ( + "one-api/common" +) + +const ( + APITypeOpenAI = iota + APITypeClaude + APITypePaLM + APITypeBaidu + APITypeZhipu + APITypeAli + APITypeXunfei + APITypeAIProxyLibrary + APITypeTencent + APITypeGemini +) + +func ChannelType2APIType(channelType int) int { + apiType := APITypeOpenAI + switch channelType { + case common.ChannelTypeAnthropic: + apiType = APITypeClaude + case common.ChannelTypeBaidu: + apiType = APITypeBaidu + case common.ChannelTypePaLM: + apiType = APITypePaLM + case common.ChannelTypeZhipu: + apiType = APITypeZhipu + case common.ChannelTypeAli: + apiType = APITypeAli + case common.ChannelTypeXunfei: + apiType = APITypeXunfei + case common.ChannelTypeAIProxyLibrary: + apiType = APITypeAIProxyLibrary + case common.ChannelTypeTencent: + apiType = APITypeTencent + case common.ChannelTypeGemini: + apiType = APITypeGemini + } + return apiType +} + +//func GetAdaptor(apiType int) channel.Adaptor { +// switch apiType { +// case APITypeOpenAI: +// return &openai.Adaptor{} +// case APITypeClaude: +// return &anthropic.Adaptor{} +// case APITypePaLM: +// return &google.Adaptor{} +// case APITypeZhipu: +// return &baidu.Adaptor{} +// case APITypeBaidu: +// return &baidu.Adaptor{} +// case APITypeAli: +// return &ali.Adaptor{} +// case APITypeXunfei: +// return &xunfei.Adaptor{} +// case APITypeAIProxyLibrary: +// return &aiproxy.Adaptor{} +// case APITypeTencent: +// return &tencent.Adaptor{} +// case APITypeGemini: +// return &google.Adaptor{} +// } +// return nil +//} diff --git a/relay/constant/common.go b/relay/constant/common.go new file mode 100644 index 00000000..b6606cc6 --- /dev/null +++ b/relay/constant/common.go @@ -0,0 +1,3 @@ +package constant + +var StopFinishReason = "stop" diff --git a/relay/constant/main.go b/relay/constant/main.go deleted file mode 100644 index b3aeaaff..00000000 --- a/relay/constant/main.go +++ /dev/null @@ -1,16 +0,0 @@ -package constant - -const ( - RelayModeUnknown = iota - RelayModeChatCompletions - RelayModeCompletions - RelayModeEmbeddings - RelayModeModerations - RelayModeImagesGenerations - RelayModeEdits - RelayModeAudioSpeech - RelayModeAudioTranscription - RelayModeAudioTranslation -) - -var StopFinishReason = "stop" diff --git a/relay/constant/relay_mode.go b/relay/constant/relay_mode.go new file mode 100644 index 00000000..5e2fe574 --- /dev/null +++ b/relay/constant/relay_mode.go @@ -0,0 +1,42 @@ +package constant + +import "strings" + +const ( + RelayModeUnknown = iota + RelayModeChatCompletions + RelayModeCompletions + RelayModeEmbeddings + RelayModeModerations + RelayModeImagesGenerations + RelayModeEdits + RelayModeAudioSpeech + RelayModeAudioTranscription + RelayModeAudioTranslation +) + +func Path2RelayMode(path string) int { + relayMode := RelayModeUnknown + if strings.HasPrefix(path, "/v1/chat/completions") { + relayMode = RelayModeChatCompletions + } else if strings.HasPrefix(path, "/v1/completions") { + relayMode = RelayModeCompletions + } else if strings.HasPrefix(path, "/v1/embeddings") { + relayMode = RelayModeEmbeddings + } else if strings.HasSuffix(path, "embeddings") { + relayMode = RelayModeEmbeddings + } else if strings.HasPrefix(path, "/v1/moderations") { + relayMode = RelayModeModerations + } else if strings.HasPrefix(path, "/v1/images/generations") { + relayMode = RelayModeImagesGenerations + } else if strings.HasPrefix(path, "/v1/edits") { + relayMode = RelayModeEdits + } else if strings.HasPrefix(path, "/v1/audio/speech") { + relayMode = RelayModeAudioSpeech + } else if strings.HasPrefix(path, "/v1/audio/transcriptions") { + relayMode = RelayModeAudioTranscription + } else if strings.HasPrefix(path, "/v1/audio/translations") { + relayMode = RelayModeAudioTranslation + } + return relayMode +} diff --git a/relay/controller/audio.go b/relay/controller/audio.go index 08d9af2a..822d7e39 100644 --- a/relay/controller/audio.go +++ b/relay/controller/audio.go @@ -11,6 +11,8 @@ import ( "io" "net/http" "one-api/common" + "one-api/common/config" + "one-api/common/logger" "one-api/model" "one-api/relay/channel/openai" "one-api/relay/constant" @@ -53,7 +55,7 @@ func RelayAudioHelper(c *gin.Context, relayMode int) *openai.ErrorWithStatusCode preConsumedQuota = int(float64(len(ttsRequest.Input)) * ratio) quota = preConsumedQuota default: - preConsumedQuota = int(float64(common.PreConsumedQuota) * ratio) + preConsumedQuota = int(float64(config.PreConsumedQuota) * ratio) } userQuota, err := model.CacheGetUserQuota(userId) if err != nil { @@ -102,7 +104,7 @@ func RelayAudioHelper(c *gin.Context, relayMode int) *openai.ErrorWithStatusCode fullRequestURL := util.GetFullRequestURL(baseURL, requestURL, channelType) if relayMode == constant.RelayModeAudioTranscription && channelType == common.ChannelTypeAzure { // https://learn.microsoft.com/en-us/azure/ai-services/openai/whisper-quickstart?tabs=command-line#rest-api - apiVersion := util.GetAPIVersion(c) + apiVersion := util.GetAzureAPIVersion(c) fullRequestURL = fmt.Sprintf("%s/openai/deployments/%s/audio/transcriptions?api-version=%s", baseURL, audioModel, apiVersion) } @@ -191,7 +193,7 @@ func RelayAudioHelper(c *gin.Context, relayMode int) *openai.ErrorWithStatusCode // negative means add quota back for token & user err := model.PostConsumeTokenQuota(tokenId, -preConsumedQuota) if err != nil { - common.LogError(ctx, fmt.Sprintf("error rollback pre-consumed quota: %s", err.Error())) + logger.Error(ctx, fmt.Sprintf("error rollback pre-consumed quota: %s", err.Error())) } }() }(c.Request.Context()) diff --git a/relay/controller/image.go b/relay/controller/image.go index be5fc3dd..9502a4d7 100644 --- a/relay/controller/image.go +++ b/relay/controller/image.go @@ -9,6 +9,7 @@ import ( "io" "net/http" "one-api/common" + "one-api/common/logger" "one-api/model" "one-api/relay/channel/openai" "one-api/relay/util" @@ -112,7 +113,7 @@ func RelayImageHelper(c *gin.Context, relayMode int) *openai.ErrorWithStatusCode fullRequestURL := util.GetFullRequestURL(baseURL, requestURL, channelType) if channelType == common.ChannelTypeAzure { // https://learn.microsoft.com/en-us/azure/ai-services/openai/dall-e-quickstart?tabs=dalle3%2Ccommand-line&pivots=rest-api - apiVersion := util.GetAPIVersion(c) + apiVersion := util.GetAzureAPIVersion(c) // https://{resource_name}.openai.azure.com/openai/deployments/dall-e-3/images/generations?api-version=2023-06-01-preview fullRequestURL = fmt.Sprintf("%s/openai/deployments/%s/images/generations?api-version=%s", baseURL, imageModel, apiVersion) } @@ -175,11 +176,11 @@ func RelayImageHelper(c *gin.Context, relayMode int) *openai.ErrorWithStatusCode } err := model.PostConsumeTokenQuota(tokenId, quota) if err != nil { - common.SysError("error consuming token remain quota: " + err.Error()) + logger.SysError("error consuming token remain quota: " + err.Error()) } err = model.CacheUpdateUserQuota(userId) if err != nil { - common.SysError("error update user quota cache: " + err.Error()) + logger.SysError("error update user quota cache: " + err.Error()) } if quota != 0 { tokenName := c.GetString("token_name") diff --git a/relay/controller/text.go b/relay/controller/text.go index b17ff950..68354628 100644 --- a/relay/controller/text.go +++ b/relay/controller/text.go @@ -1,206 +1,47 @@ package controller import ( - "bytes" "context" - "encoding/json" "errors" "fmt" "github.com/gin-gonic/gin" - "io" "math" "net/http" "one-api/common" + "one-api/common/config" + "one-api/common/logger" "one-api/model" - "one-api/relay/channel/aiproxy" - "one-api/relay/channel/ali" - "one-api/relay/channel/anthropic" - "one-api/relay/channel/baidu" - "one-api/relay/channel/google" "one-api/relay/channel/openai" - "one-api/relay/channel/tencent" - "one-api/relay/channel/xunfei" - "one-api/relay/channel/zhipu" "one-api/relay/constant" "one-api/relay/util" "strings" ) -const ( - APITypeOpenAI = iota - APITypeClaude - APITypePaLM - APITypeBaidu - APITypeZhipu - APITypeAli - APITypeXunfei - APITypeAIProxyLibrary - APITypeTencent - APITypeGemini -) - func RelayTextHelper(c *gin.Context, relayMode int) *openai.ErrorWithStatusCode { - channelType := c.GetInt("channel") - channelId := c.GetInt("channel_id") - tokenId := c.GetInt("token_id") - userId := c.GetInt("id") - group := c.GetString("group") + ctx := c.Request.Context() + meta := util.GetRelayMeta(c) var textRequest openai.GeneralOpenAIRequest err := common.UnmarshalBodyReusable(c, &textRequest) if err != nil { return openai.ErrorWrapper(err, "bind_request_body_failed", http.StatusBadRequest) } - if textRequest.MaxTokens < 0 || textRequest.MaxTokens > math.MaxInt32/2 { - return openai.ErrorWrapper(errors.New("max_tokens is invalid"), "invalid_max_tokens", http.StatusBadRequest) - } if relayMode == constant.RelayModeModerations && textRequest.Model == "" { textRequest.Model = "text-moderation-latest" } if relayMode == constant.RelayModeEmbeddings && textRequest.Model == "" { textRequest.Model = c.Param("model") } - // request validation - if textRequest.Model == "" { - return openai.ErrorWrapper(errors.New("model is required"), "required_field_missing", http.StatusBadRequest) + err = util.ValidateTextRequest(&textRequest, relayMode) + if err != nil { + return openai.ErrorWrapper(err, "invalid_text_request", http.StatusBadRequest) } - switch relayMode { - case constant.RelayModeCompletions: - if textRequest.Prompt == "" { - return openai.ErrorWrapper(errors.New("field prompt is required"), "required_field_missing", http.StatusBadRequest) - } - case constant.RelayModeChatCompletions: - if textRequest.Messages == nil || len(textRequest.Messages) == 0 { - return openai.ErrorWrapper(errors.New("field messages is required"), "required_field_missing", http.StatusBadRequest) - } - case constant.RelayModeEmbeddings: - case constant.RelayModeModerations: - if textRequest.Input == "" { - return openai.ErrorWrapper(errors.New("field input is required"), "required_field_missing", http.StatusBadRequest) - } - case constant.RelayModeEdits: - if textRequest.Instruction == "" { - return openai.ErrorWrapper(errors.New("field instruction is required"), "required_field_missing", http.StatusBadRequest) - } - } - // map model name - modelMapping := c.GetString("model_mapping") - isModelMapped := false - if modelMapping != "" && modelMapping != "{}" { - modelMap := make(map[string]string) - err := json.Unmarshal([]byte(modelMapping), &modelMap) - if err != nil { - return openai.ErrorWrapper(err, "unmarshal_model_mapping_failed", http.StatusInternalServerError) - } - if modelMap[textRequest.Model] != "" { - textRequest.Model = modelMap[textRequest.Model] - isModelMapped = true - } - } - apiType := APITypeOpenAI - switch channelType { - case common.ChannelTypeAnthropic: - apiType = APITypeClaude - case common.ChannelTypeBaidu: - apiType = APITypeBaidu - case common.ChannelTypePaLM: - apiType = APITypePaLM - case common.ChannelTypeZhipu: - apiType = APITypeZhipu - case common.ChannelTypeAli: - apiType = APITypeAli - case common.ChannelTypeXunfei: - apiType = APITypeXunfei - case common.ChannelTypeAIProxyLibrary: - apiType = APITypeAIProxyLibrary - case common.ChannelTypeTencent: - apiType = APITypeTencent - case common.ChannelTypeGemini: - apiType = APITypeGemini - } - baseURL := common.ChannelBaseURLs[channelType] - requestURL := c.Request.URL.String() - if c.GetString("base_url") != "" { - baseURL = c.GetString("base_url") - } - fullRequestURL := util.GetFullRequestURL(baseURL, requestURL, channelType) - switch apiType { - case APITypeOpenAI: - if channelType == common.ChannelTypeAzure { - // https://learn.microsoft.com/en-us/azure/cognitive-services/openai/chatgpt-quickstart?pivots=rest-api&tabs=command-line#rest-api - apiVersion := util.GetAPIVersion(c) - requestURL := strings.Split(requestURL, "?")[0] - requestURL = fmt.Sprintf("%s?api-version=%s", requestURL, apiVersion) - baseURL = c.GetString("base_url") - task := strings.TrimPrefix(requestURL, "/v1/") - model_ := textRequest.Model - model_ = strings.Replace(model_, ".", "", -1) - // https://github.com/songquanpeng/one-api/issues/67 - model_ = strings.TrimSuffix(model_, "-0301") - model_ = strings.TrimSuffix(model_, "-0314") - model_ = strings.TrimSuffix(model_, "-0613") - - requestURL = fmt.Sprintf("/openai/deployments/%s/%s", model_, task) - fullRequestURL = util.GetFullRequestURL(baseURL, requestURL, channelType) - } - case APITypeClaude: - fullRequestURL = "https://api.anthropic.com/v1/complete" - if baseURL != "" { - fullRequestURL = fmt.Sprintf("%s/v1/complete", baseURL) - } - case APITypeBaidu: - switch textRequest.Model { - case "ERNIE-Bot": - fullRequestURL = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions" - case "ERNIE-Bot-turbo": - fullRequestURL = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/eb-instant" - case "ERNIE-Bot-4": - fullRequestURL = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions_pro" - case "BLOOMZ-7B": - fullRequestURL = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/bloomz_7b1" - case "Embedding-V1": - fullRequestURL = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/embeddings/embedding-v1" - } - apiKey := c.Request.Header.Get("Authorization") - apiKey = strings.TrimPrefix(apiKey, "Bearer ") - var err error - if apiKey, err = baidu.GetAccessToken(apiKey); err != nil { - return openai.ErrorWrapper(err, "invalid_baidu_config", http.StatusInternalServerError) - } - fullRequestURL += "?access_token=" + apiKey - case APITypePaLM: - fullRequestURL = "https://generativelanguage.googleapis.com/v1beta2/models/chat-bison-001:generateMessage" - if baseURL != "" { - fullRequestURL = fmt.Sprintf("%s/v1beta2/models/chat-bison-001:generateMessage", baseURL) - } - case APITypeGemini: - requestBaseURL := "https://generativelanguage.googleapis.com" - if baseURL != "" { - requestBaseURL = baseURL - } - version := "v1" - if c.GetString("api_version") != "" { - version = c.GetString("api_version") - } - action := "generateContent" - if textRequest.Stream { - action = "streamGenerateContent" - } - fullRequestURL = fmt.Sprintf("%s/%s/models/%s:%s", requestBaseURL, version, textRequest.Model, action) - case APITypeZhipu: - method := "invoke" - if textRequest.Stream { - method = "sse-invoke" - } - fullRequestURL = fmt.Sprintf("https://open.bigmodel.cn/api/paas/v3/model-api/%s/%s", textRequest.Model, method) - case APITypeAli: - fullRequestURL = "https://dashscope.aliyuncs.com/api/v1/services/aigc/text-generation/generation" - if relayMode == constant.RelayModeEmbeddings { - fullRequestURL = "https://dashscope.aliyuncs.com/api/v1/services/embeddings/text-embedding/text-embedding" - } - case APITypeTencent: - fullRequestURL = "https://hunyuan.cloud.tencent.com/hyllm/v1/chat/completions" - case APITypeAIProxyLibrary: - fullRequestURL = fmt.Sprintf("%s/api/library/ask", baseURL) + var isModelMapped bool + textRequest.Model, isModelMapped = util.GetMappedModelName(textRequest.Model, meta.ModelMapping) + apiType := constant.ChannelType2APIType(meta.ChannelType) + fullRequestURL, err := GetRequestURL(c.Request.URL.String(), apiType, relayMode, meta, &textRequest) + if err != nil { + logger.Error(ctx, fmt.Sprintf("util.GetRequestURL failed: %s", err.Error())) + return openai.ErrorWrapper(fmt.Errorf("util.GetRequestURL failed"), "get_request_url_failed", http.StatusInternalServerError) } var promptTokens int var completionTokens int @@ -212,22 +53,22 @@ func RelayTextHelper(c *gin.Context, relayMode int) *openai.ErrorWithStatusCode case constant.RelayModeModerations: promptTokens = openai.CountTokenInput(textRequest.Input, textRequest.Model) } - preConsumedTokens := common.PreConsumedQuota + preConsumedTokens := config.PreConsumedQuota if textRequest.MaxTokens != 0 { preConsumedTokens = promptTokens + textRequest.MaxTokens } modelRatio := common.GetModelRatio(textRequest.Model) - groupRatio := common.GetGroupRatio(group) + groupRatio := common.GetGroupRatio(meta.Group) ratio := modelRatio * groupRatio preConsumedQuota := int(float64(preConsumedTokens) * ratio) - userQuota, err := model.CacheGetUserQuota(userId) + userQuota, err := model.CacheGetUserQuota(meta.UserId) if err != nil { return openai.ErrorWrapper(err, "get_user_quota_failed", http.StatusInternalServerError) } if userQuota-preConsumedQuota < 0 { return openai.ErrorWrapper(errors.New("user quota is not enough"), "insufficient_user_quota", http.StatusForbidden) } - err = model.CacheDecreaseUserQuota(userId, preConsumedQuota) + err = model.CacheDecreaseUserQuota(meta.UserId, preConsumedQuota) if err != nil { return openai.ErrorWrapper(err, "decrease_user_quota_failed", http.StatusInternalServerError) } @@ -235,165 +76,28 @@ func RelayTextHelper(c *gin.Context, relayMode int) *openai.ErrorWithStatusCode // in this case, we do not pre-consume quota // because the user has enough quota preConsumedQuota = 0 - common.LogInfo(c.Request.Context(), fmt.Sprintf("user %d has enough quota %d, trusted and no need to pre-consume", userId, userQuota)) + logger.Info(c.Request.Context(), fmt.Sprintf("user %d has enough quota %d, trusted and no need to pre-consume", meta.UserId, userQuota)) } if preConsumedQuota > 0 { - err := model.PreConsumeTokenQuota(tokenId, preConsumedQuota) + err := model.PreConsumeTokenQuota(meta.TokenId, preConsumedQuota) if err != nil { return openai.ErrorWrapper(err, "pre_consume_token_quota_failed", http.StatusForbidden) } } - var requestBody io.Reader - if isModelMapped { - jsonStr, err := json.Marshal(textRequest) - if err != nil { - return openai.ErrorWrapper(err, "marshal_text_request_failed", http.StatusInternalServerError) - } - requestBody = bytes.NewBuffer(jsonStr) - } else { - requestBody = c.Request.Body + requestBody, err := GetRequestBody(c, textRequest, isModelMapped, apiType, relayMode) + if err != nil { + return openai.ErrorWrapper(err, "get_request_body_failed", http.StatusInternalServerError) } - switch apiType { - case APITypeClaude: - claudeRequest := anthropic.ConvertRequest(textRequest) - jsonStr, err := json.Marshal(claudeRequest) - if err != nil { - return openai.ErrorWrapper(err, "marshal_text_request_failed", http.StatusInternalServerError) - } - requestBody = bytes.NewBuffer(jsonStr) - case APITypeBaidu: - var jsonData []byte - var err error - switch relayMode { - case constant.RelayModeEmbeddings: - baiduEmbeddingRequest := baidu.ConvertEmbeddingRequest(textRequest) - jsonData, err = json.Marshal(baiduEmbeddingRequest) - default: - baiduRequest := baidu.ConvertRequest(textRequest) - jsonData, err = json.Marshal(baiduRequest) - } - if err != nil { - return openai.ErrorWrapper(err, "marshal_text_request_failed", http.StatusInternalServerError) - } - requestBody = bytes.NewBuffer(jsonData) - case APITypePaLM: - palmRequest := google.ConvertPaLMRequest(textRequest) - jsonStr, err := json.Marshal(palmRequest) - if err != nil { - return openai.ErrorWrapper(err, "marshal_text_request_failed", http.StatusInternalServerError) - } - requestBody = bytes.NewBuffer(jsonStr) - case APITypeGemini: - geminiChatRequest := google.ConvertGeminiRequest(textRequest) - jsonStr, err := json.Marshal(geminiChatRequest) - if err != nil { - return openai.ErrorWrapper(err, "marshal_text_request_failed", http.StatusInternalServerError) - } - requestBody = bytes.NewBuffer(jsonStr) - case APITypeZhipu: - zhipuRequest := zhipu.ConvertRequest(textRequest) - jsonStr, err := json.Marshal(zhipuRequest) - if err != nil { - return openai.ErrorWrapper(err, "marshal_text_request_failed", http.StatusInternalServerError) - } - requestBody = bytes.NewBuffer(jsonStr) - case APITypeAli: - var jsonStr []byte - var err error - switch relayMode { - case constant.RelayModeEmbeddings: - aliEmbeddingRequest := ali.ConvertEmbeddingRequest(textRequest) - jsonStr, err = json.Marshal(aliEmbeddingRequest) - default: - aliRequest := ali.ConvertRequest(textRequest) - jsonStr, err = json.Marshal(aliRequest) - } - if err != nil { - return openai.ErrorWrapper(err, "marshal_text_request_failed", http.StatusInternalServerError) - } - requestBody = bytes.NewBuffer(jsonStr) - case APITypeTencent: - apiKey := c.Request.Header.Get("Authorization") - apiKey = strings.TrimPrefix(apiKey, "Bearer ") - appId, secretId, secretKey, err := tencent.ParseConfig(apiKey) - if err != nil { - return openai.ErrorWrapper(err, "invalid_tencent_config", http.StatusInternalServerError) - } - tencentRequest := tencent.ConvertRequest(textRequest) - tencentRequest.AppId = appId - tencentRequest.SecretId = secretId - jsonStr, err := json.Marshal(tencentRequest) - if err != nil { - return openai.ErrorWrapper(err, "marshal_text_request_failed", http.StatusInternalServerError) - } - sign := tencent.GetSign(*tencentRequest, secretKey) - c.Request.Header.Set("Authorization", sign) - requestBody = bytes.NewBuffer(jsonStr) - case APITypeAIProxyLibrary: - aiProxyLibraryRequest := aiproxy.ConvertRequest(textRequest) - aiProxyLibraryRequest.LibraryId = c.GetString("library_id") - jsonStr, err := json.Marshal(aiProxyLibraryRequest) - if err != nil { - return openai.ErrorWrapper(err, "marshal_text_request_failed", http.StatusInternalServerError) - } - requestBody = bytes.NewBuffer(jsonStr) - } - var req *http.Request var resp *http.Response isStream := textRequest.Stream - if apiType != APITypeXunfei { // cause xunfei use websocket + if apiType != constant.APITypeXunfei { // cause xunfei use websocket req, err = http.NewRequest(c.Request.Method, fullRequestURL, requestBody) if err != nil { return openai.ErrorWrapper(err, "new_request_failed", http.StatusInternalServerError) } - apiKey := c.Request.Header.Get("Authorization") - apiKey = strings.TrimPrefix(apiKey, "Bearer ") - switch apiType { - case APITypeOpenAI: - if channelType == common.ChannelTypeAzure { - req.Header.Set("api-key", apiKey) - } else { - req.Header.Set("Authorization", c.Request.Header.Get("Authorization")) - if channelType == common.ChannelTypeOpenRouter { - req.Header.Set("HTTP-Referer", "https://github.com/songquanpeng/one-api") - req.Header.Set("X-Title", "One API") - } - } - case APITypeClaude: - req.Header.Set("x-api-key", apiKey) - anthropicVersion := c.Request.Header.Get("anthropic-version") - if anthropicVersion == "" { - anthropicVersion = "2023-06-01" - } - req.Header.Set("anthropic-version", anthropicVersion) - case APITypeZhipu: - token := zhipu.GetToken(apiKey) - req.Header.Set("Authorization", token) - case APITypeAli: - req.Header.Set("Authorization", "Bearer "+apiKey) - if textRequest.Stream { - req.Header.Set("X-DashScope-SSE", "enable") - } - if c.GetString("plugin") != "" { - req.Header.Set("X-DashScope-Plugin", c.GetString("plugin")) - } - case APITypeTencent: - req.Header.Set("Authorization", apiKey) - case APITypePaLM: - req.Header.Set("x-goog-api-key", apiKey) - case APITypeGemini: - req.Header.Set("x-goog-api-key", apiKey) - default: - req.Header.Set("Authorization", "Bearer "+apiKey) - } - req.Header.Set("Content-Type", c.Request.Header.Get("Content-Type")) - req.Header.Set("Accept", c.Request.Header.Get("Accept")) - if isStream && c.Request.Header.Get("Accept") == "" { - req.Header.Set("Accept", "text/event-stream") - } - //req.Header.Set("Connection", c.Request.Header.Get("Connection")) + SetupRequestHeaders(c, req, apiType, meta, isStream) resp, err = util.HTTPClient.Do(req) if err != nil { return openai.ErrorWrapper(err, "do_request_failed", http.StatusInternalServerError) @@ -409,29 +113,31 @@ func RelayTextHelper(c *gin.Context, relayMode int) *openai.ErrorWithStatusCode isStream = isStream || strings.HasPrefix(resp.Header.Get("Content-Type"), "text/event-stream") if resp.StatusCode != http.StatusOK { - if preConsumedQuota != 0 { - go func(ctx context.Context) { - // return pre-consumed quota - err := model.PostConsumeTokenQuota(tokenId, -preConsumedQuota) - if err != nil { - common.LogError(ctx, "error return pre-consumed quota: "+err.Error()) - } - }(c.Request.Context()) - } + util.ReturnPreConsumedQuota(ctx, preConsumedQuota, meta.TokenId) return util.RelayErrorHandler(resp) } } - var textResponse openai.SlimTextResponse - tokenName := c.GetString("token_name") + var respErr *openai.ErrorWithStatusCode + var usage *openai.Usage defer func(ctx context.Context) { - // c.Writer.Flush() + // Why we use defer here? Because if error happened, we will have to return the pre-consumed quota. + if respErr != nil { + logger.Errorf(ctx, "respErr is not nil: %+v", respErr) + util.ReturnPreConsumedQuota(ctx, preConsumedQuota, meta.TokenId) + return + } + if usage == nil { + logger.Error(ctx, "usage is nil, which is unexpected") + return + } + go func() { quota := 0 completionRatio := common.GetCompletionRatio(textRequest.Model) - promptTokens = textResponse.Usage.PromptTokens - completionTokens = textResponse.Usage.CompletionTokens + promptTokens = usage.PromptTokens + completionTokens = usage.CompletionTokens quota = int(math.Ceil((float64(promptTokens) + float64(completionTokens)*completionRatio) * ratio)) if ratio != 0 && quota <= 0 { quota = 1 @@ -443,239 +149,25 @@ func RelayTextHelper(c *gin.Context, relayMode int) *openai.ErrorWithStatusCode quota = 0 } quotaDelta := quota - preConsumedQuota - err := model.PostConsumeTokenQuota(tokenId, quotaDelta) + err := model.PostConsumeTokenQuota(meta.TokenId, quotaDelta) if err != nil { - common.LogError(ctx, "error consuming token remain quota: "+err.Error()) + logger.Error(ctx, "error consuming token remain quota: "+err.Error()) } - err = model.CacheUpdateUserQuota(userId) + err = model.CacheUpdateUserQuota(meta.UserId) if err != nil { - common.LogError(ctx, "error update user quota cache: "+err.Error()) + logger.Error(ctx, "error update user quota cache: "+err.Error()) } if quota != 0 { logContent := fmt.Sprintf("模型倍率 %.2f,分组倍率 %.2f", modelRatio, groupRatio) - model.RecordConsumeLog(ctx, userId, channelId, promptTokens, completionTokens, textRequest.Model, tokenName, quota, logContent) - model.UpdateUserUsedQuotaAndRequestCount(userId, quota) - model.UpdateChannelUsedQuota(channelId, quota) + model.RecordConsumeLog(ctx, meta.UserId, meta.ChannelId, promptTokens, completionTokens, textRequest.Model, meta.TokenName, quota, logContent) + model.UpdateUserUsedQuotaAndRequestCount(meta.UserId, quota) + model.UpdateChannelUsedQuota(meta.ChannelId, quota) } - }() - }(c.Request.Context()) - switch apiType { - case APITypeOpenAI: - if isStream { - err, responseText := openai.StreamHandler(c, resp, relayMode) - if err != nil { - return err - } - textResponse.Usage.PromptTokens = promptTokens - textResponse.Usage.CompletionTokens = openai.CountTokenText(responseText, textRequest.Model) - return nil - } else { - err, usage := openai.Handler(c, resp, promptTokens, textRequest.Model) - if err != nil { - return err - } - if usage != nil { - textResponse.Usage = *usage - } - return nil - } - case APITypeClaude: - if isStream { - err, responseText := anthropic.StreamHandler(c, resp) - if err != nil { - return err - } - textResponse.Usage.PromptTokens = promptTokens - textResponse.Usage.CompletionTokens = openai.CountTokenText(responseText, textRequest.Model) - return nil - } else { - err, usage := anthropic.Handler(c, resp, promptTokens, textRequest.Model) - if err != nil { - return err - } - if usage != nil { - textResponse.Usage = *usage - } - return nil - } - case APITypeBaidu: - if isStream { - err, usage := baidu.StreamHandler(c, resp) - if err != nil { - return err - } - if usage != nil { - textResponse.Usage = *usage - } - return nil - } else { - var err *openai.ErrorWithStatusCode - var usage *openai.Usage - switch relayMode { - case constant.RelayModeEmbeddings: - err, usage = baidu.EmbeddingHandler(c, resp) - default: - err, usage = baidu.Handler(c, resp) - } - if err != nil { - return err - } - if usage != nil { - textResponse.Usage = *usage - } - return nil - } - case APITypePaLM: - if textRequest.Stream { // PaLM2 API does not support stream - err, responseText := google.PaLMStreamHandler(c, resp) - if err != nil { - return err - } - textResponse.Usage.PromptTokens = promptTokens - textResponse.Usage.CompletionTokens = openai.CountTokenText(responseText, textRequest.Model) - return nil - } else { - err, usage := google.PaLMHandler(c, resp, promptTokens, textRequest.Model) - if err != nil { - return err - } - if usage != nil { - textResponse.Usage = *usage - } - return nil - } - case APITypeGemini: - if textRequest.Stream { - err, responseText := google.StreamHandler(c, resp) - if err != nil { - return err - } - textResponse.Usage.PromptTokens = promptTokens - textResponse.Usage.CompletionTokens = openai.CountTokenText(responseText, textRequest.Model) - return nil - } else { - err, usage := google.GeminiHandler(c, resp, promptTokens, textRequest.Model) - if err != nil { - return err - } - if usage != nil { - textResponse.Usage = *usage - } - return nil - } - case APITypeZhipu: - if isStream { - err, usage := zhipu.StreamHandler(c, resp) - if err != nil { - return err - } - if usage != nil { - textResponse.Usage = *usage - } - // zhipu's API does not return prompt tokens & completion tokens - textResponse.Usage.PromptTokens = textResponse.Usage.TotalTokens - return nil - } else { - err, usage := zhipu.Handler(c, resp) - if err != nil { - return err - } - if usage != nil { - textResponse.Usage = *usage - } - // zhipu's API does not return prompt tokens & completion tokens - textResponse.Usage.PromptTokens = textResponse.Usage.TotalTokens - return nil - } - case APITypeAli: - if isStream { - err, usage := ali.StreamHandler(c, resp) - if err != nil { - return err - } - if usage != nil { - textResponse.Usage = *usage - } - return nil - } else { - var err *openai.ErrorWithStatusCode - var usage *openai.Usage - switch relayMode { - case constant.RelayModeEmbeddings: - err, usage = ali.EmbeddingHandler(c, resp) - default: - err, usage = ali.Handler(c, resp) - } - if err != nil { - return err - } - if usage != nil { - textResponse.Usage = *usage - } - return nil - } - case APITypeXunfei: - auth := c.Request.Header.Get("Authorization") - auth = strings.TrimPrefix(auth, "Bearer ") - splits := strings.Split(auth, "|") - if len(splits) != 3 { - return openai.ErrorWrapper(errors.New("invalid auth"), "invalid_auth", http.StatusBadRequest) - } - var err *openai.ErrorWithStatusCode - var usage *openai.Usage - if isStream { - err, usage = xunfei.StreamHandler(c, textRequest, splits[0], splits[1], splits[2]) - } else { - err, usage = xunfei.Handler(c, textRequest, splits[0], splits[1], splits[2]) - } - if err != nil { - return err - } - if usage != nil { - textResponse.Usage = *usage - } - return nil - case APITypeAIProxyLibrary: - if isStream { - err, usage := aiproxy.StreamHandler(c, resp) - if err != nil { - return err - } - if usage != nil { - textResponse.Usage = *usage - } - return nil - } else { - err, usage := aiproxy.Handler(c, resp) - if err != nil { - return err - } - if usage != nil { - textResponse.Usage = *usage - } - return nil - } - case APITypeTencent: - if isStream { - err, responseText := tencent.StreamHandler(c, resp) - if err != nil { - return err - } - textResponse.Usage.PromptTokens = promptTokens - textResponse.Usage.CompletionTokens = openai.CountTokenText(responseText, textRequest.Model) - return nil - } else { - err, usage := tencent.Handler(c, resp) - if err != nil { - return err - } - if usage != nil { - textResponse.Usage = *usage - } - return nil - } - default: - return openai.ErrorWrapper(errors.New("unknown api type"), "unknown_api_type", http.StatusInternalServerError) + }(ctx) + usage, respErr = DoResponse(c, &textRequest, resp, relayMode, apiType, isStream, promptTokens) + if respErr != nil { + return respErr } + return nil } diff --git a/relay/controller/util.go b/relay/controller/util.go new file mode 100644 index 00000000..02f1b30f --- /dev/null +++ b/relay/controller/util.go @@ -0,0 +1,337 @@ +package controller + +import ( + "bytes" + "encoding/json" + "errors" + "fmt" + "github.com/gin-gonic/gin" + "io" + "net/http" + "one-api/common" + "one-api/common/helper" + "one-api/relay/channel/aiproxy" + "one-api/relay/channel/ali" + "one-api/relay/channel/anthropic" + "one-api/relay/channel/baidu" + "one-api/relay/channel/google" + "one-api/relay/channel/openai" + "one-api/relay/channel/tencent" + "one-api/relay/channel/xunfei" + "one-api/relay/channel/zhipu" + "one-api/relay/constant" + "one-api/relay/util" + "strings" +) + +func GetRequestURL(requestURL string, apiType int, relayMode int, meta *util.RelayMeta, textRequest *openai.GeneralOpenAIRequest) (string, error) { + fullRequestURL := util.GetFullRequestURL(meta.BaseURL, requestURL, meta.ChannelType) + switch apiType { + case constant.APITypeOpenAI: + if meta.ChannelType == common.ChannelTypeAzure { + // https://learn.microsoft.com/en-us/azure/cognitive-services/openai/chatgpt-quickstart?pivots=rest-api&tabs=command-line#rest-api + requestURL := strings.Split(requestURL, "?")[0] + requestURL = fmt.Sprintf("%s?api-version=%s", requestURL, meta.APIVersion) + task := strings.TrimPrefix(requestURL, "/v1/") + model_ := textRequest.Model + model_ = strings.Replace(model_, ".", "", -1) + // https://github.com/songquanpeng/one-api/issues/67 + model_ = strings.TrimSuffix(model_, "-0301") + model_ = strings.TrimSuffix(model_, "-0314") + model_ = strings.TrimSuffix(model_, "-0613") + + requestURL = fmt.Sprintf("/openai/deployments/%s/%s", model_, task) + fullRequestURL = util.GetFullRequestURL(meta.BaseURL, requestURL, meta.ChannelType) + } + case constant.APITypeClaude: + fullRequestURL = fmt.Sprintf("%s/v1/complete", meta.BaseURL) + case constant.APITypeBaidu: + switch textRequest.Model { + case "ERNIE-Bot": + fullRequestURL = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions" + case "ERNIE-Bot-turbo": + fullRequestURL = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/eb-instant" + case "ERNIE-Bot-4": + fullRequestURL = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions_pro" + case "BLOOMZ-7B": + fullRequestURL = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/bloomz_7b1" + case "Embedding-V1": + fullRequestURL = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/embeddings/embedding-v1" + } + var accessToken string + var err error + if accessToken, err = baidu.GetAccessToken(meta.APIKey); err != nil { + return "", fmt.Errorf("failed to get baidu access token: %w", err) + } + fullRequestURL += "?access_token=" + accessToken + case constant.APITypePaLM: + fullRequestURL = fmt.Sprintf("%s/v1beta2/models/chat-bison-001:generateMessage", meta.BaseURL) + case constant.APITypeGemini: + version := helper.AssignOrDefault(meta.APIVersion, "v1") + action := "generateContent" + if textRequest.Stream { + action = "streamGenerateContent" + } + fullRequestURL = fmt.Sprintf("%s/%s/models/%s:%s", meta.BaseURL, version, textRequest.Model, action) + case constant.APITypeZhipu: + method := "invoke" + if textRequest.Stream { + method = "sse-invoke" + } + fullRequestURL = fmt.Sprintf("https://open.bigmodel.cn/api/paas/v3/model-api/%s/%s", textRequest.Model, method) + case constant.APITypeAli: + fullRequestURL = "https://dashscope.aliyuncs.com/api/v1/services/aigc/text-generation/generation" + if relayMode == constant.RelayModeEmbeddings { + fullRequestURL = "https://dashscope.aliyuncs.com/api/v1/services/embeddings/text-embedding/text-embedding" + } + case constant.APITypeTencent: + fullRequestURL = "https://hunyuan.cloud.tencent.com/hyllm/v1/chat/completions" + case constant.APITypeAIProxyLibrary: + fullRequestURL = fmt.Sprintf("%s/api/library/ask", meta.BaseURL) + } + return fullRequestURL, nil +} + +func GetRequestBody(c *gin.Context, textRequest openai.GeneralOpenAIRequest, isModelMapped bool, apiType int, relayMode int) (io.Reader, error) { + var requestBody io.Reader + if isModelMapped { + jsonStr, err := json.Marshal(textRequest) + if err != nil { + return nil, err + } + requestBody = bytes.NewBuffer(jsonStr) + } else { + requestBody = c.Request.Body + } + switch apiType { + case constant.APITypeClaude: + claudeRequest := anthropic.ConvertRequest(textRequest) + jsonStr, err := json.Marshal(claudeRequest) + if err != nil { + return nil, err + } + requestBody = bytes.NewBuffer(jsonStr) + case constant.APITypeBaidu: + var jsonData []byte + var err error + switch relayMode { + case constant.RelayModeEmbeddings: + baiduEmbeddingRequest := baidu.ConvertEmbeddingRequest(textRequest) + jsonData, err = json.Marshal(baiduEmbeddingRequest) + default: + baiduRequest := baidu.ConvertRequest(textRequest) + jsonData, err = json.Marshal(baiduRequest) + } + if err != nil { + return nil, err + } + requestBody = bytes.NewBuffer(jsonData) + case constant.APITypePaLM: + palmRequest := google.ConvertPaLMRequest(textRequest) + jsonStr, err := json.Marshal(palmRequest) + if err != nil { + return nil, err + } + requestBody = bytes.NewBuffer(jsonStr) + case constant.APITypeGemini: + geminiChatRequest := google.ConvertGeminiRequest(textRequest) + jsonStr, err := json.Marshal(geminiChatRequest) + if err != nil { + return nil, err + } + requestBody = bytes.NewBuffer(jsonStr) + case constant.APITypeZhipu: + zhipuRequest := zhipu.ConvertRequest(textRequest) + jsonStr, err := json.Marshal(zhipuRequest) + if err != nil { + return nil, err + } + requestBody = bytes.NewBuffer(jsonStr) + case constant.APITypeAli: + var jsonStr []byte + var err error + switch relayMode { + case constant.RelayModeEmbeddings: + aliEmbeddingRequest := ali.ConvertEmbeddingRequest(textRequest) + jsonStr, err = json.Marshal(aliEmbeddingRequest) + default: + aliRequest := ali.ConvertRequest(textRequest) + jsonStr, err = json.Marshal(aliRequest) + } + if err != nil { + return nil, err + } + requestBody = bytes.NewBuffer(jsonStr) + case constant.APITypeTencent: + apiKey := c.Request.Header.Get("Authorization") + apiKey = strings.TrimPrefix(apiKey, "Bearer ") + appId, secretId, secretKey, err := tencent.ParseConfig(apiKey) + if err != nil { + return nil, err + } + tencentRequest := tencent.ConvertRequest(textRequest) + tencentRequest.AppId = appId + tencentRequest.SecretId = secretId + jsonStr, err := json.Marshal(tencentRequest) + if err != nil { + return nil, err + } + sign := tencent.GetSign(*tencentRequest, secretKey) + c.Request.Header.Set("Authorization", sign) + requestBody = bytes.NewBuffer(jsonStr) + case constant.APITypeAIProxyLibrary: + aiProxyLibraryRequest := aiproxy.ConvertRequest(textRequest) + aiProxyLibraryRequest.LibraryId = c.GetString("library_id") + jsonStr, err := json.Marshal(aiProxyLibraryRequest) + if err != nil { + return nil, err + } + requestBody = bytes.NewBuffer(jsonStr) + } + return requestBody, nil +} + +func SetupRequestHeaders(c *gin.Context, req *http.Request, apiType int, meta *util.RelayMeta, isStream bool) { + SetupAuthHeaders(c, req, apiType, meta, isStream) + req.Header.Set("Content-Type", c.Request.Header.Get("Content-Type")) + req.Header.Set("Accept", c.Request.Header.Get("Accept")) + if isStream && c.Request.Header.Get("Accept") == "" { + req.Header.Set("Accept", "text/event-stream") + } +} + +func SetupAuthHeaders(c *gin.Context, req *http.Request, apiType int, meta *util.RelayMeta, isStream bool) { + apiKey := meta.APIKey + switch apiType { + case constant.APITypeOpenAI: + if meta.ChannelType == common.ChannelTypeAzure { + req.Header.Set("api-key", apiKey) + } else { + req.Header.Set("Authorization", c.Request.Header.Get("Authorization")) + if meta.ChannelType == common.ChannelTypeOpenRouter { + req.Header.Set("HTTP-Referer", "https://github.com/songquanpeng/one-api") + req.Header.Set("X-Title", "One API") + } + } + case constant.APITypeClaude: + req.Header.Set("x-api-key", apiKey) + anthropicVersion := c.Request.Header.Get("anthropic-version") + if anthropicVersion == "" { + anthropicVersion = "2023-06-01" + } + req.Header.Set("anthropic-version", anthropicVersion) + case constant.APITypeZhipu: + token := zhipu.GetToken(apiKey) + req.Header.Set("Authorization", token) + case constant.APITypeAli: + req.Header.Set("Authorization", "Bearer "+apiKey) + if isStream { + req.Header.Set("X-DashScope-SSE", "enable") + } + if c.GetString("plugin") != "" { + req.Header.Set("X-DashScope-Plugin", c.GetString("plugin")) + } + case constant.APITypeTencent: + req.Header.Set("Authorization", apiKey) + case constant.APITypePaLM: + req.Header.Set("x-goog-api-key", apiKey) + case constant.APITypeGemini: + req.Header.Set("x-goog-api-key", apiKey) + default: + req.Header.Set("Authorization", "Bearer "+apiKey) + } +} + +func DoResponse(c *gin.Context, textRequest *openai.GeneralOpenAIRequest, resp *http.Response, relayMode int, apiType int, isStream bool, promptTokens int) (usage *openai.Usage, err *openai.ErrorWithStatusCode) { + var responseText string + switch apiType { + case constant.APITypeOpenAI: + if isStream { + err, responseText = openai.StreamHandler(c, resp, relayMode) + } else { + err, usage = openai.Handler(c, resp, promptTokens, textRequest.Model) + } + case constant.APITypeClaude: + if isStream { + err, responseText = anthropic.StreamHandler(c, resp) + } else { + err, usage = anthropic.Handler(c, resp, promptTokens, textRequest.Model) + } + case constant.APITypeBaidu: + if isStream { + err, usage = baidu.StreamHandler(c, resp) + } else { + switch relayMode { + case constant.RelayModeEmbeddings: + err, usage = baidu.EmbeddingHandler(c, resp) + default: + err, usage = baidu.Handler(c, resp) + } + } + case constant.APITypePaLM: + if isStream { // PaLM2 API does not support stream + err, responseText = google.PaLMStreamHandler(c, resp) + } else { + err, usage = google.PaLMHandler(c, resp, promptTokens, textRequest.Model) + } + case constant.APITypeGemini: + if isStream { + err, responseText = google.StreamHandler(c, resp) + } else { + err, usage = google.GeminiHandler(c, resp, promptTokens, textRequest.Model) + } + case constant.APITypeZhipu: + if isStream { + err, usage = zhipu.StreamHandler(c, resp) + } else { + err, usage = zhipu.Handler(c, resp) + } + case constant.APITypeAli: + if isStream { + err, usage = ali.StreamHandler(c, resp) + } else { + switch relayMode { + case constant.RelayModeEmbeddings: + err, usage = ali.EmbeddingHandler(c, resp) + default: + err, usage = ali.Handler(c, resp) + } + } + case constant.APITypeXunfei: + auth := c.Request.Header.Get("Authorization") + auth = strings.TrimPrefix(auth, "Bearer ") + splits := strings.Split(auth, "|") + if len(splits) != 3 { + return nil, openai.ErrorWrapper(errors.New("invalid auth"), "invalid_auth", http.StatusBadRequest) + } + if isStream { + err, usage = xunfei.StreamHandler(c, *textRequest, splits[0], splits[1], splits[2]) + } else { + err, usage = xunfei.Handler(c, *textRequest, splits[0], splits[1], splits[2]) + } + case constant.APITypeAIProxyLibrary: + if isStream { + err, usage = aiproxy.StreamHandler(c, resp) + } else { + err, usage = aiproxy.Handler(c, resp) + } + case constant.APITypeTencent: + if isStream { + err, responseText = tencent.StreamHandler(c, resp) + } else { + err, usage = tencent.Handler(c, resp) + } + default: + return nil, openai.ErrorWrapper(errors.New("unknown api type"), "unknown_api_type", http.StatusInternalServerError) + } + if err != nil { + return nil, err + } + if usage == nil && responseText != "" { + usage = &openai.Usage{} + usage.PromptTokens = promptTokens + usage.CompletionTokens = openai.CountTokenText(responseText, textRequest.Model) + usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens + } + return usage, nil +} diff --git a/relay/util/billing.go b/relay/util/billing.go new file mode 100644 index 00000000..35fb28a4 --- /dev/null +++ b/relay/util/billing.go @@ -0,0 +1,19 @@ +package util + +import ( + "context" + "one-api/common/logger" + "one-api/model" +) + +func ReturnPreConsumedQuota(ctx context.Context, preConsumedQuota int, tokenId int) { + if preConsumedQuota != 0 { + go func(ctx context.Context) { + // return pre-consumed quota + err := model.PostConsumeTokenQuota(tokenId, -preConsumedQuota) + if err != nil { + logger.Error(ctx, "error return pre-consumed quota: "+err.Error()) + } + }(ctx) + } +} diff --git a/relay/util/common.go b/relay/util/common.go index 9d13b12e..be31857b 100644 --- a/relay/util/common.go +++ b/relay/util/common.go @@ -7,6 +7,8 @@ import ( "io" "net/http" "one-api/common" + "one-api/common/config" + "one-api/common/logger" "one-api/model" "one-api/relay/channel/openai" "strconv" @@ -16,7 +18,7 @@ import ( ) func ShouldDisableChannel(err *openai.Error, statusCode int) bool { - if !common.AutomaticDisableChannelEnabled { + if !config.AutomaticDisableChannelEnabled { return false } if err == nil { @@ -32,7 +34,7 @@ func ShouldDisableChannel(err *openai.Error, statusCode int) bool { } func ShouldEnableChannel(err error, openAIErr *openai.Error) bool { - if !common.AutomaticEnableChannelEnabled { + if !config.AutomaticEnableChannelEnabled { return false } if err != nil { @@ -138,11 +140,11 @@ func PostConsumeQuota(ctx context.Context, tokenId int, quotaDelta int, totalQuo // quotaDelta is remaining quota to be consumed err := model.PostConsumeTokenQuota(tokenId, quotaDelta) if err != nil { - common.SysError("error consuming token remain quota: " + err.Error()) + logger.SysError("error consuming token remain quota: " + err.Error()) } err = model.CacheUpdateUserQuota(userId) if err != nil { - common.SysError("error update user quota cache: " + err.Error()) + logger.SysError("error update user quota cache: " + err.Error()) } // totalQuota is total quota consumed if totalQuota != 0 { @@ -152,11 +154,11 @@ func PostConsumeQuota(ctx context.Context, tokenId int, quotaDelta int, totalQuo model.UpdateChannelUsedQuota(channelId, totalQuota) } if totalQuota <= 0 { - common.LogError(ctx, fmt.Sprintf("totalQuota consumed is %d, something is wrong", totalQuota)) + logger.Error(ctx, fmt.Sprintf("totalQuota consumed is %d, something is wrong", totalQuota)) } } -func GetAPIVersion(c *gin.Context) string { +func GetAzureAPIVersion(c *gin.Context) string { query := c.Request.URL.Query() apiVersion := query.Get("api-version") if apiVersion == "" { diff --git a/relay/util/init.go b/relay/util/init.go index d308d900..62d44d15 100644 --- a/relay/util/init.go +++ b/relay/util/init.go @@ -2,7 +2,7 @@ package util import ( "net/http" - "one-api/common" + "one-api/common/config" "time" ) @@ -10,11 +10,11 @@ var HTTPClient *http.Client var ImpatientHTTPClient *http.Client func init() { - if common.RelayTimeout == 0 { + if config.RelayTimeout == 0 { HTTPClient = &http.Client{} } else { HTTPClient = &http.Client{ - Timeout: time.Duration(common.RelayTimeout) * time.Second, + Timeout: time.Duration(config.RelayTimeout) * time.Second, } } diff --git a/relay/util/model_mapping.go b/relay/util/model_mapping.go new file mode 100644 index 00000000..39e062a1 --- /dev/null +++ b/relay/util/model_mapping.go @@ -0,0 +1,12 @@ +package util + +func GetMappedModelName(modelName string, mapping map[string]string) (string, bool) { + if mapping == nil { + return modelName, false + } + mappedModelName := mapping[modelName] + if mappedModelName != "" { + return mappedModelName, true + } + return modelName, false +} diff --git a/relay/util/relay_meta.go b/relay/util/relay_meta.go new file mode 100644 index 00000000..19936e49 --- /dev/null +++ b/relay/util/relay_meta.go @@ -0,0 +1,44 @@ +package util + +import ( + "github.com/gin-gonic/gin" + "one-api/common" + "strings" +) + +type RelayMeta struct { + ChannelType int + ChannelId int + TokenId int + TokenName string + UserId int + Group string + ModelMapping map[string]string + BaseURL string + APIVersion string + APIKey string + Config map[string]string +} + +func GetRelayMeta(c *gin.Context) *RelayMeta { + meta := RelayMeta{ + ChannelType: c.GetInt("channel"), + ChannelId: c.GetInt("channel_id"), + TokenId: c.GetInt("token_id"), + TokenName: c.GetString("token_name"), + UserId: c.GetInt("id"), + Group: c.GetString("group"), + ModelMapping: c.GetStringMapString("model_mapping"), + BaseURL: c.GetString("base_url"), + APIVersion: c.GetString("api_version"), + APIKey: strings.TrimPrefix(c.Request.Header.Get("Authorization"), "Bearer "), + Config: nil, + } + if meta.ChannelType == common.ChannelTypeAzure { + meta.APIVersion = GetAzureAPIVersion(c) + } + if meta.BaseURL == "" { + meta.BaseURL = common.ChannelBaseURLs[meta.ChannelType] + } + return &meta +} diff --git a/relay/util/validation.go b/relay/util/validation.go new file mode 100644 index 00000000..48b42d94 --- /dev/null +++ b/relay/util/validation.go @@ -0,0 +1,37 @@ +package util + +import ( + "errors" + "math" + "one-api/relay/channel/openai" + "one-api/relay/constant" +) + +func ValidateTextRequest(textRequest *openai.GeneralOpenAIRequest, relayMode int) error { + if textRequest.MaxTokens < 0 || textRequest.MaxTokens > math.MaxInt32/2 { + return errors.New("max_tokens is invalid") + } + if textRequest.Model == "" { + return errors.New("model is required") + } + switch relayMode { + case constant.RelayModeCompletions: + if textRequest.Prompt == "" { + return errors.New("field prompt is required") + } + case constant.RelayModeChatCompletions: + if textRequest.Messages == nil || len(textRequest.Messages) == 0 { + return errors.New("field messages is required") + } + case constant.RelayModeEmbeddings: + case constant.RelayModeModerations: + if textRequest.Input == "" { + return errors.New("field input is required") + } + case constant.RelayModeEdits: + if textRequest.Instruction == "" { + return errors.New("field instruction is required") + } + } + return nil +} diff --git a/router/main.go b/router/main.go index 85127a1a..6504b312 100644 --- a/router/main.go +++ b/router/main.go @@ -5,7 +5,8 @@ import ( "fmt" "github.com/gin-gonic/gin" "net/http" - "one-api/common" + "one-api/common/config" + "one-api/common/logger" "os" "strings" ) @@ -15,9 +16,9 @@ func SetRouter(router *gin.Engine, buildFS embed.FS) { SetDashboardRouter(router) SetRelayRouter(router) frontendBaseUrl := os.Getenv("FRONTEND_BASE_URL") - if common.IsMasterNode && frontendBaseUrl != "" { + if config.IsMasterNode && frontendBaseUrl != "" { frontendBaseUrl = "" - common.SysLog("FRONTEND_BASE_URL is ignored on master node") + logger.SysLog("FRONTEND_BASE_URL is ignored on master node") } if frontendBaseUrl == "" { SetWebRouter(router, buildFS) diff --git a/router/web-router.go b/router/web-router.go index 7328c7a3..95d8fcb9 100644 --- a/router/web-router.go +++ b/router/web-router.go @@ -8,17 +8,18 @@ import ( "github.com/gin-gonic/gin" "net/http" "one-api/common" + "one-api/common/config" "one-api/controller" "one-api/middleware" "strings" ) func SetWebRouter(router *gin.Engine, buildFS embed.FS) { - indexPageData, _ := buildFS.ReadFile(fmt.Sprintf("web/build/%s/index.html", common.Theme)) + indexPageData, _ := buildFS.ReadFile(fmt.Sprintf("web/build/%s/index.html", config.Theme)) router.Use(gzip.Gzip(gzip.DefaultCompression)) router.Use(middleware.GlobalWebRateLimit()) router.Use(middleware.Cache()) - router.Use(static.Serve("/", common.EmbedFolder(buildFS, fmt.Sprintf("web/build/%s", common.Theme)))) + router.Use(static.Serve("/", common.EmbedFolder(buildFS, fmt.Sprintf("web/build/%s", config.Theme)))) router.NoRoute(func(c *gin.Context) { if strings.HasPrefix(c.Request.RequestURI, "/v1") || strings.HasPrefix(c.Request.RequestURI, "/api") { controller.RelayNotFound(c)