Skip to content

phsl / new-api

Go to a project
Commit 7633863c by Seefs
Options

feat: unify param/header overrides with retry-aware conditions and flexible header operations

parent aebc8ae2
Showing with 324 additions and 13 deletions
controller/channel-test.go
...@@ -385,7 +385,7 @@ func testChannel(channel *model.Channel, testModel string, endpointType string, ...@@ -385,7 +385,7 @@ 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.ApplyParamOverrideWithRelayInfo(jsonData, info)
if err != nil { if err != nil {
if fixedErr, ok := relaycommon.AsParamOverrideReturnError(err); ok { if fixedErr, ok := relaycommon.AsParamOverrideReturnError(err); ok {
return testResult{ return testResult{
......
relay/channel/api_request.go
...@@ -168,11 +168,19 @@ func applyHeaderOverridePlaceholders(template string, c *gin.Context, apiKey str ...@@ -168,11 +168,19 @@ func applyHeaderOverridePlaceholders(template string, c *gin.Context, apiKey str
// Passthrough rules are applied first, then normal overrides are applied, so explicit overrides win. // Passthrough rules are applied first, then normal overrides are applied, so explicit overrides win.
func processHeaderOverride(info *common.RelayInfo, c *gin.Context) (map[string]string, error) { func processHeaderOverride(info *common.RelayInfo, c *gin.Context) (map[string]string, error) {
headerOverride := make(map[string]string) headerOverride := make(map[string]string)
if info == nil {
return headerOverride, nil
}
headerOverrideSource := info.HeadersOverride
if info.UseRuntimeHeadersOverride {
headerOverrideSource = info.RuntimeHeadersOverride
}
passAll := false passAll := false
var passthroughRegex []*regexp.Regexp var passthroughRegex []*regexp.Regexp
if !info.IsChannelTest { if !info.IsChannelTest {
for k := range info.HeadersOverride { for k := range headerOverrideSource {
key := strings.TrimSpace(k) key := strings.TrimSpace(k)
if key == "" { if key == "" {
continue continue
...@@ -232,7 +240,7 @@ func processHeaderOverride(info *common.RelayInfo, c *gin.Context) (map[string]s ...@@ -232,7 +240,7 @@ func processHeaderOverride(info *common.RelayInfo, c *gin.Context) (map[string]s
} }
} }
for k, v := range info.HeadersOverride { for k, v := range headerOverrideSource {
if isHeaderPassthroughRuleKey(k) { if isHeaderPassthroughRuleKey(k) {
continue continue
} }
......
relay/channel/api_request_test.go
...@@ -79,3 +79,34 @@ func TestProcessHeaderOverride_NonTestKeepsClientHeaderPlaceholder(t *testing.T) ...@@ -79,3 +79,34 @@ func TestProcessHeaderOverride_NonTestKeepsClientHeaderPlaceholder(t *testing.T)
require.NoError(t, err) require.NoError(t, err)
require.Equal(t, "trace-123", headers["X-Upstream-Trace"]) require.Equal(t, "trace-123", headers["X-Upstream-Trace"])
} }
func TestProcessHeaderOverride_RuntimeOverrideHasPriority(t *testing.T) {
t.Parallel()
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
ctx, _ := gin.CreateTestContext(recorder)
ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
info := &relaycommon.RelayInfo{
IsChannelTest: false,
UseRuntimeHeadersOverride: true,
RuntimeHeadersOverride: map[string]any{
"X-Static": "runtime-value",
"X-Runtime": "runtime-only",
},
ChannelMeta: &relaycommon.ChannelMeta{
HeadersOverride: map[string]any{
"X-Static": "legacy-value",
"X-Legacy": "legacy-only",
},
},
}
headers, err := processHeaderOverride(info, ctx)
require.NoError(t, err)
require.Equal(t, "runtime-value", headers["X-Static"])
require.Equal(t, "runtime-only", headers["X-Runtime"])
_, ok := headers["X-Legacy"]
require.False(t, ok)
}
relay/chat_completions_via_responses.go
...@@ -70,7 +70,6 @@ func applySystemPromptIfNeeded(c *gin.Context, info *relaycommon.RelayInfo, requ ...@@ -70,7 +70,6 @@ func applySystemPromptIfNeeded(c *gin.Context, info *relaycommon.RelayInfo, requ
} }
func chatCompletionsViaResponses(c *gin.Context, info *relaycommon.RelayInfo, adaptor channel.Adaptor, request *dto.GeneralOpenAIRequest) (*dto.Usage, *types.NewAPIError) { func chatCompletionsViaResponses(c *gin.Context, info *relaycommon.RelayInfo, adaptor channel.Adaptor, request *dto.GeneralOpenAIRequest) (*dto.Usage, *types.NewAPIError) {
overrideCtx := relaycommon.BuildParamOverrideContext(info)
chatJSON, err := common.Marshal(request) chatJSON, err := common.Marshal(request)
if err != nil { if err != nil {
return nil, types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) return nil, types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
...@@ -82,7 +81,7 @@ func chatCompletionsViaResponses(c *gin.Context, info *relaycommon.RelayInfo, ad ...@@ -82,7 +81,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.ApplyParamOverrideWithRelayInfo(chatJSON, info)
if err != nil { if err != nil {
return nil, newAPIErrorFromParamOverride(err) return nil, newAPIErrorFromParamOverride(err)
} }
......
relay/claude_handler.go
...@@ -153,7 +153,7 @@ func ClaudeHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ ...@@ -153,7 +153,7 @@ func ClaudeHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ
// apply param override // apply param override
if len(info.ParamOverride) > 0 { if len(info.ParamOverride) > 0 {
jsonData, err = relaycommon.ApplyParamOverride(jsonData, info.ParamOverride, relaycommon.BuildParamOverrideContext(info)) jsonData, err = relaycommon.ApplyParamOverrideWithRelayInfo(jsonData, info)
if err != nil { if err != nil {
return newAPIErrorFromParamOverride(err) return newAPIErrorFromParamOverride(err)
} }
......
relay/common/override_test.go
...@@ -956,6 +956,254 @@ func TestApplyParamOverrideConditionFromRetryAndLastErrorContext(t *testing.T) { ...@@ -956,6 +956,254 @@ func TestApplyParamOverrideConditionFromRetryAndLastErrorContext(t *testing.T) {
assertJSONEqual(t, `{"temperature":0.1}`, string(out)) assertJSONEqual(t, `{"temperature":0.1}`, string(out))
} }
func TestApplyParamOverrideConditionFromRequestHeaders(t *testing.T) {
input := []byte(`{"temperature":0.7}`)
override := map[string]interface{}{
"operations": []interface{}{
map[string]interface{}{
"path": "temperature",
"mode": "set",
"value": 0.1,
"conditions": []interface{}{
map[string]interface{}{
"path": "request_headers.authorization",
"mode": "contains",
"value": "Bearer ",
},
},
},
},
}
ctx := map[string]interface{}{
"request_headers": map[string]interface{}{
"authorization": "Bearer token-123",
},
}
out, err := ApplyParamOverride(input, override, ctx)
if err != nil {
t.Fatalf("ApplyParamOverride returned error: %v", err)
}
assertJSONEqual(t, `{"temperature":0.1}`, string(out))
}
func TestApplyParamOverrideSetHeaderAndUseInLaterCondition(t *testing.T) {
input := []byte(`{"temperature":0.7}`)
override := map[string]interface{}{
"operations": []interface{}{
map[string]interface{}{
"mode": "set_header",
"path": "X-Debug-Mode",
"value": "enabled",
},
map[string]interface{}{
"path": "temperature",
"mode": "set",
"value": 0.1,
"conditions": []interface{}{
map[string]interface{}{
"path": "header_override_normalized.x_debug_mode",
"mode": "full",
"value": "enabled",
},
},
},
},
}
out, err := ApplyParamOverride(input, override, nil)
if err != nil {
t.Fatalf("ApplyParamOverride returned error: %v", err)
}
assertJSONEqual(t, `{"temperature":0.1}`, string(out))
}
func TestApplyParamOverrideCopyHeaderFromRequestHeaders(t *testing.T) {
input := []byte(`{"temperature":0.7}`)
override := map[string]interface{}{
"operations": []interface{}{
map[string]interface{}{
"mode": "copy_header",
"from": "Authorization",
"to": "X-Upstream-Auth",
},
map[string]interface{}{
"path": "temperature",
"mode": "set",
"value": 0.1,
"conditions": []interface{}{
map[string]interface{}{
"path": "header_override_normalized.x_upstream_auth",
"mode": "contains",
"value": "Bearer ",
},
},
},
},
}
ctx := map[string]interface{}{
"request_headers_raw": map[string]interface{}{
"Authorization": "Bearer token-123",
},
"request_headers": map[string]interface{}{
"authorization": "Bearer token-123",
},
}
out, err := ApplyParamOverride(input, override, ctx)
if err != nil {
t.Fatalf("ApplyParamOverride returned error: %v", err)
}
assertJSONEqual(t, `{"temperature":0.1}`, string(out))
}
func TestApplyParamOverrideSetHeaderKeepOrigin(t *testing.T) {
input := []byte(`{"temperature":0.7}`)
override := map[string]interface{}{
"operations": []interface{}{
map[string]interface{}{
"mode": "set_header",
"path": "X-Feature-Flag",
"value": "new-value",
"keep_origin": true,
},
},
}
ctx := map[string]interface{}{
"header_override": map[string]interface{}{
"X-Feature-Flag": "legacy-value",
},
"header_override_normalized": map[string]interface{}{
"x_feature_flag": "legacy-value",
},
}
_, err := ApplyParamOverride(input, override, ctx)
if err != nil {
t.Fatalf("ApplyParamOverride returned error: %v", err)
}
headers, ok := ctx["header_override"].(map[string]interface{})
if !ok {
t.Fatalf("expected header_override context map")
}
if headers["X-Feature-Flag"] != "legacy-value" {
t.Fatalf("expected keep_origin to preserve old value, got: %v", headers["X-Feature-Flag"])
}
}
func TestApplyParamOverrideConditionsObjectShorthand(t *testing.T) {
input := []byte(`{"temperature":0.7}`)
override := map[string]interface{}{
"operations": []interface{}{
map[string]interface{}{
"path": "temperature",
"mode": "set",
"value": 0.1,
"logic": "AND",
"conditions": map[string]interface{}{
"is_retry": true,
"last_error.status_code": 400.0,
},
},
},
}
ctx := map[string]interface{}{
"is_retry": true,
"last_error": map[string]interface{}{
"status_code": 400.0,
},
}
out, err := ApplyParamOverride(input, override, ctx)
if err != nil {
t.Fatalf("ApplyParamOverride returned error: %v", err)
}
assertJSONEqual(t, `{"temperature":0.1}`, string(out))
}
func TestApplyParamOverrideWithRelayInfoSyncRuntimeHeaders(t *testing.T) {
info := &RelayInfo{
ChannelMeta: &ChannelMeta{
ParamOverride: map[string]interface{}{
"operations": []interface{}{
map[string]interface{}{
"mode": "set_header",
"path": "X-Injected-By-Param-Override",
"value": "enabled",
},
map[string]interface{}{
"mode": "delete_header",
"path": "X-Delete-Me",
},
},
},
HeadersOverride: map[string]interface{}{
"X-Delete-Me": "legacy",
"X-Keep-Me": "keep",
},
},
}
input := []byte(`{"temperature":0.7}`)
out, err := ApplyParamOverrideWithRelayInfo(input, info)
if err != nil {
t.Fatalf("ApplyParamOverrideWithRelayInfo returned error: %v", err)
}
assertJSONEqual(t, `{"temperature":0.7}`, string(out))
if !info.UseRuntimeHeadersOverride {
t.Fatalf("expected runtime header override to be enabled")
}
if info.RuntimeHeadersOverride["X-Keep-Me"] != "keep" {
t.Fatalf("expected X-Keep-Me header to be preserved, got: %v", info.RuntimeHeadersOverride["X-Keep-Me"])
}
if info.RuntimeHeadersOverride["X-Injected-By-Param-Override"] != "enabled" {
t.Fatalf("expected X-Injected-By-Param-Override header to be set, got: %v", info.RuntimeHeadersOverride["X-Injected-By-Param-Override"])
}
if _, exists := info.RuntimeHeadersOverride["X-Delete-Me"]; exists {
t.Fatalf("expected X-Delete-Me header to be deleted")
}
}
func TestApplyParamOverrideWithRelayInfoMoveAndCopyHeaders(t *testing.T) {
info := &RelayInfo{
ChannelMeta: &ChannelMeta{
ParamOverride: map[string]interface{}{
"operations": []interface{}{
map[string]interface{}{
"mode": "move_header",
"from": "X-Legacy-Trace",
"to": "X-Trace",
},
map[string]interface{}{
"mode": "copy_header",
"from": "X-Trace",
"to": "X-Trace-Backup",
},
},
},
HeadersOverride: map[string]interface{}{
"X-Legacy-Trace": "trace-123",
},
},
}
input := []byte(`{"temperature":0.7}`)
_, err := ApplyParamOverrideWithRelayInfo(input, info)
if err != nil {
t.Fatalf("ApplyParamOverrideWithRelayInfo returned error: %v", err)
}
if _, exists := info.RuntimeHeadersOverride["X-Legacy-Trace"]; exists {
t.Fatalf("expected source header to be removed after move")
}
if info.RuntimeHeadersOverride["X-Trace"] != "trace-123" {
t.Fatalf("expected X-Trace to be set, got: %v", info.RuntimeHeadersOverride["X-Trace"])
}
if info.RuntimeHeadersOverride["X-Trace-Backup"] != "trace-123" {
t.Fatalf("expected X-Trace-Backup to be copied, got: %v", info.RuntimeHeadersOverride["X-Trace-Backup"])
}
}
func assertJSONEqual(t *testing.T, want, got string) { func assertJSONEqual(t *testing.T, want, got string) {
t.Helper() t.Helper()
......
relay/common/relay_info.go
...@@ -101,6 +101,7 @@ type RelayInfo struct { ...@@ -101,6 +101,7 @@ type RelayInfo struct {
RelayMode int RelayMode int
OriginModelName string OriginModelName string
RequestURLPath string RequestURLPath string
RequestHeaders map[string]string
ShouldIncludeUsage bool ShouldIncludeUsage bool
DisablePing bool // 是否禁止向下游发送自定义 Ping DisablePing bool // 是否禁止向下游发送自定义 Ping
ClientWs *websocket.Conn ClientWs *websocket.Conn
...@@ -142,6 +143,8 @@ type RelayInfo struct { ...@@ -142,6 +143,8 @@ type RelayInfo struct {
IsChannelTest bool // channel test request IsChannelTest bool // channel test request
RetryIndex int RetryIndex int
LastError *types.NewAPIError LastError *types.NewAPIError
RuntimeHeadersOverride map[string]interface{}
UseRuntimeHeadersOverride bool
PriceData types.PriceData PriceData types.PriceData
...@@ -458,6 +461,7 @@ func genBaseRelayInfo(c *gin.Context, request dto.Request) *RelayInfo { ...@@ -458,6 +461,7 @@ func genBaseRelayInfo(c *gin.Context, request dto.Request) *RelayInfo {
isFirstResponse: true, isFirstResponse: true,
RelayMode: relayconstant.Path2RelayMode(c.Request.URL.Path), RelayMode: relayconstant.Path2RelayMode(c.Request.URL.Path),
RequestURLPath: c.Request.URL.String(), RequestURLPath: c.Request.URL.String(),
RequestHeaders: cloneRequestHeaders(c),
IsStream: isStream, IsStream: isStream,
StartTime: startTime, StartTime: startTime,
...@@ -490,6 +494,27 @@ func genBaseRelayInfo(c *gin.Context, request dto.Request) *RelayInfo { ...@@ -490,6 +494,27 @@ func genBaseRelayInfo(c *gin.Context, request dto.Request) *RelayInfo {
return info return info
} }
func cloneRequestHeaders(c *gin.Context) map[string]string {
if c == nil || c.Request == nil {
return nil
}
if len(c.Request.Header) == 0 {
return nil
}
headers := make(map[string]string, len(c.Request.Header))
for key := range c.Request.Header {
value := strings.TrimSpace(c.Request.Header.Get(key))
if value == "" {
continue
}
headers[key] = value
}
if len(headers) == 0 {
return nil
}
return headers
}
func GenRelayInfo(c *gin.Context, relayFormat types.RelayFormat, request dto.Request, ws *websocket.Conn) (*RelayInfo, error) { func GenRelayInfo(c *gin.Context, relayFormat types.RelayFormat, request dto.Request, ws *websocket.Conn) (*RelayInfo, error) {
var info *RelayInfo var info *RelayInfo
var err error var err error
......
relay/compatible_handler.go
...@@ -172,7 +172,7 @@ func TextHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *types ...@@ -172,7 +172,7 @@ func TextHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *types
// apply param override // apply param override
if len(info.ParamOverride) > 0 { if len(info.ParamOverride) > 0 {
jsonData, err = relaycommon.ApplyParamOverride(jsonData, info.ParamOverride, relaycommon.BuildParamOverrideContext(info)) jsonData, err = relaycommon.ApplyParamOverrideWithRelayInfo(jsonData, info)
if err != nil { if err != nil {
return newAPIErrorFromParamOverride(err) return newAPIErrorFromParamOverride(err)
} }
......
relay/embedding_handler.go
...@@ -51,7 +51,7 @@ func EmbeddingHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError * ...@@ -51,7 +51,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.ApplyParamOverrideWithRelayInfo(jsonData, info)
if err != nil { if err != nil {
return newAPIErrorFromParamOverride(err) return newAPIErrorFromParamOverride(err)
} }
......
relay/gemini_handler.go
...@@ -157,7 +157,7 @@ func GeminiHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ ...@@ -157,7 +157,7 @@ func GeminiHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ
// apply param override // apply param override
if len(info.ParamOverride) > 0 { if len(info.ParamOverride) > 0 {
jsonData, err = relaycommon.ApplyParamOverride(jsonData, info.ParamOverride, relaycommon.BuildParamOverrideContext(info)) jsonData, err = relaycommon.ApplyParamOverrideWithRelayInfo(jsonData, info)
if err != nil { if err != nil {
return newAPIErrorFromParamOverride(err) return newAPIErrorFromParamOverride(err)
} }
...@@ -257,7 +257,7 @@ func GeminiEmbeddingHandler(c *gin.Context, info *relaycommon.RelayInfo) (newAPI ...@@ -257,7 +257,7 @@ 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 {
jsonData, err = relaycommon.ApplyParamOverride(jsonData, info.ParamOverride, relaycommon.BuildParamOverrideContext(info)) jsonData, err = relaycommon.ApplyParamOverrideWithRelayInfo(jsonData, info)
if err != nil { if err != nil {
return newAPIErrorFromParamOverride(err) return newAPIErrorFromParamOverride(err)
} }
......
relay/image_handler.go
...@@ -70,7 +70,7 @@ func ImageHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *type ...@@ -70,7 +70,7 @@ func ImageHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *type
// apply param override // apply param override
if len(info.ParamOverride) > 0 { if len(info.ParamOverride) > 0 {
jsonData, err = relaycommon.ApplyParamOverride(jsonData, info.ParamOverride, relaycommon.BuildParamOverrideContext(info)) jsonData, err = relaycommon.ApplyParamOverrideWithRelayInfo(jsonData, info)
if err != nil { if err != nil {
return newAPIErrorFromParamOverride(err) return newAPIErrorFromParamOverride(err)
} }
......
relay/rerank_handler.go
...@@ -61,7 +61,7 @@ func RerankHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ ...@@ -61,7 +61,7 @@ func RerankHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ
// apply param override // apply param override
if len(info.ParamOverride) > 0 { if len(info.ParamOverride) > 0 {
jsonData, err = relaycommon.ApplyParamOverride(jsonData, info.ParamOverride, relaycommon.BuildParamOverrideContext(info)) jsonData, err = relaycommon.ApplyParamOverrideWithRelayInfo(jsonData, info)
if err != nil { if err != nil {
return newAPIErrorFromParamOverride(err) return newAPIErrorFromParamOverride(err)
} }
......
relay/responses_handler.go
...@@ -96,7 +96,7 @@ func ResponsesHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError * ...@@ -96,7 +96,7 @@ func ResponsesHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *
// apply param override // apply param override
if len(info.ParamOverride) > 0 { if len(info.ParamOverride) > 0 {
jsonData, err = relaycommon.ApplyParamOverride(jsonData, info.ParamOverride, relaycommon.BuildParamOverrideContext(info)) jsonData, err = relaycommon.ApplyParamOverrideWithRelayInfo(jsonData, info)
if err != nil { if err != nil {
return newAPIErrorFromParamOverride(err) 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