Refactor codebase, introduce relaymode package, update constants and improve consistency

- Refactor constant definitions and organization
- Clean up package level variables and functions
- Introduce new `relaymode` and `apitype` packages for constant definitions
- Refactor and simplify code in several packages including `openai`, `relay/channel/baidu`, `relay/util`, `relay/controller`, `relay/channeltype`
- Add helper functions in `relay/channeltype` package to convert channel type constants to corresponding API type constants
- Remove deprecated functions such as `ResponseText2Usage` from `relay/channel/openai/helper.go`
- Modify code in `relay/util/validation.go` and related files to use new `validator.ValidateTextRequest` function
- Rename `util` package to `relaymode` and update related imports in several packages
This commit is contained in:
Laisky.Cai 2024-04-06 05:18:04 +00:00
parent 7c59a97a6b
commit 50bab08496
157 changed files with 3217 additions and 19160 deletions

1
.gitignore vendored
View File

@ -8,3 +8,4 @@ build
logs logs
data data
node_modules node_modules
cmd.md

View File

@ -80,6 +80,9 @@ var SMTPToken = ""
var GitHubClientId = "" var GitHubClientId = ""
var GitHubClientSecret = "" var GitHubClientSecret = ""
var LarkClientId = ""
var LarkClientSecret = ""
var WeChatServerAddress = "" var WeChatServerAddress = ""
var WeChatServerToken = "" var WeChatServerToken = ""
var WeChatAccountQRCodeImageURL = "" var WeChatAccountQRCodeImageURL = ""

9
common/config/key.go Normal file
View File

@ -0,0 +1,9 @@
package config
const (
KeyPrefix = "cfg_"
KeyAPIVersion = KeyPrefix + "api_version"
KeyLibraryID = KeyPrefix + "library_id"
KeyPlugin = KeyPrefix + "plugin"
)

View File

@ -4,116 +4,3 @@ import "time"
var StartTime = time.Now().Unix() // unit: second var StartTime = time.Now().Unix() // unit: second
var Version = "v0.0.0" // this hard coding will be replaced automatically when building, no need to manually change var 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"
)

View File

@ -2,16 +2,15 @@ package helper
import ( import (
"fmt" "fmt"
"github.com/google/uuid"
"html/template" "html/template"
"log" "log"
"math/rand"
"net" "net"
"os/exec" "os/exec"
"runtime" "runtime"
"strconv" "strconv"
"strings" "strings"
"time"
"github.com/songquanpeng/one-api/common/random"
) )
func OpenBrowser(url string) { func OpenBrowser(url string) {
@ -79,31 +78,6 @@ func Bytes2Size(num int64) string {
return numStr + " " + unit 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 { func Interface2String(inter interface{}) string {
switch inter := inter.(type) { switch inter := inter.(type) {
case string: case string:
@ -128,65 +102,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 { func GenRequestID() string {
return GetTimeString() + GetRandomNumberString(8) return GetTimeString() + random.GetRandomNumberString(8)
} }
func Max(a int, b int) int { func Max(a int, b int) int {

15
common/helper/time.go Normal file
View File

@ -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)
}

52
common/network/ip.go Normal file
View File

@ -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
}

19
common/network/ip_test.go Normal file
View File

@ -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)
})
}

View File

@ -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)
}

62
common/random/main.go Normal file
View File

@ -0,0 +1,62 @@
package random
import (
"math/rand"
"strings"
"time"
"github.com/google/uuid"
)
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)
}

View File

