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

fix: Improve error handling and add checks in relay.go

Browse Source 
- Refactor error handling in `relay.go` using the `errors` package
- Implement specific error handling for `specific_channel_id` in `shouldRetry` function
- Add checks for various status codes (`StatusTooManyRequests`, `Status5xx`, `Status2xx`) in `Relay` function
- Improve code quality by adding missing imports (`io`, `net/http`) in `relay.go`
This commit is contained in:
Laisky.Cai 2024-03-08 02:32:47 +00:00
parent ba9b258a4b
commit 849920b91f
1 changed files with 14 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

26
controller/relay.go
View File

@ -4,7 +4,11 @@ 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/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"
@ -15,8 +19,6 @@ import (
"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/util"
"io"
"net/http"
) )
// https://platform.openai.com/docs/api-reference/chat // https://platform.openai.com/docs/api-reference/chat
@ -57,8 +59,8 @@ func Relay(c *gin.Context) {
go processChannelRelayError(ctx, channelId, channelName, bizErr) go processChannelRelayError(ctx, channelId, channelName, bizErr)
requestId := c.GetString(logger.RequestIdKey) requestId := c.GetString(logger.RequestIdKey)
retryTimes := config.RetryTimes retryTimes := config.RetryTimes
if !shouldRetry(c, bizErr.StatusCode) { if err := shouldRetry(c, bizErr.StatusCode); err != nil {
logger.Errorf(ctx, "relay error happen, status code is %d, won't retry in this case", bizErr.StatusCode) logger.Errorf(ctx, "relay error happen, won't retry since of %v", err.Error())
retryTimes = 0 retryTimes = 0
} }
for i := retryTimes; i > 0; i-- { for i := retryTimes; i > 0; i-- {
@ -94,23 +96,23 @@ func Relay(c *gin.Context) {
} }
} }
func shouldRetry(c *gin.Context, statusCode int) bool { func shouldRetry(c *gin.Context, statusCode int) error {
if _, ok := c.Get("specific_channel_id"); ok { if v, ok := c.Get("specific_channel_id"); ok {
return false return errors.Errorf("specific channel = %v", v)
} }
if statusCode == http.StatusTooManyRequests { if statusCode == http.StatusTooManyRequests {
return true return nil
} }
if statusCode/100 == 5 { if statusCode/100 == 5 {
return true return nil
} }
if statusCode == http.StatusBadRequest { if statusCode == http.StatusBadRequest {
return false return errors.Errorf("status code = %d", statusCode)
} }
if statusCode/100 == 2 { if statusCode/100 == 2 {
return false return errors.Errorf("status code = %d", statusCode)
} }
return true return nil
} }
func processChannelRelayError(ctx context.Context, channelId int, channelName string, err *model.ErrorWithStatusCode) { func processChannelRelayError(ctx context.Context, channelId int, channelName string, err *model.ErrorWithStatusCode) {