Skip to content

添加渠道错误信息覆盖功能 #5082

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
xin-zheqi wants to merge 1 commit into from
Closed

添加渠道错误信息覆盖功能 #5082

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions constant/context_key.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,8 @@ const (
ContextKeyChannelAutoBan ContextKey = "auto_ban"
ContextKeyChannelModelMapping ContextKey = "model_mapping"
ContextKeyChannelStatusCodeMapping ContextKey = "status_code_mapping"
ContextKeyChannelErrorMessageMapping ContextKey = "error_message_mapping"
ContextKeyChannelErrorOverrideSummary ContextKey = "channel_error_override_summary"
ContextKeyChannelIsMultiKey ContextKey = "channel_is_multi_key"
ContextKeyChannelMultiKeyIndex ContextKey = "channel_multi_key_index"
ContextKeyChannelKey ContextKey = "channel_key"
Expand Down
11 changes: 11 additions & 0 deletions controller/relay.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -378,6 +378,16 @@ func processChannelError(c *gin.Context, channelError types.ChannelError, err *t
other["error_type"] = err.GetErrorType()
other["error_code"] = err.GetErrorCode()
other["status_code"] = err.StatusCode
if overrideSummary, ok := c.Get(string(constant.ContextKeyChannelErrorOverrideSummary)); ok {
if summary, ok := overrideSummary.(*service.ChannelErrorOverrideSummary); ok && summary != nil {
other["original_status_code"] = summary.OriginalStatusCode
other["final_status_code"] = summary.FinalStatusCode
other["original_error_message"] = summary.OriginalMessage
other["final_error_message"] = summary.FinalMessage
other["status_code_rewritten"] = summary.StatusCodeRewritten
other["error_message_rewritten"] = summary.MessageRewritten
}
}
other["channel_id"] = channelId
other["channel_name"] = c.GetString("channel_name")
other["channel_type"] = c.GetInt("channel_type")
Expand Down Expand Up @@ -607,6 +617,7 @@ func respondTaskError(c *gin.Context, taskErr *dto.TaskError) {
if taskErr.StatusCode == http.StatusTooManyRequests {
taskErr.Message = "当前分组上游负载已饱和，请稍后再试"
}
service.ApplyChannelTaskErrorMessageOverride(taskErr, c.GetString("error_message_mapping"))
c.JSON(taskErr.StatusCode, taskErr)
}

Expand Down
5 changes: 4 additions & 1 deletion docker-compose.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,10 @@ version: '3.4' # For compatibility with older Docker versions

services:
new-api:
image: calciumion/new-api:latest
build:
context: .
dockerfile: Dockerfile
image: new-api:latest
container_name: new-api
restart: always
command: --log-dir /app/logs
Expand Down
1 change: 1 addition & 0 deletions middleware/distributor.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -420,6 +420,7 @@ func SetupContextForSelectedChannel(c *gin.Context, channel *model.Channel, mode
common.SetContextKey(c, constant.ContextKeyChannelAutoBan, channel.GetAutoBan())
common.SetContextKey(c, constant.ContextKeyChannelModelMapping, channel.GetModelMapping())
common.SetContextKey(c, constant.ContextKeyChannelStatusCodeMapping, channel.GetStatusCodeMapping())
common.SetContextKey(c, constant.ContextKeyChannelErrorMessageMapping, channel.GetErrorMessageMapping())

key, index, newAPIError := channel.GetNextEnabledKey()
if newAPIError != nil {
Expand Down
26 changes: 17 additions & 9 deletions model/channel.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,15 +41,16 @@ type Channel struct {
UsedQuota int64 `json:"used_quota" gorm:"bigint;default:0"`
ModelMapping *string `json:"model_mapping" gorm:"type:text"`
//MaxInputTokens *int `json:"max_input_tokens" gorm:"default:0"`
StatusCodeMapping *string `json:"status_code_mapping" gorm:"type:varchar(1024);default:''"`
Priority *int64 `json:"priority" gorm:"bigint;default:0"`
AutoBan *int `json:"auto_ban" gorm:"default:1"`
OtherInfo string `json:"other_info"`
Tag *string `json:"tag" gorm:"index"`
Setting *string `json:"setting" gorm:"type:text"` // 渠道额外设置
ParamOverride *string `json:"param_override" gorm:"type:text"`
HeaderOverride *string `json:"header_override" gorm:"type:text"`
Remark *string `json:"remark" gorm:"type:varchar(255)" validate:"max=255"`
StatusCodeMapping *string `json:"status_code_mapping" gorm:"type:varchar(1024);default:''"`
ErrorMessageMapping *string `json:"error_message_mapping" gorm:"type:text"`
Priority *int64 `json:"priority" gorm:"bigint;default:0"`
AutoBan *int `json:"auto_ban" gorm:"default:1"`
OtherInfo string `json:"other_info"`
Tag *string `json:"tag" gorm:"index"`
Setting *string `json:"setting" gorm:"type:text"` // 渠道额外设置
ParamOverride *string `json:"param_override" gorm:"type:text"`
HeaderOverride *string `json:"header_override" gorm:"type:text"`
Remark *string `json:"remark" gorm:"type:varchar(255)" validate:"max=255"`
// add after v0.8.5
ChannelInfo ChannelInfo `json:"channel_info" gorm:"type:json"`

Expand Down Expand Up @@ -513,6 +514,13 @@ func (channel *Channel) GetStatusCodeMapping() string {
return *channel.StatusCodeMapping
}

func (channel *Channel) GetErrorMessageMapping() string {
if channel.ErrorMessageMapping == nil {
return ""
}
return *channel.ErrorMessageMapping
}

func (channel *Channel) Insert() error {
var err error
err = DB.Create(channel).Error
Expand Down
6 changes: 4 additions & 2 deletions relay/audio_handler.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ import (
"net/http"

"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/dto"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relay/helper"
Expand Down Expand Up @@ -49,22 +50,23 @@ func AudioHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *type
return types.NewOpenAIError(err, types.ErrorCodeDoRequestFailed, http.StatusInternalServerError)
}
statusCodeMappingStr := c.GetString("status_code_mapping")
errorMessageMappingStr := c.GetString("error_message_mapping")

var httpResp *http.Response
if resp != nil {
httpResp = resp.(*http.Response)
if httpResp.StatusCode != http.StatusOK {
newAPIError = service.RelayErrorHandler(c.Request.Context(), httpResp, false)
// reset status code 重置状态码
service.ResetStatusCode(newAPIError, statusCodeMappingStr)
common.SetContextKey(c, constant.ContextKeyChannelErrorOverrideSummary, service.ApplyChannelErrorOverrides(newAPIError, statusCodeMappingStr, errorMessageMappingStr))
return newAPIError
}
}

usage, newAPIError := adaptor.DoResponse(c, httpResp, info)
if newAPIError != nil {
// reset status code 重置状态码
service.ResetStatusCode(newAPIError, statusCodeMappingStr)
common.SetContextKey(c, constant.ContextKeyChannelErrorOverrideSummary, service.ApplyChannelErrorOverrides(newAPIError, statusCodeMappingStr, errorMessageMappingStr))
return newAPIError
}
if usage.(*dto.Usage).CompletionTokenDetails.AudioTokens > 0 || usage.(*dto.Usage).PromptTokensDetails.AudioTokens > 0 {
Expand Down
2 changes: 1 addition & 1 deletion relay/channel/gemini/relay-gemini.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1519,7 +1519,7 @@ func GeminiChatHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.R
)
}

service.ResetStatusCode(newAPIError, c.GetString("status_code_mapping"))
common.SetContextKey(c, constant.ContextKeyChannelErrorOverrideSummary, service.ApplyChannelErrorOverrides(newAPIError, c.GetString("status_code_mapping"), c.GetString("error_message_mapping")))

switch info.RelayFormat {
case types.RelayFormatClaude:
Expand Down
7 changes: 4 additions & 3 deletions relay/chat_completions_via_responses.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -143,27 +143,28 @@ func chatCompletionsViaResponses(c *gin.Context, info *relaycommon.RelayInfo, ad
}

statusCodeMappingStr := c.GetString("status_code_mapping")
errorMessageMappingStr := c.GetString("error_message_mapping")

httpResp = resp.(*http.Response)
info.IsStream = info.IsStream || strings.HasPrefix(httpResp.Header.Get("Content-Type"), "text/event-stream")
if httpResp.StatusCode != http.StatusOK {
newApiErr := service.RelayErrorHandler(c.Request.Context(), httpResp, false)
service.ResetStatusCode(newApiErr, statusCodeMappingStr)
common.SetContextKey(c, constant.ContextKeyChannelErrorOverrideSummary, service.ApplyChannelErrorOverrides(newApiErr, statusCodeMappingStr, errorMessageMappingStr))
return nil, newApiErr
}

if info.IsStream {
usage, newApiErr := openaichannel.OaiResponsesToChatStreamHandler(c, info, httpResp)
if newApiErr != nil {
service.ResetStatusCode(newApiErr, statusCodeMappingStr)
common.SetContextKey(c, constant.ContextKeyChannelErrorOverrideSummary, service.ApplyChannelErrorOverrides(newApiErr, statusCodeMappingStr, errorMessageMappingStr))
return nil, newApiErr
}
return usage, nil
}

usage, newApiErr := openaichannel.OaiResponsesToChatHandler(c, info, httpResp)
if newApiErr != nil {
service.ResetStatusCode(newApiErr, statusCodeMappingStr)
common.SetContextKey(c, constant.ContextKeyChannelErrorOverrideSummary, service.ApplyChannelErrorOverrides(newApiErr, statusCodeMappingStr, errorMessageMappingStr))
return nil, newApiErr
}
return usage, nil
Expand Down
5 changes: 3 additions & 2 deletions relay/claude_handler.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -189,6 +189,7 @@ func ClaudeHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ
}

statusCodeMappingStr := c.GetString("status_code_mapping")
errorMessageMappingStr := c.GetString("error_message_mapping")
var httpResp *http.Response
resp, err := adaptor.DoRequest(c, info, requestBody)
if err != nil {
Expand All @@ -201,15 +202,15 @@ func ClaudeHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ
if httpResp.StatusCode != http.StatusOK {
newAPIError = service.RelayErrorHandler(c.Request.Context(), httpResp, false)
// reset status code 重置状态码
service.ResetStatusCode(newAPIError, statusCodeMappingStr)
common.SetContextKey(c, constant.ContextKeyChannelErrorOverrideSummary, service.ApplyChannelErrorOverrides(newAPIError, statusCodeMappingStr, errorMessageMappingStr))
return newAPIError
}
}

usage, newAPIError := adaptor.DoResponse(c, httpResp, info)
if newAPIError != nil {
// reset status code 重置状态码
service.ResetStatusCode(newAPIError, statusCodeMappingStr)
common.SetContextKey(c, constant.ContextKeyChannelErrorOverrideSummary, service.ApplyChannelErrorOverrides(newAPIError, statusCodeMappingStr, errorMessageMappingStr))
return newAPIError
}

Expand Down
5 changes: 3 additions & 2 deletions relay/compatible_handler.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -192,22 +192,23 @@ func TextHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *types
}

statusCodeMappingStr := c.GetString("status_code_mapping")
errorMessageMappingStr := c.GetString("error_message_mapping")

if resp != nil {
httpResp = resp.(*http.Response)
info.IsStream = info.IsStream || strings.HasPrefix(httpResp.Header.Get("Content-Type"), "text/event-stream")
if httpResp.StatusCode != http.StatusOK {
newApiErr := service.RelayErrorHandler(c.Request.Context(), httpResp, false)
// reset status code 重置状态码
service.ResetStatusCode(newApiErr, statusCodeMappingStr)
common.SetContextKey(c, constant.ContextKeyChannelErrorOverrideSummary, service.ApplyChannelErrorOverrides(newApiErr, statusCodeMappingStr, errorMessageMappingStr))
return newApiErr
}
}

usage, newApiErr := adaptor.DoResponse(c, httpResp, info)
if newApiErr != nil {
// reset status code 重置状态码
service.ResetStatusCode(newApiErr, statusCodeMappingStr)
common.SetContextKey(c, constant.ContextKeyChannelErrorOverrideSummary, service.ApplyChannelErrorOverrides(newApiErr, statusCodeMappingStr, errorMessageMappingStr))
return newApiErr
}

Expand Down
6 changes: 4 additions & 2 deletions relay/embedding_handler.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ import (
"net/http"

"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/logger"
relaycommon "github.com/QuantumNous/new-api/relay/common"
Expand Down Expand Up @@ -67,6 +68,7 @@ func EmbeddingHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *
info.UpstreamRequestBodySize = size
var requestBody io.Reader = body
statusCodeMappingStr := c.GetString("status_code_mapping")
errorMessageMappingStr := c.GetString("error_message_mapping")
resp, err := adaptor.DoRequest(c, info, requestBody)
if err != nil {
return types.NewOpenAIError(err, types.ErrorCodeDoRequestFailed, http.StatusInternalServerError)
Expand All @@ -78,15 +80,15 @@ func EmbeddingHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *
if httpResp.StatusCode != http.StatusOK {
newAPIError = service.RelayErrorHandler(c.Request.Context(), httpResp, false)
// reset status code 重置状态码
service.ResetStatusCode(newAPIError, statusCodeMappingStr)
common.SetContextKey(c, constant.ContextKeyChannelErrorOverrideSummary, service.ApplyChannelErrorOverrides(newAPIError, statusCodeMappingStr, errorMessageMappingStr))
return newAPIError
}
}

usage, newAPIError := adaptor.DoResponse(c, httpResp, info)
if newAPIError != nil {
// reset status code 重置状态码
service.ResetStatusCode(newAPIError, statusCodeMappingStr)
common.SetContextKey(c, constant.ContextKeyChannelErrorOverrideSummary, service.ApplyChannelErrorOverrides(newAPIError, statusCodeMappingStr, errorMessageMappingStr))
return newAPIError
}
service.PostTextConsumeQuota(c, info, usage.(*dto.Usage), nil)
Expand Down
10 changes: 6 additions & 4 deletions relay/gemini_handler.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -181,6 +181,7 @@ func GeminiHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ
}

statusCodeMappingStr := c.GetString("status_code_mapping")
errorMessageMappingStr := c.GetString("error_message_mapping")

var httpResp *http.Response
if resp != nil {
Expand All @@ -189,14 +190,14 @@ func GeminiHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ
if httpResp.StatusCode != http.StatusOK {
newAPIError = service.RelayErrorHandler(c.Request.Context(), httpResp, false)
// reset status code 重置状态码
service.ResetStatusCode(newAPIError, statusCodeMappingStr)
common.SetContextKey(c, constant.ContextKeyChannelErrorOverrideSummary, service.ApplyChannelErrorOverrides(newAPIError, statusCodeMappingStr, errorMessageMappingStr))
return newAPIError
}
}

usage, openaiErr := adaptor.DoResponse(c, resp.(*http.Response), info)
if openaiErr != nil {
service.ResetStatusCode(openaiErr, statusCodeMappingStr)
common.SetContextKey(c, constant.ContextKeyChannelErrorOverrideSummary, service.ApplyChannelErrorOverrides(openaiErr, statusCodeMappingStr, errorMessageMappingStr))
return openaiErr
}

Expand Down Expand Up @@ -285,19 +286,20 @@ func GeminiEmbeddingHandler(c *gin.Context, info *relaycommon.RelayInfo) (newAPI
}

statusCodeMappingStr := c.GetString("status_code_mapping")
errorMessageMappingStr := c.GetString("error_message_mapping")
var httpResp *http.Response
if resp != nil {
httpResp = resp.(*http.Response)
if httpResp.StatusCode != http.StatusOK {
newAPIError = service.RelayErrorHandler(c.Request.Context(), httpResp, false)
service.ResetStatusCode(newAPIError, statusCodeMappingStr)
common.SetContextKey(c, constant.ContextKeyChannelErrorOverrideSummary, service.ApplyChannelErrorOverrides(newAPIError, statusCodeMappingStr, errorMessageMappingStr))
return newAPIError
}
}

usage, openaiErr := adaptor.DoResponse(c, resp.(*http.Response), info)
if openaiErr != nil {
service.ResetStatusCode(openaiErr, statusCodeMappingStr)
common.SetContextKey(c, constant.ContextKeyChannelErrorOverrideSummary, service.ApplyChannelErrorOverrides(openaiErr, statusCodeMappingStr, errorMessageMappingStr))
return openaiErr
}

Expand Down
5 changes: 3 additions & 2 deletions relay/image_handler.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -89,6 +89,7 @@ func ImageHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *type
}

statusCodeMappingStr := c.GetString("status_code_mapping")
errorMessageMappingStr := c.GetString("error_message_mapping")

resp, err := adaptor.DoRequest(c, info, requestBody)
if err != nil {
Expand All @@ -105,7 +106,7 @@ func ImageHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *type
} else {
newAPIError = service.RelayErrorHandler(c.Request.Context(), httpResp, false)
// reset status code 重置状态码
service.ResetStatusCode(newAPIError, statusCodeMappingStr)
common.SetContextKey(c, constant.ContextKeyChannelErrorOverrideSummary, service.ApplyChannelErrorOverrides(newAPIError, statusCodeMappingStr, errorMessageMappingStr))
return newAPIError
}
}
Expand All @@ -114,7 +115,7 @@ func ImageHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *type
usage, newAPIError := adaptor.DoResponse(c, httpResp, info)
if newAPIError != nil {
// reset status code 重置状态码
service.ResetStatusCode(newAPIError, statusCodeMappingStr)
common.SetContextKey(c, constant.ContextKeyChannelErrorOverrideSummary, service.ApplyChannelErrorOverrides(newAPIError, statusCodeMappingStr, errorMessageMappingStr))
return newAPIError
}

Expand Down
6 changes: 4 additions & 2 deletions relay/rerank_handler.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ import (
"net/http"

"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/logger"
relaycommon "github.com/QuantumNous/new-api/relay/common"
Expand Down Expand Up @@ -84,21 +85,22 @@ func RerankHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ
}

statusCodeMappingStr := c.GetString("status_code_mapping")
errorMessageMappingStr := c.GetString("error_message_mapping")
var httpResp *http.Response
if resp != nil {
httpResp = resp.(*http.Response)
if httpResp.StatusCode != http.StatusOK {
newAPIError = service.RelayErrorHandler(c.Request.Context(), httpResp, false)
// reset status code 重置状态码
service.ResetStatusCode(newAPIError, statusCodeMappingStr)
common.SetContextKey(c, constant.ContextKeyChannelErrorOverrideSummary, service.ApplyChannelErrorOverrides(newAPIError, statusCodeMappingStr, errorMessageMappingStr))
return newAPIError
}
}

usage, newAPIError := adaptor.DoResponse(c, httpResp, info)
if newAPIError != nil {
// reset status code 重置状态码
service.ResetStatusCode(newAPIError, statusCodeMappingStr)
common.SetContextKey(c, constant.ContextKeyChannelErrorOverrideSummary, service.ApplyChannelErrorOverrides(newAPIError, statusCodeMappingStr, errorMessageMappingStr))
return newAPIError
}
service.PostTextConsumeQuota(c, info, usage.(*dto.Usage), nil)
Expand Down
6 changes: 4 additions & 2 deletions relay/responses_handler.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ import (
"strings"

"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
appconstant "github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/logger"
Expand Down Expand Up @@ -120,22 +121,23 @@ func ResponsesHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *
}

statusCodeMappingStr := c.GetString("status_code_mapping")
errorMessageMappingStr := c.GetString("error_message_mapping")

if resp != nil {
httpResp = resp.(*http.Response)

if httpResp.StatusCode != http.StatusOK {
newAPIError = service.RelayErrorHandler(c.Request.Context(), httpResp, false)
// reset status code 重置状态码
service.ResetStatusCode(newAPIError, statusCodeMappingStr)
common.SetContextKey(c, constant.ContextKeyChannelErrorOverrideSummary, service.ApplyChannelErrorOverrides(newAPIError, statusCodeMappingStr, errorMessageMappingStr))
return newAPIError
}
}

usage, newAPIError := adaptor.DoResponse(c, httpResp, info)
if newAPIError != nil {
// reset status code 重置状态码
service.ResetStatusCode(newAPIError, statusCodeMappingStr)
common.SetContextKey(c, constant.ContextKeyChannelErrorOverrideSummary, service.ApplyChannelErrorOverrides(newAPIError, statusCodeMappingStr, errorMessageMappingStr))
return newAPIError
}

Expand Down
Loading