@ -1,4 +1,4 @@
package controller package auth
import ( import (
"bytes" "bytes"
@ -7,10 +7,10 @@ import (
"github.com/Laisky/errors/v2" "github.com/Laisky/errors/v2"
"github.com/gin-contrib/sessions" "github.com/gin-contrib/sessions"
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
"github.com/songquanpeng/one-api/common"
"github.com/songquanpeng/one-api/common/config" "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/logger"
"github.com/songquanpeng/one-api/common/random"
"github.com/songquanpeng/one-api/controller"
"github.com/songquanpeng/one-api/model" "github.com/songquanpeng/one-api/model"
"net/http" "net/http"
"strconv" "strconv"
@ -133,8 +133,8 @@ func GitHubOAuth(c *gin.Context) {
user.DisplayName = "GitHub User" user.DisplayName = "GitHub User"
} }
user.Email = githubUser.Email user.Email = githubUser.Email
user.Role = common.RoleCommonUser user.Role = model.RoleCommonUser
user.Status = common.UserStatusEnabled user.Status = model.UserStatusEnabled
if err := user.Insert(0); err != nil { if err := user.Insert(0); err != nil {
c.JSON(http.StatusOK, gin.H{ 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{ c.JSON(http.StatusOK, gin.H{
"message": "用户已被封禁", "message": "用户已被封禁",
"success": false, "success": false,
}) })
return return
} }
setupLogin(&user, c) controller.SetupLogin(&user, c)
} }
func GitHubBind(c *gin.Context) { func GitHubBind(c *gin.Context) {
@ -219,7 +219,7 @@ func GitHubBind(c *gin.Context) {
func GenerateOAuthCode(c *gin.Context) { func GenerateOAuthCode(c *gin.Context) {
session := sessions.Default(c) session := sessions.Default(c)
state := helper.GetRandomString(12) state := random.GetRandomString(12)
session.Set("oauth_state", state) session.Set("oauth_state", state)
err := session.Save() err := session.Save()
if err != nil { if err != nil {

201
controller/auth/lark.go Normal file
View File

@ -0,0 +1,201 @@
package auth
import (
"bytes"
"encoding/json"
"fmt"
"net/http"
"strconv"
"time"
"github.com/Laisky/errors/v2"
"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"
)
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
}

View File

@ -1,12 +1,12 @@
package controller package auth
import ( import (
"encoding/json" "encoding/json"
"fmt" "fmt"
"github.com/Laisky/errors/v2" "github.com/Laisky/errors/v2"
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
"github.com/songquanpeng/one-api/common"
"github.com/songquanpeng/one-api/common/config" "github.com/songquanpeng/one-api/common/config"
"github.com/songquanpeng/one-api/controller"
"github.com/songquanpeng/one-api/model" "github.com/songquanpeng/one-api/model"
"net/http" "net/http"
"strconv" "strconv"
@ -83,8 +83,8 @@ func WeChatAuth(c *gin.Context) {
if config.RegisterEnabled { if config.RegisterEnabled {
user.Username = "wechat_" + strconv.Itoa(model.GetMaxUserId()+1) user.Username = "wechat_" + strconv.Itoa(model.GetMaxUserId()+1)
user.DisplayName = "WeChat User" user.DisplayName = "WeChat User"
user.Role = common.RoleCommonUser user.Role = model.RoleCommonUser
user.Status = common.UserStatusEnabled user.Status = model.UserStatusEnabled
if err := user.Insert(0); err != nil { if err := user.Insert(0); err != nil {
c.JSON(http.StatusOK, gin.H{ 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{ c.JSON(http.StatusOK, gin.H{
"message": "用户已被封禁", "message": "用户已被封禁",
"success": false, "success": false,
}) })
return return
} }
setupLogin(&user, c) controller.SetupLogin(&user, c)
} }
func WeChatBind(c *gin.Context) { func WeChatBind(c *gin.Context) {

View File

@ -3,18 +3,19 @@ package controller
import ( import (
"encoding/json" "encoding/json"
"fmt" "fmt"
"github.com/Laisky/errors/v2"
"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"
"io" "io"
"net/http" "net/http"
"strconv" "strconv"
"time" "time"
"github.com/Laisky/errors/v2"
"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/channeltype"
"github.com/songquanpeng/one-api/relay/client"
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
) )
@ -96,7 +97,7 @@ func GetResponseBody(method, url string, channel *model.Channel, headers http.He
for k := range headers { for k := range headers {
req.Header.Add(k, headers.Get(k)) req.Header.Add(k, headers.Get(k))
} }
res, err := util.HTTPClient.Do(req) res, err := client.HTTPClient.Do(req)
if err != nil { if err != nil {
return nil, err return nil, err
} }
@ -204,28 +205,28 @@ func updateChannelAIGC2DBalance(channel *model.Channel) (float64, error) {
} }
func updateChannelBalance(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() == "" { if channel.GetBaseURL() == "" {
channel.BaseURL = &baseURL channel.BaseURL = &baseURL
} }
switch channel.Type { switch channel.Type {
case common.ChannelTypeOpenAI: case channeltype.OpenAI:
if channel.GetBaseURL() != "" { if channel.GetBaseURL() != "" {
baseURL = channel.GetBaseURL() baseURL = channel.GetBaseURL()
} }
case common.ChannelTypeAzure: case channeltype.Azure:
return 0, errors.New("尚未实现") return 0, errors.New("尚未实现")
case common.ChannelTypeCustom: case channeltype.Custom:
baseURL = channel.GetBaseURL() baseURL = channel.GetBaseURL()
case common.ChannelTypeCloseAI: case channeltype.CloseAI:
return updateChannelCloseAIBalance(channel) return updateChannelCloseAIBalance(channel)
case common.ChannelTypeOpenAISB: case channeltype.OpenAISB:
return updateChannelOpenAISBBalance(channel) return updateChannelOpenAISBBalance(channel)
case common.ChannelTypeAIProxy: case channeltype.AIProxy:
return updateChannelAIProxyBalance(channel) return updateChannelAIProxyBalance(channel)
case common.ChannelTypeAPI2GPT: case channeltype.API2GPT:
return updateChannelAPI2GPTBalance(channel) return updateChannelAPI2GPTBalance(channel)
case common.ChannelTypeAIGC2D: case channeltype.AIGC2D:
return updateChannelAIGC2DBalance(channel) return updateChannelAIGC2DBalance(channel)
default: default:
return 0, errors.New("尚未实现") return 0, errors.New("尚未实现")
@ -301,11 +302,11 @@ func updateAllChannelsBalance() error {
return err return err
} }
for _, channel := range channels { for _, channel := range channels {
if channel.Status != common.ChannelStatusEnabled { if channel.Status != model.ChannelStatusEnabled {
continue continue
} }
// TODO: support Azure // TODO: support Azure
if channel.Type != common.ChannelTypeOpenAI && channel.Type != common.ChannelTypeCustom { if channel.Type != channeltype.OpenAI && channel.Type != channeltype.Custom {
continue continue
} }
balance, err := updateChannelBalance(channel) balance, err := updateChannelBalance(channel)

View File

@ -4,18 +4,6 @@ import (
"bytes" "bytes"
"encoding/json" "encoding/json"
"fmt" "fmt"
"github.com/Laisky/errors/v2"
"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"
relaymodel "github.com/songquanpeng/one-api/relay/model"
"github.com/songquanpeng/one-api/relay/util"
"io" "io"
"net/http" "net/http"
"net/http/httptest" "net/http/httptest"
@ -25,6 +13,20 @@ import (
"sync" "sync"
"time" "time"
"github.com/Laisky/errors/v2"
"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"
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/relaymode"
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
) )
@ -56,9 +58,9 @@ func testChannel(channel *model.Channel) (err error, openaiErr *relaymodel.Error
c.Set("channel", channel.Type) c.Set("channel", channel.Type)
c.Set("base_url", channel.GetBaseURL()) c.Set("base_url", channel.GetBaseURL())
middleware.SetupContextForSelectedChannel(c, channel, "") middleware.SetupContextForSelectedChannel(c, channel, "")
meta := util.GetRelayMeta(c) meta := meta.GetByContext(c)
apiType := constant.ChannelType2APIType(channel.Type) apiType := channeltype.ToAPIType(channel.Type)
adaptor := helper.GetAdaptor(apiType) adaptor := relay.GetAdaptor(apiType)
if adaptor == nil { if adaptor == nil {
return errors.Errorf("invalid api type: %d, adaptor is nil", apiType), nil return errors.Errorf("invalid api type: %d, adaptor is nil", apiType), nil
} }
@ -73,7 +75,7 @@ func testChannel(channel *model.Channel) (err error, openaiErr *relaymodel.Error
request := buildTestRequest() request := buildTestRequest()
request.Model = modelName request.Model = modelName
meta.OriginModelName, meta.ActualModelName = modelName, 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 { if err != nil {
return err, nil return err, nil
} }
@ -88,8 +90,8 @@ func testChannel(channel *model.Channel) (err error, openaiErr *relaymodel.Error
return err, nil return err, nil
} }
if resp.StatusCode != http.StatusOK { if resp.StatusCode != http.StatusOK {
err := util.RelayErrorHandler(resp) err := controller.RelayErrorHandler(resp)
return errors.Errorf("status code %d: %s", resp.StatusCode, err.Error.Message), &err.Error return fmt.Errorf("status code %d: %s", resp.StatusCode, err.Error.Message), &err.Error
} }
usage, respErr := adaptor.DoResponse(c, resp, meta) usage, respErr := adaptor.DoResponse(c, resp, meta)
if respErr != nil { if respErr != nil {
@ -171,7 +173,7 @@ func testChannels(notify bool, scope string) error {
} }
go func() { go func() {
for _, channel := range channels { for _, channel := range channels {
isChannelEnabled := channel.Status == common.ChannelStatusEnabled isChannelEnabled := channel.Status == model.ChannelStatusEnabled
tik := time.Now() tik := time.Now()
err, openaiErr := testChannel(channel) err, openaiErr := testChannel(channel)
tok := time.Now() tok := time.Now()
@ -184,10 +186,10 @@ func testChannels(notify bool, scope string) error {
_ = message.Notify(message.ByAll, fmt.Sprintf("渠道 %s %d测试超时", channel.Name, channel.Id), "", err.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()) 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) monitor.EnableChannel(channel.Id, channel.Name)
} }
channel.UpdateResponseTime(milliseconds) channel.UpdateResponseTime(milliseconds)

View File

@ -2,13 +2,13 @@ package controller
import ( import (
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
"github.com/songquanpeng/one-api/common" billingratio "github.com/songquanpeng/one-api/relay/billing/ratio"
"net/http" "net/http"
) )
func GetGroups(c *gin.Context) { func GetGroups(c *gin.Context) {
groupNames := make([]string, 0) groupNames := make([]string, 0)
for groupName := range common.GroupRatio { for groupName := range billingratio.GroupRatio {
groupNames = append(groupNames, groupName) groupNames = append(groupNames, groupName)
} }
c.JSON(http.StatusOK, gin.H{ c.JSON(http.StatusOK, gin.H{

View File

@ -23,6 +23,7 @@ func GetStatus(c *gin.Context) {
"email_verification": config.EmailVerificationEnabled, "email_verification": config.EmailVerificationEnabled,
"github_oauth": config.GitHubOAuthEnabled, "github_oauth": config.GitHubOAuthEnabled,
"github_client_id": config.GitHubClientId, "github_client_id": config.GitHubClientId,
"lark_client_id": config.LarkClientId,
"system_name": config.SystemName, "system_name": config.SystemName,
"logo": config.Logo, "logo": config.Logo,
"footer_html": config.Footer, "footer_html": config.Footer,

View File

@ -3,13 +3,15 @@ package controller
import ( import (
"fmt" "fmt"
"github.com/gin-gonic/gin" "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" relay "github.com/songquanpeng/one-api/relay"
"github.com/songquanpeng/one-api/relay/constant" "github.com/songquanpeng/one-api/relay/adaptor/openai"
"github.com/songquanpeng/one-api/relay/helper" "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" relaymodel "github.com/songquanpeng/one-api/relay/model"
"github.com/songquanpeng/one-api/relay/util"
"net/http" "net/http"
"strings"
) )
// https://platform.openai.com/docs/api-reference/models/list // https://platform.openai.com/docs/api-reference/models/list
@ -39,8 +41,8 @@ type OpenAIModels struct {
Parent *string `json:"parent"` Parent *string `json:"parent"`
} }
var openAIModels []OpenAIModels var models []OpenAIModels
var openAIModelsMap map[string]OpenAIModels var modelsMap map[string]OpenAIModels
var channelId2Models map[int][]string var channelId2Models map[int][]string
func init() { func init() {
@ -60,11 +62,11 @@ func init() {
IsBlocking: false, IsBlocking: false,
}) })
// https://platform.openai.com/docs/models/model-endpoint-compatibility // https://platform.openai.com/docs/models/model-endpoint-compatibility
for i := 0; i < constant.APITypeDummy; i++ { for i := 0; i < apitype.Dummy; i++ {
if i == constant.APITypeAIProxyLibrary { if i == apitype.AIProxyLibrary {
continue continue
} }
adaptor := helper.GetAdaptor(i) adaptor := relay.GetAdaptor(i)
if adaptor == nil { if adaptor == nil {
continue continue
} }
@ -72,7 +74,7 @@ func init() {
channelName := adaptor.GetChannelName() channelName := adaptor.GetChannelName()
modelNames := adaptor.GetModelList() modelNames := adaptor.GetModelList()
for _, modelName := range modelNames { for _, modelName := range modelNames {
openAIModels = append(openAIModels, OpenAIModels{ models = append(models, OpenAIModels{
Id: modelName, Id: modelName,
Object: "model", Object: "model",
Created: 1626777600, Created: 1626777600,
@ -84,12 +86,12 @@ func init() {
} }
} }
for _, channelType := range openai.CompatibleChannels { for _, channelType := range openai.CompatibleChannels {
if channelType == common.ChannelTypeAzure { if channelType == channeltype.Azure {
continue continue
} }
channelName, channelModelList := openai.GetCompatibleChannelMeta(channelType) channelName, channelModelList := openai.GetCompatibleChannelMeta(channelType)
for _, modelName := range channelModelList { for _, modelName := range channelModelList {
openAIModels = append(openAIModels, OpenAIModels{ models = append(models, OpenAIModels{
Id: modelName, Id: modelName,
Object: "model", Object: "model",
Created: 1626777600, Created: 1626777600,
@ -100,18 +102,18 @@ func init() {
}) })
} }
} }
openAIModelsMap = make(map[string]OpenAIModels) modelsMap = make(map[string]OpenAIModels)
for _, model := range openAIModels { for _, model := range models {
openAIModelsMap[model.Id] = model modelsMap[model.Id] = model
} }
channelId2Models = make(map[int][]string) channelId2Models = make(map[int][]string)
for i := 1; i < common.ChannelTypeDummy; i++ { for i := 1; i < channeltype.Dummy; i++ {
adaptor := helper.GetAdaptor(constant.ChannelType2APIType(i)) adaptor := relay.GetAdaptor(channeltype.ToAPIType(i))
if adaptor == nil { if adaptor == nil {
continue continue
} }
meta := &util.RelayMeta{ meta := &meta.Meta{
ChannelType: i, ChannelType: i,
} }
adaptor.Init(meta) adaptor.Init(meta)
@ -127,16 +129,55 @@ func DashboardListModels(c *gin.Context) {
}) })
} }
func ListModels(c *gin.Context) { func ListAllModels(c *gin.Context) {
c.JSON(200, gin.H{ c.JSON(200, gin.H{
"object": "list", "object": "list",
"data": openAIModels, "data": models,
})
}
func ListModels(c *gin.Context) {
ctx := c.Request.Context()
var availableModels []string
if c.GetString("available_models") != "" {
availableModels = strings.Split(c.GetString("available_models"), ",")
} else {
userId := c.GetInt("id")
userGroup, _ := model.CacheGetUserGroup(userId)
availableModels, _ = model.CacheGetGroupModels(ctx, userGroup)
}
modelSet := make(map[string]bool)
for _, availableModel := range availableModels {
modelSet[availableModel] = true
}
availableOpenAIModels := make([]OpenAIModels, 0)
for _, model := range models {
if _, ok := modelSet[model.Id]; ok {
modelSet[model.Id] = false
availableOpenAIModels = append(availableOpenAIModels, model)
}
}
for modelName, ok := range modelSet {
if ok {
availableOpenAIModels = append(availableOpenAIModels, OpenAIModels{
Id: modelName,
Object: "model",
Created: 1626777600,
OwnedBy: "custom",
Root: modelName,
Parent: nil,
})
}
}
c.JSON(200, gin.H{
"object": "list",
"data": availableOpenAIModels,
}) })
} }
func RetrieveModel(c *gin.Context) { func RetrieveModel(c *gin.Context) {
modelId := c.Param("model") modelId := c.Param("model")
if model, ok := openAIModelsMap[modelId]; ok { if model, ok := modelsMap[modelId]; ok {
c.JSON(200, model) c.JSON(200, model)
} else { } else {
Error := relaymodel.Error{ Error := relaymodel.Error{
@ -150,3 +191,30 @@ func RetrieveModel(c *gin.Context) {
}) })
} }
} }
func GetUserAvailableModels(c *gin.Context) {
ctx := c.Request.Context()
id := c.GetInt("id")
userGroup, err := model.CacheGetUserGroup(id)
if err != nil {
c.JSON(http.StatusOK, gin.H{
"success": false,
"message": err.Error(),
})
return
}
models, err := model.CacheGetGroupModels(ctx, userGroup)
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": models,
})
return
}

View File

@ -4,6 +4,7 @@ import (
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
"github.com/songquanpeng/one-api/common/config" "github.com/songquanpeng/one-api/common/config"
"github.com/songquanpeng/one-api/common/helper" "github.com/songquanpeng/one-api/common/helper"
"github.com/songquanpeng/one-api/common/random"
"github.com/songquanpeng/one-api/model" "github.com/songquanpeng/one-api/model"
"net/http" "net/http"
"strconv" "strconv"
@ -106,7 +107,7 @@ func AddRedemption(c *gin.Context) {
} }
var keys []string var keys []string
for i := 0; i < redemption.Count; i++ { for i := 0; i < redemption.Count; i++ {
key := helper.GetUUID() key := random.GetUUID()
cleanRedemption := model.Redemption{ cleanRedemption := model.Redemption{
UserId: c.GetInt("id"), UserId: c.GetInt("id"),
Name: redemption.Name, Name: redemption.Name,

View File

@ -4,9 +4,6 @@ import (
"bytes" "bytes"
"context" "context"
"fmt" "fmt"
"io"
"net/http"
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
"github.com/pkg/errors" "github.com/pkg/errors"
"github.com/songquanpeng/one-api/common" "github.com/songquanpeng/one-api/common"
@ -16,24 +13,25 @@ import (
"github.com/songquanpeng/one-api/middleware" "github.com/songquanpeng/one-api/middleware"
dbmodel "github.com/songquanpeng/one-api/model" dbmodel "github.com/songquanpeng/one-api/model"
"github.com/songquanpeng/one-api/monitor" "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/controller"
"github.com/songquanpeng/one-api/relay/model" "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 // 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 var err *model.ErrorWithStatusCode
switch relayMode { switch relayMode {
case constant.RelayModeImagesGenerations: case relaymode.ImagesGenerations:
err = controller.RelayImageHelper(c, relayMode) err = controller.RelayImageHelper(c, relayMode)
case constant.RelayModeAudioSpeech: case relaymode.AudioSpeech:
fallthrough fallthrough
case constant.RelayModeAudioTranslation: case relaymode.AudioTranslation:
fallthrough fallthrough
case constant.RelayModeAudioTranscription: case relaymode.AudioTranscription:
err = controller.RelayAudioHelper(c, relayMode) err = controller.RelayAudioHelper(c, relayMode)
default: default:
err = controller.RelayTextHelper(c) err = controller.RelayTextHelper(c)
@ -43,13 +41,13 @@ func relay(c *gin.Context, relayMode int) *model.ErrorWithStatusCode {
func Relay(c *gin.Context) { func Relay(c *gin.Context) {
ctx := c.Request.Context() ctx := c.Request.Context()
relayMode := constant.Path2RelayMode(c.Request.URL.Path) relayMode := relaymode.GetByPath(c.Request.URL.Path)
if config.DebugEnabled { if config.DebugEnabled {
requestBody, _ := common.GetRequestBody(c) requestBody, _ := common.GetRequestBody(c)
logger.Debugf(ctx, "request body: %s", string(requestBody)) logger.Debugf(ctx, "request body: %s", string(requestBody))
} }
channelId := c.GetInt("channel_id") channelId := c.GetInt("channel_id")
bizErr := relay(c, relayMode) bizErr := relayHelper(c, relayMode)
if bizErr == nil { if bizErr == nil {
monitor.Emit(channelId, true) monitor.Emit(channelId, true)
return return
@ -78,7 +76,7 @@ func Relay(c *gin.Context) {
middleware.SetupContextForSelectedChannel(c, channel, originalModel) middleware.SetupContextForSelectedChannel(c, channel, originalModel)
requestBody, err := common.GetRequestBody(c) requestBody, err := common.GetRequestBody(c)
c.Request.Body = io.NopCloser(bytes.NewBuffer(requestBody)) c.Request.Body = io.NopCloser(bytes.NewBuffer(requestBody))
bizErr = relay(c, relayMode) bizErr = relayHelper(c, relayMode)
if bizErr == nil { if bizErr == nil {
return return
} }
@ -117,7 +115,7 @@ func shouldRetry(c *gin.Context, statusCode int) error {
func processChannelRelayError(ctx context.Context, channelId int, channelName string, err *model.ErrorWithStatusCode) { func processChannelRelayError(ctx context.Context, channelId int, channelName string, err *model.ErrorWithStatusCode) {
logger.Errorf(ctx, "relay error (channel #%d): %s", channelId, err.Message) logger.Errorf(ctx, "relay error (channel #%d): %s", channelId, err.Message)
// https://platform.openai.com/docs/guides/error-codes/api-errors // 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) monitor.DisableChannel(channelId, channelName, err.Message)
} else { } else {
monitor.Emit(channelId, false) monitor.Emit(channelId, false)

View File

@ -6,9 +6,12 @@ import (
"strconv" "strconv"
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
"github.com/jinzhu/copier"
"github.com/songquanpeng/one-api/common" "github.com/songquanpeng/one-api/common"
"github.com/songquanpeng/one-api/common/config" "github.com/songquanpeng/one-api/common/config"
"github.com/songquanpeng/one-api/common/helper" "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" "github.com/songquanpeng/one-api/model"
) )
@ -106,9 +109,24 @@ 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) { func AddToken(c *gin.Context) {
token := model.Token{} token := new(model.Token)
err := c.ShouldBindJSON(&token) err := c.ShouldBindJSON(token)
if err != nil { if err != nil {
c.JSON(http.StatusOK, gin.H{ c.JSON(http.StatusOK, gin.H{
"success": false, "success": false,
@ -116,22 +134,27 @@ func AddToken(c *gin.Context) {
}) })
return return
} }
if len(token.Name) > 30 {
err = validateToken(c, token)
if err != nil {
c.JSON(http.StatusOK, gin.H{ c.JSON(http.StatusOK, gin.H{
"success": false, "success": false,
"message": "令牌名称过长", "message": fmt.Sprintf("参数错误:%s", err.Error()),
}) })
return return
} }
cleanToken := model.Token{ cleanToken := model.Token{
UserId: c.GetInt("id"), UserId: c.GetInt("id"),
Name: token.Name, Name: token.Name,
Key: helper.GenerateKey(), Key: random.GenerateKey(),
CreatedTime: helper.GetTimestamp(), CreatedTime: helper.GetTimestamp(),
AccessedTime: helper.GetTimestamp(), AccessedTime: helper.GetTimestamp(),
ExpiredTime: token.ExpiredTime, ExpiredTime: token.ExpiredTime,
RemainQuota: token.RemainQuota, RemainQuota: token.RemainQuota,
UnlimitedQuota: token.UnlimitedQuota, UnlimitedQuota: token.UnlimitedQuota,
Models: token.Models,
Subnet: token.Subnet,
} }
err = cleanToken.Insert() err = cleanToken.Insert()
if err != nil { if err != nil {
@ -168,12 +191,7 @@ func DeleteToken(c *gin.Context) {
} }
type updateTokenDto struct { type updateTokenDto struct {
Id int `json:"id"` model.Token
Status int `json:"status" gorm:"default:1"`
Name *string `json:"name" gorm:"index" `
ExpiredTime *int64 `json:"expired_time" gorm:"bigint;default:-1"` // -1 means never expired
RemainQuota *int `json:"remain_quota" gorm:"default:0"`
UnlimitedQuota *bool `json:"unlimited_quota" gorm:"default:false"`
// AddUsedQuota add or subtract used quota // AddUsedQuota add or subtract used quota
AddUsedQuota int `json:"add_used_quota" gorm:"-"` AddUsedQuota int `json:"add_used_quota" gorm:"-"`
AddReason string `json:"add_reason" gorm:"-"` AddReason string `json:"add_reason" gorm:"-"`
@ -183,43 +201,51 @@ func UpdateToken(c *gin.Context) {
userId := c.GetInt("id") userId := c.GetInt("id")
statusOnly := c.Query("status_only") statusOnly := c.Query("status_only")
tokenPatch := new(updateTokenDto) tokenPatch := new(updateTokenDto)
if err := c.ShouldBindJSON(tokenPatch); err != nil { err := c.ShouldBindJSON(tokenPatch)
c.JSON(http.StatusOK, gin.H{
"success": false,
"message": "parse request join: " + err.Error(),
})
return
}
if tokenPatch.Name != nil &&
(len(*tokenPatch.Name) > 30 || len(*tokenPatch.Name) == 0) {
c.JSON(http.StatusOK, gin.H{
"success": false,
"message": "令牌名称错误,长度应在 1-30 之间",
})
return
}
tokenInDB, err := model.GetTokenByIds(tokenPatch.Id, userId)
if err != nil { if err != nil {
c.JSON(http.StatusOK, gin.H{ c.JSON(http.StatusOK, gin.H{
"success": false, "success": false,
"message": fmt.Sprintf("get token by id %d: %s", tokenPatch.Id, err.Error()), "message": err.Error(),
}) })
return return
} }
if tokenPatch.Status == common.TokenStatusEnabled { token := new(model.Token)
if tokenInDB.Status == common.TokenStatusExpired && tokenInDB.ExpiredTime <= helper.GetTimestamp() && tokenInDB.ExpiredTime != -1 { if err = copier.Copy(token, tokenPatch); err != nil {
c.JSON(http.StatusOK, gin.H{
"success": false,
"message": err.Error(),
})
return
}
err = validateToken(c, token)
if err != nil {
c.JSON(http.StatusOK, gin.H{
"success": false,
"message": fmt.Sprintf("参数错误:%s", err.Error()),
})
return
}
cleanToken, err := model.GetTokenByIds(token.Id, userId)
if err != nil {
c.JSON(http.StatusOK, gin.H{
"success": false,
"message": err.Error(),
})
return
}
if token.Status == model.TokenStatusEnabled {
if cleanToken.Status == model.TokenStatusExpired && cleanToken.ExpiredTime <= helper.GetTimestamp() && cleanToken.ExpiredTime != -1 {
c.JSON(http.StatusOK, gin.H{ c.JSON(http.StatusOK, gin.H{
"success": false, "success": false,
"message": "令牌已过期,无法启用,请先修改令牌过期时间,或者设置为永不过期", "message": "令牌已过期,无法启用,请先修改令牌过期时间,或者设置为永不过期",
}) })
return return
} }
if tokenInDB.Status == common.TokenStatusExhausted && if cleanToken.Status == model.TokenStatusExhausted && cleanToken.RemainQuota <= 0 && !cleanToken.UnlimitedQuota {
tokenInDB.RemainQuota <= 0 &&
!tokenInDB.UnlimitedQuota {
c.JSON(http.StatusOK, gin.H{ c.JSON(http.StatusOK, gin.H{
"success": false, "success": false,
"message": "令牌可用额度已用尽,无法启用,请先修改令牌剩余额度,或者设置为无限额度", "message": "令牌可用额度已用尽,无法启用,请先修改令牌剩余额度,或者设置为无限额度",
@ -228,42 +254,33 @@ func UpdateToken(c *gin.Context) {
} }
} }
if statusOnly != "" { if statusOnly != "" {
tokenInDB.Status = tokenPatch.Status cleanToken.Status = token.Status
} else { } else {
// If you add more fields, please also update tokenPatch.Update() // If you add more fields, please also update token.Update()
if tokenPatch.Name != nil { cleanToken.Name = token.Name
tokenInDB.Name = *tokenPatch.Name cleanToken.ExpiredTime = token.ExpiredTime
cleanToken.RemainQuota = token.RemainQuota
cleanToken.UnlimitedQuota = token.UnlimitedQuota
cleanToken.Models = token.Models
cleanToken.Subnet = token.Subnet
} }
if tokenPatch.ExpiredTime != nil {
tokenInDB.ExpiredTime = *tokenPatch.ExpiredTime
}
if tokenPatch.RemainQuota != nil {
tokenInDB.RemainQuota = int64(*tokenPatch.RemainQuota)
}
if tokenPatch.UnlimitedQuota != nil {
tokenInDB.UnlimitedQuota = *tokenPatch.UnlimitedQuota
}
}
tokenInDB.RemainQuota -= int64(tokenPatch.AddUsedQuota)
tokenInDB.UsedQuota += int64(tokenPatch.AddUsedQuota)
if tokenPatch.AddUsedQuota != 0 { if tokenPatch.AddUsedQuota != 0 {
model.RecordLog(userId, model.LogTypeConsume, fmt.Sprintf("外部(%s)消耗 %s", tokenPatch.AddReason, common.LogQuota(int64(tokenPatch.AddUsedQuota)))) model.RecordLog(userId, model.LogTypeConsume, fmt.Sprintf("外部(%s)消耗 %s", tokenPatch.AddReason, common.LogQuota(int64(tokenPatch.AddUsedQuota))))
} }
if err = tokenInDB.Update(); err != nil { err = cleanToken.Update()
if err != nil {
c.JSON(http.StatusOK, gin.H{ c.JSON(http.StatusOK, gin.H{
"success": false, "success": false,
"message": "update token: " + err.Error(), "message": err.Error(),
}) })
return return
} }
c.JSON(http.StatusOK, gin.H{ c.JSON(http.StatusOK, gin.H{
"success": true, "success": true,
"message": "", "message": "",
"data": tokenInDB, "data": cleanToken,
}) })
return return
} }

View File

@ -5,7 +5,7 @@ import (
"fmt" "fmt"
"github.com/songquanpeng/one-api/common" "github.com/songquanpeng/one-api/common"
"github.com/songquanpeng/one-api/common/config" "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" "github.com/songquanpeng/one-api/model"
"net/http" "net/http"
"strconv" "strconv"
@ -58,11 +58,11 @@ func Login(c *gin.Context) {
}) })
return return
} }
setupLogin(&user, c) SetupLogin(&user, c)
} }
// setup session & cookies and then return user info // 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 := sessions.Default(c)
session.Set("id", user.Id) session.Set("id", user.Id)
session.Set("username", user.Username) session.Set("username", user.Username)
@ -245,7 +245,7 @@ func GetUser(c *gin.Context) {
return return
} }
myRole := c.GetInt("role") myRole := c.GetInt("role")
if myRole <= user.Role && myRole != common.RoleRootUser { if myRole <= user.Role && myRole != model.RoleRootUser {
c.JSON(http.StatusOK, gin.H{ c.JSON(http.StatusOK, gin.H{
"success": false, "success": false,
"message": "无权获取同级或更高等级用户的信息", "message": "无权获取同级或更高等级用户的信息",
@ -293,7 +293,7 @@ func GenerateAccessToken(c *gin.Context) {
}) })
return return
} }
user.AccessToken = helper.GetUUID() user.AccessToken = random.GetUUID()
if model.DB.Where("access_token = ?", user.AccessToken).First(user).RowsAffected != 0 { if model.DB.Where("access_token = ?", user.AccessToken).First(user).RowsAffected != 0 {
c.JSON(http.StatusOK, gin.H{ c.JSON(http.StatusOK, gin.H{
@ -330,7 +330,7 @@ func GetAffCode(c *gin.Context) {
return return
} }
if user.AffCode == "" { if user.AffCode == "" {
user.AffCode = helper.GetRandomString(4) user.AffCode = random.GetRandomString(4)
if err := user.Update(false); err != nil { if err := user.Update(false); err != nil {
c.JSON(http.StatusOK, gin.H{ c.JSON(http.StatusOK, gin.H{
"success": false, "success": false,
@ -404,14 +404,14 @@ func UpdateUser(c *gin.Context) {
return return
} }
myRole := c.GetInt("role") myRole := c.GetInt("role")
if myRole <= originUser.Role && myRole != common.RoleRootUser { if myRole <= originUser.Role && myRole != model.RoleRootUser {
c.JSON(http.StatusOK, gin.H{ c.JSON(http.StatusOK, gin.H{
"success": false, "success": false,
"message": "无权更新同权限等级或更高权限等级的用户信息", "message": "无权更新同权限等级或更高权限等级的用户信息",
}) })
return return
} }
if myRole <= updatedUser.Role && myRole != common.RoleRootUser { if myRole <= updatedUser.Role && myRole != model.RoleRootUser {
c.JSON(http.StatusOK, gin.H{ c.JSON(http.StatusOK, gin.H{
"success": false, "success": false,
"message": "无权将其他用户权限等级提升到大于等于自己的权限等级", "message": "无权将其他用户权限等级提升到大于等于自己的权限等级",
@ -525,7 +525,7 @@ func DeleteSelf(c *gin.Context) {
id := c.GetInt("id") id := c.GetInt("id")
user, _ := model.GetUserById(id, false) user, _ := model.GetUserById(id, false)
if user.Role == common.RoleRootUser { if user.Role == model.RoleRootUser {
c.JSON(http.StatusOK, gin.H{ c.JSON(http.StatusOK, gin.H{
"success": false, "success": false,
"message": "不能删除超级管理员账户", "message": "不能删除超级管理员账户",
@ -627,7 +627,7 @@ func ManageUser(c *gin.Context) {
return return
} }
myRole := c.GetInt("role") myRole := c.GetInt("role")
if myRole <= user.Role && myRole != common.RoleRootUser { if myRole <= user.Role && myRole != model.RoleRootUser {
c.JSON(http.StatusOK, gin.H{ c.JSON(http.StatusOK, gin.H{
"success": false, "success": false,
"message": "无权更新同权限等级或更高权限等级的用户信息", "message": "无权更新同权限等级或更高权限等级的用户信息",
@ -636,8 +636,8 @@ func ManageUser(c *gin.Context) {
} }
switch req.Action { switch req.Action {
case "disable": case "disable":
user.Status = common.UserStatusDisabled user.Status = model.UserStatusDisabled
if user.Role == common.RoleRootUser { if user.Role == model.RoleRootUser {
c.JSON(http.StatusOK, gin.H{ c.JSON(http.StatusOK, gin.H{
"success": false, "success": false,
"message": "无法禁用超级管理员用户", "message": "无法禁用超级管理员用户",
@ -645,9 +645,9 @@ func ManageUser(c *gin.Context) {
return return
} }
case "enable": case "enable":
user.Status = common.UserStatusEnabled user.Status = model.UserStatusEnabled
case "delete": case "delete":
if user.Role == common.RoleRootUser { if user.Role == model.RoleRootUser {
c.JSON(http.StatusOK, gin.H{ c.JSON(http.StatusOK, gin.H{
"success": false, "success": false,
"message": "无法删除超级管理员用户", "message": "无法删除超级管理员用户",
@ -662,37 +662,37 @@ func ManageUser(c *gin.Context) {
return return
} }
case "promote": case "promote":
if myRole != common.RoleRootUser { if myRole != model.RoleRootUser {
c.JSON(http.StatusOK, gin.H{ c.JSON(http.StatusOK, gin.H{
"success": false, "success": false,
"message": "普通管理员用户无法提升其他用户为管理员", "message": "普通管理员用户无法提升其他用户为管理员",
}) })
return return
} }
if user.Role >= common.RoleAdminUser { if user.Role >= model.RoleAdminUser {
c.JSON(http.StatusOK, gin.H{ c.JSON(http.StatusOK, gin.H{
"success": false, "success": false,
"message": "该用户已经是管理员", "message": "该用户已经是管理员",
}) })
return return
} }
user.Role = common.RoleAdminUser user.Role = model.RoleAdminUser
case "demote": case "demote":
if user.Role == common.RoleRootUser { if user.Role == model.RoleRootUser {
c.JSON(http.StatusOK, gin.H{ c.JSON(http.StatusOK, gin.H{
"success": false, "success": false,
"message": "无法降级超级管理员用户", "message": "无法降级超级管理员用户",
}) })
return return
} }
if user.Role == common.RoleCommonUser { if user.Role == model.RoleCommonUser {
c.JSON(http.StatusOK, gin.H{ c.JSON(http.StatusOK, gin.H{
"success": false, "success": false,
"message": "该用户已经是普通用户", "message": "该用户已经是普通用户",
}) })
return return
} }
user.Role = common.RoleCommonUser user.Role = model.RoleCommonUser
} }
if err := user.Update(false); err != nil { if err := user.Update(false); err != nil {
@ -746,7 +746,7 @@ func EmailBind(c *gin.Context) {
}) })
return return
} }
if user.Role == common.RoleRootUser { if user.Role == model.RoleRootUser {
config.RootUserEmail = email config.RootUserEmail = email
} }
c.JSON(http.StatusOK, gin.H{ c.JSON(http.StatusOK, gin.H{
@ -786,3 +786,38 @@ func TopUp(c *gin.Context) {
}) })
return 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
}

53
docs/API.md Normal file
View File

@ -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 补齐。

10
go.mod
View File

@ -13,8 +13,10 @@ require (
github.com/go-playground/validator/v10 v10.19.0 github.com/go-playground/validator/v10 v10.19.0
github.com/go-redis/redis/v8 v8.11.5 github.com/go-redis/redis/v8 v8.11.5
github.com/google/uuid v1.6.0 github.com/google/uuid v1.6.0
github.com/jinzhu/copier v0.4.0
github.com/pkg/errors v0.9.1 github.com/pkg/errors v0.9.1
github.com/pkoukk/tiktoken-go v0.1.6 github.com/pkoukk/tiktoken-go v0.1.6
github.com/smartystreets/goconvey v1.8.1
github.com/stretchr/testify v1.8.4 github.com/stretchr/testify v1.8.4
golang.org/x/crypto v0.21.0 golang.org/x/crypto v0.21.0
golang.org/x/image v0.15.0 golang.org/x/image v0.15.0
@ -48,16 +50,17 @@ require (
github.com/go-sql-driver/mysql v1.7.0 // indirect github.com/go-sql-driver/mysql v1.7.0 // indirect
github.com/goccy/go-json v0.10.2 // indirect github.com/goccy/go-json v0.10.2 // indirect
github.com/google/go-cpy v0.0.0-20211218193943-a9c933c06932 // indirect github.com/google/go-cpy v0.0.0-20211218193943-a9c933c06932 // indirect
github.com/gopherjs/gopherjs v1.17.2 // indirect
github.com/gorilla/context v1.1.1 // indirect github.com/gorilla/context v1.1.1 // indirect
github.com/gorilla/securecookie v1.1.1 // indirect github.com/gorilla/securecookie v1.1.1 // indirect
github.com/gorilla/sessions v1.2.1 // indirect github.com/gorilla/sessions v1.2.1 // indirect
github.com/jackc/pgpassfile v1.0.0 // indirect github.com/jackc/pgpassfile v1.0.0 // indirect
github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a // indirect github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a // indirect
github.com/jackc/pgx/v5 v5.5.4 // indirect github.com/jackc/pgx/v5 v5.4.3 // indirect
github.com/jackc/puddle/v2 v2.2.1 // indirect
github.com/jinzhu/inflection v1.0.0 // indirect github.com/jinzhu/inflection v1.0.0 // indirect
github.com/jinzhu/now v1.1.5 // indirect github.com/jinzhu/now v1.1.5 // indirect
github.com/json-iterator/go v1.1.12 // 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.7 // indirect github.com/klauspost/cpuid/v2 v2.2.7 // indirect
github.com/leodido/go-urn v1.4.0 // indirect github.com/leodido/go-urn v1.4.0 // indirect
github.com/mattn/go-isatty v0.0.20 // indirect github.com/mattn/go-isatty v0.0.20 // indirect
@ -66,6 +69,7 @@ require (
github.com/modern-go/reflect2 v1.0.2 // indirect github.com/modern-go/reflect2 v1.0.2 // indirect
github.com/pelletier/go-toml/v2 v2.1.1 // indirect github.com/pelletier/go-toml/v2 v2.1.1 // indirect
github.com/pmezard/go-difflib v1.0.0 // indirect github.com/pmezard/go-difflib v1.0.0 // indirect
github.com/smarty/assertions v1.15.0 // indirect
github.com/tailscale/hujson v0.0.0-20221223112325-20486734a56a // indirect github.com/tailscale/hujson v0.0.0-20221223112325-20486734a56a // indirect
github.com/twitchyliquid64/golang-asm v0.15.1 // indirect github.com/twitchyliquid64/golang-asm v0.15.1 // indirect
github.com/ugorji/go/codec v1.2.12 // indirect github.com/ugorji/go/codec v1.2.12 // indirect
@ -78,7 +82,7 @@ require (
golang.org/x/sys v0.18.0 // indirect golang.org/x/sys v0.18.0 // indirect
golang.org/x/term v0.18.0 // indirect golang.org/x/term v0.18.0 // indirect
golang.org/x/text v0.14.0 // indirect golang.org/x/text v0.14.0 // indirect
golang.org/x/tools v0.6.0 // indirect golang.org/x/tools v0.7.0 // indirect
google.golang.org/protobuf v1.33.0 // indirect google.golang.org/protobuf v1.33.0 // indirect
gopkg.in/yaml.v3 v3.0.1 // indirect gopkg.in/yaml.v3 v3.0.1 // indirect
) )

24
go.sum
View File

@ -96,6 +96,8 @@ github.com/google/go-cpy v0.0.0-20211218193943-a9c933c06932/go.mod h1:cC6EdPbj/1
github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg= github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=
github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= github.com/google/uuid v1.6.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 h1:AWwleXJkX/nhcU9bZSnZoi3h/qGYqQAGhq6zZe/aQW8=
github.com/gorilla/context v1.1.1/go.mod h1:kBGZzfjB9CEq2AlWe17Uuf7NDRt0dE0s8S51q0aT7Yg= github.com/gorilla/context v1.1.1/go.mod h1:kBGZzfjB9CEq2AlWe17Uuf7NDRt0dE0s8S51q0aT7Yg=
github.com/gorilla/securecookie v1.1.1 h1:miw7JPhV+b/lAHSXz4qd/nN9jRiAFV5FwjeKyCS8BvQ= github.com/gorilla/securecookie v1.1.1 h1:miw7JPhV+b/lAHSXz4qd/nN9jRiAFV5FwjeKyCS8BvQ=
@ -107,16 +109,18 @@ github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsI
github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg= github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg=
github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a h1:bbPeKD0xmW/Y25WS6cokEszi5g+S0QxI/d45PkRi7Nk= github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a h1:bbPeKD0xmW/Y25WS6cokEszi5g+S0QxI/d45PkRi7Nk=
github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a/go.mod h1:5TJZWKEWniPve33vlWYSoGYefn3gLQRzjfDlhSJ9ZKM= github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a/go.mod h1:5TJZWKEWniPve33vlWYSoGYefn3gLQRzjfDlhSJ9ZKM=
github.com/jackc/pgx/v5 v5.5.4 h1:Xp2aQS8uXButQdnCMWNmvx6UysWQQC+u1EoizjguY+8= github.com/jackc/pgx/v5 v5.4.3 h1:cxFyXhxlvAifxnkKKdlxv8XqUf59tDlYjnV5YYfsJJY=
github.com/jackc/pgx/v5 v5.5.4/go.mod h1:ez9gk+OAat140fv9ErkZDYFWmXLfV+++K0uAOiwgm1A= github.com/jackc/pgx/v5 v5.4.3/go.mod h1:Ig06C2Vu0t5qXC60W8sqIthScaEnFvojjj9dSljmHRA=
github.com/jackc/puddle/v2 v2.2.1 h1:RhxXJtFG022u4ibrCSMSiu5aOq1i77R3OHKNJj77OAk= github.com/jinzhu/copier v0.4.0 h1:w3ciUoD19shMCRargcpm0cm91ytaBhDvuRpz1ODO/U8=
github.com/jackc/puddle/v2 v2.2.1/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4= github.com/jinzhu/copier v0.4.0/go.mod h1:DfbEm0FYsaqBcKcFuvmOZb218JkPGtvSHsKg8S8hyyg=
github.com/jinzhu/inflection v1.0.0 h1:K317FqzuhWc8YvSVlFMCCUb36O/S9MCKRDI7QkRKD/E= github.com/jinzhu/inflection v1.0.0 h1:K317FqzuhWc8YvSVlFMCCUb36O/S9MCKRDI7QkRKD/E=
github.com/jinzhu/inflection v1.0.0/go.mod h1:h+uFLlag+Qp1Va5pdKtLDYj+kHp5pxUVkryuEj+Srlc= github.com/jinzhu/inflection v1.0.0/go.mod h1:h+uFLlag+Qp1Va5pdKtLDYj+kHp5pxUVkryuEj+Srlc=
github.com/jinzhu/now v1.1.5 h1:/o9tlHleP7gOFmsnYNz3RGnqzefHA47wQpKrrdTIwXQ= github.com/jinzhu/now v1.1.5 h1:/o9tlHleP7gOFmsnYNz3RGnqzefHA47wQpKrrdTIwXQ=
github.com/jinzhu/now v1.1.5/go.mod h1:d3SSVoowX0Lcu0IBviAWJpolVfI5UJVZZ7cO71lE/z8= github.com/jinzhu/now v1.1.5/go.mod h1:d3SSVoowX0Lcu0IBviAWJpolVfI5UJVZZ7cO71lE/z8=
github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM= 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/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.0.9/go.mod h1:FInQzS24/EEf25PyTYn52gqo7WaD8xa0213Md/qVLRg=
github.com/klauspost/cpuid/v2 v2.2.7 h1:ZWSB3igEs+d0qvnxR/ZBzXVmxkgt8DdzP6m9pfuVLDM= github.com/klauspost/cpuid/v2 v2.2.7 h1:ZWSB3igEs+d0qvnxR/ZBzXVmxkgt8DdzP6m9pfuVLDM=
github.com/klauspost/cpuid/v2 v2.2.7/go.mod h1:Lcz8mBdAVJIBVzewtcLocK12l3Y+JytZYpaMropDUws= github.com/klauspost/cpuid/v2 v2.2.7/go.mod h1:Lcz8mBdAVJIBVzewtcLocK12l3Y+JytZYpaMropDUws=
@ -164,6 +168,10 @@ github.com/prashantv/gostub v1.1.0/go.mod h1:A5zLQHz7ieHGG7is6LLXLz7I8+3LZzsrV0P
github.com/rogpeppe/go-internal v1.6.1/go.mod h1:xXDCJY+GAPziupqXw64V24skbSoqbTEfhy4qGm1nDQc= 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 h1:FCbCCtXNOY3UtUuHUYaghJg4y7Fd14rXifAYUAtL9R8=
github.com/rogpeppe/go-internal v1.8.0/go.mod h1:WmiCO8CzOY8rg0OYDC4/i/2WRWAB6poM+XZ2dLUbcbE= 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/spaolacci/murmur3 v0.0.0-20180118202830-f09979ecbc72 h1:qLC7fQah7D6K1B0ujays3HV9gkFtllcxhzImRR7ArPQ= github.com/spaolacci/murmur3 v0.0.0-20180118202830-f09979ecbc72 h1:qLC7fQah7D6K1B0ujays3HV9gkFtllcxhzImRR7ArPQ=
github.com/spaolacci/murmur3 v0.0.0-20180118202830-f09979ecbc72/go.mod h1:JwIasOWyU6f++ZhiEuf87xNszmSA2myDM2Kzu9HwQUA= github.com/spaolacci/murmur3 v0.0.0-20180118202830-f09979ecbc72/go.mod h1:JwIasOWyU6f++ZhiEuf87xNszmSA2myDM2Kzu9HwQUA=
github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
@ -206,8 +214,8 @@ golang.org/x/image v0.15.0/go.mod h1:HUYqC05R2ZcZ3ejNQsIHQDQiwWM4JBqmm6MKANTp4LE
golang.org/x/lint v0.0.0-20210508222113-6edffad5e616 h1:VLliZ0d+/avPrXXH+OakdXhpJuEoBZuwh1m2j7U6Iug= golang.org/x/lint v0.0.0-20210508222113-6edffad5e616 h1:VLliZ0d+/avPrXXH+OakdXhpJuEoBZuwh1m2j7U6Iug=
golang.org/x/lint v0.0.0-20210508222113-6edffad5e616/go.mod h1:3xt1FjdF8hUf6vQPIChWIBhFzV8gjjsPE/fR3IyQdNY= golang.org/x/lint v0.0.0-20210508222113-6edffad5e616/go.mod h1:3xt1FjdF8hUf6vQPIChWIBhFzV8gjjsPE/fR3IyQdNY=
golang.org/x/mod v0.1.1-0.20191105210325-c90efee705ee/go.mod h1:QqPTAvyqsEbceGzBzNggFXnrqF1CaUcvgkdR5Ot7KZg= golang.org/x/mod v0.1.1-0.20191105210325-c90efee705ee/go.mod h1:QqPTAvyqsEbceGzBzNggFXnrqF1CaUcvgkdR5Ot7KZg=
golang.org/x/mod v0.8.0 h1:LUYupSeNrTNCGzR/hVBk2NHZO4hXcVaW1k4Qx7rjPx8= golang.org/x/mod v0.9.0 h1:KENHtAZL2y3NLMYZeHY9DW8HW8V+kQyJsY/V9JlKvCs=
golang.org/x/mod v0.8.0/go.mod h1:iBbtSCu2XBx23ZKBPSOrRkjjQPZFPuis4dIYUhu/chs= golang.org/x/mod v0.9.0/go.mod h1:iBbtSCu2XBx23ZKBPSOrRkjjQPZFPuis4dIYUhu/chs=
golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg= golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=
golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s= golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg= golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg=
@ -238,8 +246,8 @@ golang.org/x/time v0.3.0 h1:rg5rLMjNzMS1RkNLzCG38eapWhnYLFYXDXj2gOlr8j4=
golang.org/x/time v0.3.0/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ= golang.org/x/time v0.3.0/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
golang.org/x/tools v0.0.0-20200130002326-2f3ba24bd6e7/go.mod h1:TB2adYChydJhpapKDTa4BR/hXlZSLoq2Wpct/0txZ28= golang.org/x/tools v0.0.0-20200130002326-2f3ba24bd6e7/go.mod h1:TB2adYChydJhpapKDTa4BR/hXlZSLoq2Wpct/0txZ28=
golang.org/x/tools v0.6.0 h1:BOw41kyTf3PuCW1pVQf8+Cyg8pMlkYB1oo9iJ6D/lKM= golang.org/x/tools v0.7.0 h1:W4OVu8VVOaIO0yzWMNdepAulS7YfoS3Zabrm8DOXXU4=
golang.org/x/tools v0.6.0/go.mod h1:Xwgl3UAJ/d3gWutnCtw505GrjyAbvKui8lOU390QaIU= golang.org/x/tools v0.7.0/go.mod h1:4pg6aUX35JBAogB10C9AtvVL+qowtN4pT3CGSQex14s=
golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw= google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw=

View File

@ -16,7 +16,7 @@ import (
"github.com/songquanpeng/one-api/controller" "github.com/songquanpeng/one-api/controller"
"github.com/songquanpeng/one-api/middleware" "github.com/songquanpeng/one-api/middleware"
"github.com/songquanpeng/one-api/model" "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" "github.com/songquanpeng/one-api/router"
) )

View File

@ -1,15 +1,15 @@
package middleware package middleware
import ( import (
"net/http" "fmt"
"strings"
"github.com/gin-contrib/sessions" "github.com/gin-contrib/sessions"
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
"github.com/songquanpeng/one-api/common"
"github.com/songquanpeng/one-api/common/blacklist" "github.com/songquanpeng/one-api/common/blacklist"
"github.com/songquanpeng/one-api/common/logger" "github.com/songquanpeng/one-api/common/logger"
"github.com/songquanpeng/one-api/common/network"
"github.com/songquanpeng/one-api/model" "github.com/songquanpeng/one-api/model"
"net/http"
"strings"
) )
func authHelper(c *gin.Context, minRole int) { func authHelper(c *gin.Context, minRole int) {
@ -47,7 +47,7 @@ func authHelper(c *gin.Context, minRole int) {
return 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{ c.JSON(http.StatusOK, gin.H{
"success": false, "success": false,
"message": "用户已被封禁", "message": "用户已被封禁",
@ -74,24 +74,25 @@ func authHelper(c *gin.Context, minRole int) {
func UserAuth() func(c *gin.Context) { func UserAuth() func(c *gin.Context) {
return func(c *gin.Context) { return func(c *gin.Context) {
authHelper(c, common.RoleCommonUser) authHelper(c, model.RoleCommonUser)
} }
} }
func AdminAuth() func(c *gin.Context) { func AdminAuth() func(c *gin.Context) {
return func(c *gin.Context) { return func(c *gin.Context) {
authHelper(c, common.RoleAdminUser) authHelper(c, model.RoleAdminUser)
} }
} }
func RootAuth() func(c *gin.Context) { func RootAuth() func(c *gin.Context) {
return func(c *gin.Context) { return func(c *gin.Context) {
authHelper(c, common.RoleRootUser) authHelper(c, model.RoleRootUser)
} }
} }
func TokenAuth() func(c *gin.Context) { func TokenAuth() func(c *gin.Context) {
return func(c *gin.Context) { return func(c *gin.Context) {
ctx := c.Request.Context()
key := c.Request.Header.Get("Authorization") key := c.Request.Header.Get("Authorization")
key = strings.TrimPrefix(key, "Bearer ") key = strings.TrimPrefix(key, "Bearer ")
key = strings.TrimPrefix(strings.TrimPrefix(key, "sk-"), "laisky-") key = strings.TrimPrefix(strings.TrimPrefix(key, "sk-"), "laisky-")
@ -102,6 +103,12 @@ func TokenAuth() func(c *gin.Context) {
abortWithMessage(c, http.StatusUnauthorized, err.Error()) abortWithMessage(c, http.StatusUnauthorized, err.Error())
return 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) userEnabled, err := model.CacheIsUserEnabled(token.UserId)
if err != nil { if err != nil {
abortWithMessage(c, http.StatusInternalServerError, err.Error()) abortWithMessage(c, http.StatusInternalServerError, err.Error())
@ -111,6 +118,19 @@ func TokenAuth() func(c *gin.Context) {
abortWithMessage(c, http.StatusForbidden, "用户已被封禁") abortWithMessage(c, http.StatusForbidden, "用户已被封禁")
return return
} }
requestModel, err := getRequestModel(c)
if err != nil && !strings.HasPrefix(c.Request.URL.Path, "/v1/models") {
abortWithMessage(c, http.StatusBadRequest, err.Error())
return
}
c.Set("request_model", requestModel)
if token.Models != nil && *token.Models != "" {
c.Set("available_models", *token.Models)
if requestModel != "" && !isModelInList(requestModel, *token.Models) {
abortWithMessage(c, http.StatusForbidden, fmt.Sprintf("该令牌无权使用模型:%s", requestModel))
return
}
}
c.Set("id", token.UserId) c.Set("id", token.UserId)
c.Set("token_id", token.Id) c.Set("token_id", token.Id)
c.Set("token_name", token.Name) c.Set("token_name", token.Name)

View File

@ -2,14 +2,16 @@ package middleware
import ( import (
"fmt" "fmt"
"github.com/songquanpeng/one-api/common"
"github.com/songquanpeng/one-api/common/logger"
"github.com/songquanpeng/one-api/model"
"net/http" "net/http"
"strconv" "strconv"
"strings" "strings"
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
"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/billing/ratio"
"github.com/songquanpeng/one-api/relay/channeltype"
) )
type ModelRequest struct { type ModelRequest struct {
@ -35,42 +37,16 @@ func Distribute() func(c *gin.Context) {
abortWithMessage(c, http.StatusBadRequest, "无效的渠道 Id") abortWithMessage(c, http.StatusBadRequest, "无效的渠道 Id")
return return
} }
if channel.Status != common.ChannelStatusEnabled { if channel.Status != model.ChannelStatusEnabled {
abortWithMessage(c, http.StatusForbidden, "该渠道已被禁用") abortWithMessage(c, http.StatusForbidden, "该渠道已被禁用")
return return
} }
} else { } else {
// Select a channel for the user requestModel := c.GetString("request_model")
var modelRequest ModelRequest var err error
err := common.UnmarshalBodyReusable(c, &modelRequest) channel, err = model.CacheGetRandomSatisfiedChannel(userGroup, requestModel, false)
if err != nil { if err != nil {
abortWithMessage(c, http.StatusBadRequest, fmt.Sprintf("无效的请求: %+v", err)) message := fmt.Sprintf("当前分组 %s 下对于模型 %s 无可用渠道", userGroup, requestModel)
return
}
if strings.HasPrefix(c.Request.URL.Path, "/v1/moderations") {
if modelRequest.Model == "" {
modelRequest.Model = "text-moderation-stable"
}
}
if strings.HasSuffix(c.Request.URL.Path, "embeddings") {
if modelRequest.Model == "" {
modelRequest.Model = c.Param("model")
}
}
if strings.HasPrefix(c.Request.URL.Path, "/v1/images/generations") {
if modelRequest.Model == "" {
modelRequest.Model = "dall-e-2"
}
}
if strings.HasPrefix(c.Request.URL.Path, "/v1/audio/transcriptions") || strings.HasPrefix(c.Request.URL.Path, "/v1/audio/translations") {
if modelRequest.Model == "" {
modelRequest.Model = "whisper-1"
}
}
requestModel = modelRequest.Model
channel, err = model.CacheGetRandomSatisfiedChannel(userGroup, modelRequest.Model, false)
if err != nil {
message := fmt.Sprintf("当前分组 %s 下对于模型 %s 无可用渠道", userGroup, modelRequest.Model)
if channel != nil { if channel != nil {
logger.SysError(fmt.Sprintf("渠道不存在:%d", channel.Id)) logger.SysError(fmt.Sprintf("渠道不存在:%d", channel.Id))
message = "数据库一致性已被破坏,请联系管理员" message = "数据库一致性已被破坏,请联系管理员"
@ -85,17 +61,18 @@ func Distribute() func(c *gin.Context) {
} }
func SetupContextForSelectedChannel(c *gin.Context, channel *model.Channel, modelName string) { func SetupContextForSelectedChannel(c *gin.Context, channel *model.Channel, modelName string) {
// one channel could relates to multiple groups,
// and each groud has individual ratio,
// set minimal group ratio as channel_ratio // set minimal group ratio as channel_ratio
var minimalRatio float64 = -1 var minimalRatio float64 = -1
for _, grp := range strings.Split(channel.Group, ",") { for _, grp := range strings.Split(channel.Group, ",") {
v := common.GetGroupRatio(grp) v := ratio.GetGroupRatio(grp)
if minimalRatio < 0 || v < minimalRatio { if minimalRatio < 0 || v < minimalRatio {
minimalRatio = v minimalRatio = v
} }
} }
logger.Info(c.Request.Context(), fmt.Sprintf("set channel %s ratio to %f", channel.Name, minimalRatio)) logger.Info(c.Request.Context(), fmt.Sprintf("set channel %s ratio to %f", channel.Name, minimalRatio))
c.Set("channel_ratio", minimalRatio) c.Set("channel_ratio", minimalRatio)
c.Set("channel", channel.Type) c.Set("channel", channel.Type)
c.Set("channel_id", channel.Id) c.Set("channel_id", channel.Id)
c.Set("channel_name", channel.Name) c.Set("channel_name", channel.Name)
@ -105,19 +82,19 @@ func SetupContextForSelectedChannel(c *gin.Context, channel *model.Channel, mode
c.Set("base_url", channel.GetBaseURL()) c.Set("base_url", channel.GetBaseURL())
// this is for backward compatibility // this is for backward compatibility
switch channel.Type { switch channel.Type {
case common.ChannelTypeAzure: case channeltype.Azure:
c.Set(common.ConfigKeyAPIVersion, channel.Other) c.Set(config.KeyAPIVersion, channel.Other)
case common.ChannelTypeXunfei: case channeltype.Xunfei:
c.Set(common.ConfigKeyAPIVersion, channel.Other) c.Set(config.KeyAPIVersion, channel.Other)
case common.ChannelTypeGemini: case channeltype.Gemini:
c.Set(common.ConfigKeyAPIVersion, channel.Other) c.Set(config.KeyAPIVersion, channel.Other)
case common.ChannelTypeAIProxyLibrary: case channeltype.AIProxyLibrary:
c.Set(common.ConfigKeyLibraryID, channel.Other) c.Set(config.KeyLibraryID, channel.Other)
case common.ChannelTypeAli: case channeltype.Ali:
c.Set(common.ConfigKeyPlugin, channel.Other) c.Set(config.KeyPlugin, channel.Other)
} }
cfg, _ := channel.LoadConfig() cfg, _ := channel.LoadConfig()
for k, v := range cfg { for k, v := range cfg {
c.Set(common.ConfigKeyPrefix+k, v) c.Set(config.KeyPrefix+k, v)
} }
} }

View File

@ -1,9 +1,12 @@
package middleware package middleware
import ( import (
"fmt"
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
"github.com/songquanpeng/one-api/common"
"github.com/songquanpeng/one-api/common/helper" "github.com/songquanpeng/one-api/common/helper"
"github.com/songquanpeng/one-api/common/logger" "github.com/songquanpeng/one-api/common/logger"
"strings"
) )
func abortWithMessage(c *gin.Context, statusCode int, message string) { func abortWithMessage(c *gin.Context, statusCode int, message string) {
@ -16,3 +19,42 @@ func abortWithMessage(c *gin.Context, statusCode int, message string) {
c.Abort() c.Abort()
logger.Error(c.Request.Context(), message) logger.Error(c.Request.Context(), message)
} }
func getRequestModel(c *gin.Context) (string, error) {
var modelRequest ModelRequest
err := common.UnmarshalBodyReusable(c, &modelRequest)
if err != nil {
return "", fmt.Errorf("common.UnmarshalBodyReusable failed: %w", err)
}
if strings.HasPrefix(c.Request.URL.Path, "/v1/moderations") {
if modelRequest.Model == "" {
modelRequest.Model = "text-moderation-stable"
}
}
if strings.HasSuffix(c.Request.URL.Path, "embeddings") {
if modelRequest.Model == "" {
modelRequest.Model = c.Param("model")
}
}
if strings.HasPrefix(c.Request.URL.Path, "/v1/images/generations") {
if modelRequest.Model == "" {
modelRequest.Model = "dall-e-2"
}
}
if strings.HasPrefix(c.Request.URL.Path, "/v1/audio/transcriptions") || strings.HasPrefix(c.Request.URL.Path, "/v1/audio/translations") {
if modelRequest.Model == "" {
modelRequest.Model = "whisper-1"
}
}
return modelRequest.Model, nil
}
func isModelInList(modelName string, models string) bool {
modelList := strings.Split(models, ",")
for _, model := range modelList {
if modelName == model {
return true
}
}
return false
}

View File

@ -1,8 +1,10 @@
package model package model
import ( import (
"context"
"github.com/songquanpeng/one-api/common" "github.com/songquanpeng/one-api/common"
"gorm.io/gorm" "gorm.io/gorm"
"sort"
"strings" "strings"
) )
@ -55,7 +57,7 @@ func (channel *Channel) AddAbilities() error {
Group: group, Group: group,
Model: model, Model: model,
ChannelId: channel.Id, ChannelId: channel.Id,
Enabled: channel.Status == common.ChannelStatusEnabled, Enabled: channel.Status == ChannelStatusEnabled,
Priority: channel.Priority, Priority: channel.Priority,
} }
abilities = append(abilities, ability) abilities = append(abilities, ability)
@ -88,3 +90,19 @@ func (channel *Channel) UpdateAbilities() error {
func UpdateAbilityStatus(channelId int, status bool) error { func UpdateAbilityStatus(channelId int, status bool) error {
return DB.Model(&Ability{}).Where("channel_id = ?", channelId).Select("enabled").Update("enabled", status).Error return DB.Model(&Ability{}).Where("channel_id = ?", channelId).Select("enabled").Update("enabled", status).Error
} }
func GetGroupModels(ctx context.Context, group string) ([]string, error) {
groupCol := "`group`"
trueVal := "1"
if common.UsingPostgreSQL {
groupCol = `"group"`
trueVal = "true"
}
var models []string
err := DB.Model(&Ability{}).Distinct("model").Where(groupCol+" = ? and enabled = "+trueVal, group).Pluck("model", &models).Error
if err != nil {
return nil, err
}
sort.Strings(models)
return models, err
}

View File

@ -8,6 +8,7 @@ import (
"github.com/songquanpeng/one-api/common" "github.com/songquanpeng/one-api/common"
"github.com/songquanpeng/one-api/common/config" "github.com/songquanpeng/one-api/common/config"
"github.com/songquanpeng/one-api/common/logger" "github.com/songquanpeng/one-api/common/logger"
"github.com/songquanpeng/one-api/common/random"
"math/rand" "math/rand"
"sort" "sort"
"strconv" "strconv"
@ -21,6 +22,7 @@ var (
UserId2GroupCacheSeconds = config.SyncFrequency UserId2GroupCacheSeconds = config.SyncFrequency
UserId2QuotaCacheSeconds = config.SyncFrequency UserId2QuotaCacheSeconds = config.SyncFrequency
UserId2StatusCacheSeconds = config.SyncFrequency UserId2StatusCacheSeconds = config.SyncFrequency
GroupModelsCacheSeconds = config.SyncFrequency
) )
func CacheGetTokenByKey(key string) (*Token, error) { func CacheGetTokenByKey(key string) (*Token, error) {
@ -147,13 +149,32 @@ func CacheIsUserEnabled(userId int) (bool, error) {
return userEnabled, err return userEnabled, err
} }
func CacheGetGroupModels(ctx context.Context, group string) ([]string, error) {
if !common.RedisEnabled {
return GetGroupModels(ctx, group)
}
modelsStr, err := common.RedisGet(fmt.Sprintf("group_models:%s", group))
if err == nil {
return strings.Split(modelsStr, ","), nil
}
models, err := GetGroupModels(ctx, group)
if err != nil {
return nil, err
}
err = common.RedisSet(fmt.Sprintf("group_models:%s", group), strings.Join(models, ","), time.Duration(GroupModelsCacheSeconds)*time.Second)
if err != nil {
logger.SysError("Redis set group models error: " + err.Error())
}
return models, nil
}
var group2model2channels map[string]map[string][]*Channel var group2model2channels map[string]map[string][]*Channel
var channelSyncLock sync.RWMutex var channelSyncLock sync.RWMutex
func InitChannelCache() { func InitChannelCache() {
newChannelId2channel := make(map[int]*Channel) newChannelId2channel := make(map[int]*Channel)
var channels []*Channel var channels []*Channel
DB.Where("status = ?", common.ChannelStatusEnabled).Find(&channels) DB.Where("status = ?", ChannelStatusEnabled).Find(&channels)
for _, channel := range channels { for _, channel := range channels {
newChannelId2channel[channel.Id] = channel newChannelId2channel[channel.Id] = channel
} }
@ -228,7 +249,7 @@ func CacheGetRandomSatisfiedChannel(group string, model string, ignoreFirstPrior
idx := rand.Intn(endIdx) idx := rand.Intn(endIdx)
if ignoreFirstPriority { if ignoreFirstPriority {
if endIdx < len(channels) { // which means there are more than one priority 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 return channels[idx], nil

View File

@ -3,13 +3,19 @@ package model
import ( import (
"encoding/json" "encoding/json"
"fmt" "fmt"
"github.com/songquanpeng/one-api/common"
"github.com/songquanpeng/one-api/common/config" "github.com/songquanpeng/one-api/common/config"
"github.com/songquanpeng/one-api/common/helper" "github.com/songquanpeng/one-api/common/helper"
"github.com/songquanpeng/one-api/common/logger" "github.com/songquanpeng/one-api/common/logger"
"gorm.io/gorm" "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 { type Channel struct {
Id int `json:"id"` Id int `json:"id"`
Type int `json:"type" gorm:"default:0"` Type int `json:"type" gorm:"default:0"`
@ -39,7 +45,7 @@ func GetAllChannels(startIdx int, num int, scope string) ([]*Channel, error) {
case "all": case "all":
err = DB.Order("id desc").Find(&channels).Error err = DB.Order("id desc").Find(&channels).Error
case "disabled": 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: default:
err = DB.Order("id desc").Limit(num).Offset(startIdx).Omit("key").Find(&channels).Error 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) { func UpdateChannelStatusById(id int, status int) {
err := UpdateAbilityStatus(id, status == common.ChannelStatusEnabled) err := UpdateAbilityStatus(id, status == ChannelStatusEnabled)
if err != nil { if err != nil {
logger.SysError("failed to update ability status: " + err.Error()) logger.SysError("failed to update ability status: " + err.Error())
} }
@ -199,6 +205,6 @@ func DeleteChannelByStatus(status int64) (int64, error) {
} }
func DeleteDisabledChannel() (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 return result.RowsAffected, result.Error
} }

View File

@ -7,7 +7,6 @@ import (
"github.com/songquanpeng/one-api/common/config" "github.com/songquanpeng/one-api/common/config"
"github.com/songquanpeng/one-api/common/helper" "github.com/songquanpeng/one-api/common/helper"
"github.com/songquanpeng/one-api/common/logger" "github.com/songquanpeng/one-api/common/logger"
"gorm.io/gorm" "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) { 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)) 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 { if !config.LogConsumeEnabled {

View File

@ -12,6 +12,7 @@ import (
"github.com/songquanpeng/one-api/common/env" "github.com/songquanpeng/one-api/common/env"
"github.com/songquanpeng/one-api/common/helper" "github.com/songquanpeng/one-api/common/helper"
"github.com/songquanpeng/one-api/common/logger" "github.com/songquanpeng/one-api/common/logger"
"github.com/songquanpeng/one-api/common/random"
"gorm.io/driver/mysql" "gorm.io/driver/mysql"
"gorm.io/driver/postgres" "gorm.io/driver/postgres"
"gorm.io/driver/sqlite" "gorm.io/driver/sqlite"
@ -33,10 +34,10 @@ func CreateRootAccountIfNeed() error {
rootUser := User{ rootUser := User{
Username: "root", Username: "root",
Password: hashedPassword, Password: hashedPassword,
Role: common.RoleRootUser, Role: RoleRootUser,
Status: common.UserStatusEnabled, Status: UserStatusEnabled,
DisplayName: "Root User", DisplayName: "Root User",
AccessToken: helper.GetUUID(), AccessToken: random.GetUUID(),
Quota: 500000000000000, Quota: 500000000000000,
} }
DB.Create(&rootUser) DB.Create(&rootUser)
@ -46,7 +47,7 @@ func CreateRootAccountIfNeed() error {
Id: 1, Id: 1,
UserId: rootUser.Id, UserId: rootUser.Id,
Key: config.InitialRootToken, Key: config.InitialRootToken,
Status: common.TokenStatusEnabled, Status: TokenStatusEnabled,
Name: "Initial Root Token", Name: "Initial Root Token",
CreatedTime: helper.GetTimestamp(), CreatedTime: helper.GetTimestamp(),
AccessedTime: helper.GetTimestamp(), AccessedTime: helper.GetTimestamp(),

View File

@ -1,9 +1,9 @@
package model package model
import ( import (
"github.com/songquanpeng/one-api/common"
"github.com/songquanpeng/one-api/common/config" "github.com/songquanpeng/one-api/common/config"
"github.com/songquanpeng/one-api/common/logger" "github.com/songquanpeng/one-api/common/logger"
billingratio "github.com/songquanpeng/one-api/relay/billing/ratio"
"strconv" "strconv"
"strings" "strings"
"time" "time"
@ -66,9 +66,9 @@ func InitOptionMap() {
config.OptionMap["QuotaForInvitee"] = strconv.FormatInt(config.QuotaForInvitee, 10) config.OptionMap["QuotaForInvitee"] = strconv.FormatInt(config.QuotaForInvitee, 10)
config.OptionMap["QuotaRemindThreshold"] = strconv.FormatInt(config.QuotaRemindThreshold, 10) config.OptionMap["QuotaRemindThreshold"] = strconv.FormatInt(config.QuotaRemindThreshold, 10)
config.OptionMap["PreConsumedQuota"] = strconv.FormatInt(config.PreConsumedQuota, 10) config.OptionMap["PreConsumedQuota"] = strconv.FormatInt(config.PreConsumedQuota, 10)
config.OptionMap["ModelRatio"] = common.ModelRatio2JSONString() config.OptionMap["ModelRatio"] = billingratio.ModelRatio2JSONString()
config.OptionMap["GroupRatio"] = common.GroupRatio2JSONString() config.OptionMap["GroupRatio"] = billingratio.GroupRatio2JSONString()
config.OptionMap["CompletionRatio"] = common.CompletionRatio2JSONString() config.OptionMap["CompletionRatio"] = billingratio.CompletionRatio2JSONString()
config.OptionMap["TopUpLink"] = config.TopUpLink config.OptionMap["TopUpLink"] = config.TopUpLink
config.OptionMap["ChatLink"] = config.ChatLink config.OptionMap["ChatLink"] = config.ChatLink
config.OptionMap["QuotaPerUnit"] = strconv.FormatFloat(config.QuotaPerUnit, 'f', -1, 64) config.OptionMap["QuotaPerUnit"] = strconv.FormatFloat(config.QuotaPerUnit, 'f', -1, 64)
@ -82,7 +82,7 @@ func loadOptionsFromDatabase() {
options, _ := AllOption() options, _ := AllOption()
for _, option := range options { for _, option := range options {
if option.Key == "ModelRatio" { if option.Key == "ModelRatio" {
option.Value = common.AddNewMissingRatio(option.Value) option.Value = billingratio.AddNewMissingRatio(option.Value)
} }
err := updateOptionMap(option.Key, option.Value) err := updateOptionMap(option.Key, option.Value)
if err != nil { if err != nil {
@ -172,6 +172,10 @@ func updateOptionMap(key string, value string) (err error) {
config.GitHubClientId = value config.GitHubClientId = value
case "GitHubClientSecret": case "GitHubClientSecret":
config.GitHubClientSecret = value config.GitHubClientSecret = value
case "LarkClientId":
config.LarkClientId = value
case "LarkClientSecret":
config.LarkClientSecret = value
case "Footer": case "Footer":
config.Footer = value config.Footer = value
case "SystemName": case "SystemName":
@ -205,11 +209,11 @@ func updateOptionMap(key string, value string) (err error) {
case "RetryTimes": case "RetryTimes":
config.RetryTimes, _ = strconv.Atoi(value) config.RetryTimes, _ = strconv.Atoi(value)
case "ModelRatio": case "ModelRatio":
err = common.UpdateModelRatioByJSONString(value) err = billingratio.UpdateModelRatioByJSONString(value)
case "GroupRatio": case "GroupRatio":
err = common.UpdateGroupRatioByJSONString(value) err = billingratio.UpdateGroupRatioByJSONString(value)
case "CompletionRatio": case "CompletionRatio":
err = common.UpdateCompletionRatioByJSONString(value) err = billingratio.UpdateCompletionRatioByJSONString(value)
case "TopUpLink": case "TopUpLink":
config.TopUpLink = value config.TopUpLink = value
case "ChatLink": case "ChatLink":

View File

@ -8,6 +8,12 @@ import (
"gorm.io/gorm" "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 { type Redemption struct {
Id int `json:"id"` Id int `json:"id"`
UserId int `json:"user_id"` UserId int `json:"user_id"`
@ -61,7 +67,7 @@ func Redeem(key string, userId int) (quota int64, err error) {
if err != nil { if err != nil {
return errors.New("无效的兑换码") return errors.New("无效的兑换码")
} }
if redemption.Status != common.RedemptionCodeStatusEnabled { if redemption.Status != RedemptionCodeStatusEnabled {
return errors.New("该兑换码已被使用") return errors.New("该兑换码已被使用")
} }
err = tx.Model(&User{}).Where("id = ?", userId).Update("quota", gorm.Expr("quota + ?", redemption.Quota)).Error 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 return err
} }
redemption.RedeemedTime = helper.GetTimestamp() redemption.RedeemedTime = helper.GetTimestamp()
redemption.Status = common.RedemptionCodeStatusUsed redemption.Status = RedemptionCodeStatusUsed
err = tx.Save(redemption).Error err = tx.Save(redemption).Error
return err return err
}) })

View File

@ -12,6 +12,13 @@ import (
"gorm.io/gorm" "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 { type Token struct {
Id int `json:"id"` Id int `json:"id"`
UserId int `json:"user_id"` UserId int `json:"user_id"`
@ -24,6 +31,8 @@ type Token struct {
RemainQuota int64 `json:"remain_quota" gorm:"bigint;default:0"` RemainQuota int64 `json:"remain_quota" gorm:"bigint;default:0"`
UnlimitedQuota bool `json:"unlimited_quota" gorm:"default:false"` UnlimitedQuota bool `json:"unlimited_quota" gorm:"default:false"`
UsedQuota int64 `json:"used_quota" gorm:"bigint;default:0"` // used quota UsedQuota int64 `json:"used_quota" gorm:"bigint;default:0"` // used quota
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) { func GetAllUserTokens(userId int, startIdx int, num int, order string) ([]*Token, error) {
@ -62,18 +71,17 @@ func ValidateUserToken(key string) (token *Token, err error) {
return nil, errors.Wrap(err, "failed to get token by key") return nil, errors.Wrap(err, "failed to get token by key")
} }
if token.Status == TokenStatusExhausted {
if token.Status == common.TokenStatusExhausted { return nil, fmt.Errorf("令牌 %s#%d额度已用尽", token.Name, token.Id)
return nil, errors.New("该令牌额度已用尽") } else if token.Status == TokenStatusExpired {
} else if token.Status == common.TokenStatusExpired {
return nil, errors.New("该令牌已过期") return nil, errors.New("该令牌已过期")
} }
if token.Status != common.TokenStatusEnabled { if token.Status != TokenStatusEnabled {
return nil, errors.New("该令牌状态不可用") return nil, errors.New("该令牌状态不可用")
} }
if token.ExpiredTime != -1 && token.ExpiredTime < helper.GetTimestamp() { if token.ExpiredTime != -1 && token.ExpiredTime < helper.GetTimestamp() {
if !common.RedisEnabled { if !common.RedisEnabled {
token.Status = common.TokenStatusExpired token.Status = TokenStatusExpired
err := token.SelectUpdate() err := token.SelectUpdate()
if err != nil { if err != nil {
logger.SysError("failed to update token status" + err.Error()) logger.SysError("failed to update token status" + err.Error())
@ -84,7 +92,7 @@ func ValidateUserToken(key string) (token *Token, err error) {
if !token.UnlimitedQuota && token.RemainQuota <= 0 { if !token.UnlimitedQuota && token.RemainQuota <= 0 {
if !common.RedisEnabled { if !common.RedisEnabled {
// in this case, we can make sure the token is exhausted // in this case, we can make sure the token is exhausted
token.Status = common.TokenStatusExhausted token.Status = TokenStatusExhausted
err := token.SelectUpdate() err := token.SelectUpdate()
if err != nil { if err != nil {
logger.SysError("failed to update token status" + err.Error()) logger.SysError("failed to update token status" + err.Error())
@ -124,7 +132,7 @@ func (token *Token) Insert() error {
// Update Make sure your token's fields is completed, because this will update non-zero values // Update Make sure your token's fields is completed, because this will update non-zero values
func (token *Token) Update() error { func (token *Token) Update() error {
var err error var err error
err = DB.Model(token).Select("name", "status", "expired_time", "remain_quota", "unlimited_quota").Updates(token).Error err = DB.Model(token).Select("name", "status", "expired_time", "remain_quota", "unlimited_quota", "models", "subnet").Updates(token).Error
return err return err
} }

View File

@ -8,11 +8,24 @@ import (
"github.com/songquanpeng/one-api/common" "github.com/songquanpeng/one-api/common"
"github.com/songquanpeng/one-api/common/blacklist" "github.com/songquanpeng/one-api/common/blacklist"
"github.com/songquanpeng/one-api/common/config" "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/logger"
"github.com/songquanpeng/one-api/common/random"
"gorm.io/gorm" "gorm.io/gorm"
) )
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. // 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! // Otherwise, the sensitive information will be saved on local storage in plain text!
type User struct { type User struct {
@ -25,6 +38,7 @@ type User struct {
Email string `json:"email" gorm:"index" validate:"max=50"` Email string `json:"email" gorm:"index" validate:"max=50"`
GitHubId string `json:"github_id" gorm:"column:github_id;index"` GitHubId string `json:"github_id" gorm:"column:github_id;index"`
WeChatId string `json:"wechat_id" gorm:"column:wechat_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! 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 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"` Quota int64 `json:"quota" gorm:"bigint;default:0"`
@ -42,7 +56,7 @@ func GetMaxUserId() int {
} }
func GetAllUsers(startIdx int, num int, order string) (users []*User, err error) { func GetAllUsers(startIdx int, num int, order string) (users []*User, err error) {
query := DB.Limit(num).Offset(startIdx).Omit("password").Where("status != ?", common.UserStatusDeleted) query := DB.Limit(num).Offset(startIdx).Omit("password").Where("status != ?", UserStatusDeleted)
switch order { switch order {
case "quota": case "quota":
@ -108,8 +122,8 @@ func (user *User) Insert(inviterId int) error {
} }
} }
user.Quota = config.QuotaForNewUser user.Quota = config.QuotaForNewUser
user.AccessToken = helper.GetUUID() user.AccessToken = random.GetUUID()
user.AffCode = helper.GetRandomString(4) user.AffCode = random.GetRandomString(4)
result := DB.Create(user) result := DB.Create(user)
if result.Error != nil { if result.Error != nil {
return result.Error return result.Error
@ -138,9 +152,9 @@ func (user *User) Update(updatePassword bool) error {
return err return err
} }
} }
if user.Status == common.UserStatusDisabled { if user.Status == UserStatusDisabled {
blacklist.BanUser(user.Id) blacklist.BanUser(user.Id)
} else if user.Status == common.UserStatusEnabled { } else if user.Status == UserStatusEnabled {
blacklist.UnbanUser(user.Id) blacklist.UnbanUser(user.Id)
} }
err = DB.Model(user).Updates(user).Error err = DB.Model(user).Updates(user).Error
@ -152,8 +166,8 @@ func (user *User) Delete() error {
return errors.New("id 为空!") return errors.New("id 为空!")
} }
blacklist.BanUser(user.Id) blacklist.BanUser(user.Id)
user.Username = fmt.Sprintf("deleted_%s", helper.GetUUID()) user.Username = fmt.Sprintf("deleted_%s", random.GetUUID())
user.Status = common.UserStatusDeleted user.Status = UserStatusDeleted
err := DB.Model(user).Updates(user).Error err := DB.Model(user).Updates(user).Error
return err return err
} }
@ -177,7 +191,7 @@ func (user *User) ValidateAndFill() (err error) {
} }
} }
okay := common.ValidatePasswordAndHash(password, user.Password) okay := common.ValidatePasswordAndHash(password, user.Password)
if !okay || user.Status != common.UserStatusEnabled { if !okay || user.Status != UserStatusEnabled {
return errors.New("用户名或密码错误,或用户已被封禁") return errors.New("用户名或密码错误,或用户已被封禁")
} }
return nil return nil
@ -207,6 +221,14 @@ func (user *User) FillUserByGitHubId() error {
return nil 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 { func (user *User) FillUserByWeChatId() error {
if user.WeChatId == "" { if user.WeChatId == "" {
return errors.New("WeChat id 为空!") return errors.New("WeChat id 为空!")
@ -235,6 +257,10 @@ func IsGitHubIdAlreadyTaken(githubId string) bool {
return DB.Where("github_id = ?", githubId).Find(&User{}).RowsAffected == 1 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 { func IsUsernameAlreadyTaken(username string) bool {
return DB.Where("username = ?", username).Find(&User{}).RowsAffected == 1 return DB.Where("username = ?", username).Find(&User{}).RowsAffected == 1
} }
@ -261,7 +287,7 @@ func IsAdmin(userId int) bool {
logger.SysError("no such user " + err.Error()) logger.SysError("no such user " + err.Error())
return false return false
} }
return user.Role >= common.RoleAdminUser return user.Role >= RoleAdminUser
} }
func IsUserEnabled(userId int) (bool, error) { func IsUserEnabled(userId int) (bool, error) {
@ -273,7 +299,7 @@ func IsUserEnabled(userId int) (bool, error) {
if err != nil { if err != nil {
return false, err return false, err
} }
return user.Status == common.UserStatusEnabled, nil return user.Status == UserStatusEnabled, nil
} }
func ValidateAccessToken(token string) (user *User) { func ValidateAccessToken(token string) (user *User) {
@ -346,7 +372,7 @@ func decreaseUserQuota(id int, quota int64) (err error) {
} }
func GetRootUserEmail() (email string) { 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 return email
} }

View File

@ -2,7 +2,6 @@ package monitor
import ( import (
"fmt" "fmt"
"github.com/songquanpeng/one-api/common"
"github.com/songquanpeng/one-api/common/config" "github.com/songquanpeng/one-api/common/config"
"github.com/songquanpeng/one-api/common/logger" "github.com/songquanpeng/one-api/common/logger"
"github.com/songquanpeng/one-api/common/message" "github.com/songquanpeng/one-api/common/message"
@ -29,7 +28,7 @@ func notifyRootUser(subject string, content string) {
// DisableChannel disable & notify // DisableChannel disable & notify
func DisableChannel(channelId int, channelName string, reason string) { 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)) logger.SysLog(fmt.Sprintf("channel #%d has been disabled: %s", channelId, reason))
subject := fmt.Sprintf("渠道「%s」#%d已被禁用", channelName, channelId) subject := fmt.Sprintf("渠道「%s」#%d已被禁用", channelName, channelId)
content := fmt.Sprintf("渠道「%s」#%d已被禁用原因%s", channelName, channelId, reason) 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) { 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)) logger.SysLog(fmt.Sprintf("channel #%d has been disabled due to low success rate: %.2f", channelId, successRate*100))
subject := fmt.Sprintf("渠道 #%d 已被禁用", channelId) subject := fmt.Sprintf("渠道 #%d 已被禁用", channelId)
content := fmt.Sprintf("该渠道(#%d在最近 %d 次调用中成功率为 %.2f%%,低于阈值 %.2f%%,因此被系统自动禁用。", content := fmt.Sprintf("该渠道(#%d在最近 %d 次调用中成功率为 %.2f%%,低于阈值 %.2f%%,因此被系统自动禁用。",
@ -47,7 +46,7 @@ func MetricDisableChannel(channelId int, successRate float64) {
// EnableChannel enable & notify // EnableChannel enable & notify
func EnableChannel(channelId int, channelName string) { 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)) logger.SysLog(fmt.Sprintf("channel #%d has been enabled", channelId))
subject := fmt.Sprintf("渠道「%s」#%d已被启用", channelName, channelId) subject := fmt.Sprintf("渠道「%s」#%d已被启用", channelName, channelId)
content := fmt.Sprintf("渠道「%s」#%d已被启用", channelName, channelId) content := fmt.Sprintf("渠道「%s」#%d已被启用", channelName, channelId)

62
monitor/manage.go Normal file
View File

@ -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
}

41
relay/adaptor.go Normal file
View File

@ -0,0 +1,41 @@
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/anthropic"
"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/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
}

View File

@ -4,10 +4,10 @@ import (
"fmt" "fmt"
"github.com/Laisky/errors/v2" "github.com/Laisky/errors/v2"
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
"github.com/songquanpeng/one-api/common" "github.com/songquanpeng/one-api/common/config"
"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/model"
"github.com/songquanpeng/one-api/relay/util"
"io" "io"
"net/http" "net/http"
) )
@ -15,16 +15,16 @@ import (
type Adaptor struct { 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 return fmt.Sprintf("%s/api/library/ask", meta.BaseURL), 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 {
channel.SetupCommonRequestHeader(c, req, meta) adaptor.SetupCommonRequestHeader(c, req, meta)
req.Header.Set("Authorization", "Bearer "+meta.APIKey) req.Header.Set("Authorization", "Bearer "+meta.APIKey)
return nil 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") return nil, errors.New("request is nil")
} }
aiProxyLibraryRequest := ConvertRequest(*request) aiProxyLibraryRequest := ConvertRequest(*request)
aiProxyLibraryRequest.LibraryId = c.GetString(common.ConfigKeyLibraryID) aiProxyLibraryRequest.LibraryId = c.GetString(config.KeyLibraryID)
return aiProxyLibraryRequest, nil return aiProxyLibraryRequest, 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) {
return channel.DoRequestHelper(a, c, meta, requestBody) 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 { if meta.IsStream {
err, usage = StreamHandler(c, resp) err, usage = StreamHandler(c, resp)
} else { } else {

View File

@ -1,6 +1,6 @@
package aiproxy package aiproxy
import "github.com/songquanpeng/one-api/relay/channel/openai" import "github.com/songquanpeng/one-api/relay/adaptor/openai"
var ModelList = []string{""} var ModelList = []string{""}

View File

@ -13,7 +13,8 @@ import (
"github.com/songquanpeng/one-api/common" "github.com/songquanpeng/one-api/common"
"github.com/songquanpeng/one-api/common/helper" "github.com/songquanpeng/one-api/common/helper"
"github.com/songquanpeng/one-api/common/logger" "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/constant"
"github.com/songquanpeng/one-api/relay/model" "github.com/songquanpeng/one-api/relay/model"
) )
@ -54,7 +55,7 @@ func responseAIProxyLibrary2OpenAI(response *LibraryResponse) *openai.TextRespon
FinishReason: "stop", FinishReason: "stop",
} }
fullTextResponse := openai.TextResponse{ fullTextResponse := openai.TextResponse{
Id: fmt.Sprintf("chatcmpl-%s", helper.GetUUID()), Id: fmt.Sprintf("chatcmpl-%s", random.GetUUID()),
Object: "chat.completion", Object: "chat.completion",
Created: helper.GetTimestamp(), Created: helper.GetTimestamp(),
Choices: []openai.TextResponseChoice{choice}, Choices: []openai.TextResponseChoice{choice},
@ -67,7 +68,7 @@ func documentsAIProxyLibrary(documents []LibraryDocument) *openai.ChatCompletion
choice.Delta.Content = aiProxyDocuments2Markdown(documents) choice.Delta.Content = aiProxyDocuments2Markdown(documents)
choice.FinishReason = &constant.StopFinishReason choice.FinishReason = &constant.StopFinishReason
return &openai.ChatCompletionsStreamResponse{ return &openai.ChatCompletionsStreamResponse{
Id: fmt.Sprintf("chatcmpl-%s", helper.GetUUID()), Id: fmt.Sprintf("chatcmpl-%s", random.GetUUID()),
Object: "chat.completion.chunk", Object: "chat.completion.chunk",
Created: helper.GetTimestamp(), Created: helper.GetTimestamp(),
Model: "", Model: "",
@ -79,7 +80,7 @@ func streamResponseAIProxyLibrary2OpenAI(response *LibraryStreamResponse) *opena
var choice openai.ChatCompletionsStreamResponseChoice var choice openai.ChatCompletionsStreamResponseChoice
choice.Delta.Content = response.Content choice.Delta.Content = response.Content
return &openai.ChatCompletionsStreamResponse{ return &openai.ChatCompletionsStreamResponse{
Id: fmt.Sprintf("chatcmpl-%s", helper.GetUUID()), Id: fmt.Sprintf("chatcmpl-%s", random.GetUUID()),
Object: "chat.completion.chunk", Object: "chat.completion.chunk",
Created: helper.GetTimestamp(), Created: helper.GetTimestamp(),
Model: response.Model, Model: response.Model,

View File

@ -0,0 +1,105 @@
package ali
// import (
// "github.com/Laisky/errors/v2"
// "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"
// }

View File

@ -3,4 +3,5 @@ package ali
var ModelList = []string{ var ModelList = []string{
"qwen-turbo", "qwen-plus", "qwen-max", "qwen-max-longcontext", "qwen-turbo", "qwen-plus", "qwen-max", "qwen-max-longcontext",
"text-embedding-v1", "text-embedding-v1",
"ali-stable-diffusion-xl", "ali-stable-diffusion-v1.5", "wanx-v1",
} }

192
relay/adaptor/ali/image.go Normal file
View File

@ -0,0 +1,192 @@
package ali
import (
"encoding/base64"
"encoding/json"
"fmt"
"github.com/Laisky/errors/v2"
"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
}

154
relay/adaptor/ali/model.go Normal file
View File

@ -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
}

View File

@ -2,31 +2,31 @@ package anthropic
import ( import (
"fmt" "fmt"
"github.com/Laisky/errors/v2"
"io" "io"
"net/http" "net/http"
"github.com/Laisky/errors/v2"
"github.com/gin-gonic/gin" "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/model"
"github.com/songquanpeng/one-api/relay/util"
) )
type Adaptor struct { 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) { // https://docs.anthropic.com/claude/reference/messages_post
// https://docs.anthropic.com/claude/reference/messages_post // anthopic migrate to Message API
// anthopic migrate to Message API func (a *Adaptor) GetRequestURL(meta *meta.Meta) (string, error) {
return fmt.Sprintf("%s/v1/messages", meta.BaseURL), nil return fmt.Sprintf("%s/v1/messages", meta.BaseURL), 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 {
channel.SetupCommonRequestHeader(c, req, meta) adaptor.SetupCommonRequestHeader(c, req, meta)
req.Header.Set("x-api-key", meta.APIKey) req.Header.Set("x-api-key", meta.APIKey)
anthropicVersion := c.Request.Header.Get("anthropic-version") anthropicVersion := c.Request.Header.Get("anthropic-version")
if anthropicVersion == "" { if anthropicVersion == "" {
@ -46,11 +46,18 @@ func (a *Adaptor) ConvertRequest(c *gin.Context, relayMode int, request *model.G
return ConvertRequest(*request), nil 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) {
return channel.DoRequestHelper(a, c, meta, requestBody) 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 { if meta.IsStream {
err, usage = StreamHandler(c, resp) err, usage = StreamHandler(c, resp)
} else { } else {

View File

@ -13,7 +13,7 @@ import (
"github.com/songquanpeng/one-api/common/helper" "github.com/songquanpeng/one-api/common/helper"
"github.com/songquanpeng/one-api/common/image" "github.com/songquanpeng/one-api/common/image"
"github.com/songquanpeng/one-api/common/logger" "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" "github.com/songquanpeng/one-api/relay/model"
) )

View File

@ -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
}

View File

@ -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",
}

View File

@ -1,4 +1,4 @@
package channel package adaptor
import ( import (
"io" "io"
@ -6,10 +6,11 @@ import (
"github.com/Laisky/errors/v2" "github.com/Laisky/errors/v2"
"github.com/gin-gonic/gin" "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"
) )
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("Content-Type", c.Request.Header.Get("Content-Type"))
req.Header.Set("Accept", c.Request.Header.Get("Accept")) req.Header.Set("Accept", c.Request.Header.Get("Accept"))
if meta.IsStream && 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) fullRequestURL, err := a.GetRequestURL(meta)
if err != nil { if err != nil {
return nil, errors.Wrap(err, "get request url failed") return nil, errors.Wrap(err, "get request url failed")
@ -43,7 +44,7 @@ func DoRequestHelper(a Adaptor, c *gin.Context, meta *util.RelayMeta, requestBod
} }
func DoRequest(c *gin.Context, req *http.Request) (*http.Response, error) { 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 { if err != nil {
return nil, err return nil, err
} }

View File

@ -8,21 +8,21 @@ import (
"github.com/Laisky/errors/v2" "github.com/Laisky/errors/v2"
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
"github.com/songquanpeng/one-api/common/helper" "github.com/songquanpeng/one-api/common/helper"
channelhelper "github.com/songquanpeng/one-api/relay/channel" channelhelper "github.com/songquanpeng/one-api/relay/adaptor"
"github.com/songquanpeng/one-api/relay/channel/openai" "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/model"
"github.com/songquanpeng/one-api/relay/util"
) )
type Adaptor struct { 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, "v1beta") version := helper.AssignOrDefault(meta.APIVersion, "v1")
action := "generateContent" action := "generateContent"
if meta.IsStream { if meta.IsStream {
action = "streamGenerateContent" action = "streamGenerateContent"
@ -30,7 +30,7 @@ func (a *Adaptor) GetRequestURL(meta *util.RelayMeta) (string, error) {
return fmt.Sprintf("%s/%s/models/%s:%s?key=%s", meta.BaseURL, version, meta.ActualModelName, action, meta.APIKey), nil return fmt.Sprintf("%s/%s/models/%s:%s?key=%s", meta.BaseURL, version, meta.ActualModelName, action, meta.APIKey), 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) channelhelper.SetupCommonRequestHeader(c, req, meta)
req.Header.Set("x-goog-api-key", meta.APIKey) req.Header.Set("x-goog-api-key", meta.APIKey)
req.URL.Query().Add("key", meta.APIKey) req.URL.Query().Add("key", meta.APIKey)
@ -44,11 +44,18 @@ func (a *Adaptor) ConvertRequest(c *gin.Context, relayMode int, request *model.G
return ConvertRequest(*request), nil 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) 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 { if meta.IsStream {
var responseText string var responseText string
err, responseText = StreamHandler(c, resp) err, responseText = StreamHandler(c, resp)

View File

@ -1,21 +1,24 @@
package gemini package gemini
import ( import (
"context" "bufio"
"encoding/json" "encoding/json"
"fmt" "fmt"
"io" "io"
"net/http" "net/http"
"strings"
"github.com/gin-gonic/gin"
"github.com/songquanpeng/one-api/common" "github.com/songquanpeng/one-api/common"
"github.com/songquanpeng/one-api/common/config" "github.com/songquanpeng/one-api/common/config"
"github.com/songquanpeng/one-api/common/helper" "github.com/songquanpeng/one-api/common/helper"
"github.com/songquanpeng/one-api/common/image" "github.com/songquanpeng/one-api/common/image"
"github.com/songquanpeng/one-api/common/logger" "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/constant"
"github.com/songquanpeng/one-api/relay/model" "github.com/songquanpeng/one-api/relay/model"
"github.com/gin-gonic/gin"
) )
// https://ai.google.dev/docs/gemini_api_overview?hl=zh-cn // https://ai.google.dev/docs/gemini_api_overview?hl=zh-cn
@ -82,13 +85,7 @@ func ConvertRequest(textRequest model.GeneralOpenAIRequest) *ChatRequest {
if imageNum > VisionMaxImageNum { if imageNum > VisionMaxImageNum {
continue continue
} }
mimeType, data, err := image.GetImageFromUrl(part.ImageURL.Url) mimeType, data, _ := image.GetImageFromUrl(part.ImageURL.Url)
if err != nil {
logger.Warn(context.TODO(),
fmt.Sprintf("get image from url %s got %+v", part.ImageURL.Url, err))
continue
}
parts = append(parts, Part{ parts = append(parts, Part{
InlineData: &InlineData{ InlineData: &InlineData{
MimeType: mimeType, MimeType: mimeType,
@ -97,9 +94,6 @@ func ConvertRequest(textRequest model.GeneralOpenAIRequest) *ChatRequest {
}) })
} }
} }
logger.Info(context.TODO(),
fmt.Sprintf("send %d messages to gemini with %d images", len(parts), imageNum))
content.Parts = parts content.Parts = parts
// there's no assistant role in gemini and API shall vomit if Role is not user or model // there's no assistant role in gemini and API shall vomit if Role is not user or model
@ -163,7 +157,7 @@ type ChatPromptFeedback struct {
func responseGeminiChat2OpenAI(response *ChatResponse) *openai.TextResponse { func responseGeminiChat2OpenAI(response *ChatResponse) *openai.TextResponse {
fullTextResponse := openai.TextResponse{ fullTextResponse := openai.TextResponse{
Id: fmt.Sprintf("chatcmpl-%s", helper.GetUUID()), Id: fmt.Sprintf("chatcmpl-%s", random.GetUUID()),
Object: "chat.completion", Object: "chat.completion",
Created: helper.GetTimestamp(), Created: helper.GetTimestamp(),
Choices: make([]openai.TextResponseChoice, 0, len(response.Candidates)), Choices: make([]openai.TextResponseChoice, 0, len(response.Candidates)),
@ -196,182 +190,73 @@ func streamResponseGeminiChat2OpenAI(geminiResponse *ChatResponse) *openai.ChatC
return &response return &response
} }
// [{
// "candidates": [
// {
// "content": {
// "parts": [
// {
// "text": "```go \n\n// Package ratelimit implements tokens bucket algorithm.\npackage rate"
// }
// ],
// "role": "model"
// },
// "finishReason": "STOP",
// "index": 0,
// "safetyRatings": [
// {
// "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
// "probability": "NEGLIGIBLE"
// },
// {
// "category": "HARM_CATEGORY_HATE_SPEECH",
// "probability": "NEGLIGIBLE"
// },
// {
// "category": "HARM_CATEGORY_HARASSMENT",
// "probability": "NEGLIGIBLE"
// },
// {
// "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
// "probability": "NEGLIGIBLE"
// }
// ]
// }
// ],
// "promptFeedback": {
// "safetyRatings": [
// {
// "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
// "probability": "NEGLIGIBLE"
// },
// {
// "category": "HARM_CATEGORY_HATE_SPEECH",
// "probability": "NEGLIGIBLE"
// },
// {
// "category": "HARM_CATEGORY_HARASSMENT",
// "probability": "NEGLIGIBLE"
// },
// {
// "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
// "probability": "NEGLIGIBLE"
// }
// ]
// }
// }]
type GeminiStreamResp struct {
Candidates []struct {
Content struct {
Parts []struct {
Text string `json:"text"`
} `json:"parts"`
Role string `json:"role"`
} `json:"content"`
FinishReason string `json:"finishReason"`
Index int64 `json:"index"`
} `json:"candidates"`
}
func StreamHandler(c *gin.Context, resp *http.Response) (*model.ErrorWithStatusCode, string) { func StreamHandler(c *gin.Context, resp *http.Response) (*model.ErrorWithStatusCode, string) {
responseText := "" responseText := ""
dataChan := make(chan string)
respBody, err := io.ReadAll(resp.Body) stopChan := make(chan bool)
if err != nil { scanner := bufio.NewScanner(resp.Body)
return openai.ErrorWrapper(err, "read upstream's body", http.StatusInternalServerError), responseText scanner.Split(func(data []byte, atEOF bool) (advance int, token []byte, err error) {
if atEOF && len(data) == 0 {
return 0, nil, nil
} }
if i := strings.Index(string(data), "\n"); i >= 0 {
var respData []GeminiStreamResp return i + 1, data[0:i], nil
if err = json.Unmarshal(respBody, &respData); err != nil {
return openai.ErrorWrapper(err, "unmarshal upstream's body", http.StatusInternalServerError), responseText
} }
if atEOF {
for _, chunk := range respData { return len(data), data, nil
for _, cad := range chunk.Candidates {
for _, part := range cad.Content.Parts {
responseText += part.Text
} }
return 0, nil, nil
})
go func() {
for scanner.Scan() {
data := scanner.Text()
data = strings.TrimSpace(data)
if !strings.HasPrefix(data, "\"text\": \"") {
continue
} }
data = strings.TrimPrefix(data, "\"text\": \"")
data = strings.TrimSuffix(data, "\"")
dataChan <- data
} }
stopChan <- true
}()
common.SetEventStreamHeaders(c)
c.Stream(func(w io.Writer) bool {
select {
case data := <-dataChan:
// this is used to prevent annoying \ related format bug
data = fmt.Sprintf("{\"content\": \"%s\"}", data)
type dummyStruct struct {
Content string `json:"content"`
}
var dummy dummyStruct
err := json.Unmarshal([]byte(data), &dummy)
responseText += dummy.Content
var choice openai.ChatCompletionsStreamResponseChoice var choice openai.ChatCompletionsStreamResponseChoice
choice.Delta.Content = responseText choice.Delta.Content = dummy.Content
resp2cli, err := json.Marshal(&openai.ChatCompletionsStreamResponse{ response := openai.ChatCompletionsStreamResponse{
Id: fmt.Sprintf("chatcmpl-%s", helper.GetUUID()), Id: fmt.Sprintf("chatcmpl-%s", random.GetUUID()),
Object: "chat.completion.chunk", Object: "chat.completion.chunk",
Created: helper.GetTimestamp(), Created: helper.GetTimestamp(),
Model: "gemini-pro", Model: "gemini-pro",
Choices: []openai.ChatCompletionsStreamResponseChoice{choice}, Choices: []openai.ChatCompletionsStreamResponseChoice{choice},
})
if err != nil {
return openai.ErrorWrapper(err, "marshal upstream's body", http.StatusInternalServerError), responseText
} }
jsonResponse, err := json.Marshal(response)
c.Render(-1, common.CustomEvent{Data: "data: " + string(resp2cli)}) if err != nil {
logger.SysError("error marshalling stream response: " + err.Error())
return true
}
c.Render(-1, common.CustomEvent{Data: "data: " + string(jsonResponse)})
return true
case <-stopChan:
c.Render(-1, common.CustomEvent{Data: "data: [DONE]"}) c.Render(-1, common.CustomEvent{Data: "data: [DONE]"})
return false
// dataChan := make(chan string) }
// stopChan := make(chan bool) })
// scanner := bufio.NewScanner(resp.Body) err := resp.Body.Close()
// scanner.Split(bufio.ScanLines) if err != nil {
// // scanner.Split(func(data []byte, atEOF bool) (advance int, token []byte, err error) {
// // if atEOF && len(data) == 0 {
// // return 0, nil, nil
// // }
// // if i := strings.Index(string(data), "\n"); i >= 0 {
// // return i + 1, data[0:i], nil
// // }
// // if atEOF {
// // return len(data), data, nil
// // }
// // return 0, nil, nil
// // })
// go func() {
// var content string
// for scanner.Scan() {
// line := strings.TrimSpace(scanner.Text())
// fmt.Printf("> gemini got line: %s\n", line)
// content += line
// // if !strings.HasPrefix(data, "\"text\": \"") {
// // continue
// // }
// // data = strings.TrimPrefix(data, "\"text\": \"")
// // data = strings.TrimSuffix(data, "\"")
// // dataChan <- data
// }
// dataChan <- content
// stopChan <- true
// }()
// common.SetEventStreamHeaders(c)
// c.Stream(func(w io.Writer) bool {
// select {
// case data := <-dataChan:
// // this is used to prevent annoying \ related format bug
// data = fmt.Sprintf("{\"content\": \"%s\"}", data)
// type dummyStruct struct {
// Content string `json:"content"`
// }
// var dummy dummyStruct
// err := json.Unmarshal([]byte(data), &dummy)
// responseText += dummy.Content
// var choice openai.ChatCompletionsStreamResponseChoice
// choice.Delta.Content = dummy.Content
// response := openai.ChatCompletionsStreamResponse{
// Id: fmt.Sprintf("chatcmpl-%s", helper.GetUUID()),
// Object: "chat.completion.chunk",
// Created: helper.GetTimestamp(),
// Model: "gemini-pro",
// Choices: []openai.ChatCompletionsStreamResponseChoice{choice},
// }
// jsonResponse, err := json.Marshal(response)
// if err != nil {
// logger.SysError("error marshalling stream response: " + err.Error())
// return true
// }
// c.Render(-1, common.CustomEvent{Data: "data: " + string(jsonResponse)})
// return true
// case <-stopChan:
// c.Render(-1, common.CustomEvent{Data: "data: [DONE]"})
// return false
// }
// })
if err := resp.Body.Close(); err != nil {
return openai.ErrorWrapper(err, "close_response_body_failed", http.StatusInternalServerError), "" return openai.ErrorWrapper(err, "close_response_body_failed", http.StatusInternalServerError), ""
} }
return nil, responseText return nil, responseText
} }

View File

@ -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
}

View File

@ -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)
}

View File

@ -1,36 +1,36 @@
package ollama package ollama
import ( import (
"errors"
"fmt" "fmt"
"io" "io"
"net/http" "net/http"
"github.com/Laisky/errors/v2"
"github.com/gin-gonic/gin" "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/constant" "github.com/songquanpeng/one-api/relay/meta"
"github.com/songquanpeng/one-api/relay/model" "github.com/songquanpeng/one-api/relay/model"
"github.com/songquanpeng/one-api/relay/util" "github.com/songquanpeng/one-api/relay/relaymode"
) )
type Adaptor struct { 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 // https://github.com/ollama/ollama/blob/main/docs/api.md
fullRequestURL := fmt.Sprintf("%s/api/chat", meta.BaseURL) 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) fullRequestURL = fmt.Sprintf("%s/api/embeddings", meta.BaseURL)
} }
return fullRequestURL, nil return fullRequestURL, 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 {
channel.SetupCommonRequestHeader(c, req, meta) adaptor.SetupCommonRequestHeader(c, req, meta)
req.Header.Set("Authorization", "Bearer "+meta.APIKey) req.Header.Set("Authorization", "Bearer "+meta.APIKey)
return nil 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") return nil, errors.New("request is nil")
} }
switch relayMode { switch relayMode {
case constant.RelayModeEmbeddings: case relaymode.Embeddings:
ollamaEmbeddingRequest := ConvertEmbeddingRequest(*request) ollamaEmbeddingRequest := ConvertEmbeddingRequest(*request)
return ollamaEmbeddingRequest, nil return ollamaEmbeddingRequest, nil
default: 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) { func (a *Adaptor) ConvertImageRequest(request *model.ImageRequest) (any, error) {
return channel.DoRequestHelper(a, c, meta, requestBody) 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 { if meta.IsStream {
err, usage = StreamHandler(c, resp) err, usage = StreamHandler(c, resp)
} else { } else {
switch meta.Mode { switch meta.Mode {
case constant.RelayModeEmbeddings: case relaymode.Embeddings:
err, usage = EmbeddingHandler(c, resp) err, usage = EmbeddingHandler(c, resp)
default: default:
err, usage = Handler(c, resp) err, usage = Handler(c, resp)

View File

@ -5,15 +5,16 @@ import (
"context" "context"
"encoding/json" "encoding/json"
"fmt" "fmt"
"github.com/songquanpeng/one-api/common/helper"
"github.com/songquanpeng/one-api/common/random"
"io" "io"
"net/http" "net/http"
"strings" "strings"
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
"github.com/songquanpeng/one-api/common" "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/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/constant"
"github.com/songquanpeng/one-api/relay/model" "github.com/songquanpeng/one-api/relay/model"
) )
@ -51,7 +52,7 @@ func responseOllama2OpenAI(response *ChatResponse) *openai.TextResponse {
choice.FinishReason = "stop" choice.FinishReason = "stop"
} }
fullTextResponse := openai.TextResponse{ fullTextResponse := openai.TextResponse{
Id: fmt.Sprintf("chatcmpl-%s", helper.GetUUID()), Id: fmt.Sprintf("chatcmpl-%s", random.GetUUID()),
Object: "chat.completion", Object: "chat.completion",
Created: helper.GetTimestamp(), Created: helper.GetTimestamp(),
Choices: []openai.TextResponseChoice{choice}, Choices: []openai.TextResponseChoice{choice},
@ -72,7 +73,7 @@ func streamResponseOllama2OpenAI(ollamaResponse *ChatResponse) *openai.ChatCompl
choice.FinishReason = &constant.StopFinishReason choice.FinishReason = &constant.StopFinishReason
} }
response := openai.ChatCompletionsStreamResponse{ response := openai.ChatCompletionsStreamResponse{
Id: fmt.Sprintf("chatcmpl-%s", helper.GetUUID()), Id: fmt.Sprintf("chatcmpl-%s", random.GetUUID()),
Object: "chat.completion.chunk", Object: "chat.completion.chunk",
Created: helper.GetTimestamp(), Created: helper.GetTimestamp(),
Model: ollamaResponse.Model, Model: ollamaResponse.Model,

View File

@ -4,11 +4,12 @@ import (
"fmt" "fmt"
"github.com/Laisky/errors/v2" "github.com/Laisky/errors/v2"
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
"github.com/songquanpeng/one-api/common" "github.com/songquanpeng/one-api/relay/adaptor"
"github.com/songquanpeng/one-api/relay/channel" "github.com/songquanpeng/one-api/relay/adaptor/minimax"
"github.com/songquanpeng/one-api/relay/channel/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/model"
"github.com/songquanpeng/one-api/relay/util" "github.com/songquanpeng/one-api/relay/relaymode"
"io" "io"
"net/http" "net/http"
"strings" "strings"
@ -18,13 +19,20 @@ type Adaptor struct {
ChannelType int ChannelType int
} }
func (a *Adaptor) Init(meta *util.RelayMeta) { func (a *Adaptor) Init(meta *meta.Meta) {
a.ChannelType = meta.ChannelType 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 { 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 // 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 := strings.Split(meta.RequestURLPath, "?")[0]
requestURL = fmt.Sprintf("%s?api-version=%s", requestURL, meta.APIVersion) 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 //https://github.com/songquanpeng/one-api/issues/1191
// {your endpoint}/openai/deployments/{your azure_model}/chat/completions?api-version={api_version} // {your endpoint}/openai/deployments/{your azure_model}/chat/completions?api-version={api_version}
requestURL = fmt.Sprintf("/openai/deployments/%s/%s", model_, task) requestURL = fmt.Sprintf("/openai/deployments/%s/%s", model_, task)
return util.GetFullRequestURL(meta.BaseURL, requestURL, meta.ChannelType), nil return GetFullRequestURL(meta.BaseURL, requestURL, meta.ChannelType), nil
case common.ChannelTypeMinimax: case channeltype.Minimax:
return minimax.GetRequestURL(meta) return minimax.GetRequestURL(meta)
default: 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 { func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Request, meta *meta.Meta) error {
channel.SetupCommonRequestHeader(c, req, meta) adaptor.SetupCommonRequestHeader(c, req, meta)
if meta.ChannelType == common.ChannelTypeAzure { if meta.ChannelType == channeltype.Azure {
req.Header.Set("api-key", meta.APIKey) req.Header.Set("api-key", meta.APIKey)
return nil return nil
} }
req.Header.Set("Authorization", "Bearer "+meta.APIKey) 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("HTTP-Referer", "https://github.com/songquanpeng/one-api")
req.Header.Set("X-Title", "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 return 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) {
return channel.DoRequestHelper(a, c, meta, requestBody) 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 { if meta.IsStream {
var responseText string var responseText string
err, responseText, usage = StreamHandler(c, resp, meta.Mode) err, responseText, usage = StreamHandler(c, resp, meta.Mode)
@ -75,8 +90,13 @@ func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, meta *util.Rel
usage = ResponseText2Usage(responseText, meta.ActualModelName, meta.PromptTokens) usage = ResponseText2Usage(responseText, meta.ActualModelName, meta.PromptTokens)
} }
} else { } else {
switch meta.Mode {
case relaymode.ImagesGenerations:
err, _ = ImageHandler(c, resp)
default:
err, usage = Handler(c, resp, meta.PromptTokens, meta.ActualModelName) err, usage = Handler(c, resp, meta.PromptTokens, meta.ActualModelName)
} }
}
return return
} }

View File

@ -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
}
}

View File

@ -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
}

View File

@ -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
}

View File

@ -8,8 +8,8 @@ import (
"github.com/songquanpeng/one-api/common" "github.com/songquanpeng/one-api/common"
"github.com/songquanpeng/one-api/common/conv" "github.com/songquanpeng/one-api/common/conv"
"github.com/songquanpeng/one-api/common/logger" "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/model"
"github.com/songquanpeng/one-api/relay/relaymode"
"io" "io"
"net/http" "net/http"
"strings" "strings"
@ -46,7 +46,7 @@ func StreamHandler(c *gin.Context, resp *http.Response, relayMode int) (*model.E
data = data[6:] data = data[6:]
if !strings.HasPrefix(data, "[DONE]") { if !strings.HasPrefix(data, "[DONE]") {
switch relayMode { switch relayMode {
case constant.RelayModeChatCompletions: case relaymode.ChatCompletions:
var streamResponse ChatCompletionsStreamResponse var streamResponse ChatCompletionsStreamResponse
err := json.Unmarshal([]byte(data), &streamResponse) err := json.Unmarshal([]byte(data), &streamResponse)
if err != nil { if err != nil {
@ -59,7 +59,7 @@ func StreamHandler(c *gin.Context, resp *http.Response, relayMode int) (*model.E
if streamResponse.Usage != nil { if streamResponse.Usage != nil {
usage = streamResponse.Usage usage = streamResponse.Usage
} }
case constant.RelayModeCompletions: case relaymode.Completions:
var streamResponse CompletionsStreamResponse var streamResponse CompletionsStreamResponse
err := json.Unmarshal([]byte(data), &streamResponse) err := json.Unmarshal([]byte(data), &streamResponse)
if err != nil { if err != nil {

View File

@ -110,11 +110,16 @@ type EmbeddingResponse struct {
model.Usage `json:"usage"` 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 { type ImageResponse struct {
Created int `json:"created"` Created int64 `json:"created"`
Data []struct { Data []ImageData `json:"data"`
Url string `json:"url"` //model.Usage `json:"usage"`
}
} }
type ChatCompletionsStreamResponseChoice struct { type ChatCompletionsStreamResponseChoice struct {

View File

@ -4,10 +4,10 @@ import (
"fmt" "fmt"
"github.com/Laisky/errors/v2" "github.com/Laisky/errors/v2"
"github.com/pkoukk/tiktoken-go" "github.com/pkoukk/tiktoken-go"
"github.com/songquanpeng/one-api/common"
"github.com/songquanpeng/one-api/common/config" "github.com/songquanpeng/one-api/common/config"
"github.com/songquanpeng/one-api/common/image" "github.com/songquanpeng/one-api/common/image"
"github.com/songquanpeng/one-api/common/logger" "github.com/songquanpeng/one-api/common/logger"
billingratio "github.com/songquanpeng/one-api/relay/billing/ratio"
"github.com/songquanpeng/one-api/relay/model" "github.com/songquanpeng/one-api/relay/model"
"math" "math"
"strings" "strings"
@ -28,7 +28,7 @@ func InitTokenEncoders() {
if err != nil { if err != nil {
logger.FatalLog(fmt.Sprintf("failed to get gpt-4 token encoder: %s", err.Error())) logger.FatalLog(fmt.Sprintf("failed to get gpt-4 token encoder: %s", err.Error()))
} }
for model := range common.ModelRatio { for model := range billingratio.ModelRatio {
if strings.HasPrefix(model, "gpt-3.5") { if strings.HasPrefix(model, "gpt-3.5") {
tokenEncoderMap[model] = gpt35TokenEncoder tokenEncoderMap[model] = gpt35TokenEncoder
} else if strings.HasPrefix(model, "gpt-4") { } else if strings.HasPrefix(model, "gpt-4") {

View File

@ -4,10 +4,10 @@ import (
"fmt" "fmt"
"github.com/Laisky/errors/v2" "github.com/Laisky/errors/v2"
"github.com/gin-gonic/gin" "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/channel/openai" "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/model"
"github.com/songquanpeng/one-api/relay/util"
"io" "io"
"net/http" "net/http"
) )
@ -15,16 +15,16 @@ import (
type Adaptor struct { 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 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 { func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Request, meta *meta.Meta) error {
channel.SetupCommonRequestHeader(c, req, meta) adaptor.SetupCommonRequestHeader(c, req, meta)
req.Header.Set("x-goog-api-key", meta.APIKey) req.Header.Set("x-goog-api-key", meta.APIKey)
return nil return nil
} }
@ -36,11 +36,18 @@ func (a *Adaptor) ConvertRequest(c *gin.Context, relayMode int, request *model.G
return ConvertRequest(*request), nil 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) {
return channel.DoRequestHelper(a, c, meta, requestBody) 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 { if meta.IsStream {
var responseText string var responseText string
err, responseText = StreamHandler(c, resp) err, responseText = StreamHandler(c, resp)

View File

@ -7,7 +7,8 @@ import (
"github.com/songquanpeng/one-api/common" "github.com/songquanpeng/one-api/common"
"github.com/songquanpeng/one-api/common/helper" "github.com/songquanpeng/one-api/common/helper"
"github.com/songquanpeng/one-api/common/logger" "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/constant"
"github.com/songquanpeng/one-api/relay/model" "github.com/songquanpeng/one-api/relay/model"
"io" "io"
@ -74,7 +75,7 @@ func streamResponsePaLM2OpenAI(palmResponse *ChatResponse) *openai.ChatCompletio
func StreamHandler(c *gin.Context, resp *http.Response) (*model.ErrorWithStatusCode, string) { func StreamHandler(c *gin.Context, resp *http.Response) (*model.ErrorWithStatusCode, string) {
responseText := "" responseText := ""
responseId := fmt.Sprintf("chatcmpl-%s", helper.GetUUID()) responseId := fmt.Sprintf("chatcmpl-%s", random.GetUUID())
createdTime := helper.GetTimestamp() createdTime := helper.GetTimestamp()
dataChan := make(chan string) dataChan := make(chan string)
stopChan := make(chan bool) stopChan := make(chan bool)

View File

@ -0,0 +1,7 @@
package stepfun
var ModelList = []string{
"step-1-32k",
"step-1v-32k",
"step-1-200k",
}

View File

@ -0,0 +1,238 @@
package tencent
// import (
// "bufio"
// "crypto/hmac"
// "crypto/sha1"
// "encoding/base64"
// "encoding/json"
// "github.com/Laisky/errors/v2"
// "fmt"
// "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/constant"
// "github.com/songquanpeng/one-api/relay/model"
// "io"
// "net/http"
// "sort"
// "strconv"
// "strings"
// )
// // https://cloud.tencent.com/document/product/1729/97732
// func ConvertRequest(request model.GeneralOpenAIRequest) *ChatRequest {
// messages := make([]Message, 0, len(request.Messages))
// for i := 0; i < len(request.Messages); i++ {
// message := request.Messages[i]
// if message.Role == "system" {
// messages = append(messages, Message{
// Role: "user",
// Content: message.StringContent(),
// })
// messages = append(messages, Message{
// Role: "assistant",
// Content: "Okay",
// })
// continue
// }
// messages = append(messages, Message{
// Content: message.StringContent(),
// Role: message.Role,
// })
// }
// stream := 0
// if request.Stream {
// stream = 1
// }
// return &ChatRequest{
// Timestamp: helper.GetTimestamp(),
// Expired: helper.GetTimestamp() + 24*60*60,
// QueryID: helper.GetUUID(),
// Temperature: request.Temperature,
// TopP: request.TopP,
// Stream: stream,
// Messages: messages,
// }
// }
// func responseTencent2OpenAI(response *ChatResponse) *openai.TextResponse {
// fullTextResponse := openai.TextResponse{
// Object: "chat.completion",
// Created: helper.GetTimestamp(),
// Usage: response.Usage,
// }
// if len(response.Choices) > 0 {
// choice := openai.TextResponseChoice{
// Index: 0,
// Message: model.Message{
// Role: "assistant",
// Content: response.Choices[0].Messages.Content,
// },
// FinishReason: response.Choices[0].FinishReason,
// }
// fullTextResponse.Choices = append(fullTextResponse.Choices, choice)
// }
// return &fullTextResponse
// }
// func streamResponseTencent2OpenAI(TencentResponse *ChatResponse) *openai.ChatCompletionsStreamResponse {
// response := openai.ChatCompletionsStreamResponse{
// Object: "chat.completion.chunk",
// Created: helper.GetTimestamp(),
// Model: "tencent-hunyuan",
// }
// if len(TencentResponse.Choices) > 0 {
// var choice openai.ChatCompletionsStreamResponseChoice
// choice.Delta.Content = TencentResponse.Choices[0].Delta.Content
// if TencentResponse.Choices[0].FinishReason == "stop" {
// choice.FinishReason = &constant.StopFinishReason
// }
// response.Choices = append(response.Choices, choice)
// }
// return &response
// }
// func StreamHandler(c *gin.Context, resp *http.Response) (*model.ErrorWithStatusCode, string) {
// var responseText string
// scanner := bufio.NewScanner(resp.Body)
// scanner.Split(func(data []byte, atEOF bool) (advance int, token []byte, err error) {
// if atEOF && len(data) == 0 {
// return 0, nil, nil
// }
// if i := strings.Index(string(data), "\n"); i >= 0 {
// return i + 1, data[0:i], nil
// }
// if atEOF {
// return len(data), data, nil
// }
// return 0, nil, nil
// })
// dataChan := make(chan string)
// stopChan := make(chan bool)
// go func() {
// for scanner.Scan() {
// data := scanner.Text()
// if len(data) < 5 { // ignore blank line or wrong format
// continue
// }
// if data[:5] != "data:" {
// continue
// }
// data = data[5:]
// dataChan <- data
// }
// stopChan <- true
// }()
// common.SetEventStreamHeaders(c)
// c.Stream(func(w io.Writer) bool {
// select {
// case data := <-dataChan:
// var TencentResponse ChatResponse
// err := json.Unmarshal([]byte(data), &TencentResponse)
// if err != nil {
// logger.SysError("error unmarshalling stream response: " + err.Error())
// return true
// }
// response := streamResponseTencent2OpenAI(&TencentResponse)
// if len(response.Choices) != 0 {
// responseText += response.Choices[0].Delta.Content
// }
// jsonResponse, err := json.Marshal(response)
// if err != nil {
// logger.SysError("error marshalling stream response: " + err.Error())
// return true
// }
// c.Render(-1, common.CustomEvent{Data: "data: " + string(jsonResponse)})
// return true
// case <-stopChan:
// c.Render(-1, common.CustomEvent{Data: "data: [DONE]"})
// return false
// }
// })
// err := resp.Body.Close()
// if err != nil {
// return openai.ErrorWrapper(err, "close_response_body_failed", http.StatusInternalServerError), ""
// }
// return nil, responseText
// }
// func Handler(c *gin.Context, resp *http.Response) (*model.ErrorWithStatusCode, *model.Usage) {
// var TencentResponse ChatResponse
// 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, &TencentResponse)
// if err != nil {
// return openai.ErrorWrapper(err, "unmarshal_response_body_failed", http.StatusInternalServerError), nil
// }
// if TencentResponse.Error.Code != 0 {
// return &model.ErrorWithStatusCode{
// Error: model.Error{
// Message: TencentResponse.Error.Message,
// Code: TencentResponse.Error.Code,
// },
// StatusCode: resp.StatusCode,
// }, nil
// }
// fullTextResponse := responseTencent2OpenAI(&TencentResponse)
// fullTextResponse.Model = "hunyuan"
// 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)
// if err != nil {
// return openai.ErrorWrapper(err, "write_response_body_failed", http.StatusInternalServerError), nil
// }
// return nil, &fullTextResponse.Usage
// }
// func ParseConfig(config string) (appId int64, secretId string, secretKey string, err error) {
// parts := strings.Split(config, "|")
// if len(parts) != 3 {
// err = errors.New("invalid tencent config")
// return
// }
// appId, err = strconv.ParseInt(parts[0], 10, 64)
// secretId = parts[1]
// secretKey = parts[2]
// return
// }
// func GetSign(req ChatRequest, secretKey string) string {
// params := make([]string, 0)
// params = append(params, "app_id="+strconv.FormatInt(req.AppId, 10))
// params = append(params, "secret_id="+req.SecretId)
// params = append(params, "timestamp="+strconv.FormatInt(req.Timestamp, 10))
// params = append(params, "query_id="+req.QueryID)
// params = append(params, "temperature="+strconv.FormatFloat(req.Temperature, 'f', -1, 64))
// params = append(params, "top_p="+strconv.FormatFloat(req.TopP, 'f', -1, 64))
// params = append(params, "stream="+strconv.Itoa(req.Stream))
// params = append(params, "expired="+strconv.FormatInt(req.Expired, 10))
// var messageStr string
// for _, msg := range req.Messages {
// messageStr += fmt.Sprintf(`{"role":"%s","content":"%s"},`, msg.Role, msg.Content)
// }
// messageStr = strings.TrimSuffix(messageStr, ",")
// params = append(params, "messages=["+messageStr+"]")
// sort.Strings(params)
// url := "hunyuan.cloud.tencent.com/hyllm/v1/chat/completions?" + strings.Join(params, "&")
// mac := hmac.New(sha1.New, []byte(secretKey))
// signURL := url
// mac.Write([]byte(signURL))
// sign := mac.Sum([]byte(nil))
// return base64.StdEncoding.EncodeToString(sign)
// }

View File

@ -0,0 +1,145 @@
package zhipu
// import (
// "github.com/Laisky/errors/v2"
// "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"
// }

Some files were not shown because too many files have changed in this diff Show More