Ecosystem/ai-gateway
5
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #70 from songquanpeng/main

Browse Source 
Fork Sync: Update from parent repository
This commit is contained in:
quzard 2023-06-23 23:19:29 +08:00 committed by GitHub
commit e6faa5db4d
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 22 additions and 22 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

24
controller/relay-text.go
View File

@ -77,7 +77,7 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
preConsumedQuota := int(float64(preConsumedTokens) * ratio) preConsumedQuota := int(float64(preConsumedTokens) * ratio)
userQuota, err := model.CacheGetUserQuota(userId) userQuota, err := model.CacheGetUserQuota(userId)
if err != nil { if err != nil {
return errorWrapper(err, "get_user_quota_failed", http.StatusOK) return errorWrapper(err, "get_user_quota_failed", http.StatusInternalServerError)
} }
if userQuota > 10*preConsumedQuota { if userQuota > 10*preConsumedQuota {
// in this case, we do not pre-consume quota // in this case, we do not pre-consume quota
@ -87,12 +87,12 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
if consumeQuota && preConsumedQuota > 0 && strings.Contains(channelName, "免费") == false { if consumeQuota && preConsumedQuota > 0 && strings.Contains(channelName, "免费") == false {
err := model.PreConsumeTokenQuota(tokenId, preConsumedQuota) err := model.PreConsumeTokenQuota(tokenId, preConsumedQuota)
if err != nil { if err != nil {
return errorWrapper(err, "pre_consume_token_quota_failed", http.StatusOK) return errorWrapper(err, "pre_consume_token_quota_failed", http.StatusForbidden)
} }
} }
req, err := http.NewRequest(c.Request.Method, fullRequestURL, c.Request.Body) req, err := http.NewRequest(c.Request.Method, fullRequestURL, c.Request.Body)
if err != nil { if err != nil {
return errorWrapper(err, "new_request_failed", http.StatusOK) return errorWrapper(err, "new_request_failed", http.StatusInternalServerError)
} }
if channelType == common.ChannelTypeAzure { if channelType == common.ChannelTypeAzure {
key := c.Request.Header.Get("Authorization") key := c.Request.Header.Get("Authorization")
@ -107,15 +107,15 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
client := &http.Client{} client := &http.Client{}
resp, err := client.Do(req) resp, err := client.Do(req)
if err != nil { if err != nil {
return errorWrapper(err, "do_request_failed", http.StatusOK) return errorWrapper(err, "do_request_failed", http.StatusInternalServerError)
} }
err = req.Body.Close() err = req.Body.Close()
if err != nil { if err != nil {
return errorWrapper(err, "close_request_body_failed", http.StatusOK) return errorWrapper(err, "close_request_body_failed", http.StatusInternalServerError)
} }
err = c.Request.Body.Close() err = c.Request.Body.Close()
if err != nil { if err != nil {
return errorWrapper(err, "close_request_body_failed", http.StatusOK) return errorWrapper(err, "close_request_body_failed", http.StatusInternalServerError)
} }
var textResponse TextResponse var textResponse TextResponse
isStream := strings.HasPrefix(resp.Header.Get("Content-Type"), "text/event-stream") isStream := strings.HasPrefix(resp.Header.Get("Content-Type"), "text/event-stream")
@ -237,22 +237,22 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
}) })
err = resp.Body.Close() err = resp.Body.Close()
if err != nil { if err != nil {
return errorWrapper(err, "close_response_body_failed", http.StatusOK) return errorWrapper(err, "close_response_body_failed", http.StatusInternalServerError)
} }
return nil return nil
} else { } else {
if consumeQuota { if consumeQuota {
responseBody, err := io.ReadAll(resp.Body) responseBody, err := io.ReadAll(resp.Body)
if err != nil { if err != nil {
return errorWrapper(err, "read_response_body_failed", http.StatusOK) return errorWrapper(err, "read_response_body_failed", http.StatusInternalServerError)
} }
err = resp.Body.Close() err = resp.Body.Close()
if err != nil { if err != nil {
return errorWrapper(err, "close_response_body_failed", http.StatusOK) return errorWrapper(err, "close_response_body_failed", http.StatusInternalServerError)
} }
err = json.Unmarshal(responseBody, &textResponse) err = json.Unmarshal(responseBody, &textResponse)
if err != nil { if err != nil {
return errorWrapper(err, "unmarshal_response_body_failed", http.StatusOK) return errorWrapper(err, "unmarshal_response_body_failed", http.StatusInternalServerError)
} }
if textResponse.Error.Type != "" { if textResponse.Error.Type != "" {
return &OpenAIErrorWithStatusCode{ return &OpenAIErrorWithStatusCode{
@ -273,11 +273,11 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
c.Writer.WriteHeader(resp.StatusCode) c.Writer.WriteHeader(resp.StatusCode)
_, err = io.Copy(c.Writer, resp.Body) _, err = io.Copy(c.Writer, resp.Body)
if err != nil { if err != nil {
return errorWrapper(err, "copy_response_body_failed", http.StatusOK) return errorWrapper(err, "copy_response_body_failed", http.StatusInternalServerError)
} }
err = resp.Body.Close() err = resp.Body.Close()
if err != nil { if err != nil {
return errorWrapper(err, "close_response_body_failed", http.StatusOK) return errorWrapper(err, "close_response_body_failed", http.StatusInternalServerError)
} }
return nil return nil
} }

4
controller/relay.go
View File

@ -136,7 +136,7 @@ func RelayNotImplemented(c *gin.Context) {
Param: "", Param: "",
Code: "api_not_implemented", Code: "api_not_implemented",
} }
c.JSON(http.StatusOK, gin.H{ c.JSON(http.StatusNotImplemented, gin.H{
"error": err, "error": err,
}) })
} }
@ -148,7 +148,7 @@ func RelayNotFound(c *gin.Context) {
Param: "", Param: "",
Code: "api_not_found", Code: "api_not_found",
} }
c.JSON(http.StatusOK, gin.H{ c.JSON(http.StatusNotFound, gin.H{
"error": err, "error": err,
}) })
} }

