xinyin025/one-api
1
0
Fork 0
mirror of https://github.com/songquanpeng/one-api.git synced 2025-11-08 09:43:42 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Merge commit '3d149fedf45472eff92910324974c762fc37dad6'

Browse Source
This commit is contained in:
Laisky.Cai
2024-04-21 15:05:13 +00:00
parent c735b9b4d3 3d149fedf4
commit 7047d9605e
45 changed files with 649 additions and 223 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
controller/auth/wechat.go
View File

@@ -3,14 +3,16 @@ package auth
import (
"encoding/json"
"fmt"
"github.com/Laisky/errors/v2"
"github.com/Laisky/one-api/common/config"
"github.com/Laisky/one-api/controller"
"github.com/Laisky/one-api/model"
"github.com/gin-gonic/gin"
"net/http"
"strconv"
"time"
"github.com/Laisky/errors/v2"
"github.com/Laisky/one-api/common/config"
"github.com/Laisky/one-api/common/ctxkey"
"github.com/Laisky/one-api/controller"
"github.com/Laisky/one-api/model"
"github.com/gin-gonic/gin"
)
type wechatLoginResponse struct {
@@ -136,7 +138,7 @@ func WeChatBind(c *gin.Context) {
})
return
}
id := c.GetInt("id")
id := c.GetInt(ctxkey.Id)
user := model.User{
Id: id,
}

9
controller/billing.go
View File

