diff --git a/.gitignore b/.gitignore index 974fcf63..2a8ae16e 100644 --- a/.gitignore +++ b/.gitignore @@ -8,3 +8,4 @@ build logs data /web/node_modules +cmd.md \ No newline at end of file diff --git a/README.md b/README.md index 2dcdbd4f..d5c939be 100644 --- a/README.md +++ b/README.md @@ -81,11 +81,12 @@ _✨ 通过标准的 OpenAI API 格式访问所有的大模型，开箱即用 + [x] [Groq](https://wow.groq.com/) + [x] [Ollama](https://github.com/ollama/ollama) + [x] [零一万物](https://platform.lingyiwanwu.com/) + + [x] [阶跃星辰](https://platform.stepfun.com/) 2. 支持配置镜像以及众多[第三方代理服务](https://iamazing.cn/page/openai-api-third-party-services)。 3. 支持通过**负载均衡**的方式访问多个渠道。 4. 支持 **stream 模式**，可以通过流式传输实现打字机效果。 5. 支持**多机部署**，[详见此处](#多机部署)。 -6. 支持**令牌管理**，设置令牌的过期时间和额度。 +6. 支持**令牌管理**，设置令牌的过期时间、额度、允许的 IP 范围以及允许的模型访问。 7. 支持**兑换码管理**，支持批量生成和导出兑换码，可使用兑换码为账户进行充值。 8. 支持**渠道管理**，批量创建渠道。 9. 支持**用户分组**以及**渠道分组**，支持为不同分组设置不同的倍率。 @@ -101,10 +102,11 @@ _✨ 通过标准的 OpenAI API 格式访问所有的大模型，开箱即用 19. 支持丰富的**自定义**设置， 1. 支持自定义系统名称，logo 以及页脚。 2. 支持自定义首页和关于页面，可以选择使用 HTML & Markdown 代码进行自定义，或者使用一个单独的网页通过 iframe 嵌入。 -20. 支持通过系统访问令牌访问管理 API（bearer token，用以替代 cookie，你可以自行抓包来查看 API 的用法）。 +20. 支持通过系统访问令牌调用管理 API，进而**在无需二开的情况下扩展和自定义** One API 的功能，详情请参考此处 [API 文档](./docs/API.md)。。 21. 支持 Cloudflare Turnstile 用户校验。 22. 支持用户管理，支持**多种用户登录注册方式**： + 邮箱登录注册（支持注册邮箱白名单）以及通过邮箱进行密码重置。 + + 支持使用飞书进行授权登录。 + [GitHub 开放授权](https://github.com/settings/applications/new)。 + 微信公众号授权（需要额外部署 [WeChat Server](https://github.com/songquanpeng/wechat-server)）。 23. 支持主题切换，设置环境变量 `THEME` 即可，默认为 `default`，欢迎 PR 更多主题，具体参考[此处](./web/README.md)。 diff --git a/common/config/config.go b/common/config/config.go index 3524183a..9fd7cba0 100644 --- a/common/config/config.go +++ b/common/config/config.go @@ -66,6 +66,9 @@ var SMTPToken = "" var GitHubClientId = "" var GitHubClientSecret = "" +var LarkClientId = "" +var LarkClientSecret = "" + var WeChatServerAddress = "" var WeChatServerToken = "" var WeChatAccountQRCodeImageURL = "" diff --git a/common/config/key.go b/common/config/key.go new file mode 100644 index 00000000..4b503c2d --- /dev/null +++ b/common/config/key.go @@ -0,0 +1,9 @@ +package config + +const ( + KeyPrefix = "cfg_" + + KeyAPIVersion = KeyPrefix + "api_version" + KeyLibraryID = KeyPrefix + "library_id" + KeyPlugin = KeyPrefix + "plugin" +) diff --git a/common/constants.go b/common/constants.go index 849bdce7..87221b61 100644 --- a/common/constants.go +++ b/common/constants.go @@ -4,116 +4,3 @@ 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 - -const ( - RoleGuestUser = 0 - RoleCommonUser = 1 - RoleAdminUser = 10 - RoleRootUser = 100 -) - -const ( - UserStatusEnabled = 1 // don't use 0, 0 is the default value! - UserStatusDisabled = 2 // also don't use 0 - UserStatusDeleted = 3 -) - -const ( - TokenStatusEnabled = 1 // don't use 0, 0 is the default value! - TokenStatusDisabled = 2 // also don't use 0 - TokenStatusExpired = 3 - TokenStatusExhausted = 4 -) - -const ( - RedemptionCodeStatusEnabled = 1 // don't use 0, 0 is the default value! - RedemptionCodeStatusDisabled = 2 // also don't use 0 - RedemptionCodeStatusUsed = 3 // also don't use 0 -) - -const ( - ChannelStatusUnknown = 0 - ChannelStatusEnabled = 1 // don't use 0, 0 is the default value! - ChannelStatusManuallyDisabled = 2 // also don't use 0 - ChannelStatusAutoDisabled = 3 -) - -const ( - ChannelTypeUnknown = iota - ChannelTypeOpenAI - ChannelTypeAPI2D - ChannelTypeAzure - ChannelTypeCloseAI - ChannelTypeOpenAISB - ChannelTypeOpenAIMax - ChannelTypeOhMyGPT - ChannelTypeCustom - ChannelTypeAILS - ChannelTypeAIProxy - ChannelTypePaLM - ChannelTypeAPI2GPT - ChannelTypeAIGC2D - ChannelTypeAnthropic - ChannelTypeBaidu - ChannelTypeZhipu - ChannelTypeAli - ChannelTypeXunfei - ChannelType360 - ChannelTypeOpenRouter - ChannelTypeAIProxyLibrary - ChannelTypeFastGPT - ChannelTypeTencent - ChannelTypeGemini - ChannelTypeMoonshot - ChannelTypeBaichuan - ChannelTypeMinimax - ChannelTypeMistral - ChannelTypeGroq - ChannelTypeOllama - ChannelTypeLingYiWanWu - - ChannelTypeDummy -) - -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 - "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 - "https://api.moonshot.cn", // 25 - "https://api.baichuan-ai.com", // 26 - "https://api.minimax.chat", // 27 - "https://api.mistral.ai", // 28 - "https://api.groq.com/openai", // 29 - "http://localhost:11434", // 30 - "https://api.lingyiwanwu.com", // 31 -} - -const ( - ConfigKeyPrefix = "cfg_" - - ConfigKeyAPIVersion = ConfigKeyPrefix + "api_version" - ConfigKeyLibraryID = ConfigKeyPrefix + "library_id" - ConfigKeyPlugin = ConfigKeyPrefix + "plugin" -) diff --git a/common/helper/helper.go b/common/helper/helper.go index db41ac74..cf2e1635 100644 --- a/common/helper/helper.go +++ b/common/helper/helper.go @@ -2,16 +2,14 @@ package helper import ( "fmt" - "github.com/google/uuid" + "github.com/songquanpeng/one-api/common/random" "html/template" "log" - "math/rand" "net" "os/exec" "runtime" "strconv" "strings" - "time" ) func OpenBrowser(url string) { @@ -79,31 +77,6 @@ func Bytes2Size(num int64) string { 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 := inter.(type) { case string: @@ -128,65 +101,8 @@ func IntMax(a int, b int) int { } } -func GetUUID() string { - code := uuid.New().String() - code = strings.Replace(code, "-", "", -1) - return code -} - -const keyChars = "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ" -const keyNumbers = "0123456789" - -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 GetRandomNumberString(length int) string { - rand.Seed(time.Now().UnixNano()) - key := make([]byte, length) - for i := 0; i < length; i++ { - key[i] = keyNumbers[rand.Intn(len(keyNumbers))] - } - 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 GenRequestID() string { - return GetTimeString() + GetRandomNumberString(8) + return GetTimeString() + random.GetRandomNumberString(8) } func Max(a int, b int) int { diff --git a/common/helper/time.go b/common/helper/time.go new file mode 100644 index 00000000..302746db --- /dev/null +++ b/common/helper/time.go @@ -0,0 +1,15 @@ +package helper + +import ( + "fmt" + "time" +) + +func GetTimestamp() int64 { + return time.Now().Unix() +} + +func GetTimeString() string { + now := time.Now() + return fmt.Sprintf("%s%d", now.Format("20060102150405"), now.UnixNano()%1e9) +} diff --git a/common/network/ip.go b/common/network/ip.go new file mode 100644 index 00000000..0fbe5e6f --- /dev/null +++ b/common/network/ip.go @@ -0,0 +1,52 @@ +package network + +import ( + "context" + "fmt" + "github.com/songquanpeng/one-api/common/logger" + "net" + "strings" +) + +func splitSubnets(subnets string) []string { + res := strings.Split(subnets, ",") + for i := 0; i < len(res); i++ { + res[i] = strings.TrimSpace(res[i]) + } + return res +} + +func isValidSubnet(subnet string) error { + _, _, err := net.ParseCIDR(subnet) + if err != nil { + return fmt.Errorf("failed to parse subnet: %w", err) + } + return nil +} + +func isIpInSubnet(ctx context.Context, ip string, subnet string) bool { + _, ipNet, err := net.ParseCIDR(subnet) + if err != nil { + logger.Errorf(ctx, "failed to parse subnet: %s", err.Error()) + return false + } + return ipNet.Contains(net.ParseIP(ip)) +} + +func IsValidSubnets(subnets string) error { + for _, subnet := range splitSubnets(subnets) { + if err := isValidSubnet(subnet); err != nil { + return err + } + } + return nil +} + +func IsIpInSubnets(ctx context.Context, ip string, subnets string) bool { + for _, subnet := range splitSubnets(subnets) { + if isIpInSubnet(ctx, ip, subnet) { + return true + } + } + return false +} diff --git a/common/network/ip_test.go b/common/network/ip_test.go new file mode 100644 index 00000000..6c593458 --- /dev/null +++ b/common/network/ip_test.go @@ -0,0 +1,19 @@ +package network + +import ( + "context" + "testing" + + . "github.com/smartystreets/goconvey/convey" +) + +func TestIsIpInSubnet(t *testing.T) { + ctx := context.Background() + ip1 := "192.168.0.5" + ip2 := "125.216.250.89" + subnet := "192.168.0.0/24" + Convey("TestIsIpInSubnet", t, func() { + So(isIpInSubnet(ctx, ip1, subnet), ShouldBeTrue) + So(isIpInSubnet(ctx, ip2, subnet), ShouldBeFalse) + }) +} diff --git a/common/random.go b/common/random.go deleted file mode 100644 index 44bd2856..00000000 --- a/common/random.go +++ /dev/null @@ -1,8 +0,0 @@ -package common - -import "math/rand" - -// RandRange returns a random number between min and max (max is not included) -func RandRange(min, max int) int { - return min + rand.Intn(max-min) -} diff --git a/common/random/main.go b/common/random/main.go new file mode 100644 index 00000000..dbb772cd --- /dev/null +++ b/common/random/main.go @@ -0,0 +1,61 @@ +package random + +import ( + "github.com/google/uuid" + "math/rand" + "strings" + "time" +) + +func GetUUID() string { + code := uuid.New().String() + code = strings.Replace(code, "-", "", -1) + return code +} + +const keyChars = "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ" +const keyNumbers = "0123456789" + +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 GetRandomNumberString(length int) string { + rand.Seed(time.Now().UnixNano()) + key := make([]byte, length) + for i := 0; i < length; i++ { + key[i] = keyNumbers[rand.Intn(len(keyNumbers))] + } + return string(key) +} + +// RandRange returns a random number between min and max (max is not included) +func RandRange(min, max int) int { + return min + rand.Intn(max-min) +} diff --git a/controller/github.go b/controller/auth/github.go similarity index 94% rename from controller/github.go rename to controller/auth/github.go index 7d7fa106..15542655 100644 --- a/controller/github.go +++ b/controller/auth/github.go @@ -1,4 +1,4 @@ -package controller +package auth import ( "bytes" @@ -7,10 +7,10 @@ import ( "fmt" "github.com/gin-contrib/sessions" "github.com/gin-gonic/gin" - "github.com/songquanpeng/one-api/common" "github.com/songquanpeng/one-api/common/config" - "github.com/songquanpeng/one-api/common/helper" "github.com/songquanpeng/one-api/common/logger" + "github.com/songquanpeng/one-api/common/random" + "github.com/songquanpeng/one-api/controller" "github.com/songquanpeng/one-api/model" "net/http" "strconv" @@ -133,8 +133,8 @@ func GitHubOAuth(c *gin.Context) { user.DisplayName = "GitHub User" } user.Email = githubUser.Email - user.Role = common.RoleCommonUser - user.Status = common.UserStatusEnabled + user.Role = model.RoleCommonUser + user.Status = model.UserStatusEnabled if err := user.Insert(0); err != nil { c.JSON(http.StatusOK, gin.H{ @@ -152,14 +152,14 @@ func GitHubOAuth(c *gin.Context) { } } - if user.Status != common.UserStatusEnabled { + if user.Status != model.UserStatusEnabled { c.JSON(http.StatusOK, gin.H{ "message": "用户已被封禁", "success": false, }) return } - setupLogin(&user, c) + controller.SetupLogin(&user, c) } func GitHubBind(c *gin.Context) { @@ -219,7 +219,7 @@ func GitHubBind(c *gin.Context) { func GenerateOAuthCode(c *gin.Context) { session := sessions.Default(c) - state := helper.GetRandomString(12) + state := random.GetRandomString(12) session.Set("oauth_state", state) err := session.Save() if err != nil { diff --git a/controller/auth/lark.go b/controller/auth/lark.go new file mode 100644 index 00000000..eb06dde9 --- /dev/null +++ b/controller/auth/lark.go @@ -0,0 +1,200 @@ +package auth + +import ( + "bytes" + "encoding/json" + "errors" + "fmt" + "github.com/gin-contrib/sessions" + "github.com/gin-gonic/gin" + "github.com/songquanpeng/one-api/common/config" + "github.com/songquanpeng/one-api/common/logger" + "github.com/songquanpeng/one-api/controller" + "github.com/songquanpeng/one-api/model" + "net/http" + "strconv" + "time" +) + +type LarkOAuthResponse struct { + AccessToken string `json:"access_token"` +} + +type LarkUser struct { + Name string `json:"name"` + OpenID string `json:"open_id"` +} + +func getLarkUserInfoByCode(code string) (*LarkUser, error) { + if code == "" { + return nil, errors.New("无效的参数") + } + values := map[string]string{ + "client_id": config.LarkClientId, + "client_secret": config.LarkClientSecret, + "code": code, + "grant_type": "authorization_code", + "redirect_uri": fmt.Sprintf("%s/oauth/lark", config.ServerAddress), + } + jsonData, err := json.Marshal(values) + if err != nil { + return nil, err + } + req, err := http.NewRequest("POST", "https://passport.feishu.cn/suite/passport/oauth/token", bytes.NewBuffer(jsonData)) + if err != nil { + return nil, err + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json") + client := http.Client{ + Timeout: 5 * time.Second, + } + res, err := client.Do(req) + if err != nil { + logger.SysLog(err.Error()) + return nil, errors.New("无法连接至飞书服务器，请稍后重试！") + } + defer res.Body.Close() + var oAuthResponse LarkOAuthResponse + err = json.NewDecoder(res.Body).Decode(&oAuthResponse) + if err != nil { + return nil, err + } + req, err = http.NewRequest("GET", "https://passport.feishu.cn/suite/passport/oauth/userinfo", nil) + if err != nil { + return nil, err + } + req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", oAuthResponse.AccessToken)) + res2, err := client.Do(req) + if err != nil { + logger.SysLog(err.Error()) + return nil, errors.New("无法连接至飞书服务器，请稍后重试！") + } + var larkUser LarkUser + err = json.NewDecoder(res2.Body).Decode(&larkUser) + if err != nil { + return nil, err + } + return &larkUser, nil +} + +func LarkOAuth(c *gin.Context) { + session := sessions.Default(c) + state := c.Query("state") + if state == "" || session.Get("oauth_state") == nil || state != session.Get("oauth_state").(string) { + c.JSON(http.StatusForbidden, gin.H{ + "success": false, + "message": "state is empty or not same", + }) + return + } + username := session.Get("username") + if username != nil { + LarkBind(c) + return + } + code := c.Query("code") + larkUser, err := getLarkUserInfoByCode(code) + if err != nil { + c.JSON(http.StatusOK, gin.H{ + "success": false, + "message": err.Error(), + }) + return + } + user := model.User{ + LarkId: larkUser.OpenID, + } + if model.IsLarkIdAlreadyTaken(user.LarkId) { + err := user.FillUserByLarkId() + if err != nil { + c.JSON(http.StatusOK, gin.H{ + "success": false, + "message": err.Error(), + }) + return + } + } else { + if config.RegisterEnabled { + user.Username = "lark_" + strconv.Itoa(model.GetMaxUserId()+1) + if larkUser.Name != "" { + user.DisplayName = larkUser.Name + } else { + user.DisplayName = "Lark User" + } + user.Role = model.RoleCommonUser + user.Status = model.UserStatusEnabled + + if err := user.Insert(0); err != nil { + c.JSON(http.StatusOK, gin.H{ + "success": false, + "message": err.Error(), + }) + return + } + } else { + c.JSON(http.StatusOK, gin.H{ + "success": false, + "message": "管理员关闭了新用户注册", + }) + return + } + } + + if user.Status != model.UserStatusEnabled { + c.JSON(http.StatusOK, gin.H{ + "message": "用户已被封禁", + "success": false, + }) + return + } + controller.SetupLogin(&user, c) +} + +func LarkBind(c *gin.Context) { + code := c.Query("code") + larkUser, err := getLarkUserInfoByCode(code) + if err != nil { + c.JSON(http.StatusOK, gin.H{ + "success": false, + "message": err.Error(), + }) + return + } + user := model.User{ + LarkId: larkUser.OpenID, + } + if model.IsLarkIdAlreadyTaken(user.LarkId) { + c.JSON(http.StatusOK, gin.H{ + "success": false, + "message": "该飞书账户已被绑定", + }) + return + } + session := sessions.Default(c) + id := session.Get("id") + // id := c.GetInt("id") // critical bug! + user.Id = id.(int) + err = user.FillUserById() + if err != nil { + c.JSON(http.StatusOK, gin.H{ + "success": false, + "message": err.Error(), + }) + return + } + user.LarkId = larkUser.OpenID + err = user.Update(false) + if err != nil { + c.JSON(http.StatusOK, gin.H{ + "success": false, + "message": err.Error(), + }) + return + } + c.JSON(http.StatusOK, gin.H{ + "success": true, + "message": "bind", + }) + return +} diff --git a/controller/wechat.go b/controller/auth/wechat.go similarity index 93% rename from controller/wechat.go rename to controller/auth/wechat.go index 74be5604..a64746c9 100644 --- a/controller/wechat.go +++ b/controller/auth/wechat.go @@ -1,12 +1,12 @@ -package controller +package auth import ( "encoding/json" "errors" "fmt" "github.com/gin-gonic/gin" - "github.com/songquanpeng/one-api/common" "github.com/songquanpeng/one-api/common/config" + "github.com/songquanpeng/one-api/controller" "github.com/songquanpeng/one-api/model" "net/http" "strconv" @@ -83,8 +83,8 @@ func WeChatAuth(c *gin.Context) { if config.RegisterEnabled { user.Username = "wechat_" + strconv.Itoa(model.GetMaxUserId()+1) user.DisplayName = "WeChat User" - user.Role = common.RoleCommonUser - user.Status = common.UserStatusEnabled + user.Role = model.RoleCommonUser + user.Status = model.UserStatusEnabled if err := user.Insert(0); err != nil { c.JSON(http.StatusOK, gin.H{ @@ -102,14 +102,14 @@ func WeChatAuth(c *gin.Context) { } } - if user.Status != common.UserStatusEnabled { + if user.Status != model.UserStatusEnabled { c.JSON(http.StatusOK, gin.H{ "message": "用户已被封禁", "success": false, }) return } - setupLogin(&user, c) + controller.SetupLogin(&user, c) } func WeChatBind(c *gin.Context) { diff --git a/controller/channel-billing.go b/controller/channel-billing.go index 03c97349..b7ac61fd 100644 --- a/controller/channel-billing.go +++ b/controller/channel-billing.go @@ -4,12 +4,12 @@ import ( "encoding/json" "errors" "fmt" - "github.com/songquanpeng/one-api/common" "github.com/songquanpeng/one-api/common/config" "github.com/songquanpeng/one-api/common/logger" "github.com/songquanpeng/one-api/model" "github.com/songquanpeng/one-api/monitor" - "github.com/songquanpeng/one-api/relay/util" + "github.com/songquanpeng/one-api/relay/channeltype" + "github.com/songquanpeng/one-api/relay/client" "io" "net/http" "strconv" @@ -96,7 +96,7 @@ func GetResponseBody(method, url string, channel *model.Channel, headers http.He for k := range headers { req.Header.Add(k, headers.Get(k)) } - res, err := util.HTTPClient.Do(req) + res, err := client.HTTPClient.Do(req) if err != nil { return nil, err } @@ -204,28 +204,28 @@ func updateChannelAIGC2DBalance(channel *model.Channel) (float64, error) { } func updateChannelBalance(channel *model.Channel) (float64, error) { - baseURL := common.ChannelBaseURLs[channel.Type] + baseURL := channeltype.ChannelBaseURLs[channel.Type] if channel.GetBaseURL() == "" { channel.BaseURL = &baseURL } switch channel.Type { - case common.ChannelTypeOpenAI: + case channeltype.OpenAI: if channel.GetBaseURL() != "" { baseURL = channel.GetBaseURL() } - case common.ChannelTypeAzure: + case channeltype.Azure: return 0, errors.New("尚未实现") - case common.ChannelTypeCustom: + case channeltype.Custom: baseURL = channel.GetBaseURL() - case common.ChannelTypeCloseAI: + case channeltype.CloseAI: return updateChannelCloseAIBalance(channel) - case common.ChannelTypeOpenAISB: + case channeltype.OpenAISB: return updateChannelOpenAISBBalance(channel) - case common.ChannelTypeAIProxy: + case channeltype.AIProxy: return updateChannelAIProxyBalance(channel) - case common.ChannelTypeAPI2GPT: + case channeltype.API2GPT: return updateChannelAPI2GPTBalance(channel) - case common.ChannelTypeAIGC2D: + case channeltype.AIGC2D: return updateChannelAIGC2DBalance(channel) default: return 0, errors.New("尚未实现") @@ -301,11 +301,11 @@ func updateAllChannelsBalance() error { return err } for _, channel := range channels { - if channel.Status != common.ChannelStatusEnabled { + if channel.Status != model.ChannelStatusEnabled { continue } // TODO: support Azure - if channel.Type != common.ChannelTypeOpenAI && channel.Type != common.ChannelTypeCustom { + if channel.Type != channeltype.OpenAI && channel.Type != channeltype.Custom { continue } balance, err := updateChannelBalance(channel) diff --git a/controller/channel-test.go b/controller/channel-test.go index 95f4d769..ddbe0b4a 100644 --- a/controller/channel-test.go +++ b/controller/channel-test.go @@ -5,17 +5,18 @@ import ( "encoding/json" "errors" "fmt" - "github.com/songquanpeng/one-api/common" "github.com/songquanpeng/one-api/common/config" "github.com/songquanpeng/one-api/common/logger" "github.com/songquanpeng/one-api/common/message" "github.com/songquanpeng/one-api/middleware" "github.com/songquanpeng/one-api/model" "github.com/songquanpeng/one-api/monitor" - "github.com/songquanpeng/one-api/relay/constant" - "github.com/songquanpeng/one-api/relay/helper" + relay "github.com/songquanpeng/one-api/relay" + "github.com/songquanpeng/one-api/relay/channeltype" + "github.com/songquanpeng/one-api/relay/controller" + "github.com/songquanpeng/one-api/relay/meta" relaymodel "github.com/songquanpeng/one-api/relay/model" - "github.com/songquanpeng/one-api/relay/util" + "github.com/songquanpeng/one-api/relay/relaymode" "io" "net/http" "net/http/httptest" @@ -56,9 +57,9 @@ func testChannel(channel *model.Channel) (err error, openaiErr *relaymodel.Error c.Set("channel", channel.Type) c.Set("base_url", channel.GetBaseURL()) middleware.SetupContextForSelectedChannel(c, channel, "") - meta := util.GetRelayMeta(c) - apiType := constant.ChannelType2APIType(channel.Type) - adaptor := helper.GetAdaptor(apiType) + meta := meta.GetByContext(c) + apiType := channeltype.ToAPIType(channel.Type) + adaptor := relay.GetAdaptor(apiType) if adaptor == nil { return fmt.Errorf("invalid api type: %d, adaptor is nil", apiType), nil } @@ -73,7 +74,7 @@ func testChannel(channel *model.Channel) (err error, openaiErr *relaymodel.Error request := buildTestRequest() request.Model = modelName meta.OriginModelName, meta.ActualModelName = modelName, modelName - convertedRequest, err := adaptor.ConvertRequest(c, constant.RelayModeChatCompletions, request) + convertedRequest, err := adaptor.ConvertRequest(c, relaymode.ChatCompletions, request) if err != nil { return err, nil } @@ -88,7 +89,7 @@ func testChannel(channel *model.Channel) (err error, openaiErr *relaymodel.Error return err, nil } if resp.StatusCode != http.StatusOK { - err := util.RelayErrorHandler(resp) + err := controller.RelayErrorHandler(resp) return fmt.Errorf("status code %d: %s", resp.StatusCode, err.Error.Message), &err.Error } usage, respErr := adaptor.DoResponse(c, resp, meta) @@ -171,7 +172,7 @@ func testChannels(notify bool, scope string) error { } go func() { for _, channel := range channels { - isChannelEnabled := channel.Status == common.ChannelStatusEnabled + isChannelEnabled := channel.Status == model.ChannelStatusEnabled tik := time.Now() err, openaiErr := testChannel(channel) tok := time.Now() @@ -184,10 +185,10 @@ func testChannels(notify bool, scope string) error { _ = message.Notify(message.ByAll, fmt.Sprintf("渠道 %s （%d）测试超时", channel.Name, channel.Id), "", err.Error()) } } - if isChannelEnabled && util.ShouldDisableChannel(openaiErr, -1) { + if isChannelEnabled && monitor.ShouldDisableChannel(openaiErr, -1) { monitor.DisableChannel(channel.Id, channel.Name, err.Error()) } - if !isChannelEnabled && util.ShouldEnableChannel(err, openaiErr) { + if !isChannelEnabled && monitor.ShouldEnableChannel(err, openaiErr) { monitor.EnableChannel(channel.Id, channel.Name) } channel.UpdateResponseTime(milliseconds) diff --git a/controller/group.go b/controller/group.go index 128a3527..6f02394f 100644 --- a/controller/group.go +++ b/controller/group.go @@ -2,13 +2,13 @@ package controller import ( "github.com/gin-gonic/gin" - "github.com/songquanpeng/one-api/common" + billingratio "github.com/songquanpeng/one-api/relay/billing/ratio" "net/http" ) func GetGroups(c *gin.Context) { groupNames := make([]string, 0) - for groupName := range common.GroupRatio { + for groupName := range billingratio.GroupRatio { groupNames = append(groupNames, groupName) } c.JSON(http.StatusOK, gin.H{ diff --git a/controller/misc.go b/controller/misc.go index 544fd6b7..42ca748d 100644 --- a/controller/misc.go +++ b/controller/misc.go @@ -23,6 +23,7 @@ func GetStatus(c *gin.Context) { "email_verification": config.EmailVerificationEnabled, "github_oauth": config.GitHubOAuthEnabled, "github_client_id": config.GitHubClientId, + "lark_client_id": config.LarkClientId, "system_name": config.SystemName, "logo": config.Logo, "footer_html": config.Footer, diff --git a/controller/model.go b/controller/model.go index 53649391..77e2e94e 100644 --- a/controller/model.go +++ b/controller/model.go @@ -3,13 +3,13 @@ package controller import ( "fmt" "github.com/gin-gonic/gin" - "github.com/songquanpeng/one-api/common" "github.com/songquanpeng/one-api/model" - "github.com/songquanpeng/one-api/relay/channel/openai" - "github.com/songquanpeng/one-api/relay/constant" - "github.com/songquanpeng/one-api/relay/helper" + relay "github.com/songquanpeng/one-api/relay" + "github.com/songquanpeng/one-api/relay/adaptor/openai" + "github.com/songquanpeng/one-api/relay/apitype" + "github.com/songquanpeng/one-api/relay/channeltype" + "github.com/songquanpeng/one-api/relay/meta" relaymodel "github.com/songquanpeng/one-api/relay/model" - "github.com/songquanpeng/one-api/relay/util" "net/http" "strings" ) @@ -41,8 +41,8 @@ type OpenAIModels struct { Parent *string `json:"parent"` } -var openAIModels []OpenAIModels -var openAIModelsMap map[string]OpenAIModels +var models []OpenAIModels +var modelsMap map[string]OpenAIModels var channelId2Models map[int][]string func init() { @@ -62,15 +62,15 @@ func init() { IsBlocking: false, }) // https://platform.openai.com/docs/models/model-endpoint-compatibility - for i := 0; i < constant.APITypeDummy; i++ { - if i == constant.APITypeAIProxyLibrary { + for i := 0; i < apitype.Dummy; i++ { + if i == apitype.AIProxyLibrary { continue } - adaptor := helper.GetAdaptor(i) + adaptor := relay.GetAdaptor(i) channelName := adaptor.GetChannelName() modelNames := adaptor.GetModelList() for _, modelName := range modelNames { - openAIModels = append(openAIModels, OpenAIModels{ + models = append(models, OpenAIModels{ Id: modelName, Object: "model", Created: 1626777600, @@ -82,12 +82,12 @@ func init() { } } for _, channelType := range openai.CompatibleChannels { - if channelType == common.ChannelTypeAzure { + if channelType == channeltype.Azure { continue } channelName, channelModelList := openai.GetCompatibleChannelMeta(channelType) for _, modelName := range channelModelList { - openAIModels = append(openAIModels, OpenAIModels{ + models = append(models, OpenAIModels{ Id: modelName, Object: "model", Created: 1626777600, @@ -98,14 +98,14 @@ func init() { }) } } - openAIModelsMap = make(map[string]OpenAIModels) - for _, model := range openAIModels { - openAIModelsMap[model.Id] = model + modelsMap = make(map[string]OpenAIModels) + for _, model := range models { + modelsMap[model.Id] = model } channelId2Models = make(map[int][]string) - for i := 1; i < common.ChannelTypeDummy; i++ { - adaptor := helper.GetAdaptor(constant.ChannelType2APIType(i)) - meta := &util.RelayMeta{ + for i := 1; i < channeltype.Dummy; i++ { + adaptor := relay.GetAdaptor(channeltype.ToAPIType(i)) + meta := &meta.Meta{ ChannelType: i, } adaptor.Init(meta) @@ -121,6 +121,13 @@ func DashboardListModels(c *gin.Context) { }) } +func ListAllModels(c *gin.Context) { + c.JSON(200, gin.H{ + "object": "list", + "data": models, + }) +} + func ListModels(c *gin.Context) { ctx := c.Request.Context() var availableModels []string @@ -135,8 +142,8 @@ func ListModels(c *gin.Context) { for _, availableModel := range availableModels { modelSet[availableModel] = true } - var availableOpenAIModels []OpenAIModels - for _, model := range openAIModels { + availableOpenAIModels := make([]OpenAIModels, 0) + for _, model := range models { if _, ok := modelSet[model.Id]; ok { modelSet[model.Id] = false availableOpenAIModels = append(availableOpenAIModels, model) @@ -162,7 +169,7 @@ func ListModels(c *gin.Context) { func RetrieveModel(c *gin.Context) { modelId := c.Param("model") - if model, ok := openAIModelsMap[modelId]; ok { + if model, ok := modelsMap[modelId]; ok { c.JSON(200, model) } else { Error := relaymodel.Error{ diff --git a/controller/redemption.go b/controller/redemption.go index 31c9348d..8d2b3f38 100644 --- a/controller/redemption.go +++ b/controller/redemption.go @@ -4,6 +4,7 @@ import ( "github.com/gin-gonic/gin" "github.com/songquanpeng/one-api/common/config" "github.com/songquanpeng/one-api/common/helper" + "github.com/songquanpeng/one-api/common/random" "github.com/songquanpeng/one-api/model" "net/http" "strconv" @@ -106,7 +107,7 @@ func AddRedemption(c *gin.Context) { } var keys []string for i := 0; i < redemption.Count; i++ { - key := helper.GetUUID() + key := random.GetUUID() cleanRedemption := model.Redemption{ UserId: c.GetInt("id"), Name: redemption.Name, diff --git a/controller/relay.go b/controller/relay.go index b34768df..56359a1c 100644 --- a/controller/relay.go +++ b/controller/relay.go @@ -12,26 +12,25 @@ import ( "github.com/songquanpeng/one-api/middleware" dbmodel "github.com/songquanpeng/one-api/model" "github.com/songquanpeng/one-api/monitor" - "github.com/songquanpeng/one-api/relay/constant" "github.com/songquanpeng/one-api/relay/controller" "github.com/songquanpeng/one-api/relay/model" - "github.com/songquanpeng/one-api/relay/util" + "github.com/songquanpeng/one-api/relay/relaymode" "io" "net/http" ) // https://platform.openai.com/docs/api-reference/chat -func relay(c *gin.Context, relayMode int) *model.ErrorWithStatusCode { +func relayHelper(c *gin.Context, relayMode int) *model.ErrorWithStatusCode { var err *model.ErrorWithStatusCode switch relayMode { - case constant.RelayModeImagesGenerations: + case relaymode.ImagesGenerations: err = controller.RelayImageHelper(c, relayMode) - case constant.RelayModeAudioSpeech: + case relaymode.AudioSpeech: fallthrough - case constant.RelayModeAudioTranslation: + case relaymode.AudioTranslation: fallthrough - case constant.RelayModeAudioTranscription: + case relaymode.AudioTranscription: err = controller.RelayAudioHelper(c, relayMode) default: err = controller.RelayTextHelper(c) @@ -41,13 +40,13 @@ func relay(c *gin.Context, relayMode int) *model.ErrorWithStatusCode { func Relay(c *gin.Context) { ctx := c.Request.Context() - relayMode := constant.Path2RelayMode(c.Request.URL.Path) + relayMode := relaymode.GetByPath(c.Request.URL.Path) if config.DebugEnabled { requestBody, _ := common.GetRequestBody(c) logger.Debugf(ctx, "request body: %s", string(requestBody)) } channelId := c.GetInt("channel_id") - bizErr := relay(c, relayMode) + bizErr := relayHelper(c, relayMode) if bizErr == nil { monitor.Emit(channelId, true) return @@ -76,7 +75,7 @@ func Relay(c *gin.Context) { middleware.SetupContextForSelectedChannel(c, channel, originalModel) requestBody, err := common.GetRequestBody(c) c.Request.Body = io.NopCloser(bytes.NewBuffer(requestBody)) - bizErr = relay(c, relayMode) + bizErr = relayHelper(c, relayMode) if bizErr == nil { return } @@ -118,7 +117,7 @@ func shouldRetry(c *gin.Context, statusCode int) bool { func processChannelRelayError(ctx context.Context, channelId int, channelName string, err *model.ErrorWithStatusCode) { logger.Errorf(ctx, "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) { + if monitor.ShouldDisableChannel(&err.Error, err.StatusCode) { monitor.DisableChannel(channelId, channelName, err.Message) } else { monitor.Emit(channelId, false) diff --git a/controller/token.go b/controller/token.go index 13b90de0..557b5ce1 100644 --- a/controller/token.go +++ b/controller/token.go @@ -1,10 +1,12 @@ package controller import ( + "fmt" "github.com/gin-gonic/gin" - "github.com/songquanpeng/one-api/common" "github.com/songquanpeng/one-api/common/config" "github.com/songquanpeng/one-api/common/helper" + "github.com/songquanpeng/one-api/common/network" + "github.com/songquanpeng/one-api/common/random" "github.com/songquanpeng/one-api/model" "net/http" "strconv" @@ -104,6 +106,19 @@ func GetTokenStatus(c *gin.Context) { }) } +func validateToken(c *gin.Context, token model.Token) error { + if len(token.Name) > 30 { + return fmt.Errorf("令牌名称过长") + } + if token.Subnet != nil && *token.Subnet != "" { + err := network.IsValidSubnets(*token.Subnet) + if err != nil { + return fmt.Errorf("无效的网段：%s", err.Error()) + } + } + return nil +} + func AddToken(c *gin.Context) { token := model.Token{} err := c.ShouldBindJSON(&token) @@ -114,23 +129,26 @@ func AddToken(c *gin.Context) { }) return } - if len(token.Name) > 30 { + err = validateToken(c, token) + if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, - "message": "令牌名称过长", + "message": fmt.Sprintf("参数错误：%s", err.Error()), }) return } + cleanToken := model.Token{ UserId: c.GetInt("id"), Name: token.Name, - Key: helper.GenerateKey(), + Key: random.GenerateKey(), CreatedTime: helper.GetTimestamp(), AccessedTime: helper.GetTimestamp(), ExpiredTime: token.ExpiredTime, RemainQuota: token.RemainQuota, UnlimitedQuota: token.UnlimitedQuota, Models: token.Models, + Subnet: token.Subnet, } err = cleanToken.Insert() if err != nil { @@ -178,10 +196,11 @@ func UpdateToken(c *gin.Context) { }) return } - if len(token.Name) > 30 { + err = validateToken(c, token) + if err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, - "message": "令牌名称过长", + "message": fmt.Sprintf("参数错误：%s", err.Error()), }) return } @@ -193,15 +212,15 @@ func UpdateToken(c *gin.Context) { }) return } - if token.Status == common.TokenStatusEnabled { - if cleanToken.Status == common.TokenStatusExpired && cleanToken.ExpiredTime <= helper.GetTimestamp() && cleanToken.ExpiredTime != -1 { + if token.Status == model.TokenStatusEnabled { + if cleanToken.Status == model.TokenStatusExpired && cleanToken.ExpiredTime <= helper.GetTimestamp() && cleanToken.ExpiredTime != -1 { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "令牌已过期，无法启用，请先修改令牌过期时间，或者设置为永不过期", }) return } - if cleanToken.Status == common.TokenStatusExhausted && cleanToken.RemainQuota <= 0 && !cleanToken.UnlimitedQuota { + if cleanToken.Status == model.TokenStatusExhausted && cleanToken.RemainQuota <= 0 && !cleanToken.UnlimitedQuota { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "令牌可用额度已用尽，无法启用，请先修改令牌剩余额度，或者设置为无限额度", @@ -218,6 +237,7 @@ func UpdateToken(c *gin.Context) { cleanToken.RemainQuota = token.RemainQuota cleanToken.UnlimitedQuota = token.UnlimitedQuota cleanToken.Models = token.Models + cleanToken.Subnet = token.Subnet } err = cleanToken.Update() if err != nil { diff --git a/controller/user.go b/controller/user.go index 8b614e5d..44b4f793 100644 --- a/controller/user.go +++ b/controller/user.go @@ -5,7 +5,7 @@ import ( "fmt" "github.com/songquanpeng/one-api/common" "github.com/songquanpeng/one-api/common/config" - "github.com/songquanpeng/one-api/common/helper" + "github.com/songquanpeng/one-api/common/random" "github.com/songquanpeng/one-api/model" "net/http" "strconv" @@ -58,11 +58,11 @@ func Login(c *gin.Context) { }) return } - setupLogin(&user, c) + SetupLogin(&user, c) } // setup session & cookies and then return user info -func setupLogin(user *model.User, c *gin.Context) { +func SetupLogin(user *model.User, c *gin.Context) { session := sessions.Default(c) session.Set("id", user.Id) session.Set("username", user.Username) @@ -180,27 +180,27 @@ func Register(c *gin.Context) { } func GetAllUsers(c *gin.Context) { - p, _ := strconv.Atoi(c.Query("p")) - if p < 0 { - p = 0 - } - - order := c.DefaultQuery("order", "") - users, err := model.GetAllUsers(p*config.ItemsPerPage, config.ItemsPerPage, order) - - if err != nil { - c.JSON(http.StatusOK, gin.H{ - "success": false, - "message": err.Error(), - }) - return - } - - c.JSON(http.StatusOK, gin.H{ - "success": true, - "message": "", - "data": users, - }) + p, _ := strconv.Atoi(c.Query("p")) + if p < 0 { + p = 0 + } + + order := c.DefaultQuery("order", "") + users, err := model.GetAllUsers(p*config.ItemsPerPage, config.ItemsPerPage, order) + + if err != nil { + c.JSON(http.StatusOK, gin.H{ + "success": false, + "message": err.Error(), + }) + return + } + + c.JSON(http.StatusOK, gin.H{ + "success": true, + "message": "", + "data": users, + }) } func SearchUsers(c *gin.Context) { @@ -239,7 +239,7 @@ func GetUser(c *gin.Context) { return } myRole := c.GetInt("role") - if myRole <= user.Role && myRole != common.RoleRootUser { + if myRole <= user.Role && myRole != model.RoleRootUser { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "无权获取同级或更高等级用户的信息", @@ -287,7 +287,7 @@ func GenerateAccessToken(c *gin.Context) { }) return } - user.AccessToken = helper.GetUUID() + user.AccessToken = random.GetUUID() if model.DB.Where("access_token = ?", user.AccessToken).First(user).RowsAffected != 0 { c.JSON(http.StatusOK, gin.H{ @@ -324,7 +324,7 @@ func GetAffCode(c *gin.Context) { return } if user.AffCode == "" { - user.AffCode = helper.GetRandomString(4) + user.AffCode = random.GetRandomString(4) if err := user.Update(false); err != nil { c.JSON(http.StatusOK, gin.H{ "success": false, @@ -388,14 +388,14 @@ func UpdateUser(c *gin.Context) { return } myRole := c.GetInt("role") - if myRole <= originUser.Role && myRole != common.RoleRootUser { + if myRole <= originUser.Role && myRole != model.RoleRootUser { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "无权更新同权限等级或更高权限等级的用户信息", }) return } - if myRole <= updatedUser.Role && myRole != common.RoleRootUser { + if myRole <= updatedUser.Role && myRole != model.RoleRootUser { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "无权将其他用户权限等级提升到大于等于自己的权限等级", @@ -509,7 +509,7 @@ func DeleteSelf(c *gin.Context) { id := c.GetInt("id") user, _ := model.GetUserById(id, false) - if user.Role == common.RoleRootUser { + if user.Role == model.RoleRootUser { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "不能删除超级管理员账户", @@ -611,7 +611,7 @@ func ManageUser(c *gin.Context) { return } myRole := c.GetInt("role") - if myRole <= user.Role && myRole != common.RoleRootUser { + if myRole <= user.Role && myRole != model.RoleRootUser { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "无权更新同权限等级或更高权限等级的用户信息", @@ -620,8 +620,8 @@ func ManageUser(c *gin.Context) { } switch req.Action { case "disable": - user.Status = common.UserStatusDisabled - if user.Role == common.RoleRootUser { + user.Status = model.UserStatusDisabled + if user.Role == model.RoleRootUser { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "无法禁用超级管理员用户", @@ -629,9 +629,9 @@ func ManageUser(c *gin.Context) { return } case "enable": - user.Status = common.UserStatusEnabled + user.Status = model.UserStatusEnabled case "delete": - if user.Role == common.RoleRootUser { + if user.Role == model.RoleRootUser { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "无法删除超级管理员用户", @@ -646,37 +646,37 @@ func ManageUser(c *gin.Context) { return } case "promote": - if myRole != common.RoleRootUser { + if myRole != model.RoleRootUser { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "普通管理员用户无法提升其他用户为管理员", }) return } - if user.Role >= common.RoleAdminUser { + if user.Role >= model.RoleAdminUser { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "该用户已经是管理员", }) return } - user.Role = common.RoleAdminUser + user.Role = model.RoleAdminUser case "demote": - if user.Role == common.RoleRootUser { + if user.Role == model.RoleRootUser { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "无法降级超级管理员用户", }) return } - if user.Role == common.RoleCommonUser { + if user.Role == model.RoleCommonUser { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "该用户已经是普通用户", }) return } - user.Role = common.RoleCommonUser + user.Role = model.RoleCommonUser } if err := user.Update(false); err != nil { @@ -730,7 +730,7 @@ func EmailBind(c *gin.Context) { }) return } - if user.Role == common.RoleRootUser { + if user.Role == model.RoleRootUser { config.RootUserEmail = email } c.JSON(http.StatusOK, gin.H{ @@ -770,3 +770,38 @@ func TopUp(c *gin.Context) { }) return } + +type adminTopUpRequest struct { + UserId int `json:"user_id"` + Quota int `json:"quota"` + Remark string `json:"remark"` +} + +func AdminTopUp(c *gin.Context) { + req := adminTopUpRequest{} + err := c.ShouldBindJSON(&req) + if err != nil { + c.JSON(http.StatusOK, gin.H{ + "success": false, + "message": err.Error(), + }) + return + } + err = model.IncreaseUserQuota(req.UserId, int64(req.Quota)) + if err != nil { + c.JSON(http.StatusOK, gin.H{ + "success": false, + "message": err.Error(), + }) + return + } + if req.Remark == "" { + req.Remark = fmt.Sprintf("通过 API 充值 %s", common.LogQuota(int64(req.Quota))) + } + model.RecordTopupLog(req.UserId, req.Remark, req.Quota) + c.JSON(http.StatusOK, gin.H{ + "success": true, + "message": "", + }) + return +} diff --git a/docs/API.md b/docs/API.md new file mode 100644 index 00000000..0b7ddf5a --- /dev/null +++ b/docs/API.md @@ -0,0 +1,53 @@ +# 使用 API 操控 & 扩展 One API +> 欢迎提交 PR 在此放上你的拓展项目。 + +例如，虽然 One API 本身没有直接支持支付，但是你可以通过系统扩展的 API 来实现支付功能。 + +又或者你想自定义渠道管理策略，也可以通过 API 来实现渠道的禁用与启用。 + +## 鉴权 +One API 支持两种鉴权方式：Cookie 和 Token，对于 Token，参照下图获取： + +![image](https://github.com/songquanpeng/songquanpeng.github.io/assets/39998050/c15281a7-83ed-47cb-a1f6-913cb6bf4a7c) + +之后，将 Token 作为请求头的 Authorization 字段的值即可，例如下面使用 Token 调用测试渠道的 API： +![image](https://github.com/songquanpeng/songquanpeng.github.io/assets/39998050/1273b7ae-cb60-4c0d-93a6-b1cbc039c4f8) + +## 请求格式与响应格式 +One API 使用 JSON 格式进行请求和响应。 + +对于响应体，一般格式如下： +```json +{ + "message": "请求信息", + "success": true, + "data": {} +} +``` + +## API 列表 +> 当前 API 列表不全，请自行通过浏览器抓取前端请求 + +如果现有的 API 没有办法满足你的需求，欢迎提交 issue 讨论。 + +### 获取当前登录用户信息 +**GET** `/api/user/self` + +### 为给定用户充值额度 +**POST** `/api/topup` +```json +{ + "user_id": 1, + "quota": 100000, + "remark": "充值 100000 额度" +} +``` + +## 其他 +### 充值链接上的附加参数 +One API 会在用户点击充值按钮的时候，将用户的信息和充值信息附加在链接上，例如： +`https://example.com?username=root&user_id=1&transaction_id=4b3eed80-55d5-443f-bd44-fb18c648c837` + +你可以通过解析链接上的参数来获取用户信息和充值信息，然后调用 API 来为用户充值。 + +注意，不是所有主题都支持该功能，欢迎 PR 补齐。 \ No newline at end of file diff --git a/go.mod b/go.mod index f9ed96d3..6ace51f2 100644 --- a/go.mod +++ b/go.mod @@ -15,6 +15,7 @@ require ( github.com/google/uuid v1.3.0 github.com/gorilla/websocket v1.5.0 github.com/pkoukk/tiktoken-go v0.1.5 + github.com/smartystreets/goconvey v1.8.1 github.com/stretchr/testify v1.8.3 golang.org/x/crypto v0.17.0 golang.org/x/image v0.14.0 @@ -37,6 +38,7 @@ require ( github.com/go-playground/universal-translator v0.18.1 // indirect github.com/go-sql-driver/mysql v1.6.0 // indirect github.com/goccy/go-json v0.10.2 // indirect + github.com/gopherjs/gopherjs v1.17.2 // indirect github.com/gorilla/context v1.1.1 // indirect github.com/gorilla/securecookie v1.1.1 // indirect github.com/gorilla/sessions v1.2.1 // indirect @@ -47,6 +49,7 @@ require ( github.com/jinzhu/inflection v1.0.0 // indirect github.com/jinzhu/now v1.1.5 // indirect github.com/json-iterator/go v1.1.12 // indirect + github.com/jtolds/gls v4.20.0+incompatible // indirect github.com/klauspost/cpuid/v2 v2.2.4 // indirect github.com/leodido/go-urn v1.2.4 // indirect github.com/mattn/go-isatty v0.0.19 // indirect @@ -55,6 +58,7 @@ require ( github.com/modern-go/reflect2 v1.0.2 // indirect github.com/pelletier/go-toml/v2 v2.0.8 // indirect github.com/pmezard/go-difflib v1.0.0 // indirect + github.com/smarty/assertions v1.15.0 // indirect github.com/twitchyliquid64/golang-asm v0.15.1 // indirect github.com/ugorji/go/codec v1.2.11 // indirect golang.org/x/arch v0.3.0 // indirect diff --git a/go.sum b/go.sum index 9cf056e5..3ead2711 100644 --- a/go.sum +++ b/go.sum @@ -56,11 +56,13 @@ github.com/golang-jwt/jwt v3.2.2+incompatible h1:IfV12K8xAKAnZqdXVzCZ+TOjboZ2keL github.com/golang-jwt/jwt v3.2.2+incompatible/go.mod h1:8pz2t5EyA70fFQQSrl6XZXzqecmYZeUEB8OUGHkxJ+I= github.com/golang/protobuf v1.3.3/go.mod h1:vzj43D7+SQXF/4pzW/hwtAqwc6iTitCiVSaWz5lYuqw= github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaSAoJOfIk= -github.com/google/go-cmp v0.5.5 h1:Khx7svrCpmxxtHBq5j2mp/xVjsi8hQMfNLvJFAlrGgU= github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= +github.com/google/go-cmp v0.5.6 h1:BKbKCqvP6I+rmFHt06ZmyQtvB8xAkWdhFyr0ZUNZcxQ= github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg= github.com/google/uuid v1.3.0 h1:t6JiXgmwXMjEs8VusXIJk2BXHsn+wx8BZdTaoZ5fu7I= github.com/google/uuid v1.3.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/gopherjs/gopherjs v1.17.2 h1:fQnZVsXk8uxXIStYb0N4bGk7jeyTalG/wsZjQ25dO0g= +github.com/gopherjs/gopherjs v1.17.2/go.mod h1:pRRIvn/QzFLrKfvEz3qUuEhtE/zLCWfreZ6J5gM2i+k= github.com/gorilla/context v1.1.1 h1:AWwleXJkX/nhcU9bZSnZoi3h/qGYqQAGhq6zZe/aQW8= github.com/gorilla/context v1.1.1/go.mod h1:kBGZzfjB9CEq2AlWe17Uuf7NDRt0dE0s8S51q0aT7Yg= github.com/gorilla/securecookie v1.1.1 h1:miw7JPhV+b/lAHSXz4qd/nN9jRiAFV5FwjeKyCS8BvQ= @@ -85,6 +87,8 @@ github.com/jinzhu/now v1.1.5/go.mod h1:d3SSVoowX0Lcu0IBviAWJpolVfI5UJVZZ7cO71lE/ github.com/json-iterator/go v1.1.9/go.mod h1:KdQUCv79m/52Kvf8AW2vK1V8akMuk1QjK/uOdHXbAo4= github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM= github.com/json-iterator/go v1.1.12/go.mod h1:e30LSqwooZae/UwlEbR2852Gd8hjQvJoHmT4TnhNGBo= +github.com/jtolds/gls v4.20.0+incompatible h1:xdiiI2gbIgH/gLH7ADydsJ1uDOEzR8yvV7C0MuV77Wo= +github.com/jtolds/gls v4.20.0+incompatible/go.mod h1:QJZ7F/aHp+rZTRtaJ1ow/lLfFfVYBRgL+9YlvaHOwJU= github.com/klauspost/cpuid/v2 v2.0.9/go.mod h1:FInQzS24/EEf25PyTYn52gqo7WaD8xa0213Md/qVLRg= github.com/klauspost/cpuid/v2 v2.2.4 h1:acbojRNwl3o09bUq+yDCtZFc1aiwaAAxtcn8YkZXnvk= github.com/klauspost/cpuid/v2 v2.2.4/go.mod h1:RVVoqg1df56z8g3pUjL/3lE5UfnlrJX8tyFgg4nqhuY= @@ -127,6 +131,10 @@ github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZN github.com/rogpeppe/go-internal v1.6.1/go.mod h1:xXDCJY+GAPziupqXw64V24skbSoqbTEfhy4qGm1nDQc= github.com/rogpeppe/go-internal v1.8.0 h1:FCbCCtXNOY3UtUuHUYaghJg4y7Fd14rXifAYUAtL9R8= github.com/rogpeppe/go-internal v1.8.0/go.mod h1:WmiCO8CzOY8rg0OYDC4/i/2WRWAB6poM+XZ2dLUbcbE= +github.com/smarty/assertions v1.15.0 h1:cR//PqUBUiQRakZWqBiFFQ9wb8emQGDb0HeGdqGByCY= +github.com/smarty/assertions v1.15.0/go.mod h1:yABtdzeQs6l1brC900WlRNwj6ZR55d7B+E8C6HtKdec= +github.com/smartystreets/goconvey v1.8.1 h1:qGjIddxOk4grTu9JPOU31tVfq3cNdBlNa5sSznIX1xY= +github.com/smartystreets/goconvey v1.8.1/go.mod h1:+/u4qLyY6x1jReYOp7GOM2FSt8aP9CzCZL03bI28W60= github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw= github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo= @@ -177,8 +185,8 @@ golang.org/x/text v0.3.6/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= golang.org/x/text v0.14.0 h1:ScX5w1eTa3QqT8oi6+ziP7dTV1S2+ALU0bI+0zXKWiQ= golang.org/x/text v0.14.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU= golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= -golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543 h1:E7g+9GITq07hpfrRu66IVDexMakfv52eLZ2CXBWiKr4= golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= +golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1 h1:go1bK/D/BFZV2I8cIQd1NKEZ+0owSTG1fDTci4IqFcE= google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw= google.golang.org/protobuf v1.28.0/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I= google.golang.org/protobuf v1.33.0 h1:uNO2rsAINq/JlFpSdYEKIZ0uKD/R9cpdv0T+yoGwGmI= diff --git a/main.go b/main.go index b20c6daf..a0621711 100644 --- a/main.go +++ b/main.go @@ -12,7 +12,7 @@ import ( "github.com/songquanpeng/one-api/controller" "github.com/songquanpeng/one-api/middleware" "github.com/songquanpeng/one-api/model" - "github.com/songquanpeng/one-api/relay/channel/openai" + "github.com/songquanpeng/one-api/relay/adaptor/openai" "github.com/songquanpeng/one-api/router" "os" "strconv" diff --git a/middleware/auth.go b/middleware/auth.go index 29701524..01b2cce3 100644 --- a/middleware/auth.go +++ b/middleware/auth.go @@ -4,8 +4,8 @@ import ( "fmt" "github.com/gin-contrib/sessions" "github.com/gin-gonic/gin" - "github.com/songquanpeng/one-api/common" "github.com/songquanpeng/one-api/common/blacklist" + "github.com/songquanpeng/one-api/common/network" "github.com/songquanpeng/one-api/model" "net/http" "strings" @@ -44,7 +44,7 @@ func authHelper(c *gin.Context, minRole int) { return } } - if status.(int) == common.UserStatusDisabled || blacklist.IsUserBanned(id.(int)) { + if status.(int) == model.UserStatusDisabled || blacklist.IsUserBanned(id.(int)) { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "用户已被封禁", @@ -71,24 +71,25 @@ func authHelper(c *gin.Context, minRole int) { func UserAuth() func(c *gin.Context) { return func(c *gin.Context) { - authHelper(c, common.RoleCommonUser) + authHelper(c, model.RoleCommonUser) } } func AdminAuth() func(c *gin.Context) { return func(c *gin.Context) { - authHelper(c, common.RoleAdminUser) + authHelper(c, model.RoleAdminUser) } } func RootAuth() func(c *gin.Context) { return func(c *gin.Context) { - authHelper(c, common.RoleRootUser) + authHelper(c, model.RoleRootUser) } } func TokenAuth() func(c *gin.Context) { return func(c *gin.Context) { + ctx := c.Request.Context() key := c.Request.Header.Get("Authorization") key = strings.TrimPrefix(key, "Bearer ") key = strings.TrimPrefix(key, "sk-") @@ -99,6 +100,12 @@ func TokenAuth() func(c *gin.Context) { abortWithMessage(c, http.StatusUnauthorized, err.Error()) return } + if token.Subnet != nil && *token.Subnet != "" { + if !network.IsIpInSubnets(ctx, c.ClientIP(), *token.Subnet) { + abortWithMessage(c, http.StatusForbidden, fmt.Sprintf("该令牌只能在指定网段使用：%s，当前 ip：%s", *token.Subnet, c.ClientIP())) + return + } + } userEnabled, err := model.CacheIsUserEnabled(token.UserId) if err != nil { abortWithMessage(c, http.StatusInternalServerError, err.Error()) @@ -109,7 +116,7 @@ func TokenAuth() func(c *gin.Context) { return } requestModel, err := getRequestModel(c) - if err != nil { + if err != nil && !strings.HasPrefix(c.Request.URL.Path, "/v1/models") { abortWithMessage(c, http.StatusBadRequest, err.Error()) return } diff --git a/middleware/distributor.go b/middleware/distributor.go index 04489a2b..6e0d2718 100644 --- a/middleware/distributor.go +++ b/middleware/distributor.go @@ -3,9 +3,10 @@ package middleware import ( "fmt" "github.com/gin-gonic/gin" - "github.com/songquanpeng/one-api/common" + "github.com/songquanpeng/one-api/common/config" "github.com/songquanpeng/one-api/common/logger" "github.com/songquanpeng/one-api/model" + "github.com/songquanpeng/one-api/relay/channeltype" "net/http" "strconv" ) @@ -33,7 +34,7 @@ func Distribute() func(c *gin.Context) { abortWithMessage(c, http.StatusBadRequest, "无效的渠道 Id") return } - if channel.Status != common.ChannelStatusEnabled { + if channel.Status != model.ChannelStatusEnabled { abortWithMessage(c, http.StatusForbidden, "该渠道已被禁用") return } @@ -66,19 +67,19 @@ func SetupContextForSelectedChannel(c *gin.Context, channel *model.Channel, mode c.Set("base_url", channel.GetBaseURL()) // this is for backward compatibility switch channel.Type { - case common.ChannelTypeAzure: - c.Set(common.ConfigKeyAPIVersion, channel.Other) - case common.ChannelTypeXunfei: - c.Set(common.ConfigKeyAPIVersion, channel.Other) - case common.ChannelTypeGemini: - c.Set(common.ConfigKeyAPIVersion, channel.Other) - case common.ChannelTypeAIProxyLibrary: - c.Set(common.ConfigKeyLibraryID, channel.Other) - case common.ChannelTypeAli: - c.Set(common.ConfigKeyPlugin, channel.Other) + case channeltype.Azure: + c.Set(config.KeyAPIVersion, channel.Other) + case channeltype.Xunfei: + c.Set(config.KeyAPIVersion, channel.Other) + case channeltype.Gemini: + c.Set(config.KeyAPIVersion, channel.Other) + case channeltype.AIProxyLibrary: + c.Set(config.KeyLibraryID, channel.Other) + case channeltype.Ali: + c.Set(config.KeyPlugin, channel.Other) } cfg, _ := channel.LoadConfig() for k, v := range cfg { - c.Set(common.ConfigKeyPrefix+k, v) + c.Set(config.KeyPrefix+k, v) } } diff --git a/model/ability.go b/model/ability.go index 4a48bc51..2db72518 100644 --- a/model/ability.go +++ b/model/ability.go @@ -57,7 +57,7 @@ func (channel *Channel) AddAbilities() error { Group: group, Model: model, ChannelId: channel.Id, - Enabled: channel.Status == common.ChannelStatusEnabled, + Enabled: channel.Status == ChannelStatusEnabled, Priority: channel.Priority, } abilities = append(abilities, ability) diff --git a/model/cache.go b/model/cache.go index cfc5445a..cfb0f8a4 100644 --- a/model/cache.go +++ b/model/cache.go @@ -8,6 +8,7 @@ import ( "github.com/songquanpeng/one-api/common" "github.com/songquanpeng/one-api/common/config" "github.com/songquanpeng/one-api/common/logger" + "github.com/songquanpeng/one-api/common/random" "math/rand" "sort" "strconv" @@ -172,7 +173,7 @@ var channelSyncLock sync.RWMutex func InitChannelCache() { newChannelId2channel := make(map[int]*Channel) var channels []*Channel - DB.Where("status = ?", common.ChannelStatusEnabled).Find(&channels) + DB.Where("status = ?", ChannelStatusEnabled).Find(&channels) for _, channel := range channels { newChannelId2channel[channel.Id] = channel } @@ -247,7 +248,7 @@ func CacheGetRandomSatisfiedChannel(group string, model string, ignoreFirstPrior idx := rand.Intn(endIdx) if ignoreFirstPriority { if endIdx < len(channels) { // which means there are more than one priority - idx = common.RandRange(endIdx, len(channels)) + idx = random.RandRange(endIdx, len(channels)) } } return channels[idx], nil diff --git a/model/channel.go b/model/channel.go index fc4905b1..e667f7e7 100644 --- a/model/channel.go +++ b/model/channel.go @@ -3,13 +3,19 @@ package model import ( "encoding/json" "fmt" - "github.com/songquanpeng/one-api/common" "github.com/songquanpeng/one-api/common/config" "github.com/songquanpeng/one-api/common/helper" "github.com/songquanpeng/one-api/common/logger" "gorm.io/gorm" ) +const ( + ChannelStatusUnknown = 0 + ChannelStatusEnabled = 1 // don't use 0, 0 is the default value! + ChannelStatusManuallyDisabled = 2 // also don't use 0 + ChannelStatusAutoDisabled = 3 +) + type Channel struct { Id int `json:"id"` Type int `json:"type" gorm:"default:0"` @@ -39,7 +45,7 @@ func GetAllChannels(startIdx int, num int, scope string) ([]*Channel, error) { case "all": err = DB.Order("id desc").Find(&channels).Error case "disabled": - err = DB.Order("id desc").Where("status = ? or status = ?", common.ChannelStatusAutoDisabled, common.ChannelStatusManuallyDisabled).Find(&channels).Error + err = DB.Order("id desc").Where("status = ? or status = ?", ChannelStatusAutoDisabled, ChannelStatusManuallyDisabled).Find(&channels).Error default: err = DB.Order("id desc").Limit(num).Offset(startIdx).Omit("key").Find(&channels).Error } @@ -168,7 +174,7 @@ func (channel *Channel) LoadConfig() (map[string]string, error) { } func UpdateChannelStatusById(id int, status int) { - err := UpdateAbilityStatus(id, status == common.ChannelStatusEnabled) + err := UpdateAbilityStatus(id, status == ChannelStatusEnabled) if err != nil { logger.SysError("failed to update ability status: " + err.Error()) } @@ -199,6 +205,6 @@ func DeleteChannelByStatus(status int64) (int64, error) { } func DeleteDisabledChannel() (int64, error) { - result := DB.Where("status = ? or status = ?", common.ChannelStatusAutoDisabled, common.ChannelStatusManuallyDisabled).Delete(&Channel{}) + result := DB.Where("status = ? or status = ?", ChannelStatusAutoDisabled, ChannelStatusManuallyDisabled).Delete(&Channel{}) return result.RowsAffected, result.Error } diff --git a/model/log.go b/model/log.go index 4409f73e..6fba776a 100644 --- a/model/log.go +++ b/model/log.go @@ -7,7 +7,6 @@ import ( "github.com/songquanpeng/one-api/common/config" "github.com/songquanpeng/one-api/common/helper" "github.com/songquanpeng/one-api/common/logger" - "gorm.io/gorm" ) @@ -51,6 +50,21 @@ func RecordLog(userId int, logType int, content string) { } } +func RecordTopupLog(userId int, content string, quota int) { + log := &Log{ + UserId: userId, + Username: GetUsernameById(userId), + CreatedAt: helper.GetTimestamp(), + Type: LogTypeTopup, + Content: content, + Quota: quota, + } + err := LOG_DB.Create(log).Error + if err != nil { + 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 int64, content string) { 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 { diff --git a/model/main.go b/model/main.go index ff542faf..4b5323c4 100644 --- a/model/main.go +++ b/model/main.go @@ -7,6 +7,7 @@ import ( "github.com/songquanpeng/one-api/common/env" "github.com/songquanpeng/one-api/common/helper" "github.com/songquanpeng/one-api/common/logger" + "github.com/songquanpeng/one-api/common/random" "gorm.io/driver/mysql" "gorm.io/driver/postgres" "gorm.io/driver/sqlite" @@ -31,10 +32,10 @@ func CreateRootAccountIfNeed() error { rootUser := User{ Username: "root", Password: hashedPassword, - Role: common.RoleRootUser, - Status: common.UserStatusEnabled, + Role: RoleRootUser, + Status: UserStatusEnabled, DisplayName: "Root User", - AccessToken: helper.GetUUID(), + AccessToken: random.GetUUID(), Quota: 500000000000000, } DB.Create(&rootUser) @@ -44,7 +45,7 @@ func CreateRootAccountIfNeed() error { Id: 1, UserId: rootUser.Id, Key: config.InitialRootToken, - Status: common.TokenStatusEnabled, + Status: TokenStatusEnabled, Name: "Initial Root Token", CreatedTime: helper.GetTimestamp(), AccessedTime: helper.GetTimestamp(), diff --git a/model/option.go b/model/option.go index 1d1c28b4..bed8d4c3 100644 --- a/model/option.go +++ b/model/option.go @@ -1,9 +1,9 @@ package model import ( - "github.com/songquanpeng/one-api/common" "github.com/songquanpeng/one-api/common/config" "github.com/songquanpeng/one-api/common/logger" + billingratio "github.com/songquanpeng/one-api/relay/billing/ratio" "strconv" "strings" "time" @@ -66,9 +66,9 @@ func InitOptionMap() { config.OptionMap["QuotaForInvitee"] = strconv.FormatInt(config.QuotaForInvitee, 10) config.OptionMap["QuotaRemindThreshold"] = strconv.FormatInt(config.QuotaRemindThreshold, 10) config.OptionMap["PreConsumedQuota"] = strconv.FormatInt(config.PreConsumedQuota, 10) - config.OptionMap["ModelRatio"] = common.ModelRatio2JSONString() - config.OptionMap["GroupRatio"] = common.GroupRatio2JSONString() - config.OptionMap["CompletionRatio"] = common.CompletionRatio2JSONString() + config.OptionMap["ModelRatio"] = billingratio.ModelRatio2JSONString() + config.OptionMap["GroupRatio"] = billingratio.GroupRatio2JSONString() + config.OptionMap["CompletionRatio"] = billingratio.CompletionRatio2JSONString() config.OptionMap["TopUpLink"] = config.TopUpLink config.OptionMap["ChatLink"] = config.ChatLink config.OptionMap["QuotaPerUnit"] = strconv.FormatFloat(config.QuotaPerUnit, 'f', -1, 64) @@ -82,7 +82,7 @@ func loadOptionsFromDatabase() { options, _ := AllOption() for _, option := range options { if option.Key == "ModelRatio" { - option.Value = common.AddNewMissingRatio(option.Value) + option.Value = billingratio.AddNewMissingRatio(option.Value) } err := updateOptionMap(option.Key, option.Value) if err != nil { @@ -172,6 +172,10 @@ func updateOptionMap(key string, value string) (err error) { config.GitHubClientId = value case "GitHubClientSecret": config.GitHubClientSecret = value + case "LarkClientId": + config.LarkClientId = value + case "LarkClientSecret": + config.LarkClientSecret = value case "Footer": config.Footer = value case "SystemName": @@ -205,11 +209,11 @@ func updateOptionMap(key string, value string) (err error) { case "RetryTimes": config.RetryTimes, _ = strconv.Atoi(value) case "ModelRatio": - err = common.UpdateModelRatioByJSONString(value) + err = billingratio.UpdateModelRatioByJSONString(value) case "GroupRatio": - err = common.UpdateGroupRatioByJSONString(value) + err = billingratio.UpdateGroupRatioByJSONString(value) case "CompletionRatio": - err = common.UpdateCompletionRatioByJSONString(value) + err = billingratio.UpdateCompletionRatioByJSONString(value) case "TopUpLink": config.TopUpLink = value case "ChatLink": diff --git a/model/redemption.go b/model/redemption.go index 79a5b8a9..45871a71 100644 --- a/model/redemption.go +++ b/model/redemption.go @@ -8,6 +8,12 @@ import ( "gorm.io/gorm" ) +const ( + RedemptionCodeStatusEnabled = 1 // don't use 0, 0 is the default value! + RedemptionCodeStatusDisabled = 2 // also don't use 0 + RedemptionCodeStatusUsed = 3 // also don't use 0 +) + type Redemption struct { Id int `json:"id"` UserId int `json:"user_id"` @@ -61,7 +67,7 @@ func Redeem(key string, userId int) (quota int64, err error) { if err != nil { return errors.New("无效的兑换码") } - if redemption.Status != common.RedemptionCodeStatusEnabled { + if redemption.Status != RedemptionCodeStatusEnabled { return errors.New("该兑换码已被使用") } err = tx.Model(&User{}).Where("id = ?", userId).Update("quota", gorm.Expr("quota + ?", redemption.Quota)).Error @@ -69,7 +75,7 @@ func Redeem(key string, userId int) (quota int64, err error) { return err } redemption.RedeemedTime = helper.GetTimestamp() - redemption.Status = common.RedemptionCodeStatusUsed + redemption.Status = RedemptionCodeStatusUsed err = tx.Save(redemption).Error return err }) diff --git a/model/token.go b/model/token.go index fef80fcf..96e6b491 100644 --- a/model/token.go +++ b/model/token.go @@ -11,6 +11,13 @@ import ( "gorm.io/gorm" ) +const ( + TokenStatusEnabled = 1 // don't use 0, 0 is the default value! + TokenStatusDisabled = 2 // also don't use 0 + TokenStatusExpired = 3 + TokenStatusExhausted = 4 +) + type Token struct { Id int `json:"id"` UserId int `json:"user_id"` @@ -23,7 +30,8 @@ type Token struct { RemainQuota int64 `json:"remain_quota" gorm:"bigint;default:0"` UnlimitedQuota bool `json:"unlimited_quota" gorm:"default:false"` UsedQuota int64 `json:"used_quota" gorm:"bigint;default:0"` // used quota - Models *string `json:"models" gorm:"default:''"` + Models *string `json:"models" gorm:"default:''"` // allowed models + Subnet *string `json:"subnet" gorm:"default:''"` // allowed subnet } func GetAllUserTokens(userId int, startIdx int, num int, order string) ([]*Token, error) { @@ -61,17 +69,17 @@ func ValidateUserToken(key string) (token *Token, err error) { } return nil, errors.New("令牌验证失败") } - if token.Status == common.TokenStatusExhausted { - return nil, errors.New("该令牌额度已用尽") - } else if token.Status == common.TokenStatusExpired { + if token.Status == TokenStatusExhausted { + return nil, fmt.Errorf("令牌 %s（#%d）额度已用尽", token.Name, token.Id) + } else if token.Status == TokenStatusExpired { return nil, errors.New("该令牌已过期") } - if token.Status != common.TokenStatusEnabled { + if token.Status != TokenStatusEnabled { return nil, errors.New("该令牌状态不可用") } if token.ExpiredTime != -1 && token.ExpiredTime < helper.GetTimestamp() { if !common.RedisEnabled { - token.Status = common.TokenStatusExpired + token.Status = TokenStatusExpired err := token.SelectUpdate() if err != nil { logger.SysError("failed to update token status" + err.Error()) @@ -82,7 +90,7 @@ func ValidateUserToken(key string) (token *Token, err error) { if !token.UnlimitedQuota && token.RemainQuota <= 0 { if !common.RedisEnabled { // in this case, we can make sure the token is exhausted - token.Status = common.TokenStatusExhausted + token.Status = TokenStatusExhausted err := token.SelectUpdate() if err != nil { logger.SysError("failed to update token status" + err.Error()) @@ -122,7 +130,7 @@ func (token *Token) Insert() error { // Update Make sure your token's fields is completed, because this will update non-zero values func (token *Token) Update() error { var err error - err = DB.Model(token).Select("name", "status", "expired_time", "remain_quota", "unlimited_quota", "models").Updates(token).Error + err = DB.Model(token).Select("name", "status", "expired_time", "remain_quota", "unlimited_quota", "models", "subnet").Updates(token).Error return err } diff --git a/model/user.go b/model/user.go index 5e729b5e..1dc633b1 100644 --- a/model/user.go +++ b/model/user.go @@ -6,12 +6,25 @@ import ( "github.com/songquanpeng/one-api/common" "github.com/songquanpeng/one-api/common/blacklist" "github.com/songquanpeng/one-api/common/config" - "github.com/songquanpeng/one-api/common/helper" "github.com/songquanpeng/one-api/common/logger" + "github.com/songquanpeng/one-api/common/random" "gorm.io/gorm" "strings" ) +const ( + RoleGuestUser = 0 + RoleCommonUser = 1 + RoleAdminUser = 10 + RoleRootUser = 100 +) + +const ( + UserStatusEnabled = 1 // don't use 0, 0 is the default value! + UserStatusDisabled = 2 // also don't use 0 + UserStatusDeleted = 3 +) + // User if you add sensitive fields, don't forget to clean them in setupLogin function. // Otherwise, the sensitive information will be saved on local storage in plain text! type User struct { @@ -24,6 +37,7 @@ type User struct { Email string `json:"email" gorm:"index" validate:"max=50"` GitHubId string `json:"github_id" gorm:"column:github_id;index"` WeChatId string `json:"wechat_id" gorm:"column:wechat_id;index"` + LarkId string `json:"lark_id" gorm:"column:lark_id;index"` VerificationCode string `json:"verification_code" gorm:"-:all"` // this field is only for Email verification, don't save it to database! AccessToken string `json:"access_token" gorm:"type:char(32);column:access_token;uniqueIndex"` // this token is for system management Quota int64 `json:"quota" gorm:"bigint;default:0"` @@ -41,21 +55,21 @@ func GetMaxUserId() int { } func GetAllUsers(startIdx int, num int, order string) (users []*User, err error) { - query := DB.Limit(num).Offset(startIdx).Omit("password").Where("status != ?", common.UserStatusDeleted) - - switch order { - case "quota": - query = query.Order("quota desc") - case "used_quota": - query = query.Order("used_quota desc") - case "request_count": - query = query.Order("request_count desc") - default: - query = query.Order("id desc") - } - - err = query.Find(&users).Error - return users, err + query := DB.Limit(num).Offset(startIdx).Omit("password").Where("status != ?", UserStatusDeleted) + + switch order { + case "quota": + query = query.Order("quota desc") + case "used_quota": + query = query.Order("used_quota desc") + case "request_count": + query = query.Order("request_count desc") + default: + query = query.Order("id desc") + } + + err = query.Find(&users).Error + return users, err } func SearchUsers(keyword string) (users []*User, err error) { @@ -107,8 +121,8 @@ func (user *User) Insert(inviterId int) error { } } user.Quota = config.QuotaForNewUser - user.AccessToken = helper.GetUUID() - user.AffCode = helper.GetRandomString(4) + user.AccessToken = random.GetUUID() + user.AffCode = random.GetRandomString(4) result := DB.Create(user) if result.Error != nil { return result.Error @@ -137,9 +151,9 @@ func (user *User) Update(updatePassword bool) error { return err } } - if user.Status == common.UserStatusDisabled { + if user.Status == UserStatusDisabled { blacklist.BanUser(user.Id) - } else if user.Status == common.UserStatusEnabled { + } else if user.Status == UserStatusEnabled { blacklist.UnbanUser(user.Id) } err = DB.Model(user).Updates(user).Error @@ -151,8 +165,8 @@ func (user *User) Delete() error { return errors.New("id 为空！") } blacklist.BanUser(user.Id) - user.Username = fmt.Sprintf("deleted_%s", helper.GetUUID()) - user.Status = common.UserStatusDeleted + user.Username = fmt.Sprintf("deleted_%s", random.GetUUID()) + user.Status = UserStatusDeleted err := DB.Model(user).Updates(user).Error return err } @@ -176,7 +190,7 @@ func (user *User) ValidateAndFill() (err error) { } } okay := common.ValidatePasswordAndHash(password, user.Password) - if !okay || user.Status != common.UserStatusEnabled { + if !okay || user.Status != UserStatusEnabled { return errors.New("用户名或密码错误，或用户已被封禁") } return nil @@ -206,6 +220,14 @@ func (user *User) FillUserByGitHubId() error { return nil } +func (user *User) FillUserByLarkId() error { + if user.LarkId == "" { + return errors.New("lark id 为空！") + } + DB.Where(User{LarkId: user.LarkId}).First(user) + return nil +} + func (user *User) FillUserByWeChatId() error { if user.WeChatId == "" { return errors.New("WeChat id 为空！") @@ -234,6 +256,10 @@ func IsGitHubIdAlreadyTaken(githubId string) bool { return DB.Where("github_id = ?", githubId).Find(&User{}).RowsAffected == 1 } +func IsLarkIdAlreadyTaken(githubId string) bool { + return DB.Where("lark_id = ?", githubId).Find(&User{}).RowsAffected == 1 +} + func IsUsernameAlreadyTaken(username string) bool { return DB.Where("username = ?", username).Find(&User{}).RowsAffected == 1 } @@ -260,7 +286,7 @@ func IsAdmin(userId int) bool { logger.SysError("no such user " + err.Error()) return false } - return user.Role >= common.RoleAdminUser + return user.Role >= RoleAdminUser } func IsUserEnabled(userId int) (bool, error) { @@ -272,7 +298,7 @@ func IsUserEnabled(userId int) (bool, error) { if err != nil { return false, err } - return user.Status == common.UserStatusEnabled, nil + return user.Status == UserStatusEnabled, nil } func ValidateAccessToken(token string) (user *User) { @@ -345,7 +371,7 @@ func decreaseUserQuota(id int, quota int64) (err error) { } func GetRootUserEmail() (email string) { - DB.Model(&User{}).Where("role = ?", common.RoleRootUser).Select("email").Find(&email) + DB.Model(&User{}).Where("role = ?", RoleRootUser).Select("email").Find(&email) return email } diff --git a/monitor/channel.go b/monitor/channel.go index ad82d2f5..7e5dc58a 100644 --- a/monitor/channel.go +++ b/monitor/channel.go @@ -2,7 +2,6 @@ package monitor import ( "fmt" - "github.com/songquanpeng/one-api/common" "github.com/songquanpeng/one-api/common/config" "github.com/songquanpeng/one-api/common/logger" "github.com/songquanpeng/one-api/common/message" @@ -29,7 +28,7 @@ func notifyRootUser(subject string, content string) { // DisableChannel disable & notify func DisableChannel(channelId int, channelName string, reason string) { - model.UpdateChannelStatusById(channelId, common.ChannelStatusAutoDisabled) + model.UpdateChannelStatusById(channelId, model.ChannelStatusAutoDisabled) logger.SysLog(fmt.Sprintf("channel #%d has been disabled: %s", channelId, reason)) subject := fmt.Sprintf("渠道「%s」（#%d）已被禁用", channelName, channelId) content := fmt.Sprintf("渠道「%s」（#%d）已被禁用，原因：%s", channelName, channelId, reason) @@ -37,7 +36,7 @@ func DisableChannel(channelId int, channelName string, reason string) { } func MetricDisableChannel(channelId int, successRate float64) { - model.UpdateChannelStatusById(channelId, common.ChannelStatusAutoDisabled) + model.UpdateChannelStatusById(channelId, model.ChannelStatusAutoDisabled) logger.SysLog(fmt.Sprintf("channel #%d has been disabled due to low success rate: %.2f", channelId, successRate*100)) subject := fmt.Sprintf("渠道 #%d 已被禁用", channelId) content := fmt.Sprintf("该渠道（#%d）在最近 %d 次调用中成功率为 %.2f%%，低于阈值 %.2f%%，因此被系统自动禁用。", @@ -47,7 +46,7 @@ func MetricDisableChannel(channelId int, successRate float64) { // EnableChannel enable & notify func EnableChannel(channelId int, channelName string) { - model.UpdateChannelStatusById(channelId, common.ChannelStatusEnabled) + model.UpdateChannelStatusById(channelId, model.ChannelStatusEnabled) logger.SysLog(fmt.Sprintf("channel #%d has been enabled", channelId)) subject := fmt.Sprintf("渠道「%s」（#%d）已被启用", channelName, channelId) content := fmt.Sprintf("渠道「%s」（#%d）已被启用", channelName, channelId) diff --git a/monitor/manage.go b/monitor/manage.go new file mode 100644 index 00000000..946e78af --- /dev/null +++ b/monitor/manage.go @@ -0,0 +1,62 @@ +package monitor + +import ( + "github.com/songquanpeng/one-api/common/config" + "github.com/songquanpeng/one-api/relay/model" + "net/http" + "strings" +) + +func ShouldDisableChannel(err *model.Error, statusCode int) bool { + if !config.AutomaticDisableChannelEnabled { + return false + } + if err == nil { + return false + } + if statusCode == http.StatusUnauthorized { + return true + } + switch err.Type { + case "insufficient_quota": + return true + // https://docs.anthropic.com/claude/reference/errors + case "authentication_error": + return true + case "permission_error": + return true + case "forbidden": + return true + } + if err.Code == "invalid_api_key" || err.Code == "account_deactivated" { + return true + } + if strings.HasPrefix(err.Message, "Your credit balance is too low") { // anthropic + return true + } else if strings.HasPrefix(err.Message, "This organization has been disabled.") { + return true + } + //if strings.Contains(err.Message, "quota") { + // return true + //} + if strings.Contains(err.Message, "credit") { + return true + } + if strings.Contains(err.Message, "balance") { + return true + } + return false +} + +func ShouldEnableChannel(err error, openAIErr *model.Error) bool { + if !config.AutomaticEnableChannelEnabled { + return false + } + if err != nil { + return false + } + if openAIErr != nil { + return false + } + return true +} diff --git a/relay/adaptor.go b/relay/adaptor.go new file mode 100644 index 00000000..c90bd708 --- /dev/null +++ b/relay/adaptor.go @@ -0,0 +1,45 @@ +package relay + +import ( + "github.com/songquanpeng/one-api/relay/adaptor" + "github.com/songquanpeng/one-api/relay/adaptor/aiproxy" + "github.com/songquanpeng/one-api/relay/adaptor/ali" + "github.com/songquanpeng/one-api/relay/adaptor/anthropic" + "github.com/songquanpeng/one-api/relay/adaptor/baidu" + "github.com/songquanpeng/one-api/relay/adaptor/gemini" + "github.com/songquanpeng/one-api/relay/adaptor/ollama" + "github.com/songquanpeng/one-api/relay/adaptor/openai" + "github.com/songquanpeng/one-api/relay/adaptor/palm" + "github.com/songquanpeng/one-api/relay/adaptor/tencent" + "github.com/songquanpeng/one-api/relay/adaptor/xunfei" + "github.com/songquanpeng/one-api/relay/adaptor/zhipu" + "github.com/songquanpeng/one-api/relay/apitype" +) + +func GetAdaptor(apiType int) adaptor.Adaptor { + switch apiType { + case apitype.AIProxyLibrary: + return &aiproxy.Adaptor{} + case apitype.Ali: + return &ali.Adaptor{} + case apitype.Anthropic: + return &anthropic.Adaptor{} + case apitype.Baidu: + return &baidu.Adaptor{} + case apitype.Gemini: + return &gemini.Adaptor{} + case apitype.OpenAI: + return &openai.Adaptor{} + case apitype.PaLM: + return &palm.Adaptor{} + case apitype.Tencent: + return &tencent.Adaptor{} + case apitype.Xunfei: + return &xunfei.Adaptor{} + case apitype.Zhipu: + return &zhipu.Adaptor{} + case apitype.Ollama: + return &ollama.Adaptor{} + } + return nil +} diff --git a/relay/channel/ai360/constants.go b/relay/adaptor/ai360/constants.go similarity index 100% rename from relay/channel/ai360/constants.go rename to relay/adaptor/ai360/constants.go diff --git a/relay/channel/aiproxy/adaptor.go b/relay/adaptor/aiproxy/adaptor.go similarity index 53% rename from relay/channel/aiproxy/adaptor.go rename to relay/adaptor/aiproxy/adaptor.go index 2b4e3022..7ad6225a 100644 --- a/relay/channel/aiproxy/adaptor.go +++ b/relay/adaptor/aiproxy/adaptor.go @@ -4,10 +4,10 @@ import ( "errors" "fmt" "github.com/gin-gonic/gin" - "github.com/songquanpeng/one-api/common" - "github.com/songquanpeng/one-api/relay/channel" + "github.com/songquanpeng/one-api/common/config" + "github.com/songquanpeng/one-api/relay/adaptor" + "github.com/songquanpeng/one-api/relay/meta" "github.com/songquanpeng/one-api/relay/model" - "github.com/songquanpeng/one-api/relay/util" "io" "net/http" ) @@ -15,16 +15,16 @@ import ( type Adaptor struct { } -func (a *Adaptor) Init(meta *util.RelayMeta) { +func (a *Adaptor) Init(meta *meta.Meta) { } -func (a *Adaptor) GetRequestURL(meta *util.RelayMeta) (string, error) { +func (a *Adaptor) GetRequestURL(meta *meta.Meta) (string, error) { return fmt.Sprintf("%s/api/library/ask", meta.BaseURL), nil } -func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Request, meta *util.RelayMeta) error { - channel.SetupCommonRequestHeader(c, req, meta) +func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Request, meta *meta.Meta) error { + adaptor.SetupCommonRequestHeader(c, req, meta) req.Header.Set("Authorization", "Bearer "+meta.APIKey) return nil } @@ -34,15 +34,22 @@ func (a *Adaptor) ConvertRequest(c *gin.Context, relayMode int, request *model.G return nil, errors.New("request is nil") } aiProxyLibraryRequest := ConvertRequest(*request) - aiProxyLibraryRequest.LibraryId = c.GetString(common.ConfigKeyLibraryID) + aiProxyLibraryRequest.LibraryId = c.GetString(config.KeyLibraryID) return aiProxyLibraryRequest, nil } -func (a *Adaptor) DoRequest(c *gin.Context, meta *util.RelayMeta, requestBody io.Reader) (*http.Response, error) { - return channel.DoRequestHelper(a, c, meta, requestBody) +func (a *Adaptor) ConvertImageRequest(request *model.ImageRequest) (any, error) { + if request == nil { + return nil, errors.New("request is nil") + } + return request, nil } -func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, meta *util.RelayMeta) (usage *model.Usage, err *model.ErrorWithStatusCode) { +func (a *Adaptor) DoRequest(c *gin.Context, meta *meta.Meta, requestBody io.Reader) (*http.Response, error) { + return adaptor.DoRequestHelper(a, c, meta, requestBody) +} + +func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, meta *meta.Meta) (usage *model.Usage, err *model.ErrorWithStatusCode) { if meta.IsStream { err, usage = StreamHandler(c, resp) } else { diff --git a/relay/channel/aiproxy/constants.go b/relay/adaptor/aiproxy/constants.go similarity index 60% rename from relay/channel/aiproxy/constants.go rename to relay/adaptor/aiproxy/constants.go index c4df51c4..818d2709 100644 --- a/relay/channel/aiproxy/constants.go +++ b/relay/adaptor/aiproxy/constants.go @@ -1,6 +1,6 @@ package aiproxy -import "github.com/songquanpeng/one-api/relay/channel/openai" +import "github.com/songquanpeng/one-api/relay/adaptor/openai" var ModelList = []string{""} diff --git a/relay/channel/aiproxy/main.go b/relay/adaptor/aiproxy/main.go similarity index 95% rename from relay/channel/aiproxy/main.go rename to relay/adaptor/aiproxy/main.go index 96972407..01a568f6 100644 --- a/relay/channel/aiproxy/main.go +++ b/relay/adaptor/aiproxy/main.go @@ -8,7 +8,8 @@ import ( "github.com/songquanpeng/one-api/common" "github.com/songquanpeng/one-api/common/helper" "github.com/songquanpeng/one-api/common/logger" - "github.com/songquanpeng/one-api/relay/channel/openai" + "github.com/songquanpeng/one-api/common/random" + "github.com/songquanpeng/one-api/relay/adaptor/openai" "github.com/songquanpeng/one-api/relay/constant" "github.com/songquanpeng/one-api/relay/model" "io" @@ -53,7 +54,7 @@ func responseAIProxyLibrary2OpenAI(response *LibraryResponse) *openai.TextRespon FinishReason: "stop", } fullTextResponse := openai.TextResponse{ - Id: fmt.Sprintf("chatcmpl-%s", helper.GetUUID()), + Id: fmt.Sprintf("chatcmpl-%s", random.GetUUID()), Object: "chat.completion", Created: helper.GetTimestamp(), Choices: []openai.TextResponseChoice{choice}, @@ -66,7 +67,7 @@ func documentsAIProxyLibrary(documents []LibraryDocument) *openai.ChatCompletion choice.Delta.Content = aiProxyDocuments2Markdown(documents) choice.FinishReason = &constant.StopFinishReason return &openai.ChatCompletionsStreamResponse{ - Id: fmt.Sprintf("chatcmpl-%s", helper.GetUUID()), + Id: fmt.Sprintf("chatcmpl-%s", random.GetUUID()), Object: "chat.completion.chunk", Created: helper.GetTimestamp(), Model: "", @@ -78,7 +79,7 @@ func streamResponseAIProxyLibrary2OpenAI(response *LibraryStreamResponse) *opena var choice openai.ChatCompletionsStreamResponseChoice choice.Delta.Content = response.Content return &openai.ChatCompletionsStreamResponse{ - Id: fmt.Sprintf("chatcmpl-%s", helper.GetUUID()), + Id: fmt.Sprintf("chatcmpl-%s", random.GetUUID()), Object: "chat.completion.chunk", Created: helper.GetTimestamp(), Model: response.Model, diff --git a/relay/channel/aiproxy/model.go b/relay/adaptor/aiproxy/model.go similarity index 100% rename from relay/channel/aiproxy/model.go rename to relay/adaptor/aiproxy/model.go diff --git a/relay/adaptor/ali/adaptor.go b/relay/adaptor/ali/adaptor.go new file mode 100644 index 00000000..21b5e8b8 --- /dev/null +++ b/relay/adaptor/ali/adaptor.go @@ -0,0 +1,105 @@ +package ali + +import ( + "errors" + "fmt" + "github.com/gin-gonic/gin" + "github.com/songquanpeng/one-api/common/config" + "github.com/songquanpeng/one-api/relay/adaptor" + "github.com/songquanpeng/one-api/relay/meta" + "github.com/songquanpeng/one-api/relay/model" + "github.com/songquanpeng/one-api/relay/relaymode" + "io" + "net/http" +) + +// https://help.aliyun.com/zh/dashscope/developer-reference/api-details + +type Adaptor struct { +} + +func (a *Adaptor) Init(meta *meta.Meta) { + +} + +func (a *Adaptor) GetRequestURL(meta *meta.Meta) (string, error) { + fullRequestURL := "" + switch meta.Mode { + case relaymode.Embeddings: + fullRequestURL = fmt.Sprintf("%s/api/v1/services/embeddings/text-embedding/text-embedding", meta.BaseURL) + case relaymode.ImagesGenerations: + fullRequestURL = fmt.Sprintf("%s/api/v1/services/aigc/text2image/image-synthesis", meta.BaseURL) + default: + fullRequestURL = fmt.Sprintf("%s/api/v1/services/aigc/text-generation/generation", meta.BaseURL) + } + + return fullRequestURL, nil +} + +func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Request, meta *meta.Meta) error { + adaptor.SetupCommonRequestHeader(c, req, meta) + if meta.IsStream { + req.Header.Set("Accept", "text/event-stream") + req.Header.Set("X-DashScope-SSE", "enable") + } + req.Header.Set("Authorization", "Bearer "+meta.APIKey) + + if meta.Mode == relaymode.ImagesGenerations { + req.Header.Set("X-DashScope-Async", "enable") + } + if c.GetString(config.KeyPlugin) != "" { + req.Header.Set("X-DashScope-Plugin", c.GetString(config.KeyPlugin)) + } + return nil +} + +func (a *Adaptor) ConvertRequest(c *gin.Context, relayMode int, request *model.GeneralOpenAIRequest) (any, error) { + if request == nil { + return nil, errors.New("request is nil") + } + switch relayMode { + case relaymode.Embeddings: + aliEmbeddingRequest := ConvertEmbeddingRequest(*request) + return aliEmbeddingRequest, nil + default: + aliRequest := ConvertRequest(*request) + return aliRequest, nil + } +} + +func (a *Adaptor) ConvertImageRequest(request *model.ImageRequest) (any, error) { + if request == nil { + return nil, errors.New("request is nil") + } + + aliRequest := ConvertImageRequest(*request) + return aliRequest, nil +} + +func (a *Adaptor) DoRequest(c *gin.Context, meta *meta.Meta, requestBody io.Reader) (*http.Response, error) { + return adaptor.DoRequestHelper(a, c, meta, requestBody) +} + +func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, meta *meta.Meta) (usage *model.Usage, err *model.ErrorWithStatusCode) { + if meta.IsStream { + err, usage = StreamHandler(c, resp) + } else { + switch meta.Mode { + case relaymode.Embeddings: + err, usage = EmbeddingHandler(c, resp) + case relaymode.ImagesGenerations: + err, usage = ImageHandler(c, resp) + default: + err, usage = Handler(c, resp) + } + } + return +} + +func (a *Adaptor) GetModelList() []string { + return ModelList +} + +func (a *Adaptor) GetChannelName() string { + return "ali" +} diff --git a/relay/channel/ali/constants.go b/relay/adaptor/ali/constants.go similarity index 65% rename from relay/channel/ali/constants.go rename to relay/adaptor/ali/constants.go index 16bcfca4..3f24ce2e 100644 --- a/relay/channel/ali/constants.go +++ b/relay/adaptor/ali/constants.go @@ -3,4 +3,5 @@ package ali var ModelList = []string{ "qwen-turbo", "qwen-plus", "qwen-max", "qwen-max-longcontext", "text-embedding-v1", + "ali-stable-diffusion-xl", "ali-stable-diffusion-v1.5", "wanx-v1", } diff --git a/relay/adaptor/ali/image.go b/relay/adaptor/ali/image.go new file mode 100644 index 00000000..8261803d --- /dev/null +++ b/relay/adaptor/ali/image.go @@ -0,0 +1,192 @@ +package ali + +import ( + "encoding/base64" + "encoding/json" + "errors" + "fmt" + "github.com/gin-gonic/gin" + "github.com/songquanpeng/one-api/common/helper" + "github.com/songquanpeng/one-api/common/logger" + "github.com/songquanpeng/one-api/relay/adaptor/openai" + "github.com/songquanpeng/one-api/relay/model" + "io" + "net/http" + "strings" + "time" +) + +func ImageHandler(c *gin.Context, resp *http.Response) (*model.ErrorWithStatusCode, *model.Usage) { + apiKey := c.Request.Header.Get("Authorization") + apiKey = strings.TrimPrefix(apiKey, "Bearer ") + responseFormat := c.GetString("response_format") + + var aliTaskResponse TaskResponse + responseBody, err := io.ReadAll(resp.Body) + if err != nil { + return openai.ErrorWrapper(err, "read_response_body_failed", http.StatusInternalServerError), nil + } + err = resp.Body.Close() + if err != nil { + return openai.ErrorWrapper(err, "close_response_body_failed", http.StatusInternalServerError), nil + } + err = json.Unmarshal(responseBody, &aliTaskResponse) + if err != nil { + return openai.ErrorWrapper(err, "unmarshal_response_body_failed", http.StatusInternalServerError), nil + } + + if aliTaskResponse.Message != "" { + logger.SysError("aliAsyncTask err: " + string(responseBody)) + return openai.ErrorWrapper(errors.New(aliTaskResponse.Message), "ali_async_task_failed", http.StatusInternalServerError), nil + } + + aliResponse, _, err := asyncTaskWait(aliTaskResponse.Output.TaskId, apiKey) + if err != nil { + return openai.ErrorWrapper(err, "ali_async_task_wait_failed", http.StatusInternalServerError), nil + } + + if aliResponse.Output.TaskStatus != "SUCCEEDED" { + return &model.ErrorWithStatusCode{ + Error: model.Error{ + Message: aliResponse.Output.Message, + Type: "ali_error", + Param: "", + Code: aliResponse.Output.Code, + }, + StatusCode: resp.StatusCode, + }, nil + } + + fullTextResponse := responseAli2OpenAIImage(aliResponse, responseFormat) + jsonResponse, err := json.Marshal(fullTextResponse) + if err != nil { + return openai.ErrorWrapper(err, "marshal_response_body_failed", http.StatusInternalServerError), nil + } + c.Writer.Header().Set("Content-Type", "application/json") + c.Writer.WriteHeader(resp.StatusCode) + _, err = c.Writer.Write(jsonResponse) + return nil, nil +} + +func asyncTask(taskID string, key string) (*TaskResponse, error, []byte) { + url := fmt.Sprintf("https://dashscope.aliyuncs.com/api/v1/tasks/%s", taskID) + + var aliResponse TaskResponse + + req, err := http.NewRequest("GET", url, nil) + if err != nil { + return &aliResponse, err, nil + } + + req.Header.Set("Authorization", "Bearer "+key) + + client := &http.Client{} + resp, err := client.Do(req) + if err != nil { + logger.SysError("aliAsyncTask client.Do err: " + err.Error()) + return &aliResponse, err, nil + } + defer resp.Body.Close() + + responseBody, err := io.ReadAll(resp.Body) + + var response TaskResponse + err = json.Unmarshal(responseBody, &response) + if err != nil { + logger.SysError("aliAsyncTask NewDecoder err: " + err.Error()) + return &aliResponse, err, nil + } + + return &response, nil, responseBody +} + +func asyncTaskWait(taskID string, key string) (*TaskResponse, []byte, error) { + waitSeconds := 2 + step := 0 + maxStep := 20 + + var taskResponse TaskResponse + var responseBody []byte + + for { + step++ + rsp, err, body := asyncTask(taskID, key) + responseBody = body + if err != nil { + return &taskResponse, responseBody, err + } + + if rsp.Output.TaskStatus == "" { + return &taskResponse, responseBody, nil + } + + switch rsp.Output.TaskStatus { + case "FAILED": + fallthrough + case "CANCELED": + fallthrough + case "SUCCEEDED": + fallthrough + case "UNKNOWN": + return rsp, responseBody, nil + } + if step >= maxStep { + break + } + time.Sleep(time.Duration(waitSeconds) * time.Second) + } + + return nil, nil, fmt.Errorf("aliAsyncTaskWait timeout") +} + +func responseAli2OpenAIImage(response *TaskResponse, responseFormat string) *openai.ImageResponse { + imageResponse := openai.ImageResponse{ + Created: helper.GetTimestamp(), + } + + for _, data := range response.Output.Results { + var b64Json string + if responseFormat == "b64_json" { + // 读取 data.Url 的图片数据并转存到 b64Json + imageData, err := getImageData(data.Url) + if err != nil { + // 处理获取图片数据失败的情况 + logger.SysError("getImageData Error getting image data: " + err.Error()) + continue + } + + // 将图片数据转为 Base64 编码的字符串 + b64Json = Base64Encode(imageData) + } else { + // 如果 responseFormat 不是 "b64_json"，则直接使用 data.B64Image + b64Json = data.B64Image + } + + imageResponse.Data = append(imageResponse.Data, openai.ImageData{ + Url: data.Url, + B64Json: b64Json, + RevisedPrompt: "", + }) + } + return &imageResponse +} + +func getImageData(url string) ([]byte, error) { + response, err := http.Get(url) + if err != nil { + return nil, err + } + defer response.Body.Close() + + imageData, err := io.ReadAll(response.Body) + if err != nil { + return nil, err + } + + return imageData, nil +} + +func Base64Encode(data []byte) string { + b64Json := base64.StdEncoding.EncodeToString(data) + return b64Json +} diff --git a/relay/channel/ali/main.go b/relay/adaptor/ali/main.go similarity index 95% rename from relay/channel/ali/main.go rename to relay/adaptor/ali/main.go index dd1707ee..0462c26b 100644 --- a/relay/channel/ali/main.go +++ b/relay/adaptor/ali/main.go @@ -7,7 +7,7 @@ import ( "github.com/songquanpeng/one-api/common" "github.com/songquanpeng/one-api/common/helper" "github.com/songquanpeng/one-api/common/logger" - "github.com/songquanpeng/one-api/relay/channel/openai" + "github.com/songquanpeng/one-api/relay/adaptor/openai" "github.com/songquanpeng/one-api/relay/model" "io" "net/http" @@ -66,6 +66,17 @@ func ConvertEmbeddingRequest(request model.GeneralOpenAIRequest) *EmbeddingReque } } +func ConvertImageRequest(request model.ImageRequest) *ImageRequest { + var imageRequest ImageRequest + imageRequest.Input.Prompt = request.Prompt + imageRequest.Model = request.Model + imageRequest.Parameters.Size = strings.Replace(request.Size, "x", "*", -1) + imageRequest.Parameters.N = request.N + imageRequest.ResponseFormat = request.ResponseFormat + + return &imageRequest +} + func EmbeddingHandler(c *gin.Context, resp *http.Response) (*model.ErrorWithStatusCode, *model.Usage) { var aliResponse EmbeddingResponse err := json.NewDecoder(resp.Body).Decode(&aliResponse) diff --git a/relay/adaptor/ali/model.go b/relay/adaptor/ali/model.go new file mode 100644 index 00000000..450b5f52 --- /dev/null +++ b/relay/adaptor/ali/model.go @@ -0,0 +1,154 @@ +package ali + +import ( + "github.com/songquanpeng/one-api/relay/adaptor/openai" + "github.com/songquanpeng/one-api/relay/model" +) + +type Message struct { + Content string `json:"content"` + Role string `json:"role"` +} + +type Input struct { + //Prompt string `json:"prompt"` + Messages []Message `json:"messages"` +} + +type Parameters struct { + TopP float64 `json:"top_p,omitempty"` + TopK int `json:"top_k,omitempty"` + Seed uint64 `json:"seed,omitempty"` + EnableSearch bool `json:"enable_search,omitempty"` + IncrementalOutput bool `json:"incremental_output,omitempty"` + MaxTokens int `json:"max_tokens,omitempty"` + Temperature float64 `json:"temperature,omitempty"` + ResultFormat string `json:"result_format,omitempty"` + Tools []model.Tool `json:"tools,omitempty"` +} + +type ChatRequest struct { + Model string `json:"model"` + Input Input `json:"input"` + Parameters Parameters `json:"parameters,omitempty"` +} + +type ImageRequest struct { + Model string `json:"model"` + Input struct { + Prompt string `json:"prompt"` + NegativePrompt string `json:"negative_prompt,omitempty"` + } `json:"input"` + Parameters struct { + Size string `json:"size,omitempty"` + N int `json:"n,omitempty"` + Steps string `json:"steps,omitempty"` + Scale string `json:"scale,omitempty"` + } `json:"parameters,omitempty"` + ResponseFormat string `json:"response_format,omitempty"` +} + +type TaskResponse struct { + StatusCode int `json:"status_code,omitempty"` + RequestId string `json:"request_id,omitempty"` + Code string `json:"code,omitempty"` + Message string `json:"message,omitempty"` + Output struct { + TaskId string `json:"task_id,omitempty"` + TaskStatus string `json:"task_status,omitempty"` + Code string `json:"code,omitempty"` + Message string `json:"message,omitempty"` + Results []struct { + B64Image string `json:"b64_image,omitempty"` + Url string `json:"url,omitempty"` + Code string `json:"code,omitempty"` + Message string `json:"message,omitempty"` + } `json:"results,omitempty"` + TaskMetrics struct { + Total int `json:"TOTAL,omitempty"` + Succeeded int `json:"SUCCEEDED,omitempty"` + Failed int `json:"FAILED,omitempty"` + } `json:"task_metrics,omitempty"` + } `json:"output,omitempty"` + Usage Usage `json:"usage"` +} + +type Header struct { + Action string `json:"action,omitempty"` + Streaming string `json:"streaming,omitempty"` + TaskID string `json:"task_id,omitempty"` + Event string `json:"event,omitempty"` + ErrorCode string `json:"error_code,omitempty"` + ErrorMessage string `json:"error_message,omitempty"` + Attributes any `json:"attributes,omitempty"` +} + +type Payload struct { + Model string `json:"model,omitempty"` + Task string `json:"task,omitempty"` + TaskGroup string `json:"task_group,omitempty"` + Function string `json:"function,omitempty"` + Parameters struct { + SampleRate int `json:"sample_rate,omitempty"` + Rate float64 `json:"rate,omitempty"` + Format string `json:"format,omitempty"` + } `json:"parameters,omitempty"` + Input struct { + Text string `json:"text,omitempty"` + } `json:"input,omitempty"` + Usage struct { + Characters int `json:"characters,omitempty"` + } `json:"usage,omitempty"` +} + +type WSSMessage struct { + Header Header `json:"header,omitempty"` + Payload Payload `json:"payload,omitempty"` +} + +type EmbeddingRequest struct { + Model string `json:"model"` + Input struct { + Texts []string `json:"texts"` + } `json:"input"` + Parameters *struct { + TextType string `json:"text_type,omitempty"` + } `json:"parameters,omitempty"` +} + +type Embedding struct { + Embedding []float64 `json:"embedding"` + TextIndex int `json:"text_index"` +} + +type EmbeddingResponse struct { + Output struct { + Embeddings []Embedding `json:"embeddings"` + } `json:"output"` + Usage Usage `json:"usage"` + Error +} + +type Error struct { + Code string `json:"code"` + Message string `json:"message"` + RequestId string `json:"request_id"` +} + +type Usage struct { + InputTokens int `json:"input_tokens"` + OutputTokens int `json:"output_tokens"` + TotalTokens int `json:"total_tokens"` +} + +type Output struct { + //Text string `json:"text"` + //FinishReason string `json:"finish_reason"` + Choices []openai.TextResponseChoice `json:"choices"` +} + +type ChatResponse struct { + Output Output `json:"output"` + Usage Usage `json:"usage"` + Error +} diff --git a/relay/channel/anthropic/adaptor.go b/relay/adaptor/anthropic/adaptor.go similarity index 62% rename from relay/channel/anthropic/adaptor.go rename to relay/adaptor/anthropic/adaptor.go index a165b35c..b1136e84 100644 --- a/relay/channel/anthropic/adaptor.go +++ b/relay/adaptor/anthropic/adaptor.go @@ -4,9 +4,9 @@ import ( "errors" "fmt" "github.com/gin-gonic/gin" - "github.com/songquanpeng/one-api/relay/channel" + "github.com/songquanpeng/one-api/relay/adaptor" + "github.com/songquanpeng/one-api/relay/meta" "github.com/songquanpeng/one-api/relay/model" - "github.com/songquanpeng/one-api/relay/util" "io" "net/http" ) @@ -14,16 +14,16 @@ import ( type Adaptor struct { } -func (a *Adaptor) Init(meta *util.RelayMeta) { +func (a *Adaptor) Init(meta *meta.Meta) { } -func (a *Adaptor) GetRequestURL(meta *util.RelayMeta) (string, error) { +func (a *Adaptor) GetRequestURL(meta *meta.Meta) (string, error) { return fmt.Sprintf("%s/v1/messages", meta.BaseURL), nil } -func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Request, meta *util.RelayMeta) error { - channel.SetupCommonRequestHeader(c, req, meta) +func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Request, meta *meta.Meta) error { + adaptor.SetupCommonRequestHeader(c, req, meta) req.Header.Set("x-api-key", meta.APIKey) anthropicVersion := c.Request.Header.Get("anthropic-version") if anthropicVersion == "" { @@ -41,11 +41,18 @@ func (a *Adaptor) ConvertRequest(c *gin.Context, relayMode int, request *model.G return ConvertRequest(*request), nil } -func (a *Adaptor) DoRequest(c *gin.Context, meta *util.RelayMeta, requestBody io.Reader) (*http.Response, error) { - return channel.DoRequestHelper(a, c, meta, requestBody) +func (a *Adaptor) ConvertImageRequest(request *model.ImageRequest) (any, error) { + if request == nil { + return nil, errors.New("request is nil") + } + return request, nil } -func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, meta *util.RelayMeta) (usage *model.Usage, err *model.ErrorWithStatusCode) { +func (a *Adaptor) DoRequest(c *gin.Context, meta *meta.Meta, requestBody io.Reader) (*http.Response, error) { + return adaptor.DoRequestHelper(a, c, meta, requestBody) +} + +func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, meta *meta.Meta) (usage *model.Usage, err *model.ErrorWithStatusCode) { if meta.IsStream { err, usage = StreamHandler(c, resp) } else { diff --git a/relay/channel/anthropic/constants.go b/relay/adaptor/anthropic/constants.go similarity index 100% rename from relay/channel/anthropic/constants.go rename to relay/adaptor/anthropic/constants.go diff --git a/relay/channel/anthropic/main.go b/relay/adaptor/anthropic/main.go similarity index 99% rename from relay/channel/anthropic/main.go rename to relay/adaptor/anthropic/main.go index 04e65d99..6bb82d01 100644 --- a/relay/channel/anthropic/main.go +++ b/relay/adaptor/anthropic/main.go @@ -9,7 +9,7 @@ import ( "github.com/songquanpeng/one-api/common/helper" "github.com/songquanpeng/one-api/common/image" "github.com/songquanpeng/one-api/common/logger" - "github.com/songquanpeng/one-api/relay/channel/openai" + "github.com/songquanpeng/one-api/relay/adaptor/openai" "github.com/songquanpeng/one-api/relay/model" "io" "net/http" diff --git a/relay/channel/anthropic/model.go b/relay/adaptor/anthropic/model.go similarity index 100% rename from relay/channel/anthropic/model.go rename to relay/adaptor/anthropic/model.go diff --git a/relay/adaptor/azure/helper.go b/relay/adaptor/azure/helper.go new file mode 100644 index 00000000..dd207f37 --- /dev/null +++ b/relay/adaptor/azure/helper.go @@ -0,0 +1,15 @@ +package azure + +import ( + "github.com/gin-gonic/gin" + "github.com/songquanpeng/one-api/common/config" +) + +func GetAPIVersion(c *gin.Context) string { + query := c.Request.URL.Query() + apiVersion := query.Get("api-version") + if apiVersion == "" { + apiVersion = c.GetString(config.KeyAPIVersion) + } + return apiVersion +} diff --git a/relay/channel/baichuan/constants.go b/relay/adaptor/baichuan/constants.go similarity index 100% rename from relay/channel/baichuan/constants.go rename to relay/adaptor/baichuan/constants.go diff --git a/relay/channel/baidu/adaptor.go b/relay/adaptor/baidu/adaptor.go similarity index 63% rename from relay/channel/baidu/adaptor.go rename to relay/adaptor/baidu/adaptor.go index 2d2e24f6..15306b95 100644 --- a/relay/channel/baidu/adaptor.go +++ b/relay/adaptor/baidu/adaptor.go @@ -3,25 +3,25 @@ package baidu import ( "errors" "fmt" + "github.com/songquanpeng/one-api/relay/meta" + "github.com/songquanpeng/one-api/relay/relaymode" "io" "net/http" "strings" "github.com/gin-gonic/gin" - "github.com/songquanpeng/one-api/relay/channel" - "github.com/songquanpeng/one-api/relay/constant" + "github.com/songquanpeng/one-api/relay/adaptor" "github.com/songquanpeng/one-api/relay/model" - "github.com/songquanpeng/one-api/relay/util" ) type Adaptor struct { } -func (a *Adaptor) Init(meta *util.RelayMeta) { +func (a *Adaptor) Init(meta *meta.Meta) { } -func (a *Adaptor) GetRequestURL(meta *util.RelayMeta) (string, error) { +func (a *Adaptor) GetRequestURL(meta *meta.Meta) (string, error) { // https://cloud.baidu.com/doc/WENXINWORKSHOP/s/clntwmv7t suffix := "chat/" if strings.HasPrefix(meta.ActualModelName, "Embedding") { @@ -38,16 +38,34 @@ func (a *Adaptor) GetRequestURL(meta *util.RelayMeta) (string, error) { suffix += "completions_pro" case "ERNIE-Bot-4": suffix += "completions_pro" - case "ERNIE-3.5-8K": - suffix += "completions" - case "ERNIE-Bot-8K": - suffix += "ernie_bot_8k" case "ERNIE-Bot": suffix += "completions" - case "ERNIE-Speed": - suffix += "ernie_speed" case "ERNIE-Bot-turbo": suffix += "eb-instant" + case "ERNIE-Speed": + suffix += "ernie_speed" + case "ERNIE-4.0-8K": + suffix += "completions_pro" + case "ERNIE-3.5-8K": + suffix += "completions" + case "ERNIE-3.5-8K-0205": + suffix += "ernie-3.5-8k-0205" + case "ERNIE-3.5-8K-1222": + suffix += "ernie-3.5-8k-1222" + case "ERNIE-Bot-8K": + suffix += "ernie_bot_8k" + case "ERNIE-3.5-4K-0205": + suffix += "ernie-3.5-4k-0205" + case "ERNIE-Speed-8K": + suffix += "ernie_speed" + case "ERNIE-Speed-128K": + suffix += "ernie-speed-128k" + case "ERNIE-Lite-8K-0922": + suffix += "eb-instant" + case "ERNIE-Lite-8K-0308": + suffix += "ernie-lite-8k" + case "ERNIE-Tiny-8K": + suffix += "ernie-tiny-8k" case "BLOOMZ-7B": suffix += "bloomz_7b1" case "Embedding-V1": @@ -59,7 +77,7 @@ func (a *Adaptor) GetRequestURL(meta *util.RelayMeta) (string, error) { case "tao-8k": suffix += "tao_8k" default: - suffix += meta.ActualModelName + suffix += strings.ToLower(meta.ActualModelName) } fullRequestURL := fmt.Sprintf("%s/rpc/2.0/ai_custom/v1/wenxinworkshop/%s", meta.BaseURL, suffix) var accessToken string @@ -71,8 +89,8 @@ func (a *Adaptor) GetRequestURL(meta *util.RelayMeta) (string, error) { return fullRequestURL, nil } -func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Request, meta *util.RelayMeta) error { - channel.SetupCommonRequestHeader(c, req, meta) +func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Request, meta *meta.Meta) error { + adaptor.SetupCommonRequestHeader(c, req, meta) req.Header.Set("Authorization", "Bearer "+meta.APIKey) return nil } @@ -82,7 +100,7 @@ func (a *Adaptor) ConvertRequest(c *gin.Context, relayMode int, request *model.G return nil, errors.New("request is nil") } switch relayMode { - case constant.RelayModeEmbeddings: + case relaymode.Embeddings: baiduEmbeddingRequest := ConvertEmbeddingRequest(*request) return baiduEmbeddingRequest, nil default: @@ -91,16 +109,23 @@ func (a *Adaptor) ConvertRequest(c *gin.Context, relayMode int, request *model.G } } -func (a *Adaptor) DoRequest(c *gin.Context, meta *util.RelayMeta, requestBody io.Reader) (*http.Response, error) { - return channel.DoRequestHelper(a, c, meta, requestBody) +func (a *Adaptor) ConvertImageRequest(request *model.ImageRequest) (any, error) { + if request == nil { + return nil, errors.New("request is nil") + } + return request, nil } -func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, meta *util.RelayMeta) (usage *model.Usage, err *model.ErrorWithStatusCode) { +func (a *Adaptor) DoRequest(c *gin.Context, meta *meta.Meta, requestBody io.Reader) (*http.Response, error) { + return adaptor.DoRequestHelper(a, c, meta, requestBody) +} + +func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, meta *meta.Meta) (usage *model.Usage, err *model.ErrorWithStatusCode) { if meta.IsStream { err, usage = StreamHandler(c, resp) } else { switch meta.Mode { - case constant.RelayModeEmbeddings: + case relaymode.Embeddings: err, usage = EmbeddingHandler(c, resp) default: err, usage = Handler(c, resp) diff --git a/relay/adaptor/baidu/constants.go b/relay/adaptor/baidu/constants.go new file mode 100644 index 00000000..f952adc6 --- /dev/null +++ b/relay/adaptor/baidu/constants.go @@ -0,0 +1,20 @@ +package baidu + +var ModelList = []string{ + "ERNIE-4.0-8K", + "ERNIE-3.5-8K", + "ERNIE-3.5-8K-0205", + "ERNIE-3.5-8K-1222", + "ERNIE-Bot-8K", + "ERNIE-3.5-4K-0205", + "ERNIE-Speed-8K", + "ERNIE-Speed-128K", + "ERNIE-Lite-8K-0922", + "ERNIE-Lite-8K-0308", + "ERNIE-Tiny-8K", + "BLOOMZ-7B", + "Embedding-V1", + "bge-large-zh", + "bge-large-en", + "tao-8k", +} diff --git a/relay/channel/baidu/main.go b/relay/adaptor/baidu/main.go similarity index 98% rename from relay/channel/baidu/main.go rename to relay/adaptor/baidu/main.go index 9ca9e47d..6df5ce84 100644 --- a/relay/channel/baidu/main.go +++ b/relay/adaptor/baidu/main.go @@ -8,10 +8,10 @@ import ( "github.com/gin-gonic/gin" "github.com/songquanpeng/one-api/common" "github.com/songquanpeng/one-api/common/logger" - "github.com/songquanpeng/one-api/relay/channel/openai" + "github.com/songquanpeng/one-api/relay/adaptor/openai" + "github.com/songquanpeng/one-api/relay/client" "github.com/songquanpeng/one-api/relay/constant" "github.com/songquanpeng/one-api/relay/model" - "github.com/songquanpeng/one-api/relay/util" "io" "net/http" "strings" @@ -305,7 +305,7 @@ func getBaiduAccessTokenHelper(apiKey string) (*AccessToken, error) { } req.Header.Add("Content-Type", "application/json") req.Header.Add("Accept", "application/json") - res, err := util.ImpatientHTTPClient.Do(req) + res, err := client.ImpatientHTTPClient.Do(req) if err != nil { return nil, err } diff --git a/relay/channel/baidu/model.go b/relay/adaptor/baidu/model.go similarity index 100% rename from relay/channel/baidu/model.go rename to relay/adaptor/baidu/model.go diff --git a/relay/channel/common.go b/relay/adaptor/common.go similarity index 80% rename from relay/channel/common.go rename to relay/adaptor/common.go index c6e1abf2..82a5160e 100644 --- a/relay/channel/common.go +++ b/relay/adaptor/common.go @@ -1,15 +1,16 @@ -package channel +package adaptor import ( "errors" "fmt" "github.com/gin-gonic/gin" - "github.com/songquanpeng/one-api/relay/util" + "github.com/songquanpeng/one-api/relay/client" + "github.com/songquanpeng/one-api/relay/meta" "io" "net/http" ) -func SetupCommonRequestHeader(c *gin.Context, req *http.Request, meta *util.RelayMeta) { +func SetupCommonRequestHeader(c *gin.Context, req *http.Request, meta *meta.Meta) { req.Header.Set("Content-Type", c.Request.Header.Get("Content-Type")) req.Header.Set("Accept", c.Request.Header.Get("Accept")) if meta.IsStream && c.Request.Header.Get("Accept") == "" { @@ -17,7 +18,7 @@ func SetupCommonRequestHeader(c *gin.Context, req *http.Request, meta *util.Rela } } -func DoRequestHelper(a Adaptor, c *gin.Context, meta *util.RelayMeta, requestBody io.Reader) (*http.Response, error) { +func DoRequestHelper(a Adaptor, c *gin.Context, meta *meta.Meta, requestBody io.Reader) (*http.Response, error) { fullRequestURL, err := a.GetRequestURL(meta) if err != nil { return nil, fmt.Errorf("get request url failed: %w", err) @@ -38,7 +39,7 @@ func DoRequestHelper(a Adaptor, c *gin.Context, meta *util.RelayMeta, requestBod } func DoRequest(c *gin.Context, req *http.Request) (*http.Response, error) { - resp, err := util.HTTPClient.Do(req) + resp, err := client.HTTPClient.Do(req) if err != nil { return nil, err } diff --git a/relay/channel/gemini/adaptor.go b/relay/adaptor/gemini/adaptor.go similarity index 68% rename from relay/channel/gemini/adaptor.go rename to relay/adaptor/gemini/adaptor.go index f3305e5d..45124752 100644 --- a/relay/channel/gemini/adaptor.go +++ b/relay/adaptor/gemini/adaptor.go @@ -5,10 +5,10 @@ import ( "fmt" "github.com/gin-gonic/gin" "github.com/songquanpeng/one-api/common/helper" - channelhelper "github.com/songquanpeng/one-api/relay/channel" - "github.com/songquanpeng/one-api/relay/channel/openai" + channelhelper "github.com/songquanpeng/one-api/relay/adaptor" + "github.com/songquanpeng/one-api/relay/adaptor/openai" + "github.com/songquanpeng/one-api/relay/meta" "github.com/songquanpeng/one-api/relay/model" - "github.com/songquanpeng/one-api/relay/util" "io" "net/http" ) @@ -16,11 +16,11 @@ import ( type Adaptor struct { } -func (a *Adaptor) Init(meta *util.RelayMeta) { +func (a *Adaptor) Init(meta *meta.Meta) { } -func (a *Adaptor) GetRequestURL(meta *util.RelayMeta) (string, error) { +func (a *Adaptor) GetRequestURL(meta *meta.Meta) (string, error) { version := helper.AssignOrDefault(meta.APIVersion, "v1") action := "generateContent" if meta.IsStream { @@ -29,7 +29,7 @@ func (a *Adaptor) GetRequestURL(meta *util.RelayMeta) (string, error) { return fmt.Sprintf("%s/%s/models/%s:%s", meta.BaseURL, version, meta.ActualModelName, action), nil } -func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Request, meta *util.RelayMeta) error { +func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Request, meta *meta.Meta) error { channelhelper.SetupCommonRequestHeader(c, req, meta) req.Header.Set("x-goog-api-key", meta.APIKey) return nil @@ -42,11 +42,18 @@ func (a *Adaptor) ConvertRequest(c *gin.Context, relayMode int, request *model.G return ConvertRequest(*request), nil } -func (a *Adaptor) DoRequest(c *gin.Context, meta *util.RelayMeta, requestBody io.Reader) (*http.Response, error) { +func (a *Adaptor) ConvertImageRequest(request *model.ImageRequest) (any, error) { + if request == nil { + return nil, errors.New("request is nil") + } + return request, nil +} + +func (a *Adaptor) DoRequest(c *gin.Context, meta *meta.Meta, requestBody io.Reader) (*http.Response, error) { return channelhelper.DoRequestHelper(a, c, meta, requestBody) } -func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, meta *util.RelayMeta) (usage *model.Usage, err *model.ErrorWithStatusCode) { +func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, meta *meta.Meta) (usage *model.Usage, err *model.ErrorWithStatusCode) { if meta.IsStream { var responseText string err, responseText = StreamHandler(c, resp) diff --git a/relay/channel/gemini/constants.go b/relay/adaptor/gemini/constants.go similarity index 100% rename from relay/channel/gemini/constants.go rename to relay/adaptor/gemini/constants.go diff --git a/relay/channel/gemini/main.go b/relay/adaptor/gemini/main.go similarity index 97% rename from relay/channel/gemini/main.go rename to relay/adaptor/gemini/main.go index c24694c8..6bf0c6d7 100644 --- a/relay/channel/gemini/main.go +++ b/relay/adaptor/gemini/main.go @@ -9,7 +9,8 @@ import ( "github.com/songquanpeng/one-api/common/helper" "github.com/songquanpeng/one-api/common/image" "github.com/songquanpeng/one-api/common/logger" - "github.com/songquanpeng/one-api/relay/channel/openai" + "github.com/songquanpeng/one-api/common/random" + "github.com/songquanpeng/one-api/relay/adaptor/openai" "github.com/songquanpeng/one-api/relay/constant" "github.com/songquanpeng/one-api/relay/model" "io" @@ -155,7 +156,7 @@ type ChatPromptFeedback struct { func responseGeminiChat2OpenAI(response *ChatResponse) *openai.TextResponse { fullTextResponse := openai.TextResponse{ - Id: fmt.Sprintf("chatcmpl-%s", helper.GetUUID()), + Id: fmt.Sprintf("chatcmpl-%s", random.GetUUID()), Object: "chat.completion", Created: helper.GetTimestamp(), Choices: make([]openai.TextResponseChoice, 0, len(response.Candidates)), @@ -233,7 +234,7 @@ func StreamHandler(c *gin.Context, resp *http.Response) (*model.ErrorWithStatusC var choice openai.ChatCompletionsStreamResponseChoice choice.Delta.Content = dummy.Content response := openai.ChatCompletionsStreamResponse{ - Id: fmt.Sprintf("chatcmpl-%s", helper.GetUUID()), + Id: fmt.Sprintf("chatcmpl-%s", random.GetUUID()), Object: "chat.completion.chunk", Created: helper.GetTimestamp(), Model: "gemini-pro", diff --git a/relay/channel/gemini/model.go b/relay/adaptor/gemini/model.go similarity index 100% rename from relay/channel/gemini/model.go rename to relay/adaptor/gemini/model.go diff --git a/relay/channel/groq/constants.go b/relay/adaptor/groq/constants.go similarity index 100% rename from relay/channel/groq/constants.go rename to relay/adaptor/groq/constants.go diff --git a/relay/adaptor/interface.go b/relay/adaptor/interface.go new file mode 100644 index 00000000..01b2e2cb --- /dev/null +++ b/relay/adaptor/interface.go @@ -0,0 +1,21 @@ +package adaptor + +import ( + "github.com/gin-gonic/gin" + "github.com/songquanpeng/one-api/relay/meta" + "github.com/songquanpeng/one-api/relay/model" + "io" + "net/http" +) + +type Adaptor interface { + Init(meta *meta.Meta) + GetRequestURL(meta *meta.Meta) (string, error) + SetupRequestHeader(c *gin.Context, req *http.Request, meta *meta.Meta) error + ConvertRequest(c *gin.Context, relayMode int, request *model.GeneralOpenAIRequest) (any, error) + ConvertImageRequest(request *model.ImageRequest) (any, error) + DoRequest(c *gin.Context, meta *meta.Meta, requestBody io.Reader) (*http.Response, error) + DoResponse(c *gin.Context, resp *http.Response, meta *meta.Meta) (usage *model.Usage, err *model.ErrorWithStatusCode) + GetModelList() []string + GetChannelName() string +} diff --git a/relay/channel/lingyiwanwu/constants.go b/relay/adaptor/lingyiwanwu/constants.go similarity index 100% rename from relay/channel/lingyiwanwu/constants.go rename to relay/adaptor/lingyiwanwu/constants.go diff --git a/relay/channel/minimax/constants.go b/relay/adaptor/minimax/constants.go similarity index 100% rename from relay/channel/minimax/constants.go rename to relay/adaptor/minimax/constants.go diff --git a/relay/adaptor/minimax/main.go b/relay/adaptor/minimax/main.go new file mode 100644 index 00000000..fc9b5d26 --- /dev/null +++ b/relay/adaptor/minimax/main.go @@ -0,0 +1,14 @@ +package minimax + +import ( + "fmt" + "github.com/songquanpeng/one-api/relay/meta" + "github.com/songquanpeng/one-api/relay/relaymode" +) + +func GetRequestURL(meta *meta.Meta) (string, error) { + if meta.Mode == relaymode.ChatCompletions { + return fmt.Sprintf("%s/v1/text/chatcompletion_v2", meta.BaseURL), nil + } + return "", fmt.Errorf("unsupported relay mode %d for minimax", meta.Mode) +} diff --git a/relay/channel/mistral/constants.go b/relay/adaptor/mistral/constants.go similarity index 100% rename from relay/channel/mistral/constants.go rename to relay/adaptor/mistral/constants.go diff --git a/relay/channel/moonshot/constants.go b/relay/adaptor/moonshot/constants.go similarity index 100% rename from relay/channel/moonshot/constants.go rename to relay/adaptor/moonshot/constants.go diff --git a/relay/channel/ollama/adaptor.go b/relay/adaptor/ollama/adaptor.go similarity index 59% rename from relay/channel/ollama/adaptor.go rename to relay/adaptor/ollama/adaptor.go index e2ae7d2b..66702c5d 100644 --- a/relay/channel/ollama/adaptor.go +++ b/relay/adaptor/ollama/adaptor.go @@ -3,34 +3,34 @@ package ollama import ( "errors" "fmt" + "github.com/songquanpeng/one-api/relay/meta" + "github.com/songquanpeng/one-api/relay/relaymode" "io" "net/http" "github.com/gin-gonic/gin" - "github.com/songquanpeng/one-api/relay/channel" - "github.com/songquanpeng/one-api/relay/constant" + "github.com/songquanpeng/one-api/relay/adaptor" "github.com/songquanpeng/one-api/relay/model" - "github.com/songquanpeng/one-api/relay/util" ) type Adaptor struct { } -func (a *Adaptor) Init(meta *util.RelayMeta) { +func (a *Adaptor) Init(meta *meta.Meta) { } -func (a *Adaptor) GetRequestURL(meta *util.RelayMeta) (string, error) { +func (a *Adaptor) GetRequestURL(meta *meta.Meta) (string, error) { // https://github.com/ollama/ollama/blob/main/docs/api.md fullRequestURL := fmt.Sprintf("%s/api/chat", meta.BaseURL) - if meta.Mode == constant.RelayModeEmbeddings { + if meta.Mode == relaymode.Embeddings { fullRequestURL = fmt.Sprintf("%s/api/embeddings", meta.BaseURL) } return fullRequestURL, nil } -func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Request, meta *util.RelayMeta) error { - channel.SetupCommonRequestHeader(c, req, meta) +func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Request, meta *meta.Meta) error { + adaptor.SetupCommonRequestHeader(c, req, meta) req.Header.Set("Authorization", "Bearer "+meta.APIKey) return nil } @@ -40,7 +40,7 @@ func (a *Adaptor) ConvertRequest(c *gin.Context, relayMode int, request *model.G return nil, errors.New("request is nil") } switch relayMode { - case constant.RelayModeEmbeddings: + case relaymode.Embeddings: ollamaEmbeddingRequest := ConvertEmbeddingRequest(*request) return ollamaEmbeddingRequest, nil default: @@ -48,16 +48,23 @@ func (a *Adaptor) ConvertRequest(c *gin.Context, relayMode int, request *model.G } } -func (a *Adaptor) DoRequest(c *gin.Context, meta *util.RelayMeta, requestBody io.Reader) (*http.Response, error) { - return channel.DoRequestHelper(a, c, meta, requestBody) +func (a *Adaptor) ConvertImageRequest(request *model.ImageRequest) (any, error) { + if request == nil { + return nil, errors.New("request is nil") + } + return request, nil } -func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, meta *util.RelayMeta) (usage *model.Usage, err *model.ErrorWithStatusCode) { +func (a *Adaptor) DoRequest(c *gin.Context, meta *meta.Meta, requestBody io.Reader) (*http.Response, error) { + return adaptor.DoRequestHelper(a, c, meta, requestBody) +} + +func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, meta *meta.Meta) (usage *model.Usage, err *model.ErrorWithStatusCode) { if meta.IsStream { err, usage = StreamHandler(c, resp) } else { switch meta.Mode { - case constant.RelayModeEmbeddings: + case relaymode.Embeddings: err, usage = EmbeddingHandler(c, resp) default: err, usage = Handler(c, resp) diff --git a/relay/channel/ollama/constants.go b/relay/adaptor/ollama/constants.go similarity index 100% rename from relay/channel/ollama/constants.go rename to relay/adaptor/ollama/constants.go diff --git a/relay/channel/ollama/main.go b/relay/adaptor/ollama/main.go similarity index 97% rename from relay/channel/ollama/main.go rename to relay/adaptor/ollama/main.go index 821a335b..a7e4c058 100644 --- a/relay/channel/ollama/main.go +++ b/relay/adaptor/ollama/main.go @@ -5,15 +5,16 @@ import ( "context" "encoding/json" "fmt" + "github.com/songquanpeng/one-api/common/helper" + "github.com/songquanpeng/one-api/common/random" "io" "net/http" "strings" "github.com/gin-gonic/gin" "github.com/songquanpeng/one-api/common" - "github.com/songquanpeng/one-api/common/helper" "github.com/songquanpeng/one-api/common/logger" - "github.com/songquanpeng/one-api/relay/channel/openai" + "github.com/songquanpeng/one-api/relay/adaptor/openai" "github.com/songquanpeng/one-api/relay/constant" "github.com/songquanpeng/one-api/relay/model" ) @@ -51,7 +52,7 @@ func responseOllama2OpenAI(response *ChatResponse) *openai.TextResponse { choice.FinishReason = "stop" } fullTextResponse := openai.TextResponse{ - Id: fmt.Sprintf("chatcmpl-%s", helper.GetUUID()), + Id: fmt.Sprintf("chatcmpl-%s", random.GetUUID()), Object: "chat.completion", Created: helper.GetTimestamp(), Choices: []openai.TextResponseChoice{choice}, @@ -72,7 +73,7 @@ func streamResponseOllama2OpenAI(ollamaResponse *ChatResponse) *openai.ChatCompl choice.FinishReason = &constant.StopFinishReason } response := openai.ChatCompletionsStreamResponse{ - Id: fmt.Sprintf("chatcmpl-%s", helper.GetUUID()), + Id: fmt.Sprintf("chatcmpl-%s", random.GetUUID()), Object: "chat.completion.chunk", Created: helper.GetTimestamp(), Model: ollamaResponse.Model, diff --git a/relay/channel/ollama/model.go b/relay/adaptor/ollama/model.go similarity index 100% rename from relay/channel/ollama/model.go rename to relay/adaptor/ollama/model.go diff --git a/relay/channel/openai/adaptor.go b/relay/adaptor/openai/adaptor.go similarity index 52% rename from relay/channel/openai/adaptor.go rename to relay/adaptor/openai/adaptor.go index 9be695f2..4bb2384e 100644 --- a/relay/channel/openai/adaptor.go +++ b/relay/adaptor/openai/adaptor.go @@ -4,11 +4,12 @@ import ( "errors" "fmt" "github.com/gin-gonic/gin" - "github.com/songquanpeng/one-api/common" - "github.com/songquanpeng/one-api/relay/channel" - "github.com/songquanpeng/one-api/relay/channel/minimax" + "github.com/songquanpeng/one-api/relay/adaptor" + "github.com/songquanpeng/one-api/relay/adaptor/minimax" + "github.com/songquanpeng/one-api/relay/channeltype" + "github.com/songquanpeng/one-api/relay/meta" "github.com/songquanpeng/one-api/relay/model" - "github.com/songquanpeng/one-api/relay/util" + "github.com/songquanpeng/one-api/relay/relaymode" "io" "net/http" "strings" @@ -18,13 +19,20 @@ type Adaptor struct { ChannelType int } -func (a *Adaptor) Init(meta *util.RelayMeta) { +func (a *Adaptor) Init(meta *meta.Meta) { a.ChannelType = meta.ChannelType } -func (a *Adaptor) GetRequestURL(meta *util.RelayMeta) (string, error) { +func (a *Adaptor) GetRequestURL(meta *meta.Meta) (string, error) { switch meta.ChannelType { - case common.ChannelTypeAzure: + case channeltype.Azure: + if meta.Mode == relaymode.ImagesGenerations { + // https://learn.microsoft.com/en-us/azure/ai-services/openai/dall-e-quickstart?tabs=dalle3%2Ccommand-line&pivots=rest-api + // https://{resource_name}.openai.azure.com/openai/deployments/dall-e-3/images/generations?api-version=2024-03-01-preview + fullRequestURL := fmt.Sprintf("%s/openai/deployments/%s/images/generations?api-version=%s", meta.BaseURL, meta.ActualModelName, meta.APIVersion) + return fullRequestURL, nil + } + // https://learn.microsoft.com/en-us/azure/cognitive-services/openai/chatgpt-quickstart?pivots=rest-api&tabs=command-line#rest-api requestURL := strings.Split(meta.RequestURLPath, "?")[0] requestURL = fmt.Sprintf("%s?api-version=%s", requestURL, meta.APIVersion) @@ -34,22 +42,22 @@ func (a *Adaptor) GetRequestURL(meta *util.RelayMeta) (string, error) { //https://github.com/songquanpeng/one-api/issues/1191 // {your endpoint}/openai/deployments/{your azure_model}/chat/completions?api-version={api_version} requestURL = fmt.Sprintf("/openai/deployments/%s/%s", model_, task) - return util.GetFullRequestURL(meta.BaseURL, requestURL, meta.ChannelType), nil - case common.ChannelTypeMinimax: + return GetFullRequestURL(meta.BaseURL, requestURL, meta.ChannelType), nil + case channeltype.Minimax: return minimax.GetRequestURL(meta) default: - return util.GetFullRequestURL(meta.BaseURL, meta.RequestURLPath, meta.ChannelType), nil + return GetFullRequestURL(meta.BaseURL, meta.RequestURLPath, meta.ChannelType), nil } } -func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Request, meta *util.RelayMeta) error { - channel.SetupCommonRequestHeader(c, req, meta) - if meta.ChannelType == common.ChannelTypeAzure { +func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Request, meta *meta.Meta) error { + adaptor.SetupCommonRequestHeader(c, req, meta) + if meta.ChannelType == channeltype.Azure { req.Header.Set("api-key", meta.APIKey) return nil } req.Header.Set("Authorization", "Bearer "+meta.APIKey) - if meta.ChannelType == common.ChannelTypeOpenRouter { + if meta.ChannelType == channeltype.OpenRouter { req.Header.Set("HTTP-Referer", "https://github.com/songquanpeng/one-api") req.Header.Set("X-Title", "One API") } @@ -63,11 +71,18 @@ func (a *Adaptor) ConvertRequest(c *gin.Context, relayMode int, request *model.G return request, nil } -func (a *Adaptor) DoRequest(c *gin.Context, meta *util.RelayMeta, requestBody io.Reader) (*http.Response, error) { - return channel.DoRequestHelper(a, c, meta, requestBody) +func (a *Adaptor) ConvertImageRequest(request *model.ImageRequest) (any, error) { + if request == nil { + return nil, errors.New("request is nil") + } + return request, nil } -func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, meta *util.RelayMeta) (usage *model.Usage, err *model.ErrorWithStatusCode) { +func (a *Adaptor) DoRequest(c *gin.Context, meta *meta.Meta, requestBody io.Reader) (*http.Response, error) { + return adaptor.DoRequestHelper(a, c, meta, requestBody) +} + +func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, meta *meta.Meta) (usage *model.Usage, err *model.ErrorWithStatusCode) { if meta.IsStream { var responseText string err, responseText, usage = StreamHandler(c, resp, meta.Mode) @@ -75,7 +90,12 @@ func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, meta *util.Rel usage = ResponseText2Usage(responseText, meta.ActualModelName, meta.PromptTokens) } } else { - err, usage = Handler(c, resp, meta.PromptTokens, meta.ActualModelName) + switch meta.Mode { + case relaymode.ImagesGenerations: + err, _ = ImageHandler(c, resp) + default: + err, usage = Handler(c, resp, meta.PromptTokens, meta.ActualModelName) + } } return } diff --git a/relay/adaptor/openai/compatible.go b/relay/adaptor/openai/compatible.go new file mode 100644 index 00000000..200eac44 --- /dev/null +++ b/relay/adaptor/openai/compatible.go @@ -0,0 +1,50 @@ +package openai + +import ( + "github.com/songquanpeng/one-api/relay/adaptor/ai360" + "github.com/songquanpeng/one-api/relay/adaptor/baichuan" + "github.com/songquanpeng/one-api/relay/adaptor/groq" + "github.com/songquanpeng/one-api/relay/adaptor/lingyiwanwu" + "github.com/songquanpeng/one-api/relay/adaptor/minimax" + "github.com/songquanpeng/one-api/relay/adaptor/mistral" + "github.com/songquanpeng/one-api/relay/adaptor/moonshot" + "github.com/songquanpeng/one-api/relay/adaptor/stepfun" + "github.com/songquanpeng/one-api/relay/channeltype" +) + +var CompatibleChannels = []int{ + channeltype.Azure, + channeltype.AI360, + channeltype.Moonshot, + channeltype.Baichuan, + channeltype.Minimax, + channeltype.Mistral, + channeltype.Groq, + channeltype.LingYiWanWu, + channeltype.StepFun, +} + +func GetCompatibleChannelMeta(channelType int) (string, []string) { + switch channelType { + case channeltype.Azure: + return "azure", ModelList + case channeltype.AI360: + return "360", ai360.ModelList + case channeltype.Moonshot: + return "moonshot", moonshot.ModelList + case channeltype.Baichuan: + return "baichuan", baichuan.ModelList + case channeltype.Minimax: + return "minimax", minimax.ModelList + case channeltype.Mistral: + return "mistralai", mistral.ModelList + case channeltype.Groq: + return "groq", groq.ModelList + case channeltype.LingYiWanWu: + return "lingyiwanwu", lingyiwanwu.ModelList + case channeltype.StepFun: + return "stepfun", stepfun.ModelList + default: + return "openai", ModelList + } +} diff --git a/relay/channel/openai/constants.go b/relay/adaptor/openai/constants.go similarity index 100% rename from relay/channel/openai/constants.go rename to relay/adaptor/openai/constants.go diff --git a/relay/adaptor/openai/helper.go b/relay/adaptor/openai/helper.go new file mode 100644 index 00000000..7d73303b --- /dev/null +++ b/relay/adaptor/openai/helper.go @@ -0,0 +1,30 @@ +package openai + +import ( + "fmt" + "github.com/songquanpeng/one-api/relay/channeltype" + "github.com/songquanpeng/one-api/relay/model" + "strings" +) + +func ResponseText2Usage(responseText string, modeName string, promptTokens int) *model.Usage { + usage := &model.Usage{} + usage.PromptTokens = promptTokens + usage.CompletionTokens = CountTokenText(responseText, modeName) + usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens + return usage +} + +func GetFullRequestURL(baseURL string, requestURL string, channelType int) string { + fullRequestURL := fmt.Sprintf("%s%s", baseURL, requestURL) + + if strings.HasPrefix(baseURL, "https://gateway.ai.cloudflare.com") { + switch channelType { + case channeltype.OpenAI: + fullRequestURL = fmt.Sprintf("%s%s", baseURL, strings.TrimPrefix(requestURL, "/v1")) + case channeltype.Azure: + fullRequestURL = fmt.Sprintf("%s%s", baseURL, strings.TrimPrefix(requestURL, "/openai/deployments")) + } + } + return fullRequestURL +} diff --git a/relay/adaptor/openai/image.go b/relay/adaptor/openai/image.go new file mode 100644 index 00000000..0f89618a --- /dev/null +++ b/relay/adaptor/openai/image.go @@ -0,0 +1,44 @@ +package openai + +import ( + "bytes" + "encoding/json" + "github.com/gin-gonic/gin" + "github.com/songquanpeng/one-api/relay/model" + "io" + "net/http" +) + +func ImageHandler(c *gin.Context, resp *http.Response) (*model.ErrorWithStatusCode, *model.Usage) { + var imageResponse ImageResponse + responseBody, err := io.ReadAll(resp.Body) + + if err != nil { + return ErrorWrapper(err, "read_response_body_failed", http.StatusInternalServerError), nil + } + err = resp.Body.Close() + if err != nil { + return ErrorWrapper(err, "close_response_body_failed", http.StatusInternalServerError), nil + } + err = json.Unmarshal(responseBody, &imageResponse) + if err != nil { + return ErrorWrapper(err, "unmarshal_response_body_failed", http.StatusInternalServerError), nil + } + + resp.Body = io.NopCloser(bytes.NewBuffer(responseBody)) + + for k, v := range resp.Header { + c.Writer.Header().Set(k, v[0]) + } + c.Writer.WriteHeader(resp.StatusCode) + + _, err = io.Copy(c.Writer, resp.Body) + if err != nil { + return ErrorWrapper(err, "copy_response_body_failed", http.StatusInternalServerError), nil + } + err = resp.Body.Close() + if err != nil { + return ErrorWrapper(err, "close_response_body_failed", http.StatusInternalServerError), nil + } + return nil, nil +} diff --git a/relay/channel/openai/main.go b/relay/adaptor/openai/main.go similarity index 97% rename from relay/channel/openai/main.go rename to relay/adaptor/openai/main.go index 63cb9ae8..68d8f48f 100644 --- a/relay/channel/openai/main.go +++ b/relay/adaptor/openai/main.go @@ -8,8 +8,8 @@ import ( "github.com/songquanpeng/one-api/common" "github.com/songquanpeng/one-api/common/conv" "github.com/songquanpeng/one-api/common/logger" - "github.com/songquanpeng/one-api/relay/constant" "github.com/songquanpeng/one-api/relay/model" + "github.com/songquanpeng/one-api/relay/relaymode" "io" "net/http" "strings" @@ -46,7 +46,7 @@ func StreamHandler(c *gin.Context, resp *http.Response, relayMode int) (*model.E data = data[6:] if !strings.HasPrefix(data, "[DONE]") { switch relayMode { - case constant.RelayModeChatCompletions: + case relaymode.ChatCompletions: var streamResponse ChatCompletionsStreamResponse err := json.Unmarshal([]byte(data), &streamResponse) if err != nil { @@ -59,7 +59,7 @@ func StreamHandler(c *gin.Context, resp *http.Response, relayMode int) (*model.E if streamResponse.Usage != nil { usage = streamResponse.Usage } - case constant.RelayModeCompletions: + case relaymode.Completions: var streamResponse CompletionsStreamResponse err := json.Unmarshal([]byte(data), &streamResponse) if err != nil { diff --git a/relay/channel/openai/model.go b/relay/adaptor/openai/model.go similarity index 93% rename from relay/channel/openai/model.go rename to relay/adaptor/openai/model.go index 30d77739..ce252ff6 100644 --- a/relay/channel/openai/model.go +++ b/relay/adaptor/openai/model.go @@ -110,11 +110,16 @@ type EmbeddingResponse struct { model.Usage `json:"usage"` } +type ImageData struct { + Url string `json:"url,omitempty"` + B64Json string `json:"b64_json,omitempty"` + RevisedPrompt string `json:"revised_prompt,omitempty"` +} + type ImageResponse struct { - Created int `json:"created"` - Data []struct { - Url string `json:"url"` - } + Created int64 `json:"created"` + Data []ImageData `json:"data"` + //model.Usage `json:"usage"` } type ChatCompletionsStreamResponseChoice struct { diff --git a/relay/channel/openai/token.go b/relay/adaptor/openai/token.go similarity index 98% rename from relay/channel/openai/token.go rename to relay/adaptor/openai/token.go index 0720425f..c95a7b5e 100644 --- a/relay/channel/openai/token.go +++ b/relay/adaptor/openai/token.go @@ -4,10 +4,10 @@ import ( "errors" "fmt" "github.com/pkoukk/tiktoken-go" - "github.com/songquanpeng/one-api/common" "github.com/songquanpeng/one-api/common/config" "github.com/songquanpeng/one-api/common/image" "github.com/songquanpeng/one-api/common/logger" + billingratio "github.com/songquanpeng/one-api/relay/billing/ratio" "github.com/songquanpeng/one-api/relay/model" "math" "strings" @@ -28,7 +28,7 @@ func InitTokenEncoders() { if err != nil { logger.FatalLog(fmt.Sprintf("failed to get gpt-4 token encoder: %s", err.Error())) } - for model := range common.ModelRatio { + for model := range billingratio.ModelRatio { if strings.HasPrefix(model, "gpt-3.5") { tokenEncoderMap[model] = gpt35TokenEncoder } else if strings.HasPrefix(model, "gpt-4") { diff --git a/relay/channel/openai/util.go b/relay/adaptor/openai/util.go similarity index 100% rename from relay/channel/openai/util.go rename to relay/adaptor/openai/util.go diff --git a/relay/channel/palm/adaptor.go b/relay/adaptor/palm/adaptor.go similarity index 59% rename from relay/channel/palm/adaptor.go rename to relay/adaptor/palm/adaptor.go index efd0620c..98aa3e18 100644 --- a/relay/channel/palm/adaptor.go +++ b/relay/adaptor/palm/adaptor.go @@ -4,10 +4,10 @@ import ( "errors" "fmt" "github.com/gin-gonic/gin" - "github.com/songquanpeng/one-api/relay/channel" - "github.com/songquanpeng/one-api/relay/channel/openai" + "github.com/songquanpeng/one-api/relay/adaptor" + "github.com/songquanpeng/one-api/relay/adaptor/openai" + "github.com/songquanpeng/one-api/relay/meta" "github.com/songquanpeng/one-api/relay/model" - "github.com/songquanpeng/one-api/relay/util" "io" "net/http" ) @@ -15,16 +15,16 @@ import ( type Adaptor struct { } -func (a *Adaptor) Init(meta *util.RelayMeta) { +func (a *Adaptor) Init(meta *meta.Meta) { } -func (a *Adaptor) GetRequestURL(meta *util.RelayMeta) (string, error) { +func (a *Adaptor) GetRequestURL(meta *meta.Meta) (string, error) { return fmt.Sprintf("%s/v1beta2/models/chat-bison-001:generateMessage", meta.BaseURL), nil } -func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Request, meta *util.RelayMeta) error { - channel.SetupCommonRequestHeader(c, req, meta) +func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Request, meta *meta.Meta) error { + adaptor.SetupCommonRequestHeader(c, req, meta) req.Header.Set("x-goog-api-key", meta.APIKey) return nil } @@ -36,11 +36,18 @@ func (a *Adaptor) ConvertRequest(c *gin.Context, relayMode int, request *model.G return ConvertRequest(*request), nil } -func (a *Adaptor) DoRequest(c *gin.Context, meta *util.RelayMeta, requestBody io.Reader) (*http.Response, error) { - return channel.DoRequestHelper(a, c, meta, requestBody) +func (a *Adaptor) ConvertImageRequest(request *model.ImageRequest) (any, error) { + if request == nil { + return nil, errors.New("request is nil") + } + return request, nil } -func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, meta *util.RelayMeta) (usage *model.Usage, err *model.ErrorWithStatusCode) { +func (a *Adaptor) DoRequest(c *gin.Context, meta *meta.Meta, requestBody io.Reader) (*http.Response, error) { + return adaptor.DoRequestHelper(a, c, meta, requestBody) +} + +func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, meta *meta.Meta) (usage *model.Usage, err *model.ErrorWithStatusCode) { if meta.IsStream { var responseText string err, responseText = StreamHandler(c, resp) diff --git a/relay/channel/palm/constants.go b/relay/adaptor/palm/constants.go similarity index 100% rename from relay/channel/palm/constants.go rename to relay/adaptor/palm/constants.go diff --git a/relay/channel/palm/model.go b/relay/adaptor/palm/model.go similarity index 100% rename from relay/channel/palm/model.go rename to relay/adaptor/palm/model.go diff --git a/relay/channel/palm/palm.go b/relay/adaptor/palm/palm.go similarity index 97% rename from relay/channel/palm/palm.go rename to relay/adaptor/palm/palm.go index 56738544..1e60e7cd 100644 --- a/relay/channel/palm/palm.go +++ b/relay/adaptor/palm/palm.go @@ -7,7 +7,8 @@ import ( "github.com/songquanpeng/one-api/common" "github.com/songquanpeng/one-api/common/helper" "github.com/songquanpeng/one-api/common/logger" - "github.com/songquanpeng/one-api/relay/channel/openai" + "github.com/songquanpeng/one-api/common/random" + "github.com/songquanpeng/one-api/relay/adaptor/openai" "github.com/songquanpeng/one-api/relay/constant" "github.com/songquanpeng/one-api/relay/model" "io" @@ -74,7 +75,7 @@ func streamResponsePaLM2OpenAI(palmResponse *ChatResponse) *openai.ChatCompletio func StreamHandler(c *gin.Context, resp *http.Response) (*model.ErrorWithStatusCode, string) { responseText := "" - responseId := fmt.Sprintf("chatcmpl-%s", helper.GetUUID()) + responseId := fmt.Sprintf("chatcmpl-%s", random.GetUUID()) createdTime := helper.GetTimestamp() dataChan := make(chan string) stopChan := make(chan bool) diff --git a/relay/adaptor/stepfun/constants.go b/relay/adaptor/stepfun/constants.go new file mode 100644 index 00000000..a82e562b --- /dev/null +++ b/relay/adaptor/stepfun/constants.go @@ -0,0 +1,7 @@ +package stepfun + +var ModelList = []string{ + "step-1-32k", + "step-1v-32k", + "step-1-200k", +} diff --git a/relay/channel/tencent/adaptor.go b/relay/adaptor/tencent/adaptor.go similarity index 66% rename from relay/channel/tencent/adaptor.go rename to relay/adaptor/tencent/adaptor.go index f348674e..a97476d6 100644 --- a/relay/channel/tencent/adaptor.go +++ b/relay/adaptor/tencent/adaptor.go @@ -4,10 +4,10 @@ import ( "errors" "fmt" "github.com/gin-gonic/gin" - "github.com/songquanpeng/one-api/relay/channel" - "github.com/songquanpeng/one-api/relay/channel/openai" + "github.com/songquanpeng/one-api/relay/adaptor" + "github.com/songquanpeng/one-api/relay/adaptor/openai" + "github.com/songquanpeng/one-api/relay/meta" "github.com/songquanpeng/one-api/relay/model" - "github.com/songquanpeng/one-api/relay/util" "io" "net/http" "strings" @@ -19,16 +19,16 @@ type Adaptor struct { Sign string } -func (a *Adaptor) Init(meta *util.RelayMeta) { +func (a *Adaptor) Init(meta *meta.Meta) { } -func (a *Adaptor) GetRequestURL(meta *util.RelayMeta) (string, error) { +func (a *Adaptor) GetRequestURL(meta *meta.Meta) (string, error) { return fmt.Sprintf("%s/hyllm/v1/chat/completions", meta.BaseURL), nil } -func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Request, meta *util.RelayMeta) error { - channel.SetupCommonRequestHeader(c, req, meta) +func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Request, meta *meta.Meta) error { + adaptor.SetupCommonRequestHeader(c, req, meta) req.Header.Set("Authorization", a.Sign) req.Header.Set("X-TC-Action", meta.ActualModelName) return nil @@ -52,11 +52,18 @@ func (a *Adaptor) ConvertRequest(c *gin.Context, relayMode int, request *model.G return tencentRequest, nil } -func (a *Adaptor) DoRequest(c *gin.Context, meta *util.RelayMeta, requestBody io.Reader) (*http.Response, error) { - return channel.DoRequestHelper(a, c, meta, requestBody) +func (a *Adaptor) ConvertImageRequest(request *model.ImageRequest) (any, error) { + if request == nil { + return nil, errors.New("request is nil") + } + return request, nil } -func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, meta *util.RelayMeta) (usage *model.Usage, err *model.ErrorWithStatusCode) { +func (a *Adaptor) DoRequest(c *gin.Context, meta *meta.Meta, requestBody io.Reader) (*http.Response, error) { + return adaptor.DoRequestHelper(a, c, meta, requestBody) +} + +func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, meta *meta.Meta) (usage *model.Usage, err *model.ErrorWithStatusCode) { if meta.IsStream { var responseText string err, responseText = StreamHandler(c, resp) diff --git a/relay/channel/tencent/constants.go b/relay/adaptor/tencent/constants.go similarity index 100% rename from relay/channel/tencent/constants.go rename to relay/adaptor/tencent/constants.go diff --git a/relay/channel/tencent/main.go b/relay/adaptor/tencent/main.go similarity index 97% rename from relay/channel/tencent/main.go rename to relay/adaptor/tencent/main.go index b5a64cde..2ca5724e 100644 --- a/relay/channel/tencent/main.go +++ b/relay/adaptor/tencent/main.go @@ -13,7 +13,8 @@ import ( "github.com/songquanpeng/one-api/common/conv" "github.com/songquanpeng/one-api/common/helper" "github.com/songquanpeng/one-api/common/logger" - "github.com/songquanpeng/one-api/relay/channel/openai" + "github.com/songquanpeng/one-api/common/random" + "github.com/songquanpeng/one-api/relay/adaptor/openai" "github.com/songquanpeng/one-api/relay/constant" "github.com/songquanpeng/one-api/relay/model" "io" @@ -41,7 +42,7 @@ func ConvertRequest(request model.GeneralOpenAIRequest) *ChatRequest { return &ChatRequest{ Timestamp: helper.GetTimestamp(), Expired: helper.GetTimestamp() + 24*60*60, - QueryID: helper.GetUUID(), + QueryID: random.GetUUID(), Temperature: request.Temperature, TopP: request.TopP, Stream: stream, @@ -71,7 +72,7 @@ func responseTencent2OpenAI(response *ChatResponse) *openai.TextResponse { func streamResponseTencent2OpenAI(TencentResponse *ChatResponse) *openai.ChatCompletionsStreamResponse { response := openai.ChatCompletionsStreamResponse{ - Id: fmt.Sprintf("chatcmpl-%s", helper.GetUUID()), + Id: fmt.Sprintf("chatcmpl-%s", random.GetUUID()), Object: "chat.completion.chunk", Created: helper.GetTimestamp(), Model: "tencent-hunyuan", diff --git a/relay/channel/tencent/model.go b/relay/adaptor/tencent/model.go similarity index 100% rename from relay/channel/tencent/model.go rename to relay/adaptor/tencent/model.go diff --git a/relay/channel/xunfei/adaptor.go b/relay/adaptor/xunfei/adaptor.go similarity index 67% rename from relay/channel/xunfei/adaptor.go rename to relay/adaptor/xunfei/adaptor.go index 92d9d7d6..edcd719f 100644 --- a/relay/channel/xunfei/adaptor.go +++ b/relay/adaptor/xunfei/adaptor.go @@ -3,10 +3,10 @@ package xunfei import ( "errors" "github.com/gin-gonic/gin" - "github.com/songquanpeng/one-api/relay/channel" - "github.com/songquanpeng/one-api/relay/channel/openai" + "github.com/songquanpeng/one-api/relay/adaptor" + "github.com/songquanpeng/one-api/relay/adaptor/openai" + "github.com/songquanpeng/one-api/relay/meta" "github.com/songquanpeng/one-api/relay/model" - "github.com/songquanpeng/one-api/relay/util" "io" "net/http" "strings" @@ -16,16 +16,16 @@ type Adaptor struct { request *model.GeneralOpenAIRequest } -func (a *Adaptor) Init(meta *util.RelayMeta) { +func (a *Adaptor) Init(meta *meta.Meta) { } -func (a *Adaptor) GetRequestURL(meta *util.RelayMeta) (string, error) { +func (a *Adaptor) GetRequestURL(meta *meta.Meta) (string, error) { return "", nil } -func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Request, meta *util.RelayMeta) error { - channel.SetupCommonRequestHeader(c, req, meta) +func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Request, meta *meta.Meta) error { + adaptor.SetupCommonRequestHeader(c, req, meta) // check DoResponse for auth part return nil } @@ -38,14 +38,21 @@ func (a *Adaptor) ConvertRequest(c *gin.Context, relayMode int, request *model.G return nil, nil } -func (a *Adaptor) DoRequest(c *gin.Context, meta *util.RelayMeta, requestBody io.Reader) (*http.Response, error) { +func (a *Adaptor) ConvertImageRequest(request *model.ImageRequest) (any, error) { + if request == nil { + return nil, errors.New("request is nil") + } + return request, nil +} + +func (a *Adaptor) DoRequest(c *gin.Context, meta *meta.Meta, requestBody io.Reader) (*http.Response, error) { // xunfei's request is not http request, so we don't need to do anything here dummyResp := &http.Response{} dummyResp.StatusCode = http.StatusOK return dummyResp, nil } -func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, meta *util.RelayMeta) (usage *model.Usage, err *model.ErrorWithStatusCode) { +func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, meta *meta.Meta) (usage *model.Usage, err *model.ErrorWithStatusCode) { splits := strings.Split(meta.APIKey, "|") if len(splits) != 3 { return nil, openai.ErrorWrapper(errors.New("invalid auth"), "invalid_auth", http.StatusBadRequest) diff --git a/relay/channel/xunfei/constants.go b/relay/adaptor/xunfei/constants.go similarity index 100% rename from relay/channel/xunfei/constants.go rename to relay/adaptor/xunfei/constants.go diff --git a/relay/channel/xunfei/main.go b/relay/adaptor/xunfei/main.go similarity index 96% rename from relay/channel/xunfei/main.go rename to relay/adaptor/xunfei/main.go index 67784a56..369e6227 100644 --- a/relay/channel/xunfei/main.go +++ b/relay/adaptor/xunfei/main.go @@ -9,9 +9,11 @@ import ( "github.com/gin-gonic/gin" "github.com/gorilla/websocket" "github.com/songquanpeng/one-api/common" + "github.com/songquanpeng/one-api/common/config" "github.com/songquanpeng/one-api/common/helper" "github.com/songquanpeng/one-api/common/logger" - "github.com/songquanpeng/one-api/relay/channel/openai" + "github.com/songquanpeng/one-api/common/random" + "github.com/songquanpeng/one-api/relay/adaptor/openai" "github.com/songquanpeng/one-api/relay/constant" "github.com/songquanpeng/one-api/relay/model" "io" @@ -62,7 +64,7 @@ func getToolCalls(response *ChatResponse) []model.Tool { return toolCalls } toolCall := model.Tool{ - Id: fmt.Sprintf("call_%s", helper.GetUUID()), + Id: fmt.Sprintf("call_%s", random.GetUUID()), Type: "function", Function: *item.FunctionCall, } @@ -88,7 +90,7 @@ func responseXunfei2OpenAI(response *ChatResponse) *openai.TextResponse { FinishReason: constant.StopFinishReason, } fullTextResponse := openai.TextResponse{ - Id: fmt.Sprintf("chatcmpl-%s", helper.GetUUID()), + Id: fmt.Sprintf("chatcmpl-%s", random.GetUUID()), Object: "chat.completion", Created: helper.GetTimestamp(), Choices: []openai.TextResponseChoice{choice}, @@ -112,7 +114,7 @@ func streamResponseXunfei2OpenAI(xunfeiResponse *ChatResponse) *openai.ChatCompl choice.FinishReason = &constant.StopFinishReason } response := openai.ChatCompletionsStreamResponse{ - Id: fmt.Sprintf("chatcmpl-%s", helper.GetUUID()), + Id: fmt.Sprintf("chatcmpl-%s", random.GetUUID()), Object: "chat.completion.chunk", Created: helper.GetTimestamp(), Model: "SparkDesk", @@ -278,7 +280,7 @@ func getAPIVersion(c *gin.Context, modelName string) string { return apiVersion } - apiVersion = c.GetString(common.ConfigKeyAPIVersion) + apiVersion = c.GetString(config.KeyAPIVersion) if apiVersion != "" { return apiVersion } diff --git a/relay/channel/xunfei/model.go b/relay/adaptor/xunfei/model.go similarity index 98% rename from relay/channel/xunfei/model.go rename to relay/adaptor/xunfei/model.go index e9cc59a6..97a43154 100644 --- a/relay/channel/xunfei/model.go +++ b/relay/adaptor/xunfei/model.go @@ -28,7 +28,7 @@ type ChatRequest struct { } `json:"message"` Functions struct { Text []model.Function `json:"text,omitempty"` - } `json:"functions"` + } `json:"functions,omitempty"` } `json:"payload"` } diff --git a/relay/adaptor/zhipu/adaptor.go b/relay/adaptor/zhipu/adaptor.go new file mode 100644 index 00000000..5ebafbb3 --- /dev/null +++ b/relay/adaptor/zhipu/adaptor.go @@ -0,0 +1,145 @@ +package zhipu + +import ( + "errors" + "fmt" + "github.com/gin-gonic/gin" + "github.com/songquanpeng/one-api/relay/adaptor" + "github.com/songquanpeng/one-api/relay/adaptor/openai" + "github.com/songquanpeng/one-api/relay/meta" + "github.com/songquanpeng/one-api/relay/model" + "github.com/songquanpeng/one-api/relay/relaymode" + "io" + "math" + "net/http" + "strings" +) + +type Adaptor struct { + APIVersion string +} + +func (a *Adaptor) Init(meta *meta.Meta) { + +} + +func (a *Adaptor) SetVersionByModeName(modelName string) { + if strings.HasPrefix(modelName, "glm-") { + a.APIVersion = "v4" + } else { + a.APIVersion = "v3" + } +} + +func (a *Adaptor) GetRequestURL(meta *meta.Meta) (string, error) { + switch meta.Mode { + case relaymode.ImagesGenerations: + return fmt.Sprintf("%s/api/paas/v4/images/generations", meta.BaseURL), nil + case relaymode.Embeddings: + return fmt.Sprintf("%s/api/paas/v4/embeddings", meta.BaseURL), nil + } + a.SetVersionByModeName(meta.ActualModelName) + if a.APIVersion == "v4" { + return fmt.Sprintf("%s/api/paas/v4/chat/completions", meta.BaseURL), nil + } + method := "invoke" + if meta.IsStream { + method = "sse-invoke" + } + return fmt.Sprintf("%s/api/paas/v3/model-api/%s/%s", meta.BaseURL, meta.ActualModelName, method), nil +} + +func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Request, meta *meta.Meta) error { + adaptor.SetupCommonRequestHeader(c, req, meta) + token := GetToken(meta.APIKey) + req.Header.Set("Authorization", token) + return nil +} + +func (a *Adaptor) ConvertRequest(c *gin.Context, relayMode int, request *model.GeneralOpenAIRequest) (any, error) { + if request == nil { + return nil, errors.New("request is nil") + } + switch relayMode { + case relaymode.Embeddings: + baiduEmbeddingRequest := ConvertEmbeddingRequest(*request) + return baiduEmbeddingRequest, nil + default: + // TopP (0.0, 1.0) + request.TopP = math.Min(0.99, request.TopP) + request.TopP = math.Max(0.01, request.TopP) + + // Temperature (0.0, 1.0) + request.Temperature = math.Min(0.99, request.Temperature) + request.Temperature = math.Max(0.01, request.Temperature) + a.SetVersionByModeName(request.Model) + if a.APIVersion == "v4" { + return request, nil + } + return ConvertRequest(*request), nil + } +} + +func (a *Adaptor) ConvertImageRequest(request *model.ImageRequest) (any, error) { + if request == nil { + return nil, errors.New("request is nil") + } + newRequest := ImageRequest{ + Model: request.Model, + Prompt: request.Prompt, + UserId: request.User, + } + return newRequest, nil +} + +func (a *Adaptor) DoRequest(c *gin.Context, meta *meta.Meta, requestBody io.Reader) (*http.Response, error) { + return adaptor.DoRequestHelper(a, c, meta, requestBody) +} + +func (a *Adaptor) DoResponseV4(c *gin.Context, resp *http.Response, meta *meta.Meta) (usage *model.Usage, err *model.ErrorWithStatusCode) { + if meta.IsStream { + err, _, usage = openai.StreamHandler(c, resp, meta.Mode) + } else { + err, usage = openai.Handler(c, resp, meta.PromptTokens, meta.ActualModelName) + } + return +} + +func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, meta *meta.Meta) (usage *model.Usage, err *model.ErrorWithStatusCode) { + switch meta.Mode { + case relaymode.Embeddings: + err, usage = EmbeddingsHandler(c, resp) + return + case relaymode.ImagesGenerations: + err, usage = openai.ImageHandler(c, resp) + return + } + if a.APIVersion == "v4" { + return a.DoResponseV4(c, resp, meta) + } + if meta.IsStream { + err, usage = StreamHandler(c, resp) + } else { + if meta.Mode == relaymode.Embeddings { + err, usage = EmbeddingsHandler(c, resp) + } else { + err, usage = Handler(c, resp) + } + } + return +} + +func ConvertEmbeddingRequest(request model.GeneralOpenAIRequest) *EmbeddingRequest { + return &EmbeddingRequest{ + Model: "embedding-2", + Input: request.Input.(string), + } +} + +func (a *Adaptor) GetModelList() []string { + return ModelList +} + +func (a *Adaptor) GetChannelName() string { + return "zhipu" +} diff --git a/relay/channel/zhipu/constants.go b/relay/adaptor/zhipu/constants.go similarity index 62% rename from relay/channel/zhipu/constants.go rename to relay/adaptor/zhipu/constants.go index 1655a59d..e1192123 100644 --- a/relay/channel/zhipu/constants.go +++ b/relay/adaptor/zhipu/constants.go @@ -2,5 +2,6 @@ package zhipu var ModelList = []string{ "chatglm_turbo", "chatglm_pro", "chatglm_std", "chatglm_lite", - "glm-4", "glm-4v", "glm-3-turbo", + "glm-4", "glm-4v", "glm-3-turbo", "embedding-2", + "cogview-3", } diff --git a/relay/channel/zhipu/main.go b/relay/adaptor/zhipu/main.go similarity index 80% rename from relay/channel/zhipu/main.go rename to relay/adaptor/zhipu/main.go index a46fd537..74a1a05e 100644 --- a/relay/channel/zhipu/main.go +++ b/relay/adaptor/zhipu/main.go @@ -8,7 +8,7 @@ import ( "github.com/songquanpeng/one-api/common" "github.com/songquanpeng/one-api/common/helper" "github.com/songquanpeng/one-api/common/logger" - "github.com/songquanpeng/one-api/relay/channel/openai" + "github.com/songquanpeng/one-api/relay/adaptor/openai" "github.com/songquanpeng/one-api/relay/constant" "github.com/songquanpeng/one-api/relay/model" "io" @@ -254,3 +254,50 @@ func Handler(c *gin.Context, resp *http.Response) (*model.ErrorWithStatusCode, * _, err = c.Writer.Write(jsonResponse) return nil, &fullTextResponse.Usage } + +func EmbeddingsHandler(c *gin.Context, resp *http.Response) (*model.ErrorWithStatusCode, *model.Usage) { + var zhipuResponse EmbeddingResponse + responseBody, err := io.ReadAll(resp.Body) + if err != nil { + return openai.ErrorWrapper(err, "read_response_body_failed", http.StatusInternalServerError), nil + } + err = resp.Body.Close() + if err != nil { + return openai.ErrorWrapper(err, "close_response_body_failed", http.StatusInternalServerError), nil + } + err = json.Unmarshal(responseBody, &zhipuResponse) + if err != nil { + return openai.ErrorWrapper(err, "unmarshal_response_body_failed", http.StatusInternalServerError), nil + } + fullTextResponse := embeddingResponseZhipu2OpenAI(&zhipuResponse) + jsonResponse, err := json.Marshal(fullTextResponse) + if err != nil { + return openai.ErrorWrapper(err, "marshal_response_body_failed", http.StatusInternalServerError), nil + } + c.Writer.Header().Set("Content-Type", "application/json") + c.Writer.WriteHeader(resp.StatusCode) + _, err = c.Writer.Write(jsonResponse) + return nil, &fullTextResponse.Usage +} + +func embeddingResponseZhipu2OpenAI(response *EmbeddingResponse) *openai.EmbeddingResponse { + openAIEmbeddingResponse := openai.EmbeddingResponse{ + Object: "list", + Data: make([]openai.EmbeddingResponseItem, 0, len(response.Embeddings)), + Model: response.Model, + Usage: model.Usage{ + PromptTokens: response.PromptTokens, + CompletionTokens: response.CompletionTokens, + TotalTokens: response.Usage.TotalTokens, + }, + } + + for _, item := range response.Embeddings { + openAIEmbeddingResponse.Data = append(openAIEmbeddingResponse.Data, openai.EmbeddingResponseItem{ + Object: `embedding`, + Index: item.Index, + Embedding: item.Embedding, + }) + } + return &openAIEmbeddingResponse +} diff --git a/relay/channel/zhipu/model.go b/relay/adaptor/zhipu/model.go similarity index 65% rename from relay/channel/zhipu/model.go rename to relay/adaptor/zhipu/model.go index b63e1d6f..f91de1dc 100644 --- a/relay/channel/zhipu/model.go +++ b/relay/adaptor/zhipu/model.go @@ -44,3 +44,27 @@ type tokenData struct { Token string ExpiryTime time.Time } + +type EmbeddingRequest struct { + Model string `json:"model"` + Input string `json:"input"` +} + +type EmbeddingResponse struct { + Model string `json:"model"` + Object string `json:"object"` + Embeddings []EmbeddingData `json:"data"` + model.Usage `json:"usage"` +} + +type EmbeddingData struct { + Index int `json:"index"` + Object string `json:"object"` + Embedding []float64 `json:"embedding"` +} + +type ImageRequest struct { + Model string `json:"model"` + Prompt string `json:"prompt"` + UserId string `json:"user_id,omitempty"` +} diff --git a/relay/apitype/define.go b/relay/apitype/define.go new file mode 100644 index 00000000..82d32a50 --- /dev/null +++ b/relay/apitype/define.go @@ -0,0 +1,17 @@ +package apitype + +const ( + OpenAI = iota + Anthropic + PaLM + Baidu + Zhipu + Ali + Xunfei + AIProxyLibrary + Tencent + Gemini + Ollama + + Dummy // this one is only for count, do not add any channel after this +) diff --git a/relay/billing/billing.go b/relay/billing/billing.go new file mode 100644 index 00000000..a99d37ee --- /dev/null +++ b/relay/billing/billing.go @@ -0,0 +1,42 @@ +package billing + +import ( + "context" + "fmt" + "github.com/songquanpeng/one-api/common/logger" + "github.com/songquanpeng/one-api/model" +) + +func ReturnPreConsumedQuota(ctx context.Context, preConsumedQuota int64, 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) + } +} + +func PostConsumeQuota(ctx context.Context, tokenId int, quotaDelta int64, totalQuota int64, userId int, channelId int, modelRatio float64, groupRatio float64, modelName string, tokenName string) { + // quotaDelta is remaining quota to be consumed + err := model.PostConsumeTokenQuota(tokenId, quotaDelta) + if err != nil { + logger.SysError("error consuming token remain quota: " + err.Error()) + } + err = model.CacheUpdateUserQuota(ctx, userId) + if err != nil { + logger.SysError("error update user quota cache: " + err.Error()) + } + // totalQuota is total quota consumed + if totalQuota != 0 { + logContent := fmt.Sprintf("模型倍率 %.2f，分组倍率 %.2f", modelRatio, groupRatio) + model.RecordConsumeLog(ctx, userId, channelId, int(totalQuota), 0, modelName, tokenName, totalQuota, logContent) + model.UpdateUserUsedQuotaAndRequestCount(userId, totalQuota) + model.UpdateChannelUsedQuota(channelId, totalQuota) + } + if totalQuota <= 0 { + logger.Error(ctx, fmt.Sprintf("totalQuota consumed is %d, something is wrong", totalQuota)) + } +} diff --git a/common/group-ratio.go b/relay/billing/ratio/group.go similarity index 97% rename from common/group-ratio.go rename to relay/billing/ratio/group.go index 2de6e810..8e9c5b73 100644 --- a/common/group-ratio.go +++ b/relay/billing/ratio/group.go @@ -1,4 +1,4 @@ -package common +package ratio import ( "encoding/json" diff --git a/relay/billing/ratio/image.go b/relay/billing/ratio/image.go new file mode 100644 index 00000000..5a29cddc --- /dev/null +++ b/relay/billing/ratio/image.go @@ -0,0 +1,51 @@ +package ratio + +var ImageSizeRatios = map[string]map[string]float64{ + "dall-e-2": { + "256x256": 1, + "512x512": 1.125, + "1024x1024": 1.25, + }, + "dall-e-3": { + "1024x1024": 1, + "1024x1792": 2, + "1792x1024": 2, + }, + "ali-stable-diffusion-xl": { + "512x1024": 1, + "1024x768": 1, + "1024x1024": 1, + "576x1024": 1, + "1024x576": 1, + }, + "ali-stable-diffusion-v1.5": { + "512x1024": 1, + "1024x768": 1, + "1024x1024": 1, + "576x1024": 1, + "1024x576": 1, + }, + "wanx-v1": { + "1024x1024": 1, + "720x1280": 1, + "1280x720": 1, + }, +} + +var ImageGenerationAmounts = map[string][2]int{ + "dall-e-2": {1, 10}, + "dall-e-3": {1, 1}, // OpenAI allows n=1 currently. + "ali-stable-diffusion-xl": {1, 4}, // Ali + "ali-stable-diffusion-v1.5": {1, 4}, // Ali + "wanx-v1": {1, 4}, // Ali + "cogview-3": {1, 1}, +} + +var ImagePromptLengthLimitations = map[string]int{ + "dall-e-2": 1000, + "dall-e-3": 4000, + "ali-stable-diffusion-xl": 4000, + "ali-stable-diffusion-v1.5": 4000, + "wanx-v1": 4000, + "cogview-3": 833, +} diff --git a/common/model-ratio.go b/relay/billing/ratio/model.go similarity index 84% rename from common/model-ratio.go rename to relay/billing/ratio/model.go index aa75042e..108924a1 100644 --- a/common/model-ratio.go +++ b/relay/billing/ratio/model.go @@ -1,4 +1,4 @@ -package common +package ratio import ( "encoding/json" @@ -62,8 +62,8 @@ var ModelRatio = map[string]float64{ "text-search-ada-doc-001": 10, "text-moderation-stable": 0.1, "text-moderation-latest": 0.1, - "dall-e-2": 8, // $0.016 - $0.020 / image - "dall-e-3": 20, // $0.040 - $0.120 / image + "dall-e-2": 0.02 * USD, // $0.016 - $0.020 / image + "dall-e-3": 0.04 * USD, // $0.040 - $0.120 / image // https://www.anthropic.com/api#pricing "claude-instant-1.2": 0.8 / 1000 * USD, "claude-2.0": 8.0 / 1000 * USD, @@ -72,14 +72,22 @@ var ModelRatio = map[string]float64{ "claude-3-sonnet-20240229": 3.0 / 1000 * USD, "claude-3-opus-20240229": 15.0 / 1000 * USD, // https://cloud.baidu.com/doc/WENXINWORKSHOP/s/hlrk4akp7 - "ERNIE-Bot": 0.8572, // ￥0.012 / 1k tokens - "ERNIE-Bot-turbo": 0.5715, // ￥0.008 / 1k tokens - "ERNIE-Bot-4": 0.12 * RMB, // ￥0.12 / 1k tokens - "ERNIE-Bot-8K": 0.024 * RMB, - "Embedding-V1": 0.1429, // ￥0.002 / 1k tokens - "bge-large-zh": 0.002 * RMB, - "bge-large-en": 0.002 * RMB, - "bge-large-8k": 0.002 * RMB, + "ERNIE-4.0-8K": 0.120 * RMB, + "ERNIE-3.5-8K": 0.012 * RMB, + "ERNIE-3.5-8K-0205": 0.024 * RMB, + "ERNIE-3.5-8K-1222": 0.012 * RMB, + "ERNIE-Bot-8K": 0.024 * RMB, + "ERNIE-3.5-4K-0205": 0.012 * RMB, + "ERNIE-Speed-8K": 0.004 * RMB, + "ERNIE-Speed-128K": 0.004 * RMB, + "ERNIE-Lite-8K-0922": 0.008 * RMB, + "ERNIE-Lite-8K-0308": 0.003 * RMB, + "ERNIE-Tiny-8K": 0.001 * RMB, + "BLOOMZ-7B": 0.004 * RMB, + "Embedding-V1": 0.002 * RMB, + "bge-large-zh": 0.002 * RMB, + "bge-large-en": 0.002 * RMB, + "tao-8k": 0.002 * RMB, // https://ai.google.dev/pricing "PaLM-2": 1, "gemini-pro": 1, // $0.00025 / 1k characters -> $0.001 / 1k tokens @@ -88,18 +96,24 @@ var ModelRatio = map[string]float64{ "gemini-1.0-pro-001": 1, "gemini-1.5-pro": 1, // https://open.bigmodel.cn/pricing - "glm-4": 0.1 * RMB, - "glm-4v": 0.1 * RMB, - "glm-3-turbo": 0.005 * RMB, - "chatglm_turbo": 0.3572, // ￥0.005 / 1k tokens - "chatglm_pro": 0.7143, // ￥0.01 / 1k tokens - "chatglm_std": 0.3572, // ￥0.005 / 1k tokens - "chatglm_lite": 0.1429, // ￥0.002 / 1k tokens - "qwen-turbo": 0.5715, // ￥0.008 / 1k tokens // https://help.aliyun.com/zh/dashscope/developer-reference/tongyi-thousand-questions-metering-and-billing + "glm-4": 0.1 * RMB, + "glm-4v": 0.1 * RMB, + "glm-3-turbo": 0.005 * RMB, + "embedding-2": 0.0005 * RMB, + "chatglm_turbo": 0.3572, // ￥0.005 / 1k tokens + "chatglm_pro": 0.7143, // ￥0.01 / 1k tokens + "chatglm_std": 0.3572, // ￥0.005 / 1k tokens + "chatglm_lite": 0.1429, // ￥0.002 / 1k tokens + "cogview-3": 0.25 * RMB, + // https://help.aliyun.com/zh/dashscope/developer-reference/tongyi-thousand-questions-metering-and-billing + "qwen-turbo": 0.5715, // ￥0.008 / 1k tokens "qwen-plus": 1.4286, // ￥0.02 / 1k tokens "qwen-max": 1.4286, // ￥0.02 / 1k tokens "qwen-max-longcontext": 1.4286, // ￥0.02 / 1k tokens "text-embedding-v1": 0.05, // ￥0.0007 / 1k tokens + "ali-stable-diffusion-xl": 8, + "ali-stable-diffusion-v1.5": 8, + "wanx-v1": 8, "SparkDesk": 1.2858, // ￥0.018 / 1k tokens "SparkDesk-v1.1": 1.2858, // ￥0.018 / 1k tokens "SparkDesk-v2.1": 1.2858, // ￥0.018 / 1k tokens @@ -140,6 +154,10 @@ var ModelRatio = map[string]float64{ "yi-34b-chat-0205": 2.5 / 1000 * RMB, "yi-34b-chat-200k": 12.0 / 1000 * RMB, "yi-vl-plus": 6.0 / 1000 * RMB, + // stepfun todo + "step-1v-32k": 0.024 * RMB, + "step-1-32k": 0.024 * RMB, + "step-1-200k": 0.15 * RMB, } var CompletionRatio = map[string]float64{} diff --git a/relay/channel/ali/adaptor.go b/relay/channel/ali/adaptor.go deleted file mode 100644 index 6a3245ad..00000000 --- a/relay/channel/ali/adaptor.go +++ /dev/null @@ -1,86 +0,0 @@ -package ali - -import ( - "errors" - "fmt" - "github.com/gin-gonic/gin" - "github.com/songquanpeng/one-api/common" - "github.com/songquanpeng/one-api/relay/channel" - "github.com/songquanpeng/one-api/relay/constant" - "github.com/songquanpeng/one-api/relay/model" - "github.com/songquanpeng/one-api/relay/util" - "io" - "net/http" -) - -// https://help.aliyun.com/zh/dashscope/developer-reference/api-details - -type Adaptor struct { -} - -func (a *Adaptor) Init(meta *util.RelayMeta) { - -} - -func (a *Adaptor) GetRequestURL(meta *util.RelayMeta) (string, error) { - fullRequestURL := fmt.Sprintf("%s/api/v1/services/aigc/text-generation/generation", meta.BaseURL) - if meta.Mode == constant.RelayModeEmbeddings { - fullRequestURL = fmt.Sprintf("%s/api/v1/services/embeddings/text-embedding/text-embedding", meta.BaseURL) - } - return fullRequestURL, nil -} - -func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Request, meta *util.RelayMeta) error { - channel.SetupCommonRequestHeader(c, req, meta) - if meta.IsStream { - req.Header.Set("Accept", "text/event-stream") - } - req.Header.Set("Authorization", "Bearer "+meta.APIKey) - if meta.IsStream { - req.Header.Set("X-DashScope-SSE", "enable") - } - if c.GetString(common.ConfigKeyPlugin) != "" { - req.Header.Set("X-DashScope-Plugin", c.GetString(common.ConfigKeyPlugin)) - } - return nil -} - -func (a *Adaptor) ConvertRequest(c *gin.Context, relayMode int, request *model.GeneralOpenAIRequest) (any, error) { - if request == nil { - return nil, errors.New("request is nil") - } - switch relayMode { - case constant.RelayModeEmbeddings: - baiduEmbeddingRequest := ConvertEmbeddingRequest(*request) - return baiduEmbeddingRequest, nil - default: - baiduRequest := ConvertRequest(*request) - return baiduRequest, nil - } -} - -func (a *Adaptor) DoRequest(c *gin.Context, meta *util.RelayMeta, requestBody io.Reader) (*http.Response, error) { - return channel.DoRequestHelper(a, c, meta, requestBody) -} - -func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, meta *util.RelayMeta) (usage *model.Usage, err *model.ErrorWithStatusCode) { - if meta.IsStream { - err, usage = StreamHandler(c, resp) - } else { - switch meta.Mode { - case constant.RelayModeEmbeddings: - err, usage = EmbeddingHandler(c, resp) - default: - err, usage = Handler(c, resp) - } - } - return -} - -func (a *Adaptor) GetModelList() []string { - return ModelList -} - -func (a *Adaptor) GetChannelName() string { - return "ali" -} diff --git a/relay/channel/ali/model.go b/relay/channel/ali/model.go deleted file mode 100644 index 3b8a8372..00000000 --- a/relay/channel/ali/model.go +++ /dev/null @@ -1,81 +0,0 @@ -package ali - -import ( - "github.com/songquanpeng/one-api/relay/channel/openai" - "github.com/songquanpeng/one-api/relay/model" -) - -type Message struct { - Content string `json:"content"` - Role string `json:"role"` -} - -type Input struct { - //Prompt string `json:"prompt"` - Messages []Message `json:"messages"` -} - -type Parameters struct { - TopP float64 `json:"top_p,omitempty"` - TopK int `json:"top_k,omitempty"` - Seed uint64 `json:"seed,omitempty"` - EnableSearch bool `json:"enable_search,omitempty"` - IncrementalOutput bool `json:"incremental_output,omitempty"` - MaxTokens int `json:"max_tokens,omitempty"` - Temperature float64 `json:"temperature,omitempty"` - ResultFormat string `json:"result_format,omitempty"` - Tools []model.Tool `json:"tools,omitempty"` -} - -type ChatRequest struct { - Model string `json:"model"` - Input Input `json:"input"` - Parameters Parameters `json:"parameters,omitempty"` -} - -type EmbeddingRequest struct { - Model string `json:"model"` - Input struct { - Texts []string `json:"texts"` - } `json:"input"` - Parameters *struct { - TextType string `json:"text_type,omitempty"` - } `json:"parameters,omitempty"` -} - -type Embedding struct { - Embedding []float64 `json:"embedding"` - TextIndex int `json:"text_index"` -} - -type EmbeddingResponse struct { - Output struct { - Embeddings []Embedding `json:"embeddings"` - } `json:"output"` - Usage Usage `json:"usage"` - Error -} - -type Error struct { - Code string `json:"code"` - Message string `json:"message"` - RequestId string `json:"request_id"` -} - -type Usage struct { - InputTokens int `json:"input_tokens"` - OutputTokens int `json:"output_tokens"` - TotalTokens int `json:"total_tokens"` -} - -type Output struct { - //Text string `json:"text"` - //FinishReason string `json:"finish_reason"` - Choices []openai.TextResponseChoice `json:"choices"` -} - -type ChatResponse struct { - Output Output `json:"output"` - Usage Usage `json:"usage"` - Error -} diff --git a/relay/channel/baidu/constants.go b/relay/channel/baidu/constants.go deleted file mode 100644 index 45a4e901..00000000 --- a/relay/channel/baidu/constants.go +++ /dev/null @@ -1,13 +0,0 @@ -package baidu - -var ModelList = []string{ - "ERNIE-Bot-4", - "ERNIE-Bot-8K", - "ERNIE-Bot", - "ERNIE-Speed", - "ERNIE-Bot-turbo", - "Embedding-V1", - "bge-large-zh", - "bge-large-en", - "tao-8k", -} diff --git a/relay/channel/interface.go b/relay/channel/interface.go deleted file mode 100644 index e25db83f..00000000 --- a/relay/channel/interface.go +++ /dev/null @@ -1,20 +0,0 @@ -package channel - -import ( - "github.com/gin-gonic/gin" - "github.com/songquanpeng/one-api/relay/model" - "github.com/songquanpeng/one-api/relay/util" - "io" - "net/http" -) - -type Adaptor interface { - Init(meta *util.RelayMeta) - GetRequestURL(meta *util.RelayMeta) (string, error) - SetupRequestHeader(c *gin.Context, req *http.Request, meta *util.RelayMeta) error - ConvertRequest(c *gin.Context, relayMode int, request *model.GeneralOpenAIRequest) (any, error) - DoRequest(c *gin.Context, meta *util.RelayMeta, requestBody io.Reader) (*http.Response, error) - DoResponse(c *gin.Context, resp *http.Response, meta *util.RelayMeta) (usage *model.Usage, err *model.ErrorWithStatusCode) - GetModelList() []string - GetChannelName() string -} diff --git a/relay/channel/minimax/main.go b/relay/channel/minimax/main.go deleted file mode 100644 index a01821c2..00000000 --- a/relay/channel/minimax/main.go +++ /dev/null @@ -1,14 +0,0 @@ -package minimax - -import ( - "fmt" - "github.com/songquanpeng/one-api/relay/constant" - "github.com/songquanpeng/one-api/relay/util" -) - -func GetRequestURL(meta *util.RelayMeta) (string, error) { - if meta.Mode == constant.RelayModeChatCompletions { - return fmt.Sprintf("%s/v1/text/chatcompletion_v2", meta.BaseURL), nil - } - return "", fmt.Errorf("unsupported relay mode %d for minimax", meta.Mode) -} diff --git a/relay/channel/openai/compatible.go b/relay/channel/openai/compatible.go deleted file mode 100644 index e4951a34..00000000 --- a/relay/channel/openai/compatible.go +++ /dev/null @@ -1,46 +0,0 @@ -package openai - -import ( - "github.com/songquanpeng/one-api/common" - "github.com/songquanpeng/one-api/relay/channel/ai360" - "github.com/songquanpeng/one-api/relay/channel/baichuan" - "github.com/songquanpeng/one-api/relay/channel/groq" - "github.com/songquanpeng/one-api/relay/channel/lingyiwanwu" - "github.com/songquanpeng/one-api/relay/channel/minimax" - "github.com/songquanpeng/one-api/relay/channel/mistral" - "github.com/songquanpeng/one-api/relay/channel/moonshot" -) - -var CompatibleChannels = []int{ - common.ChannelTypeAzure, - common.ChannelType360, - common.ChannelTypeMoonshot, - common.ChannelTypeBaichuan, - common.ChannelTypeMinimax, - common.ChannelTypeMistral, - common.ChannelTypeGroq, - common.ChannelTypeLingYiWanWu, -} - -func GetCompatibleChannelMeta(channelType int) (string, []string) { - switch channelType { - case common.ChannelTypeAzure: - return "azure", ModelList - case common.ChannelType360: - return "360", ai360.ModelList - case common.ChannelTypeMoonshot: - return "moonshot", moonshot.ModelList - case common.ChannelTypeBaichuan: - return "baichuan", baichuan.ModelList - case common.ChannelTypeMinimax: - return "minimax", minimax.ModelList - case common.ChannelTypeMistral: - return "mistralai", mistral.ModelList - case common.ChannelTypeGroq: - return "groq", groq.ModelList - case common.ChannelTypeLingYiWanWu: - return "lingyiwanwu", lingyiwanwu.ModelList - default: - return "openai", ModelList - } -} diff --git a/relay/channel/openai/helper.go b/relay/channel/openai/helper.go deleted file mode 100644 index 9bca8cab..00000000 --- a/relay/channel/openai/helper.go +++ /dev/null @@ -1,11 +0,0 @@ -package openai - -import "github.com/songquanpeng/one-api/relay/model" - -func ResponseText2Usage(responseText string, modeName string, promptTokens int) *model.Usage { - usage := &model.Usage{} - usage.PromptTokens = promptTokens - usage.CompletionTokens = CountTokenText(responseText, modeName) - usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens - return usage -} diff --git a/relay/channel/zhipu/adaptor.go b/relay/channel/zhipu/adaptor.go deleted file mode 100644 index 0ca23d59..00000000 --- a/relay/channel/zhipu/adaptor.go +++ /dev/null @@ -1,101 +0,0 @@ -package zhipu - -import ( - "errors" - "fmt" - "github.com/gin-gonic/gin" - "github.com/songquanpeng/one-api/relay/channel" - "github.com/songquanpeng/one-api/relay/channel/openai" - "github.com/songquanpeng/one-api/relay/model" - "github.com/songquanpeng/one-api/relay/util" - "io" - "math" - "net/http" - "strings" -) - -type Adaptor struct { - APIVersion string -} - -func (a *Adaptor) Init(meta *util.RelayMeta) { - -} - -func (a *Adaptor) SetVersionByModeName(modelName string) { - if strings.HasPrefix(modelName, "glm-") { - a.APIVersion = "v4" - } else { - a.APIVersion = "v3" - } -} - -func (a *Adaptor) GetRequestURL(meta *util.RelayMeta) (string, error) { - a.SetVersionByModeName(meta.ActualModelName) - if a.APIVersion == "v4" { - return fmt.Sprintf("%s/api/paas/v4/chat/completions", meta.BaseURL), nil - } - method := "invoke" - if meta.IsStream { - method = "sse-invoke" - } - return fmt.Sprintf("%s/api/paas/v3/model-api/%s/%s", meta.BaseURL, meta.ActualModelName, method), nil -} - -func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Request, meta *util.RelayMeta) error { - channel.SetupCommonRequestHeader(c, req, meta) - token := GetToken(meta.APIKey) - req.Header.Set("Authorization", token) - return nil -} - -func (a *Adaptor) ConvertRequest(c *gin.Context, relayMode int, request *model.GeneralOpenAIRequest) (any, error) { - if request == nil { - return nil, errors.New("request is nil") - } - // TopP (0.0, 1.0) - request.TopP = math.Min(0.99, request.TopP) - request.TopP = math.Max(0.01, request.TopP) - - // Temperature (0.0, 1.0) - request.Temperature = math.Min(0.99, request.Temperature) - request.Temperature = math.Max(0.01, request.Temperature) - a.SetVersionByModeName(request.Model) - if a.APIVersion == "v4" { - return request, nil - } - return ConvertRequest(*request), nil -} - -func (a *Adaptor) DoRequest(c *gin.Context, meta *util.RelayMeta, requestBody io.Reader) (*http.Response, error) { - return channel.DoRequestHelper(a, c, meta, requestBody) -} - -func (a *Adaptor) DoResponseV4(c *gin.Context, resp *http.Response, meta *util.RelayMeta) (usage *model.Usage, err *model.ErrorWithStatusCode) { - if meta.IsStream { - err, _, usage = openai.StreamHandler(c, resp, meta.Mode) - } else { - err, usage = openai.Handler(c, resp, meta.PromptTokens, meta.ActualModelName) - } - return -} - -func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, meta *util.RelayMeta) (usage *model.Usage, err *model.ErrorWithStatusCode) { - if a.APIVersion == "v4" { - return a.DoResponseV4(c, resp, meta) - } - if meta.IsStream { - err, usage = StreamHandler(c, resp) - } else { - err, usage = Handler(c, resp) - } - return -} - -func (a *Adaptor) GetModelList() []string { - return ModelList -} - -func (a *Adaptor) GetChannelName() string { - return "zhipu" -} diff --git a/relay/channeltype/define.go b/relay/channeltype/define.go new file mode 100644 index 00000000..80027a80 --- /dev/null +++ b/relay/channeltype/define.go @@ -0,0 +1,39 @@ +package channeltype + +const ( + Unknown = iota + OpenAI + API2D + Azure + CloseAI + OpenAISB + OpenAIMax + OhMyGPT + Custom + Ails + AIProxy + PaLM + API2GPT + AIGC2D + Anthropic + Baidu + Zhipu + Ali + Xunfei + AI360 + OpenRouter + AIProxyLibrary + FastGPT + Tencent + Gemini + Moonshot + Baichuan + Minimax + Mistral + Groq + Ollama + LingYiWanWu + StepFun + + Dummy +) diff --git a/relay/channeltype/helper.go b/relay/channeltype/helper.go new file mode 100644 index 00000000..01c2918c --- /dev/null +++ b/relay/channeltype/helper.go @@ -0,0 +1,30 @@ +package channeltype + +import "github.com/songquanpeng/one-api/relay/apitype" + +func ToAPIType(channelType int) int { + apiType := apitype.OpenAI + switch channelType { + case Anthropic: + apiType = apitype.Anthropic + case Baidu: + apiType = apitype.Baidu + case PaLM: + apiType = apitype.PaLM + case Zhipu: + apiType = apitype.Zhipu + case Ali: + apiType = apitype.Ali + case Xunfei: + apiType = apitype.Xunfei + case AIProxyLibrary: + apiType = apitype.AIProxyLibrary + case Tencent: + apiType = apitype.Tencent + case Gemini: + apiType = apitype.Gemini + case Ollama: + apiType = apitype.Ollama + } + return apiType +} diff --git a/relay/channeltype/url.go b/relay/channeltype/url.go new file mode 100644 index 00000000..eec59116 --- /dev/null +++ b/relay/channeltype/url.go @@ -0,0 +1,43 @@ +package channeltype + +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 + "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 + "https://api.moonshot.cn", // 25 + "https://api.baichuan-ai.com", // 26 + "https://api.minimax.chat", // 27 + "https://api.mistral.ai", // 28 + "https://api.groq.com/openai", // 29 + "http://localhost:11434", // 30 + "https://api.lingyiwanwu.com", // 31 + "https://api.stepfun.com", // 32 +} + +func init() { + if len(ChannelBaseURLs) != Dummy { + panic("channel base urls length not match") + } +} diff --git a/relay/util/init.go b/relay/client/init.go similarity index 96% rename from relay/util/init.go rename to relay/client/init.go index 03dad31b..4b59cba7 100644 --- a/relay/util/init.go +++ b/relay/client/init.go @@ -1,4 +1,4 @@ -package util +package client import ( "github.com/songquanpeng/one-api/common/config" diff --git a/relay/constant/api_type.go b/relay/constant/api_type.go deleted file mode 100644 index b249f6a2..00000000 --- a/relay/constant/api_type.go +++ /dev/null @@ -1,48 +0,0 @@ -package constant - -import ( - "github.com/songquanpeng/one-api/common" -) - -const ( - APITypeOpenAI = iota - APITypeAnthropic - APITypePaLM - APITypeBaidu - APITypeZhipu - APITypeAli - APITypeXunfei - APITypeAIProxyLibrary - APITypeTencent - APITypeGemini - APITypeOllama - - APITypeDummy // this one is only for count, do not add any channel after this -) - -func ChannelType2APIType(channelType int) int { - apiType := APITypeOpenAI - switch channelType { - case common.ChannelTypeAnthropic: - apiType = APITypeAnthropic - 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 - case common.ChannelTypeOllama: - apiType = APITypeOllama - } - return apiType -} diff --git a/relay/constant/image.go b/relay/constant/image.go deleted file mode 100644 index 5e04895f..00000000 --- a/relay/constant/image.go +++ /dev/null @@ -1,24 +0,0 @@ -package constant - -var DalleSizeRatios = map[string]map[string]float64{ - "dall-e-2": { - "256x256": 1, - "512x512": 1.125, - "1024x1024": 1.25, - }, - "dall-e-3": { - "1024x1024": 1, - "1024x1792": 2, - "1792x1024": 2, - }, -} - -var DalleGenerationImageAmounts = map[string][2]int{ - "dall-e-2": {1, 10}, - "dall-e-3": {1, 1}, // OpenAI allows n=1 currently. -} - -var DalleImagePromptLengthLimitations = map[string]int{ - "dall-e-2": 1000, - "dall-e-3": 4000, -} diff --git a/relay/constant/relay_mode.go b/relay/constant/relay_mode.go deleted file mode 100644 index 5e2fe574..00000000 --- a/relay/constant/relay_mode.go +++ /dev/null @@ -1,42 +0,0 @@ -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 cd118985..9d8cfef5 100644 --- a/relay/controller/audio.go +++ b/relay/controller/audio.go @@ -12,10 +12,14 @@ import ( "github.com/songquanpeng/one-api/common/config" "github.com/songquanpeng/one-api/common/logger" "github.com/songquanpeng/one-api/model" - "github.com/songquanpeng/one-api/relay/channel/openai" - "github.com/songquanpeng/one-api/relay/constant" + "github.com/songquanpeng/one-api/relay/adaptor/azure" + "github.com/songquanpeng/one-api/relay/adaptor/openai" + "github.com/songquanpeng/one-api/relay/billing" + billingratio "github.com/songquanpeng/one-api/relay/billing/ratio" + "github.com/songquanpeng/one-api/relay/channeltype" + "github.com/songquanpeng/one-api/relay/client" relaymodel "github.com/songquanpeng/one-api/relay/model" - "github.com/songquanpeng/one-api/relay/util" + "github.com/songquanpeng/one-api/relay/relaymode" "io" "net/http" "strings" @@ -33,7 +37,7 @@ func RelayAudioHelper(c *gin.Context, relayMode int) *relaymodel.ErrorWithStatus tokenName := c.GetString("token_name") var ttsRequest openai.TextToSpeechRequest - if relayMode == constant.RelayModeAudioSpeech { + if relayMode == relaymode.AudioSpeech { // Read JSON err := common.UnmarshalBodyReusable(c, &ttsRequest) // Check if JSON is valid @@ -47,13 +51,13 @@ func RelayAudioHelper(c *gin.Context, relayMode int) *relaymodel.ErrorWithStatus } } - modelRatio := common.GetModelRatio(audioModel) - groupRatio := common.GetGroupRatio(group) + modelRatio := billingratio.GetModelRatio(audioModel) + groupRatio := billingratio.GetGroupRatio(group) ratio := modelRatio * groupRatio var quota int64 var preConsumedQuota int64 switch relayMode { - case constant.RelayModeAudioSpeech: + case relaymode.AudioSpeech: preConsumedQuota = int64(float64(len(ttsRequest.Input)) * ratio) quota = preConsumedQuota default: @@ -115,19 +119,19 @@ func RelayAudioHelper(c *gin.Context, relayMode int) *relaymodel.ErrorWithStatus } } - baseURL := common.ChannelBaseURLs[channelType] + baseURL := channeltype.ChannelBaseURLs[channelType] requestURL := c.Request.URL.String() if c.GetString("base_url") != "" { baseURL = c.GetString("base_url") } - fullRequestURL := util.GetFullRequestURL(baseURL, requestURL, channelType) - if channelType == common.ChannelTypeAzure { - apiVersion := util.GetAzureAPIVersion(c) - if relayMode == constant.RelayModeAudioTranscription { + fullRequestURL := openai.GetFullRequestURL(baseURL, requestURL, channelType) + if channelType == channeltype.Azure { + apiVersion := azure.GetAPIVersion(c) + if relayMode == relaymode.AudioTranscription { // https://learn.microsoft.com/en-us/azure/ai-services/openai/whisper-quickstart?tabs=command-line#rest-api fullRequestURL = fmt.Sprintf("%s/openai/deployments/%s/audio/transcriptions?api-version=%s", baseURL, audioModel, apiVersion) - } else if relayMode == constant.RelayModeAudioSpeech { + } else if relayMode == relaymode.AudioSpeech { // https://learn.microsoft.com/en-us/azure/ai-services/openai/text-to-speech-quickstart?tabs=command-line#rest-api fullRequestURL = fmt.Sprintf("%s/openai/deployments/%s/audio/speech?api-version=%s", baseURL, audioModel, apiVersion) } @@ -146,7 +150,7 @@ func RelayAudioHelper(c *gin.Context, relayMode int) *relaymodel.ErrorWithStatus return openai.ErrorWrapper(err, "new_request_failed", http.StatusInternalServerError) } - if (relayMode == constant.RelayModeAudioTranscription || relayMode == constant.RelayModeAudioSpeech) && channelType == common.ChannelTypeAzure { + if (relayMode == relaymode.AudioTranscription || relayMode == relaymode.AudioSpeech) && channelType == channeltype.Azure { // https://learn.microsoft.com/en-us/azure/ai-services/openai/whisper-quickstart?tabs=command-line#rest-api apiKey := c.Request.Header.Get("Authorization") apiKey = strings.TrimPrefix(apiKey, "Bearer ") @@ -158,7 +162,7 @@ func RelayAudioHelper(c *gin.Context, relayMode int) *relaymodel.ErrorWithStatus req.Header.Set("Content-Type", c.Request.Header.Get("Content-Type")) req.Header.Set("Accept", c.Request.Header.Get("Accept")) - resp, err := util.HTTPClient.Do(req) + resp, err := client.HTTPClient.Do(req) if err != nil { return openai.ErrorWrapper(err, "do_request_failed", http.StatusInternalServerError) } @@ -172,7 +176,7 @@ func RelayAudioHelper(c *gin.Context, relayMode int) *relaymodel.ErrorWithStatus return openai.ErrorWrapper(err, "close_request_body_failed", http.StatusInternalServerError) } - if relayMode != constant.RelayModeAudioSpeech { + if relayMode != relaymode.AudioSpeech { responseBody, err := io.ReadAll(resp.Body) if err != nil { return openai.ErrorWrapper(err, "read_response_body_failed", http.StatusInternalServerError) @@ -211,12 +215,12 @@ func RelayAudioHelper(c *gin.Context, relayMode int) *relaymodel.ErrorWithStatus resp.Body = io.NopCloser(bytes.NewBuffer(responseBody)) } if resp.StatusCode != http.StatusOK { - return util.RelayErrorHandler(resp) + return RelayErrorHandler(resp) } succeed = true quotaDelta := quota - preConsumedQuota defer func(ctx context.Context) { - go util.PostConsumeQuota(ctx, tokenId, quotaDelta, quota, userId, channelId, modelRatio, groupRatio, audioModel, tokenName) + go billing.PostConsumeQuota(ctx, tokenId, quotaDelta, quota, userId, channelId, modelRatio, groupRatio, audioModel, tokenName) }(c.Request.Context()) for k, v := range resp.Header { diff --git a/relay/controller/error.go b/relay/controller/error.go new file mode 100644 index 00000000..69ece3ec --- /dev/null +++ b/relay/controller/error.go @@ -0,0 +1,91 @@ +package controller + +import ( + "encoding/json" + "fmt" + "github.com/songquanpeng/one-api/common/config" + "github.com/songquanpeng/one-api/common/logger" + "github.com/songquanpeng/one-api/relay/model" + "io" + "net/http" + "strconv" +) + +type GeneralErrorResponse struct { + Error model.Error `json:"error"` + Message string `json:"message"` + Msg string `json:"msg"` + Err string `json:"err"` + ErrorMsg string `json:"error_msg"` + Header struct { + Message string `json:"message"` + } `json:"header"` + Response struct { + Error struct { + Message string `json:"message"` + } `json:"error"` + } `json:"response"` +} + +func (e GeneralErrorResponse) ToMessage() string { + if e.Error.Message != "" { + return e.Error.Message + } + if e.Message != "" { + return e.Message + } + if e.Msg != "" { + return e.Msg + } + if e.Err != "" { + return e.Err + } + if e.ErrorMsg != "" { + return e.ErrorMsg + } + if e.Header.Message != "" { + return e.Header.Message + } + if e.Response.Error.Message != "" { + return e.Response.Error.Message + } + return "" +} + +func RelayErrorHandler(resp *http.Response) (ErrorWithStatusCode *model.ErrorWithStatusCode) { + ErrorWithStatusCode = &model.ErrorWithStatusCode{ + StatusCode: resp.StatusCode, + Error: model.Error{ + Message: "", + Type: "upstream_error", + Code: "bad_response_status_code", + Param: strconv.Itoa(resp.StatusCode), + }, + } + responseBody, err := io.ReadAll(resp.Body) + if err != nil { + return + } + if config.DebugEnabled { + logger.SysLog(fmt.Sprintf("error happened, status code: %d, response: \n%s", resp.StatusCode, string(responseBody))) + } + err = resp.Body.Close() + if err != nil { + return + } + var errResponse GeneralErrorResponse + err = json.Unmarshal(responseBody, &errResponse) + if err != nil { + return + } + if errResponse.Error.Message != "" { + // OpenAI format error, so we override the default one + ErrorWithStatusCode.Error = errResponse.Error + } else { + ErrorWithStatusCode.Error.Message = errResponse.ToMessage() + } + if ErrorWithStatusCode.Error.Message == "" { + ErrorWithStatusCode.Error.Message = fmt.Sprintf("bad response status code %d", resp.StatusCode) + } + return +} diff --git a/relay/controller/helper.go b/relay/controller/helper.go index 600a8d65..f1b40bef 100644 --- a/relay/controller/helper.go +++ b/relay/controller/helper.go @@ -9,10 +9,13 @@ import ( "github.com/songquanpeng/one-api/common/config" "github.com/songquanpeng/one-api/common/logger" "github.com/songquanpeng/one-api/model" - "github.com/songquanpeng/one-api/relay/channel/openai" - "github.com/songquanpeng/one-api/relay/constant" + "github.com/songquanpeng/one-api/relay/adaptor/openai" + billingratio "github.com/songquanpeng/one-api/relay/billing/ratio" + "github.com/songquanpeng/one-api/relay/channeltype" + "github.com/songquanpeng/one-api/relay/controller/validator" + "github.com/songquanpeng/one-api/relay/meta" relaymodel "github.com/songquanpeng/one-api/relay/model" - "github.com/songquanpeng/one-api/relay/util" + "github.com/songquanpeng/one-api/relay/relaymode" "math" "net/http" ) @@ -23,21 +26,21 @@ func getAndValidateTextRequest(c *gin.Context, relayMode int) (*relaymodel.Gener if err != nil { return nil, err } - if relayMode == constant.RelayModeModerations && textRequest.Model == "" { + if relayMode == relaymode.Moderations && textRequest.Model == "" { textRequest.Model = "text-moderation-latest" } - if relayMode == constant.RelayModeEmbeddings && textRequest.Model == "" { + if relayMode == relaymode.Embeddings && textRequest.Model == "" { textRequest.Model = c.Param("model") } - err = util.ValidateTextRequest(textRequest, relayMode) + err = validator.ValidateTextRequest(textRequest, relayMode) if err != nil { return nil, err } return textRequest, nil } -func getImageRequest(c *gin.Context, relayMode int) (*openai.ImageRequest, error) { - imageRequest := &openai.ImageRequest{} +func getImageRequest(c *gin.Context, relayMode int) (*relaymodel.ImageRequest, error) { + imageRequest := &relaymodel.ImageRequest{} err := common.UnmarshalBodyReusable(c, imageRequest) if err != nil { return nil, err @@ -54,9 +57,25 @@ func getImageRequest(c *gin.Context, relayMode int) (*openai.ImageRequest, error return imageRequest, nil } -func validateImageRequest(imageRequest *openai.ImageRequest, meta *util.RelayMeta) *relaymodel.ErrorWithStatusCode { +func isValidImageSize(model string, size string) bool { + if model == "cogview-3" { + return true + } + _, ok := billingratio.ImageSizeRatios[model][size] + return ok +} + +func getImageSizeRatio(model string, size string) float64 { + ratio, ok := billingratio.ImageSizeRatios[model][size] + if !ok { + return 1 + } + return ratio +} + +func validateImageRequest(imageRequest *relaymodel.ImageRequest, meta *meta.Meta) *relaymodel.ErrorWithStatusCode { // model validation - _, hasValidSize := constant.DalleSizeRatios[imageRequest.Model][imageRequest.Size] + hasValidSize := isValidImageSize(imageRequest.Model, imageRequest.Size) if !hasValidSize { return openai.ErrorWrapper(errors.New("size not supported for this image model"), "size_not_supported", http.StatusBadRequest) } @@ -64,27 +83,24 @@ func validateImageRequest(imageRequest *openai.ImageRequest, meta *util.RelayMet if imageRequest.Prompt == "" { return openai.ErrorWrapper(errors.New("prompt is required"), "prompt_missing", http.StatusBadRequest) } - if len(imageRequest.Prompt) > constant.DalleImagePromptLengthLimitations[imageRequest.Model] { + if len(imageRequest.Prompt) > billingratio.ImagePromptLengthLimitations[imageRequest.Model] { return openai.ErrorWrapper(errors.New("prompt is too long"), "prompt_too_long", http.StatusBadRequest) } // Number of generated images validation if !isWithinRange(imageRequest.Model, imageRequest.N) { // channel not azure - if meta.ChannelType != common.ChannelTypeAzure { + if meta.ChannelType != channeltype.Azure { return openai.ErrorWrapper(errors.New("invalid value of n"), "n_not_within_range", http.StatusBadRequest) } } return nil } -func getImageCostRatio(imageRequest *openai.ImageRequest) (float64, error) { +func getImageCostRatio(imageRequest *relaymodel.ImageRequest) (float64, error) { if imageRequest == nil { return 0, errors.New("imageRequest is nil") } - imageCostRatio, hasValidSize := constant.DalleSizeRatios[imageRequest.Model][imageRequest.Size] - if !hasValidSize { - return 0, fmt.Errorf("size not supported for this image model: %s", imageRequest.Size) - } + imageCostRatio := getImageSizeRatio(imageRequest.Model, imageRequest.Size) if imageRequest.Quality == "hd" && imageRequest.Model == "dall-e-3" { if imageRequest.Size == "1024x1024" { imageCostRatio *= 2 @@ -97,11 +113,11 @@ func getImageCostRatio(imageRequest *openai.ImageRequest) (float64, error) { func getPromptTokens(textRequest *relaymodel.GeneralOpenAIRequest, relayMode int) int { switch relayMode { - case constant.RelayModeChatCompletions: + case relaymode.ChatCompletions: return openai.CountTokenMessages(textRequest.Messages, textRequest.Model) - case constant.RelayModeCompletions: + case relaymode.Completions: return openai.CountTokenInput(textRequest.Prompt, textRequest.Model) - case constant.RelayModeModerations: + case relaymode.Moderations: return openai.CountTokenInput(textRequest.Input, textRequest.Model) } return 0 @@ -115,7 +131,7 @@ func getPreConsumedQuota(textRequest *relaymodel.GeneralOpenAIRequest, promptTok return int64(float64(preConsumedTokens) * ratio) } -func preConsumeQuota(ctx context.Context, textRequest *relaymodel.GeneralOpenAIRequest, promptTokens int, ratio float64, meta *util.RelayMeta) (int64, *relaymodel.ErrorWithStatusCode) { +func preConsumeQuota(ctx context.Context, textRequest *relaymodel.GeneralOpenAIRequest, promptTokens int, ratio float64, meta *meta.Meta) (int64, *relaymodel.ErrorWithStatusCode) { preConsumedQuota := getPreConsumedQuota(textRequest, promptTokens, ratio) userQuota, err := model.CacheGetUserQuota(ctx, meta.UserId) @@ -144,13 +160,13 @@ func preConsumeQuota(ctx context.Context, textRequest *relaymodel.GeneralOpenAIR return preConsumedQuota, nil } -func postConsumeQuota(ctx context.Context, usage *relaymodel.Usage, meta *util.RelayMeta, textRequest *relaymodel.GeneralOpenAIRequest, ratio float64, preConsumedQuota int64, modelRatio float64, groupRatio float64) { +func postConsumeQuota(ctx context.Context, usage *relaymodel.Usage, meta *meta.Meta, textRequest *relaymodel.GeneralOpenAIRequest, ratio float64, preConsumedQuota int64, modelRatio float64, groupRatio float64) { if usage == nil { logger.Error(ctx, "usage is nil, which is unexpected") return } var quota int64 - completionRatio := common.GetCompletionRatio(textRequest.Model) + completionRatio := billingratio.GetCompletionRatio(textRequest.Model) promptTokens := usage.PromptTokens completionTokens := usage.CompletionTokens quota = int64(math.Ceil((float64(promptTokens) + float64(completionTokens)*completionRatio) * ratio)) @@ -177,3 +193,14 @@ func postConsumeQuota(ctx context.Context, usage *relaymodel.Usage, meta *util.R model.UpdateUserUsedQuotaAndRequestCount(meta.UserId, quota) model.UpdateChannelUsedQuota(meta.ChannelId, quota) } + +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/controller/image.go b/relay/controller/image.go index d81dadf6..80769845 100644 --- a/relay/controller/image.go +++ b/relay/controller/image.go @@ -6,33 +6,31 @@ import ( "encoding/json" "errors" "fmt" - "github.com/songquanpeng/one-api/common" + "github.com/gin-gonic/gin" "github.com/songquanpeng/one-api/common/logger" "github.com/songquanpeng/one-api/model" - "github.com/songquanpeng/one-api/relay/channel/openai" - "github.com/songquanpeng/one-api/relay/constant" + "github.com/songquanpeng/one-api/relay" + "github.com/songquanpeng/one-api/relay/adaptor/openai" + billingratio "github.com/songquanpeng/one-api/relay/billing/ratio" + "github.com/songquanpeng/one-api/relay/channeltype" + "github.com/songquanpeng/one-api/relay/meta" relaymodel "github.com/songquanpeng/one-api/relay/model" - "github.com/songquanpeng/one-api/relay/util" "io" "net/http" - "strings" - - "github.com/gin-gonic/gin" ) func isWithinRange(element string, value int) bool { - if _, ok := constant.DalleGenerationImageAmounts[element]; !ok { + if _, ok := billingratio.ImageGenerationAmounts[element]; !ok { return false } - min := constant.DalleGenerationImageAmounts[element][0] - max := constant.DalleGenerationImageAmounts[element][1] - + min := billingratio.ImageGenerationAmounts[element][0] + max := billingratio.ImageGenerationAmounts[element][1] return value >= min && value <= max } func RelayImageHelper(c *gin.Context, relayMode int) *relaymodel.ErrorWithStatusCode { ctx := c.Request.Context() - meta := util.GetRelayMeta(c) + meta := meta.GetByContext(c) imageRequest, err := getImageRequest(c, meta.Mode) if err != nil { logger.Errorf(ctx, "getImageRequest failed: %s", err.Error()) @@ -42,7 +40,7 @@ func RelayImageHelper(c *gin.Context, relayMode int) *relaymodel.ErrorWithStatus // map model name var isModelMapped bool meta.OriginModelName = imageRequest.Model - imageRequest.Model, isModelMapped = util.GetMappedModelName(imageRequest.Model, meta.ModelMapping) + imageRequest.Model, isModelMapped = getMappedModelName(imageRequest.Model, meta.ModelMapping) meta.ActualModelName = imageRequest.Model // model validation @@ -56,17 +54,8 @@ func RelayImageHelper(c *gin.Context, relayMode int) *relaymodel.ErrorWithStatus return openai.ErrorWrapper(err, "get_image_cost_ratio_failed", http.StatusInternalServerError) } - requestURL := c.Request.URL.String() - fullRequestURL := util.GetFullRequestURL(meta.BaseURL, requestURL, meta.ChannelType) - if meta.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.GetAzureAPIVersion(c) - // https://{resource_name}.openai.azure.com/openai/deployments/dall-e-3/images/generations?api-version=2024-03-01-preview - fullRequestURL = fmt.Sprintf("%s/openai/deployments/%s/images/generations?api-version=%s", meta.BaseURL, imageRequest.Model, apiVersion) - } - var requestBody io.Reader - if isModelMapped || meta.ChannelType == common.ChannelTypeAzure { // make Azure channel request body + if isModelMapped || meta.ChannelType == channeltype.Azure { // make Azure channel request body jsonStr, err := json.Marshal(imageRequest) if err != nil { return openai.ErrorWrapper(err, "marshal_image_request_failed", http.StatusInternalServerError) @@ -76,8 +65,30 @@ func RelayImageHelper(c *gin.Context, relayMode int) *relaymodel.ErrorWithStatus requestBody = c.Request.Body } - modelRatio := common.GetModelRatio(imageRequest.Model) - groupRatio := common.GetGroupRatio(meta.Group) + adaptor := relay.GetAdaptor(meta.APIType) + if adaptor == nil { + return openai.ErrorWrapper(fmt.Errorf("invalid api type: %d", meta.APIType), "invalid_api_type", http.StatusBadRequest) + } + + switch meta.ChannelType { + case channeltype.Ali: + fallthrough + case channeltype.Baidu: + fallthrough + case channeltype.Zhipu: + finalRequest, err := adaptor.ConvertImageRequest(imageRequest) + if err != nil { + return openai.ErrorWrapper(err, "convert_image_request_failed", http.StatusInternalServerError) + } + jsonStr, err := json.Marshal(finalRequest) + if err != nil { + return openai.ErrorWrapper(err, "marshal_image_request_failed", http.StatusInternalServerError) + } + requestBody = bytes.NewBuffer(jsonStr) + } + + modelRatio := billingratio.GetModelRatio(imageRequest.Model) + groupRatio := billingratio.GetGroupRatio(meta.Group) ratio := modelRatio * groupRatio userQuota, err := model.CacheGetUserQuota(ctx, meta.UserId) @@ -87,36 +98,13 @@ func RelayImageHelper(c *gin.Context, relayMode int) *relaymodel.ErrorWithStatus return openai.ErrorWrapper(errors.New("user quota is not enough"), "insufficient_user_quota", http.StatusForbidden) } - req, err := http.NewRequest(c.Request.Method, fullRequestURL, requestBody) - if err != nil { - return openai.ErrorWrapper(err, "new_request_failed", http.StatusInternalServerError) - } - token := c.Request.Header.Get("Authorization") - if meta.ChannelType == common.ChannelTypeAzure { // Azure authentication - token = strings.TrimPrefix(token, "Bearer ") - req.Header.Set("api-key", token) - } else { - req.Header.Set("Authorization", token) - } - - req.Header.Set("Content-Type", c.Request.Header.Get("Content-Type")) - req.Header.Set("Accept", c.Request.Header.Get("Accept")) - - resp, err := util.HTTPClient.Do(req) + // do request + resp, err := adaptor.DoRequest(c, meta, requestBody) if err != nil { + logger.Errorf(ctx, "DoRequest failed: %s", err.Error()) return openai.ErrorWrapper(err, "do_request_failed", http.StatusInternalServerError) } - err = req.Body.Close() - if err != nil { - return openai.ErrorWrapper(err, "close_request_body_failed", http.StatusInternalServerError) - } - err = c.Request.Body.Close() - if err != nil { - return openai.ErrorWrapper(err, "close_request_body_failed", http.StatusInternalServerError) - } - var imageResponse openai.ImageResponse - defer func(ctx context.Context) { if resp.StatusCode != http.StatusOK { return @@ -139,34 +127,12 @@ func RelayImageHelper(c *gin.Context, relayMode int) *relaymodel.ErrorWithStatus } }(c.Request.Context()) - responseBody, err := io.ReadAll(resp.Body) - - if err != nil { - return openai.ErrorWrapper(err, "read_response_body_failed", http.StatusInternalServerError) - } - err = resp.Body.Close() - if err != nil { - return openai.ErrorWrapper(err, "close_response_body_failed", http.StatusInternalServerError) - } - err = json.Unmarshal(responseBody, &imageResponse) - if err != nil { - return openai.ErrorWrapper(err, "unmarshal_response_body_failed", http.StatusInternalServerError) + // do response + _, respErr := adaptor.DoResponse(c, resp, meta) + if respErr != nil { + logger.Errorf(ctx, "respErr is not nil: %+v", respErr) + return respErr } - resp.Body = io.NopCloser(bytes.NewBuffer(responseBody)) - - for k, v := range resp.Header { - c.Writer.Header().Set(k, v[0]) - } - c.Writer.WriteHeader(resp.StatusCode) - - _, err = io.Copy(c.Writer, resp.Body) - if err != nil { - return openai.ErrorWrapper(err, "copy_response_body_failed", http.StatusInternalServerError) - } - err = resp.Body.Close() - if err != nil { - return openai.ErrorWrapper(err, "close_response_body_failed", http.StatusInternalServerError) - } return nil } diff --git a/relay/controller/text.go b/relay/controller/text.go index ba008713..0332a23f 100644 --- a/relay/controller/text.go +++ b/relay/controller/text.go @@ -5,13 +5,15 @@ import ( "encoding/json" "fmt" "github.com/gin-gonic/gin" - "github.com/songquanpeng/one-api/common" "github.com/songquanpeng/one-api/common/logger" - "github.com/songquanpeng/one-api/relay/channel/openai" - "github.com/songquanpeng/one-api/relay/constant" - "github.com/songquanpeng/one-api/relay/helper" + "github.com/songquanpeng/one-api/relay" + "github.com/songquanpeng/one-api/relay/adaptor/openai" + "github.com/songquanpeng/one-api/relay/apitype" + "github.com/songquanpeng/one-api/relay/billing" + billingratio "github.com/songquanpeng/one-api/relay/billing/ratio" + "github.com/songquanpeng/one-api/relay/channeltype" + "github.com/songquanpeng/one-api/relay/meta" "github.com/songquanpeng/one-api/relay/model" - "github.com/songquanpeng/one-api/relay/util" "io" "net/http" "strings" @@ -19,7 +21,7 @@ import ( func RelayTextHelper(c *gin.Context) *model.ErrorWithStatusCode { ctx := c.Request.Context() - meta := util.GetRelayMeta(c) + meta := meta.GetByContext(c) // get & validate textRequest textRequest, err := getAndValidateTextRequest(c, meta.Mode) if err != nil { @@ -31,11 +33,11 @@ func RelayTextHelper(c *gin.Context) *model.ErrorWithStatusCode { // map model name var isModelMapped bool meta.OriginModelName = textRequest.Model - textRequest.Model, isModelMapped = util.GetMappedModelName(textRequest.Model, meta.ModelMapping) + textRequest.Model, isModelMapped = getMappedModelName(textRequest.Model, meta.ModelMapping) meta.ActualModelName = textRequest.Model // get model ratio & group ratio - modelRatio := common.GetModelRatio(textRequest.Model) - groupRatio := common.GetGroupRatio(meta.Group) + modelRatio := billingratio.GetModelRatio(textRequest.Model) + groupRatio := billingratio.GetGroupRatio(meta.Group) ratio := modelRatio * groupRatio // pre-consume quota promptTokens := getPromptTokens(textRequest, meta.Mode) @@ -46,16 +48,16 @@ func RelayTextHelper(c *gin.Context) *model.ErrorWithStatusCode { return bizErr } - adaptor := helper.GetAdaptor(meta.APIType) + adaptor := relay.GetAdaptor(meta.APIType) if adaptor == nil { return openai.ErrorWrapper(fmt.Errorf("invalid api type: %d", meta.APIType), "invalid_api_type", http.StatusBadRequest) } // get request body var requestBody io.Reader - if meta.APIType == constant.APITypeOpenAI { + if meta.APIType == apitype.OpenAI { // no need to convert request for openai - shouldResetRequestBody := isModelMapped || meta.ChannelType == common.ChannelTypeBaichuan // frequency_penalty 0 is not acceptable for baichuan + shouldResetRequestBody := isModelMapped || meta.ChannelType == channeltype.Baichuan // frequency_penalty 0 is not acceptable for baichuan if shouldResetRequestBody { jsonStr, err := json.Marshal(textRequest) if err != nil { @@ -86,8 +88,8 @@ func RelayTextHelper(c *gin.Context) *model.ErrorWithStatusCode { } errorHappened := (resp.StatusCode != http.StatusOK) || (meta.IsStream && resp.Header.Get("Content-Type") == "application/json") if errorHappened { - util.ReturnPreConsumedQuota(ctx, preConsumedQuota, meta.TokenId) - return util.RelayErrorHandler(resp) + billing.ReturnPreConsumedQuota(ctx, preConsumedQuota, meta.TokenId) + return RelayErrorHandler(resp) } meta.IsStream = meta.IsStream || strings.HasPrefix(resp.Header.Get("Content-Type"), "text/event-stream") @@ -95,7 +97,7 @@ func RelayTextHelper(c *gin.Context) *model.ErrorWithStatusCode { usage, respErr := adaptor.DoResponse(c, resp, meta) if respErr != nil { logger.Errorf(ctx, "respErr is not nil: %+v", respErr) - util.ReturnPreConsumedQuota(ctx, preConsumedQuota, meta.TokenId) + billing.ReturnPreConsumedQuota(ctx, preConsumedQuota, meta.TokenId) return respErr } // post-consume quota diff --git a/relay/util/validation.go b/relay/controller/validator/validation.go similarity index 76% rename from relay/util/validation.go rename to relay/controller/validator/validation.go index ef8d840c..8ff520b8 100644 --- a/relay/util/validation.go +++ b/relay/controller/validator/validation.go @@ -1,9 +1,9 @@ -package util +package validator import ( "errors" - "github.com/songquanpeng/one-api/relay/constant" "github.com/songquanpeng/one-api/relay/model" + "github.com/songquanpeng/one-api/relay/relaymode" "math" ) @@ -15,20 +15,20 @@ func ValidateTextRequest(textRequest *model.GeneralOpenAIRequest, relayMode int) return errors.New("model is required") } switch relayMode { - case constant.RelayModeCompletions: + case relaymode.Completions: if textRequest.Prompt == "" { return errors.New("field prompt is required") } - case constant.RelayModeChatCompletions: + case relaymode.ChatCompletions: if textRequest.Messages == nil || len(textRequest.Messages) == 0 { return errors.New("field messages is required") } - case constant.RelayModeEmbeddings: - case constant.RelayModeModerations: + case relaymode.Embeddings: + case relaymode.Moderations: if textRequest.Input == "" { return errors.New("field input is required") } - case constant.RelayModeEdits: + case relaymode.Edits: if textRequest.Instruction == "" { return errors.New("field instruction is required") } diff --git a/relay/helper/main.go b/relay/helper/main.go deleted file mode 100644 index e7342329..00000000 --- a/relay/helper/main.go +++ /dev/null @@ -1,45 +0,0 @@ -package helper - -import ( - "github.com/songquanpeng/one-api/relay/channel" - "github.com/songquanpeng/one-api/relay/channel/aiproxy" - "github.com/songquanpeng/one-api/relay/channel/ali" - "github.com/songquanpeng/one-api/relay/channel/anthropic" - "github.com/songquanpeng/one-api/relay/channel/baidu" - "github.com/songquanpeng/one-api/relay/channel/gemini" - "github.com/songquanpeng/one-api/relay/channel/ollama" - "github.com/songquanpeng/one-api/relay/channel/openai" - "github.com/songquanpeng/one-api/relay/channel/palm" - "github.com/songquanpeng/one-api/relay/channel/tencent" - "github.com/songquanpeng/one-api/relay/channel/xunfei" - "github.com/songquanpeng/one-api/relay/channel/zhipu" - "github.com/songquanpeng/one-api/relay/constant" -) - -func GetAdaptor(apiType int) channel.Adaptor { - switch apiType { - case constant.APITypeAIProxyLibrary: - return &aiproxy.Adaptor{} - case constant.APITypeAli: - return &ali.Adaptor{} - case constant.APITypeAnthropic: - return &anthropic.Adaptor{} - case constant.APITypeBaidu: - return &baidu.Adaptor{} - case constant.APITypeGemini: - return &gemini.Adaptor{} - case constant.APITypeOpenAI: - return &openai.Adaptor{} - case constant.APITypePaLM: - return &palm.Adaptor{} - case constant.APITypeTencent: - return &tencent.Adaptor{} - case constant.APITypeXunfei: - return &xunfei.Adaptor{} - case constant.APITypeZhipu: - return &zhipu.Adaptor{} - case constant.APITypeOllama: - return &ollama.Adaptor{} - } - return nil -} diff --git a/relay/util/relay_meta.go b/relay/meta/relay_meta.go similarity index 63% rename from relay/util/relay_meta.go rename to relay/meta/relay_meta.go index 31b9d2b4..22ef1567 100644 --- a/relay/util/relay_meta.go +++ b/relay/meta/relay_meta.go @@ -1,13 +1,15 @@ -package util +package meta import ( "github.com/gin-gonic/gin" - "github.com/songquanpeng/one-api/common" - "github.com/songquanpeng/one-api/relay/constant" + "github.com/songquanpeng/one-api/common/config" + "github.com/songquanpeng/one-api/relay/adaptor/azure" + "github.com/songquanpeng/one-api/relay/channeltype" + "github.com/songquanpeng/one-api/relay/relaymode" "strings" ) -type RelayMeta struct { +type Meta struct { Mode int ChannelType int ChannelId int @@ -28,9 +30,9 @@ type RelayMeta struct { PromptTokens int // only for DoResponse } -func GetRelayMeta(c *gin.Context) *RelayMeta { - meta := RelayMeta{ - Mode: constant.Path2RelayMode(c.Request.URL.Path), +func GetByContext(c *gin.Context) *Meta { + meta := Meta{ + Mode: relaymode.GetByPath(c.Request.URL.Path), ChannelType: c.GetInt("channel"), ChannelId: c.GetInt("channel_id"), TokenId: c.GetInt("token_id"), @@ -39,17 +41,17 @@ func GetRelayMeta(c *gin.Context) *RelayMeta { Group: c.GetString("group"), ModelMapping: c.GetStringMapString("model_mapping"), BaseURL: c.GetString("base_url"), - APIVersion: c.GetString(common.ConfigKeyAPIVersion), + APIVersion: c.GetString(config.KeyAPIVersion), APIKey: strings.TrimPrefix(c.Request.Header.Get("Authorization"), "Bearer "), Config: nil, RequestURLPath: c.Request.URL.String(), } - if meta.ChannelType == common.ChannelTypeAzure { - meta.APIVersion = GetAzureAPIVersion(c) + if meta.ChannelType == channeltype.Azure { + meta.APIVersion = azure.GetAPIVersion(c) } if meta.BaseURL == "" { - meta.BaseURL = common.ChannelBaseURLs[meta.ChannelType] + meta.BaseURL = channeltype.ChannelBaseURLs[meta.ChannelType] } - meta.APIType = constant.ChannelType2APIType(meta.ChannelType) + meta.APIType = channeltype.ToAPIType(meta.ChannelType) return &meta } diff --git a/relay/model/image.go b/relay/model/image.go new file mode 100644 index 00000000..bab84256 --- /dev/null +++ b/relay/model/image.go @@ -0,0 +1,12 @@ +package model + +type ImageRequest struct { + Model string `json:"model"` + Prompt string `json:"prompt" binding:"required"` + N int `json:"n,omitempty"` + Size string `json:"size,omitempty"` + Quality string `json:"quality,omitempty"` + ResponseFormat string `json:"response_format,omitempty"` + Style string `json:"style,omitempty"` + User string `json:"user,omitempty"` +} diff --git a/relay/relaymode/define.go b/relay/relaymode/define.go new file mode 100644 index 00000000..96d09438 --- /dev/null +++ b/relay/relaymode/define.go @@ -0,0 +1,14 @@ +package relaymode + +const ( + Unknown = iota + ChatCompletions + Completions + Embeddings + Moderations + ImagesGenerations + Edits + AudioSpeech + AudioTranscription + AudioTranslation +) diff --git a/relay/relaymode/helper.go b/relay/relaymode/helper.go new file mode 100644 index 00000000..926dd42e --- /dev/null +++ b/relay/relaymode/helper.go @@ -0,0 +1,29 @@ +package relaymode + +import "strings" + +func GetByPath(path string) int { + relayMode := Unknown + if strings.HasPrefix(path, "/v1/chat/completions") { + relayMode = ChatCompletions + } else if strings.HasPrefix(path, "/v1/completions") { + relayMode = Completions + } else if strings.HasPrefix(path, "/v1/embeddings") { + relayMode = Embeddings + } else if strings.HasSuffix(path, "embeddings") { + relayMode = Embeddings + } else if strings.HasPrefix(path, "/v1/moderations") { + relayMode = Moderations + } else if strings.HasPrefix(path, "/v1/images/generations") { + relayMode = ImagesGenerations + } else if strings.HasPrefix(path, "/v1/edits") { + relayMode = Edits + } else if strings.HasPrefix(path, "/v1/audio/speech") { + relayMode = AudioSpeech + } else if strings.HasPrefix(path, "/v1/audio/transcriptions") { + relayMode = AudioTranscription + } else if strings.HasPrefix(path, "/v1/audio/translations") { + relayMode = AudioTranslation + } + return relayMode +} diff --git a/relay/util/billing.go b/relay/util/billing.go deleted file mode 100644 index 495d011e..00000000 --- a/relay/util/billing.go +++ /dev/null @@ -1,19 +0,0 @@ -package util - -import ( - "context" - "github.com/songquanpeng/one-api/common/logger" - "github.com/songquanpeng/one-api/model" -) - -func ReturnPreConsumedQuota(ctx context.Context, preConsumedQuota int64, 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 deleted file mode 100644 index 535ef680..00000000 --- a/relay/util/common.go +++ /dev/null @@ -1,187 +0,0 @@ -package util - -import ( - "context" - "encoding/json" - "fmt" - "github.com/songquanpeng/one-api/common" - "github.com/songquanpeng/one-api/common/config" - "github.com/songquanpeng/one-api/common/logger" - "github.com/songquanpeng/one-api/model" - relaymodel "github.com/songquanpeng/one-api/relay/model" - "io" - "net/http" - "strconv" - "strings" - - "github.com/gin-gonic/gin" -) - -func ShouldDisableChannel(err *relaymodel.Error, statusCode int) bool { - if !config.AutomaticDisableChannelEnabled { - return false - } - if err == nil { - return false - } - if statusCode == http.StatusUnauthorized { - return true - } - switch err.Type { - case "insufficient_quota": - return true - // https://docs.anthropic.com/claude/reference/errors - case "authentication_error": - return true - case "permission_error": - return true - case "forbidden": - return true - } - if err.Code == "invalid_api_key" || err.Code == "account_deactivated" { - return true - } - if strings.HasPrefix(err.Message, "Your credit balance is too low") { // anthropic - return true - } else if strings.HasPrefix(err.Message, "This organization has been disabled.") { - return true - } - return false -} - -func ShouldEnableChannel(err error, openAIErr *relaymodel.Error) bool { - if !config.AutomaticEnableChannelEnabled { - return false - } - if err != nil { - return false - } - if openAIErr != nil { - return false - } - return true -} - -type GeneralErrorResponse struct { - Error relaymodel.Error `json:"error"` - Message string `json:"message"` - Msg string `json:"msg"` - Err string `json:"err"` - ErrorMsg string `json:"error_msg"` - Header struct { - Message string `json:"message"` - } `json:"header"` - Response struct { - Error struct { - Message string `json:"message"` - } `json:"error"` - } `json:"response"` -} - -func (e GeneralErrorResponse) ToMessage() string { - if e.Error.Message != "" { - return e.Error.Message - } - if e.Message != "" { - return e.Message - } - if e.Msg != "" { - return e.Msg - } - if e.Err != "" { - return e.Err - } - if e.ErrorMsg != "" { - return e.ErrorMsg - } - if e.Header.Message != "" { - return e.Header.Message - } - if e.Response.Error.Message != "" { - return e.Response.Error.Message - } - return "" -} - -func RelayErrorHandler(resp *http.Response) (ErrorWithStatusCode *relaymodel.ErrorWithStatusCode) { - ErrorWithStatusCode = &relaymodel.ErrorWithStatusCode{ - StatusCode: resp.StatusCode, - Error: relaymodel.Error{ - Message: "", - Type: "upstream_error", - Code: "bad_response_status_code", - Param: strconv.Itoa(resp.StatusCode), - }, - } - responseBody, err := io.ReadAll(resp.Body) - if err != nil { - return - } - if config.DebugEnabled { - logger.SysLog(fmt.Sprintf("error happened, status code: %d, response: \n%s", resp.StatusCode, string(responseBody))) - } - err = resp.Body.Close() - if err != nil { - return - } - var errResponse GeneralErrorResponse - err = json.Unmarshal(responseBody, &errResponse) - if err != nil { - return - } - if errResponse.Error.Message != "" { - // OpenAI format error, so we override the default one - ErrorWithStatusCode.Error = errResponse.Error - } else { - ErrorWithStatusCode.Error.Message = errResponse.ToMessage() - } - if ErrorWithStatusCode.Error.Message == "" { - ErrorWithStatusCode.Error.Message = fmt.Sprintf("bad response status code %d", resp.StatusCode) - } - return -} - -func GetFullRequestURL(baseURL string, requestURL string, channelType int) string { - fullRequestURL := fmt.Sprintf("%s%s", baseURL, requestURL) - - if strings.HasPrefix(baseURL, "https://gateway.ai.cloudflare.com") { - switch channelType { - case common.ChannelTypeOpenAI: - fullRequestURL = fmt.Sprintf("%s%s", baseURL, strings.TrimPrefix(requestURL, "/v1")) - case common.ChannelTypeAzure: - fullRequestURL = fmt.Sprintf("%s%s", baseURL, strings.TrimPrefix(requestURL, "/openai/deployments")) - } - } - return fullRequestURL -} - -func PostConsumeQuota(ctx context.Context, tokenId int, quotaDelta int64, totalQuota int64, userId int, channelId int, modelRatio float64, groupRatio float64, modelName string, tokenName string) { - // quotaDelta is remaining quota to be consumed - err := model.PostConsumeTokenQuota(tokenId, quotaDelta) - if err != nil { - logger.SysError("error consuming token remain quota: " + err.Error()) - } - err = model.CacheUpdateUserQuota(ctx, userId) - if err != nil { - logger.SysError("error update user quota cache: " + err.Error()) - } - // totalQuota is total quota consumed - if totalQuota != 0 { - logContent := fmt.Sprintf("模型倍率 %.2f，分组倍率 %.2f", modelRatio, groupRatio) - model.RecordConsumeLog(ctx, userId, channelId, int(totalQuota), 0, modelName, tokenName, totalQuota, logContent) - model.UpdateUserUsedQuotaAndRequestCount(userId, totalQuota) - model.UpdateChannelUsedQuota(channelId, totalQuota) - } - if totalQuota <= 0 { - logger.Error(ctx, fmt.Sprintf("totalQuota consumed is %d, something is wrong", totalQuota)) - } -} - -func GetAzureAPIVersion(c *gin.Context) string { - query := c.Request.URL.Query() - apiVersion := query.Get("api-version") - if apiVersion == "" { - apiVersion = c.GetString(common.ConfigKeyAPIVersion) - } - return apiVersion -} diff --git a/relay/util/model_mapping.go b/relay/util/model_mapping.go deleted file mode 100644 index 39e062a1..00000000 --- a/relay/util/model_mapping.go +++ /dev/null @@ -1,12 +0,0 @@ -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/router/api-router.go b/router/api.go similarity index 93% rename from router/api-router.go rename to router/api.go index 4aa6d830..d2ada4eb 100644 --- a/router/api-router.go +++ b/router/api.go @@ -2,6 +2,7 @@ package router import ( "github.com/songquanpeng/one-api/controller" + "github.com/songquanpeng/one-api/controller/auth" "github.com/songquanpeng/one-api/middleware" "github.com/gin-contrib/gzip" @@ -21,11 +22,13 @@ func SetApiRouter(router *gin.Engine) { apiRouter.GET("/verification", middleware.CriticalRateLimit(), middleware.TurnstileCheck(), controller.SendEmailVerification) apiRouter.GET("/reset_password", middleware.CriticalRateLimit(), middleware.TurnstileCheck(), controller.SendPasswordResetEmail) apiRouter.POST("/user/reset", middleware.CriticalRateLimit(), controller.ResetPassword) - apiRouter.GET("/oauth/github", middleware.CriticalRateLimit(), controller.GitHubOAuth) - apiRouter.GET("/oauth/state", middleware.CriticalRateLimit(), controller.GenerateOAuthCode) - apiRouter.GET("/oauth/wechat", middleware.CriticalRateLimit(), controller.WeChatAuth) - apiRouter.GET("/oauth/wechat/bind", middleware.CriticalRateLimit(), middleware.UserAuth(), controller.WeChatBind) + apiRouter.GET("/oauth/github", middleware.CriticalRateLimit(), auth.GitHubOAuth) + apiRouter.GET("/oauth/lark", middleware.CriticalRateLimit(), auth.LarkOAuth) + apiRouter.GET("/oauth/state", middleware.CriticalRateLimit(), auth.GenerateOAuthCode) + apiRouter.GET("/oauth/wechat", middleware.CriticalRateLimit(), auth.WeChatAuth) + apiRouter.GET("/oauth/wechat/bind", middleware.CriticalRateLimit(), middleware.UserAuth(), auth.WeChatBind) apiRouter.GET("/oauth/email/bind", middleware.CriticalRateLimit(), middleware.UserAuth(), controller.EmailBind) + apiRouter.POST("/topup", middleware.AdminAuth(), controller.AdminTopUp) userRoute := apiRouter.Group("/user") { @@ -69,7 +72,7 @@ func SetApiRouter(router *gin.Engine) { { channelRoute.GET("/", controller.GetAllChannels) channelRoute.GET("/search", controller.SearchChannels) - channelRoute.GET("/models", controller.ListModels) + channelRoute.GET("/models", controller.ListAllModels) channelRoute.GET("/:id", controller.GetChannel) channelRoute.GET("/test", controller.TestChannels) channelRoute.GET("/test/:id", controller.TestChannel) diff --git a/router/relay-router.go b/router/relay.go similarity index 100% rename from router/relay-router.go rename to router/relay.go diff --git a/router/web-router.go b/router/web.go similarity index 100% rename from router/web-router.go rename to router/web.go diff --git a/web/README.md b/web/README.md index 29f4713e..829271e2 100644 --- a/web/README.md +++ b/web/README.md @@ -2,6 +2,9 @@ > 每个文件夹代表一个主题，欢迎提交你的主题 +> [!WARNING] +> 不是每一个主题都及时同步了所有功能，由于精力有限，优先更新默认主题，其他主题欢迎 & 期待 PR + ## 提交新的主题 > 欢迎在页面底部保留你和 One API 的版权信息以及指向链接 diff --git a/web/air/src/components/TokensTable.js b/web/air/src/components/TokensTable.js index c106b388..0853ddfb 100644 --- a/web/air/src/components/TokensTable.js +++ b/web/air/src/components/TokensTable.js @@ -319,7 +319,7 @@ const TokensTable = () => { if (nextLink) { nextUrl = nextLink + `/#/?settings={"key":"sk-${key}","url":"${serverAddress}"}`; } else { - nextUrl = `https://chat.oneapi.pro/#/?settings={"key":"sk-${key}","url":"${serverAddress}"}`; + nextUrl = `https://app.nextchat.dev/#/?settings={"key":"sk-${key}","url":"${serverAddress}"}`; } let url; diff --git a/web/berry/src/constants/ChannelConstants.js b/web/berry/src/constants/ChannelConstants.js index ec049f7d..b74c58c7 100644 --- a/web/berry/src/constants/ChannelConstants.js +++ b/web/berry/src/constants/ChannelConstants.js @@ -107,6 +107,12 @@ export const CHANNEL_OPTIONS = { value: 31, color: 'primary' }, + 32: { + key: 32, + text: '阶跃星辰', + value: 32, + color: 'primary' + }, 8: { key: 8, text: '自定义渠道', diff --git a/web/berry/src/constants/SnackbarConstants.js b/web/berry/src/constants/SnackbarConstants.js index a05c6652..19523da1 100644 --- a/web/berry/src/constants/SnackbarConstants.js +++ b/web/berry/src/constants/SnackbarConstants.js @@ -18,7 +18,7 @@ export const snackbarConstants = { }, NOTICE: { variant: 'info', - autoHideDuration: 20000 + autoHideDuration: 7000 } }, Mobile: { diff --git a/web/berry/src/utils/common.js b/web/berry/src/utils/common.js index 25e5c635..d8dabac3 100644 --- a/web/berry/src/utils/common.js +++ b/web/berry/src/utils/common.js @@ -51,9 +51,9 @@ export function showError(error) { export function showNotice(message, isHTML = false) { if (isHTML) { - enqueueSnackbar(, getSnackbarOptions('INFO')); + enqueueSnackbar(, getSnackbarOptions('NOTICE')); } else { - enqueueSnackbar(message, getSnackbarOptions('INFO')); + enqueueSnackbar(message, getSnackbarOptions('NOTICE')); } } diff --git a/web/berry/src/views/Channel/component/EditModal.js b/web/berry/src/views/Channel/component/EditModal.js index 07111c97..cbf411b9 100644 --- a/web/berry/src/views/Channel/component/EditModal.js +++ b/web/berry/src/views/Channel/component/EditModal.js @@ -340,7 +340,9 @@ const EditModal = ({ open, channelId, onCancel, onOk }) => { }, }} > - {Object.values(CHANNEL_OPTIONS).map((option) => { + {Object.values(CHANNEL_OPTIONS).sort((a, b) => { + return a.text.localeCompare(b.text) + }).map((option) => { return ( {option.text} diff --git a/web/berry/src/views/Token/component/EditModal.js b/web/berry/src/views/Token/component/EditModal.js index cf65a96a..a5e08be1 100644 --- a/web/berry/src/views/Token/component/EditModal.js +++ b/web/berry/src/views/Token/component/EditModal.js @@ -103,7 +103,7 @@ const EditModal = ({ open, tokenId, onCancel, onOk }) => { fontSize: "1.125rem", }} > - {tokenId ? "编辑Token" : "新建Token"} + {tokenId ? "编辑令牌" : "新建令牌"} diff --git a/web/berry/src/views/Token/component/TableRow.js b/web/berry/src/views/Token/component/TableRow.js index 2753764c..51ab0d4b 100644 --- a/web/berry/src/views/Token/component/TableRow.js +++ b/web/berry/src/views/Token/component/TableRow.js @@ -28,7 +28,7 @@ const COPY_OPTIONS = [ { key: 'next', text: 'ChatGPT Next', - url: 'https://chat.oneapi.pro/#/?settings={"key":"sk-{key}","url":"{serverAddress}"}', + url: 'https://app.nextchat.dev/#/?settings={"key":"sk-{key}","url":"{serverAddress}"}', encode: false }, { key: 'ama', text: 'BotGem', url: 'ama://set-api-key?server={serverAddress}&key=sk-{key}', encode: true }, diff --git a/web/default/src/App.js b/web/default/src/App.js index 13c884dc..4ece4eeb 100644 --- a/web/default/src/App.js +++ b/web/default/src/App.js @@ -24,6 +24,7 @@ import EditRedemption from './pages/Redemption/EditRedemption'; import TopUp from './pages/TopUp'; import Log from './pages/Log'; import Chat from './pages/Chat'; +import LarkOAuth from './components/LarkOAuth'; const Home = lazy(() => import('./pages/Home')); const About = lazy(() => import('./pages/About')); @@ -239,6 +240,14 @@ function App() { } /> + }> + + + } + /> { + const [searchParams, setSearchParams] = useSearchParams(); + + const [userState, userDispatch] = useContext(UserContext); + const [prompt, setPrompt] = useState('处理中...'); + const [processing, setProcessing] = useState(true); + + let navigate = useNavigate(); + + const sendCode = async (code, state, count) => { + const res = await API.get(`/api/oauth/lark?code=${code}&state=${state}`); + const { success, message, data } = res.data; + if (success) { + if (message === 'bind') { + showSuccess('绑定成功！'); + navigate('/setting'); + } else { + userDispatch({ type: 'login', payload: data }); + localStorage.setItem('user', JSON.stringify(data)); + showSuccess('登录成功！'); + navigate('/'); + } + } else { + showError(message); + if (count === 0) { + setPrompt(`操作失败，重定向至登录界面中...`); + navigate('/setting'); // in case this is failed to bind lark + return; + } + count++; + setPrompt(`出现错误，第 ${count} 次重试中...`); + await new Promise((resolve) => setTimeout(resolve, count * 2000)); + await sendCode(code, state, count); + } + }; + + useEffect(() => { + let code = searchParams.get('code'); + let state = searchParams.get('state'); + sendCode(code, state, 0).then(); + }, []); + + return ( + + + {prompt} + + + ); +}; + +export default LarkOAuth; diff --git a/web/default/src/components/LoginForm.js b/web/default/src/components/LoginForm.js index b48f64c4..1623b549 100644 --- a/web/default/src/components/LoginForm.js +++ b/web/default/src/components/LoginForm.js @@ -3,7 +3,8 @@ import { Button, Divider, Form, Grid, Header, Image, Message, Modal, Segment } f import { Link, useNavigate, useSearchParams } from 'react-router-dom'; import { UserContext } from '../context/User'; import { API, getLogo, showError, showSuccess, showWarning } from '../helpers'; -import { onGitHubOAuthClicked } from './utils'; +import { onGitHubOAuthClicked, onLarkOAuthClicked } from './utils'; +import larkIcon from '../images/lark.svg'; const LoginForm = () => { const [inputs, setInputs] = useState({ @@ -124,29 +125,50 @@ const LoginForm = () => { 点击注册 - {status.github_oauth || status.wechat_login ? ( + {status.github_oauth || status.wechat_login || status.lark_client_id ? ( <> Or - {status.github_oauth ? ( - ) } + { + status.lark_client_id && ( + + ) + }