6
middleware/auth.go
View File

@ -91,7 +91,7 @@ func TokenAuth() func(c *gin.Context) {
key = parts[0] key = parts[0]
token, err := model.ValidateUserToken(key) token, err := model.ValidateUserToken(key)
if err != nil { if err != nil {
c.JSON(http.StatusOK, gin.H{ c.JSON(http.StatusUnauthorized, gin.H{
"error": gin.H{ "error": gin.H{
"message": err.Error(), "message": err.Error(),
"type": "one_api_error", "type": "one_api_error",
@ -101,7 +101,7 @@ func TokenAuth() func(c *gin.Context) {
return return
} }
if !model.CacheIsUserEnabled(token.UserId) { if !model.CacheIsUserEnabled(token.UserId) {
c.JSON(http.StatusOK, gin.H{ c.JSON(http.StatusForbidden, gin.H{
"error": gin.H{ "error": gin.H{
"message": "用户已被封禁", "message": "用户已被封禁",
"type": "one_api_error", "type": "one_api_error",
@ -123,7 +123,7 @@ func TokenAuth() func(c *gin.Context) {
if model.IsAdmin(token.UserId) { if model.IsAdmin(token.UserId) {
c.Set("channelId", parts[1]) c.Set("channelId", parts[1])
} else { } else {
c.JSON(http.StatusOK, gin.H{ c.JSON(http.StatusForbidden, gin.H{
"error": gin.H{ "error": gin.H{
"message": "普通用户不支持指定渠道", "message": "普通用户不支持指定渠道",
"type": "one_api_error", "type": "one_api_error",

10
middleware/distributor.go
View File

@ -24,7 +24,7 @@ func Distribute() func(c *gin.Context) {
if ok { if ok {
id, err := strconv.Atoi(channelId.(string)) id, err := strconv.Atoi(channelId.(string))
if err != nil { if err != nil {
c.JSON(http.StatusOK, gin.H{ c.JSON(http.StatusBadRequest, gin.H{
"error": gin.H{ "error": gin.H{
"message": "无效的渠道 ID", "message": "无效的渠道 ID",
"type": "one_api_error", "type": "one_api_error",
@ -35,7 +35,7 @@ func Distribute() func(c *gin.Context) {
} }
channel, err = model.GetChannelById(id, true) channel, err = model.GetChannelById(id, true)
if err != nil { if err != nil {
c.JSON(200, gin.H{ c.JSON(http.StatusBadRequest, gin.H{
"error": gin.H{ "error": gin.H{
"message": "无效的渠道 ID", "message": "无效的渠道 ID",
"type": "one_api_error", "type": "one_api_error",
@ -45,7 +45,7 @@ func Distribute() func(c *gin.Context) {
return return
} }
if channel.Status != common.ChannelStatusEnabled { if channel.Status != common.ChannelStatusEnabled {
c.JSON(200, gin.H{ c.JSON(http.StatusForbidden, gin.H{
"error": gin.H{ "error": gin.H{
"message": "该渠道已被禁用", "message": "该渠道已被禁用",
"type": "one_api_error", "type": "one_api_error",
@ -59,7 +59,7 @@ func Distribute() func(c *gin.Context) {
var modelRequest ModelRequest var modelRequest ModelRequest
err := common.UnmarshalBodyReusable(c, &modelRequest) err := common.UnmarshalBodyReusable(c, &modelRequest)
if err != nil { if err != nil {
c.JSON(200, gin.H{ c.JSON(http.StatusBadRequest, gin.H{
"error": gin.H{ "error": gin.H{
"message": "无效的请求", "message": "无效的请求",
"type": "one_api_error", "type": "one_api_error",
@ -78,7 +78,7 @@ func Distribute() func(c *gin.Context) {
} }
channel, err = model.CacheGetRandomSatisfiedChannel(userGroup, modelRequest.Model) channel, err = model.CacheGetRandomSatisfiedChannel(userGroup, modelRequest.Model)
if err != nil { if err != nil {
c.JSON(200, gin.H{ c.JSON(http.StatusServiceUnavailable, gin.H{
"error": gin.H{ "error": gin.H{
"message": "无可用渠道", "message": "无可用渠道",
"type": "one_api_error", "type": "one_api_error",