@@ -2,6 +2,7 @@ package controller
import (
"github.com/Laisky/one-api/common/config"
"github.com/Laisky/one-api/common/ctxkey"
"github.com/Laisky/one-api/model"
relaymodel "github.com/Laisky/one-api/relay/model"
"github.com/gin-gonic/gin"
@@ -14,13 +15,13 @@ func GetSubscription(c *gin.Context) {
var token *model.Token
var expiredTime int64
if config.DisplayTokenStatEnabled {
tokenId := c.GetInt("token_id")
tokenId := c.GetInt(ctxkey.TokenId)
token, err = model.GetTokenById(tokenId)
expiredTime = token.ExpiredTime
remainQuota = token.RemainQuota
usedQuota = token.UsedQuota
} else {
userId := c.GetInt("id")
userId := c.GetInt(ctxkey.Id)
remainQuota, err = model.GetUserQuota(userId)
if err != nil {
usedQuota, err = model.GetUserUsedQuota(userId)
@@ -64,11 +65,11 @@ func GetUsage(c *gin.Context) {
var err error
var token *model.Token
if config.DisplayTokenStatEnabled {
tokenId := c.GetInt("token_id")
tokenId := c.GetInt(ctxkey.TokenId)
token, err = model.GetTokenById(tokenId)
quota = token.UsedQuota
} else {
userId := c.GetInt("id")
userId := c.GetInt(ctxkey.Id)
quota, err = model.GetUserUsedQuota(userId)
}
if err != nil {

15
controller/channel-test.go
View File

@@ -15,6 +15,7 @@ import (
"github.com/Laisky/errors/v2"
"github.com/Laisky/one-api/common/config"
"github.com/Laisky/one-api/common/ctxkey"
"github.com/Laisky/one-api/common/logger"
"github.com/Laisky/one-api/common/message"
"github.com/Laisky/one-api/middleware"
@@ -26,7 +27,6 @@ import (
"github.com/Laisky/one-api/relay/meta"
relaymodel "github.com/Laisky/one-api/relay/model"
"github.com/Laisky/one-api/relay/relaymode"
"github.com/gin-gonic/gin"
)
@@ -55,8 +55,8 @@ func testChannel(channel *model.Channel) (err error, openaiErr *relaymodel.Error
}
c.Request.Header.Set("Authorization", "Bearer "+channel.Key)
c.Request.Header.Set("Content-Type", "application/json")
c.Set("channel", channel.Type)
c.Set("base_url", channel.GetBaseURL())
c.Set(ctxkey.Channel, channel.Type)
c.Set(ctxkey.BaseURL, channel.GetBaseURL())
middleware.SetupContextForSelectedChannel(c, channel, "")
meta := meta.GetByContext(c)
apiType := channeltype.ToAPIType(channel.Type)
@@ -65,8 +65,12 @@ func testChannel(channel *model.Channel) (err error, openaiErr *relaymodel.Error
return errors.Errorf("invalid api type: %d, adaptor is nil", apiType), nil
}
adaptor.Init(meta)
modelName := adaptor.GetModelList()[0]
if !strings.Contains(channel.Models, modelName) {
var modelName string
modelList := adaptor.GetModelList()
if len(modelList) != 0 {
modelName = modelList[0]
}
if modelName == "" || !strings.Contains(channel.Models, modelName) {
modelNames := strings.Split(channel.Models, ",")
if len(modelNames) > 0 {
modelName = modelNames[0]
@@ -83,6 +87,7 @@ func testChannel(channel *model.Channel) (err error, openaiErr *relaymodel.Error
if err != nil {
return err, nil
}
logger.SysLog(string(jsonData))
requestBody := bytes.NewBuffer(jsonData)
c.Request.Body = io.NopCloser(requestBody)
resp, err := adaptor.DoRequest(c, meta, requestBody)

14
controller/log.go
View File

@@ -1,11 +1,13 @@
package controller
import (
"github.com/Laisky/one-api/common/config"
"github.com/Laisky/one-api/model"
"github.com/gin-gonic/gin"
"net/http"
"strconv"
"github.com/Laisky/one-api/common/config"
"github.com/Laisky/one-api/common/ctxkey"
"github.com/Laisky/one-api/model"
"github.com/gin-gonic/gin"
)
func GetAllLogs(c *gin.Context) {
@@ -41,7 +43,7 @@ func GetUserLogs(c *gin.Context) {
if p < 0 {
p = 0
}
userId := c.GetInt("id")
userId := c.GetInt(ctxkey.Id)
logType, _ := strconv.Atoi(c.Query("type"))
startTimestamp, _ := strconv.ParseInt(c.Query("start_timestamp"), 10, 64)
endTimestamp, _ := strconv.ParseInt(c.Query("end_timestamp"), 10, 64)
@@ -83,7 +85,7 @@ func SearchAllLogs(c *gin.Context) {
func SearchUserLogs(c *gin.Context) {
keyword := c.Query("keyword")
userId := c.GetInt("id")
userId := c.GetInt(ctxkey.Id)
logs, err := model.SearchUserLogs(userId, keyword)
if err != nil {
c.JSON(http.StatusOK, gin.H{
@@ -122,7 +124,7 @@ func GetLogsStat(c *gin.Context) {
}
func GetLogsSelfStat(c *gin.Context) {
username := c.GetString("username")
username := c.GetString(ctxkey.Username)
logType, _ := strconv.Atoi(c.Query("type"))
startTimestamp, _ := strconv.ParseInt(c.Query("start_timestamp"), 10, 64)
endTimestamp, _ := strconv.ParseInt(c.Query("end_timestamp"), 10, 64)

14
controller/model.go
View File

@@ -2,6 +2,10 @@ package controller
import (
"fmt"
"net/http"
"strings"
"github.com/Laisky/one-api/common/ctxkey"
"github.com/Laisky/one-api/model"
relay "github.com/Laisky/one-api/relay"
"github.com/Laisky/one-api/relay/adaptor/openai"
@@ -10,8 +14,6 @@ import (
"github.com/Laisky/one-api/relay/meta"
relaymodel "github.com/Laisky/one-api/relay/model"
"github.com/gin-gonic/gin"
"net/http"
"strings"
)
// https://platform.openai.com/docs/api-reference/models/list
@@ -139,10 +141,10 @@ func ListAllModels(c *gin.Context) {
func ListModels(c *gin.Context) {
ctx := c.Request.Context()
var availableModels []string
if c.GetString("available_models") != "" {
availableModels = strings.Split(c.GetString("available_models"), ",")
if c.GetString(ctxkey.AvailableModels) != "" {
availableModels = strings.Split(c.GetString(ctxkey.AvailableModels), ",")
} else {
userId := c.GetInt("id")
userId := c.GetInt(ctxkey.Id)
userGroup, _ := model.CacheGetUserGroup(userId)
availableModels, _ = model.CacheGetGroupModels(ctx, userGroup)
}
@@ -194,7 +196,7 @@ func RetrieveModel(c *gin.Context) {
func GetUserAvailableModels(c *gin.Context) {
ctx := c.Request.Context()
id := c.GetInt("id")
id := c.GetInt(ctxkey.Id)
userGroup, err := model.CacheGetUserGroup(id)
if err != nil {
c.JSON(http.StatusOK, gin.H{

8
controller/redemption.go
View File

@@ -1,13 +1,15 @@
package controller
import (
"net/http"
"strconv"
"github.com/Laisky/one-api/common/config"
"github.com/Laisky/one-api/common/ctxkey"
"github.com/Laisky/one-api/common/helper"
"github.com/Laisky/one-api/common/random"
"github.com/Laisky/one-api/model"
"github.com/gin-gonic/gin"
"net/http"
"strconv"
)
func GetAllRedemptions(c *gin.Context) {
@@ -109,7 +111,7 @@ func AddRedemption(c *gin.Context) {
for i := 0; i < redemption.Count; i++ {
key := random.GetUUID()
cleanRedemption := model.Redemption{
UserId: c.GetInt("id"),
UserId: c.GetInt(ctxkey.Id),
Name: redemption.Name,
Key: key,
CreatedTime: helper.GetTimestamp(),

16
controller/relay.go
View File

@@ -48,15 +48,15 @@ func Relay(c *gin.Context) {
requestBody, _ := common.GetRequestBody(c)
logger.Debugf(ctx, "request body: %s", string(requestBody))
}
channelId := c.GetInt("channel_id")
channelId := c.GetInt(ctxkey.ChannelId)
bizErr := relayHelper(c, relayMode)
if bizErr == nil {
monitor.Emit(channelId, true)
return
}
lastFailedChannelId := channelId
channelName := c.GetString("channel_name")
group := c.GetString("group")
channelName := c.GetString(ctxkey.ChannelName)
group := c.GetString(ctxkey.Group)
originalModel := c.GetString(ctxkey.OriginalModel)
go processChannelRelayError(ctx, channelId, channelName, bizErr)
requestId := c.GetString(logger.RequestIdKey)
@@ -82,9 +82,9 @@ func Relay(c *gin.Context) {
if bizErr == nil {
return
}
channelId := c.GetInt("channel_id")
channelId := c.GetInt(ctxkey.ChannelId)
lastFailedChannelId = channelId
channelName := c.GetString("channel_name")
channelName := c.GetString(ctxkey.ChannelName)
go processChannelRelayError(ctx, channelId, channelName, bizErr)
}
if bizErr != nil {
@@ -100,14 +100,14 @@ func Relay(c *gin.Context) {
// shouldRetry returns nil if should retry, otherwise returns error
func shouldRetry(c *gin.Context, statusCode int) error {
if v, ok := c.Get("specific_channel_id"); ok {
if v, ok := c.Get(ctxkey.SpecificChannelId); ok {
return errors.Errorf("specific channel = %v", v)
}
if statusCode == http.StatusBadRequest {
if statusCode == http.StatusTooManyRequests {
return errors.Errorf("status code = %d", statusCode)
}
if statusCode/100 == 2 {
if statusCode/100 == 5 {
return errors.Errorf("status code = %d", statusCode)
}

17
controller/token.go
View File

@@ -7,6 +7,7 @@ import (
"github.com/Laisky/one-api/common"
"github.com/Laisky/one-api/common/config"
"github.com/Laisky/one-api/common/ctxkey"
"github.com/Laisky/one-api/common/helper"
"github.com/Laisky/one-api/common/network"
"github.com/Laisky/one-api/common/random"
@@ -16,7 +17,7 @@ import (
)
func GetAllTokens(c *gin.Context) {
userId := c.GetInt("id")
userId := c.GetInt(ctxkey.Id)
p, _ := strconv.Atoi(c.Query("p"))
if p < 0 {
p = 0
@@ -41,7 +42,7 @@ func GetAllTokens(c *gin.Context) {
}
func SearchTokens(c *gin.Context) {
userId := c.GetInt("id")
userId := c.GetInt(ctxkey.Id)
keyword := c.Query("keyword")
tokens, err := model.SearchUserTokens(userId, keyword)
if err != nil {
@@ -61,7 +62,7 @@ func SearchTokens(c *gin.Context) {
func GetToken(c *gin.Context) {
id, err := strconv.Atoi(c.Param("id"))
userId := c.GetInt("id")
userId := c.GetInt(ctxkey.Id)
if err != nil {
c.JSON(http.StatusOK, gin.H{
"success": false,
@@ -86,8 +87,8 @@ func GetToken(c *gin.Context) {
}
func GetTokenStatus(c *gin.Context) {
tokenId := c.GetInt("token_id")
userId := c.GetInt("id")
tokenId := c.GetInt(ctxkey.TokenId)
userId := c.GetInt(ctxkey.Id)
token, err := model.GetTokenByIds(tokenId, userId)
if err != nil {
c.JSON(http.StatusOK, gin.H{
@@ -145,7 +146,7 @@ func AddToken(c *gin.Context) {
}
cleanToken := model.Token{
UserId: c.GetInt("id"),
UserId: c.GetInt(ctxkey.Id),
Name: token.Name,
Key: random.GenerateKey(),
CreatedTime: helper.GetTimestamp(),
@@ -174,7 +175,7 @@ func AddToken(c *gin.Context) {
func DeleteToken(c *gin.Context) {
id, _ := strconv.Atoi(c.Param("id"))
userId := c.GetInt("id")
userId := c.GetInt(ctxkey.Id)
err := model.DeleteTokenById(id, userId)
if err != nil {
c.JSON(http.StatusOK, gin.H{
@@ -198,7 +199,7 @@ type updateTokenDto struct {
}
func UpdateToken(c *gin.Context) {
userId := c.GetInt("id")
userId := c.GetInt(ctxkey.Id)
statusOnly := c.Query("status_only")
tokenPatch := new(updateTokenDto)
err := c.ShouldBindJSON(tokenPatch)

23
controller/user.go
View File

@@ -3,14 +3,15 @@ package controller
import (
"encoding/json"
"fmt"
"github.com/Laisky/one-api/common"
"github.com/Laisky/one-api/common/config"
"github.com/Laisky/one-api/common/random"
"github.com/Laisky/one-api/model"
"net/http"
"strconv"
"time"
"github.com/Laisky/one-api/common"
"github.com/Laisky/one-api/common/config"
"github.com/Laisky/one-api/common/ctxkey"
"github.com/Laisky/one-api/common/random"
"github.com/Laisky/one-api/model"
"github.com/gin-contrib/sessions"
"github.com/gin-gonic/gin"
)
@@ -244,7 +245,7 @@ func GetUser(c *gin.Context) {
})
return
}
myRole := c.GetInt("role")
myRole := c.GetInt(ctxkey.Role)
if myRole <= user.Role && myRole != model.RoleRootUser {
c.JSON(http.StatusOK, gin.H{
"success": false,
@@ -261,7 +262,7 @@ func GetUser(c *gin.Context) {
}
func GetUserDashboard(c *gin.Context) {
id := c.GetInt("id")
id := c.GetInt(ctxkey.Id)
now := time.Now()
startOfDay := now.Truncate(24*time.Hour).AddDate(0, 0, -6).Unix()
endOfDay := now.Truncate(24 * time.Hour).Add(24*time.Hour - time.Second).Unix()
@@ -284,7 +285,7 @@ func GetUserDashboard(c *gin.Context) {
}
func GenerateAccessToken(c *gin.Context) {
id := c.GetInt("id")
id := c.GetInt(ctxkey.Id)
user, err := model.GetUserById(id, true)
if err != nil {
c.JSON(http.StatusOK, gin.H{
@@ -320,7 +321,7 @@ func GenerateAccessToken(c *gin.Context) {
}
func GetAffCode(c *gin.Context) {
id := c.GetInt("id")
id := c.GetInt(ctxkey.Id)
user, err := model.GetUserById(id, true)
if err != nil {
c.JSON(http.StatusOK, gin.H{
@@ -358,7 +359,7 @@ func GetSelfByToken(c *gin.Context) {
}
func GetSelf(c *gin.Context) {
id := c.GetInt("id")
id := c.GetInt(ctxkey.Id)
user, err := model.GetUserById(id, false)
if err != nil {
c.JSON(http.StatusOK, gin.H{
@@ -403,7 +404,7 @@ func UpdateUser(c *gin.Context) {
})
return
}
myRole := c.GetInt("role")
myRole := c.GetInt(ctxkey.Role)
if myRole <= originUser.Role && myRole != model.RoleRootUser {
c.JSON(http.StatusOK, gin.H{
"success": false,
@@ -461,7 +462,7 @@ func UpdateSelf(c *gin.Context) {
}
cleanUser := model.User{
Id: c.GetInt("id"),
Id: c.GetInt(ctxkey.Id),
Username: user.Username,
Password: user.Password,
DisplayName: user.DisplayName,