Commit aebc8ae2 by Seefs

feat: add retry-aware param override with return_error and prune_objects

parent 37e4fccb
...@@ -366,7 +366,7 @@ func testChannel(channel *model.Channel, testModel string, endpointType string, ...@@ -366,7 +366,7 @@ func testChannel(channel *model.Channel, testModel string, endpointType string,
newAPIError: types.NewError(err, types.ErrorCodeConvertRequestFailed), newAPIError: types.NewError(err, types.ErrorCodeConvertRequestFailed),
} }
} }
jsonData, err := json.Marshal(convertedRequest) jsonData, err := common.Marshal(convertedRequest)
if err != nil { if err != nil {
return testResult{ return testResult{
context: c, context: c,
...@@ -387,6 +387,13 @@ func testChannel(channel *model.Channel, testModel string, endpointType string, ...@@ -387,6 +387,13 @@ func testChannel(channel *model.Channel, testModel string, endpointType string,
if len(info.ParamOverride) > 0 { if len(info.ParamOverride) > 0 {
jsonData, err = relaycommon.ApplyParamOverride(jsonData, info.ParamOverride, relaycommon.BuildParamOverrideContext(info)) jsonData, err = relaycommon.ApplyParamOverride(jsonData, info.ParamOverride, relaycommon.BuildParamOverrideContext(info))
if err != nil { if err != nil {
if fixedErr, ok := relaycommon.AsParamOverrideReturnError(err); ok {
return testResult{
context: c,
localErr: fixedErr,
newAPIError: relaycommon.NewAPIErrorFromParamOverride(fixedErr),
}
}
return testResult{ return testResult{
context: c, context: c,
localErr: err, localErr: err,
......
...@@ -182,8 +182,11 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) { ...@@ -182,8 +182,11 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) {
ModelName: relayInfo.OriginModelName, ModelName: relayInfo.OriginModelName,
Retry: common.GetPointer(0), Retry: common.GetPointer(0),
} }
relayInfo.RetryIndex = 0
relayInfo.LastError = nil
for ; retryParam.GetRetry() <= common.RetryTimes; retryParam.IncreaseRetry() { for ; retryParam.GetRetry() <= common.RetryTimes; retryParam.IncreaseRetry() {
relayInfo.RetryIndex = retryParam.GetRetry()
channel, channelErr := getChannel(c, relayInfo, retryParam) channel, channelErr := getChannel(c, relayInfo, retryParam)
if channelErr != nil { if channelErr != nil {
logger.LogError(c, channelErr.Error()) logger.LogError(c, channelErr.Error())
...@@ -216,10 +219,12 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) { ...@@ -216,10 +219,12 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) {
} }
if newAPIError == nil { if newAPIError == nil {
relayInfo.LastError = nil
return return
} }
newAPIError = service.NormalizeViolationFeeError(newAPIError) newAPIError = service.NormalizeViolationFeeError(newAPIError)
relayInfo.LastError = newAPIError
processChannelError(c, *types.NewChannelError(channel.Id, channel.Type, channel.Name, channel.ChannelInfo.IsMultiKey, common.GetContextKeyString(c, constant.ContextKeyChannelKey), channel.GetAutoBan()), newAPIError) processChannelError(c, *types.NewChannelError(channel.Id, channel.Type, channel.Name, channel.ChannelInfo.IsMultiKey, common.GetContextKeyString(c, constant.ContextKeyChannelKey), channel.GetAutoBan()), newAPIError)
......
...@@ -84,7 +84,7 @@ func chatCompletionsViaResponses(c *gin.Context, info *relaycommon.RelayInfo, ad ...@@ -84,7 +84,7 @@ func chatCompletionsViaResponses(c *gin.Context, info *relaycommon.RelayInfo, ad
if len(info.ParamOverride) > 0 { if len(info.ParamOverride) > 0 {
chatJSON, err = relaycommon.ApplyParamOverride(chatJSON, info.ParamOverride, overrideCtx) chatJSON, err = relaycommon.ApplyParamOverride(chatJSON, info.ParamOverride, overrideCtx)
if err != nil { if err != nil {
return nil, types.NewError(err, types.ErrorCodeChannelParamOverrideInvalid, types.ErrOptionWithSkipRetry()) return nil, newAPIErrorFromParamOverride(err)
} }
} }
......
...@@ -155,7 +155,7 @@ func ClaudeHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ ...@@ -155,7 +155,7 @@ func ClaudeHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ
if len(info.ParamOverride) > 0 { if len(info.ParamOverride) > 0 {
jsonData, err = relaycommon.ApplyParamOverride(jsonData, info.ParamOverride, relaycommon.BuildParamOverrideContext(info)) jsonData, err = relaycommon.ApplyParamOverride(jsonData, info.ParamOverride, relaycommon.BuildParamOverrideContext(info))
if err != nil { if err != nil {
return types.NewError(err, types.ErrorCodeChannelParamOverrideInvalid, types.ErrOptionWithSkipRetry()) return newAPIErrorFromParamOverride(err)
} }
} }
......
...@@ -4,6 +4,8 @@ import ( ...@@ -4,6 +4,8 @@ import (
"encoding/json" "encoding/json"
"reflect" "reflect"
"testing" "testing"
"github.com/QuantumNous/new-api/types"
) )
func TestApplyParamOverrideTrimPrefix(t *testing.T) { func TestApplyParamOverrideTrimPrefix(t *testing.T) {
...@@ -772,6 +774,188 @@ func TestApplyParamOverrideToUpper(t *testing.T) { ...@@ -772,6 +774,188 @@ func TestApplyParamOverrideToUpper(t *testing.T) {
assertJSONEqual(t, `{"model":"GPT-4"}`, string(out)) assertJSONEqual(t, `{"model":"GPT-4"}`, string(out))
} }
func TestApplyParamOverrideReturnError(t *testing.T) {
input := []byte(`{"model":"gemini-2.5-pro"}`)
override := map[string]interface{}{
"operations": []interface{}{
map[string]interface{}{
"mode": "return_error",
"value": map[string]interface{}{
"message": "forced bad request by param override",
"status_code": 422,
"code": "forced_bad_request",
"type": "invalid_request_error",
"skip_retry": true,
},
"conditions": []interface{}{
map[string]interface{}{
"path": "retry.is_retry",
"mode": "full",
"value": true,
},
},
},
},
}
ctx := map[string]interface{}{
"retry": map[string]interface{}{
"index": 1,
"is_retry": true,
},
}
_, err := ApplyParamOverride(input, override, ctx)
if err == nil {
t.Fatalf("expected error, got nil")
}
returnErr, ok := AsParamOverrideReturnError(err)
if !ok {
t.Fatalf("expected ParamOverrideReturnError, got %T: %v", err, err)
}
if returnErr.StatusCode != 422 {
t.Fatalf("expected status 422, got %d", returnErr.StatusCode)
}
if returnErr.Code != "forced_bad_request" {
t.Fatalf("expected code forced_bad_request, got %s", returnErr.Code)
}
if !returnErr.SkipRetry {
t.Fatalf("expected skip_retry true")
}
}
func TestApplyParamOverridePruneObjectsByTypeString(t *testing.T) {
input := []byte(`{
"messages":[
{"role":"assistant","content":[
{"type":"output_text","text":"a"},
{"type":"redacted_thinking","text":"secret"},
{"type":"tool_call","name":"tool_a"}
]},
{"role":"assistant","content":[
{"type":"output_text","text":"b"},
{"type":"wrapper","parts":[
{"type":"redacted_thinking","text":"secret2"},
{"type":"output_text","text":"c"}
]}
]}
]
}`)
override := map[string]interface{}{
"operations": []interface{}{
map[string]interface{}{
"mode": "prune_objects",
"value": "redacted_thinking",
},
},
}
out, err := ApplyParamOverride(input, override, nil)
if err != nil {
t.Fatalf("ApplyParamOverride returned error: %v", err)
}
assertJSONEqual(t, `{
"messages":[
{"role":"assistant","content":[
{"type":"output_text","text":"a"},
{"type":"tool_call","name":"tool_a"}
]},
{"role":"assistant","content":[
{"type":"output_text","text":"b"},
{"type":"wrapper","parts":[
{"type":"output_text","text":"c"}
]}
]}
]
}`, string(out))
}
func TestApplyParamOverridePruneObjectsWhereAndPath(t *testing.T) {
input := []byte(`{
"a":{"items":[{"type":"redacted_thinking","id":1},{"type":"output_text","id":2}]},
"b":{"items":[{"type":"redacted_thinking","id":3},{"type":"output_text","id":4}]}
}`)
override := map[string]interface{}{
"operations": []interface{}{
map[string]interface{}{
"path": "a",
"mode": "prune_objects",
"value": map[string]interface{}{
"where": map[string]interface{}{
"type": "redacted_thinking",
},
},
},
},
}
out, err := ApplyParamOverride(input, override, nil)
if err != nil {
t.Fatalf("ApplyParamOverride returned error: %v", err)
}
assertJSONEqual(t, `{
"a":{"items":[{"type":"output_text","id":2}]},
"b":{"items":[{"type":"redacted_thinking","id":3},{"type":"output_text","id":4}]}
}`, string(out))
}
func TestApplyParamOverrideNormalizeThinkingSignatureUnsupported(t *testing.T) {
input := []byte(`{"items":[{"type":"redacted_thinking"}]}`)
override := map[string]interface{}{
"operations": []interface{}{
map[string]interface{}{
"mode": "normalize_thinking_signature",
},
},
}
_, err := ApplyParamOverride(input, override, nil)
if err == nil {
t.Fatalf("expected error, got nil")
}
}
func TestApplyParamOverrideConditionFromRetryAndLastErrorContext(t *testing.T) {
info := &RelayInfo{
RetryIndex: 1,
LastError: types.WithOpenAIError(types.OpenAIError{
Message: "invalid thinking signature",
Type: "invalid_request_error",
Code: "bad_thought_signature",
}, 400),
}
ctx := BuildParamOverrideContext(info)
input := []byte(`{"temperature":0.7}`)
override := map[string]interface{}{
"operations": []interface{}{
map[string]interface{}{
"path": "temperature",
"mode": "set",
"value": 0.1,
"logic": "AND",
"conditions": []interface{}{
map[string]interface{}{
"path": "is_retry",
"mode": "full",
"value": true,
},
map[string]interface{}{
"path": "last_error.code",
"mode": "contains",
"value": "thought_signature",
},
},
},
},
}
out, err := ApplyParamOverride(input, override, ctx)
if err != nil {
t.Fatalf("ApplyParamOverride returned error: %v", err)
}
assertJSONEqual(t, `{"temperature":0.1}`, string(out))
}
func assertJSONEqual(t *testing.T, want, got string) { func assertJSONEqual(t *testing.T, want, got string) {
t.Helper() t.Helper()
......
...@@ -140,6 +140,8 @@ type RelayInfo struct { ...@@ -140,6 +140,8 @@ type RelayInfo struct {
SubscriptionAmountUsedAfterPreConsume int64 SubscriptionAmountUsedAfterPreConsume int64
IsClaudeBetaQuery bool // /v1/messages?beta=true IsClaudeBetaQuery bool // /v1/messages?beta=true
IsChannelTest bool // channel test request IsChannelTest bool // channel test request
RetryIndex int
LastError *types.NewAPIError
PriceData types.PriceData PriceData types.PriceData
......
...@@ -174,7 +174,7 @@ func TextHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *types ...@@ -174,7 +174,7 @@ func TextHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *types
if len(info.ParamOverride) > 0 { if len(info.ParamOverride) > 0 {
jsonData, err = relaycommon.ApplyParamOverride(jsonData, info.ParamOverride, relaycommon.BuildParamOverrideContext(info)) jsonData, err = relaycommon.ApplyParamOverride(jsonData, info.ParamOverride, relaycommon.BuildParamOverrideContext(info))
if err != nil { if err != nil {
return types.NewError(err, types.ErrorCodeChannelParamOverrideInvalid, types.ErrOptionWithSkipRetry()) return newAPIErrorFromParamOverride(err)
} }
} }
......
...@@ -2,7 +2,6 @@ package relay ...@@ -2,7 +2,6 @@ package relay
import ( import (
"bytes" "bytes"
"encoding/json"
"fmt" "fmt"
"net/http" "net/http"
...@@ -46,7 +45,7 @@ func EmbeddingHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError * ...@@ -46,7 +45,7 @@ func EmbeddingHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *
return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
} }
relaycommon.AppendRequestConversionFromRequest(info, convertedRequest) relaycommon.AppendRequestConversionFromRequest(info, convertedRequest)
jsonData, err := json.Marshal(convertedRequest) jsonData, err := common.Marshal(convertedRequest)
if err != nil { if err != nil {
return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
} }
...@@ -54,7 +53,7 @@ func EmbeddingHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError * ...@@ -54,7 +53,7 @@ func EmbeddingHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *
if len(info.ParamOverride) > 0 { if len(info.ParamOverride) > 0 {
jsonData, err = relaycommon.ApplyParamOverride(jsonData, info.ParamOverride, relaycommon.BuildParamOverrideContext(info)) jsonData, err = relaycommon.ApplyParamOverride(jsonData, info.ParamOverride, relaycommon.BuildParamOverrideContext(info))
if err != nil { if err != nil {
return types.NewError(err, types.ErrorCodeChannelParamOverrideInvalid, types.ErrOptionWithSkipRetry()) return newAPIErrorFromParamOverride(err)
} }
} }
......
...@@ -159,7 +159,7 @@ func GeminiHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ ...@@ -159,7 +159,7 @@ func GeminiHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ
if len(info.ParamOverride) > 0 { if len(info.ParamOverride) > 0 {
jsonData, err = relaycommon.ApplyParamOverride(jsonData, info.ParamOverride, relaycommon.BuildParamOverrideContext(info)) jsonData, err = relaycommon.ApplyParamOverride(jsonData, info.ParamOverride, relaycommon.BuildParamOverrideContext(info))
if err != nil { if err != nil {
return types.NewError(err, types.ErrorCodeChannelParamOverrideInvalid, types.ErrOptionWithSkipRetry()) return newAPIErrorFromParamOverride(err)
} }
} }
...@@ -257,14 +257,9 @@ func GeminiEmbeddingHandler(c *gin.Context, info *relaycommon.RelayInfo) (newAPI ...@@ -257,14 +257,9 @@ func GeminiEmbeddingHandler(c *gin.Context, info *relaycommon.RelayInfo) (newAPI
// apply param override // apply param override
if len(info.ParamOverride) > 0 { if len(info.ParamOverride) > 0 {
reqMap := make(map[string]interface{}) jsonData, err = relaycommon.ApplyParamOverride(jsonData, info.ParamOverride, relaycommon.BuildParamOverrideContext(info))
_ = common.Unmarshal(jsonData, &reqMap)
for key, value := range info.ParamOverride {
reqMap[key] = value
}
jsonData, err = common.Marshal(reqMap)
if err != nil { if err != nil {
return types.NewError(err, types.ErrorCodeChannelParamOverrideInvalid, types.ErrOptionWithSkipRetry()) return newAPIErrorFromParamOverride(err)
} }
} }
logger.LogDebug(c, "Gemini embedding request body: "+string(jsonData)) logger.LogDebug(c, "Gemini embedding request body: "+string(jsonData))
......
...@@ -72,7 +72,7 @@ func ImageHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *type ...@@ -72,7 +72,7 @@ func ImageHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *type
if len(info.ParamOverride) > 0 { if len(info.ParamOverride) > 0 {
jsonData, err = relaycommon.ApplyParamOverride(jsonData, info.ParamOverride, relaycommon.BuildParamOverrideContext(info)) jsonData, err = relaycommon.ApplyParamOverride(jsonData, info.ParamOverride, relaycommon.BuildParamOverrideContext(info))
if err != nil { if err != nil {
return types.NewError(err, types.ErrorCodeChannelParamOverrideInvalid, types.ErrOptionWithSkipRetry()) return newAPIErrorFromParamOverride(err)
} }
} }
......
package relay
import (
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/types"
)
func newAPIErrorFromParamOverride(err error) *types.NewAPIError {
if fixedErr, ok := relaycommon.AsParamOverrideReturnError(err); ok {
return relaycommon.NewAPIErrorFromParamOverride(fixedErr)
}
return types.NewError(err, types.ErrorCodeChannelParamOverrideInvalid, types.ErrOptionWithSkipRetry())
}
...@@ -63,7 +63,7 @@ func RerankHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ ...@@ -63,7 +63,7 @@ func RerankHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ
if len(info.ParamOverride) > 0 { if len(info.ParamOverride) > 0 {
jsonData, err = relaycommon.ApplyParamOverride(jsonData, info.ParamOverride, relaycommon.BuildParamOverrideContext(info)) jsonData, err = relaycommon.ApplyParamOverride(jsonData, info.ParamOverride, relaycommon.BuildParamOverrideContext(info))
if err != nil { if err != nil {
return types.NewError(err, types.ErrorCodeChannelParamOverrideInvalid, types.ErrOptionWithSkipRetry()) return newAPIErrorFromParamOverride(err)
} }
} }
......
...@@ -98,7 +98,7 @@ func ResponsesHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError * ...@@ -98,7 +98,7 @@ func ResponsesHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *
if len(info.ParamOverride) > 0 { if len(info.ParamOverride) > 0 {
jsonData, err = relaycommon.ApplyParamOverride(jsonData, info.ParamOverride, relaycommon.BuildParamOverrideContext(info)) jsonData, err = relaycommon.ApplyParamOverride(jsonData, info.ParamOverride, relaycommon.BuildParamOverrideContext(info))
if err != nil { if err != nil {
return types.NewError(err, types.ErrorCodeChannelParamOverrideInvalid, types.ErrOptionWithSkipRetry()) return newAPIErrorFromParamOverride(err)
} }
} }
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or sign in to comment