feat: sync deployed v1.0.0-rc.38-custom.42
Sync production source including routing, billing, Astra detection and public uptime. Update custom build instructions and license packaging.
This commit is contained in:
1 parent
aa6c3f50f4
commit
ceddfd3d6b
828 files changed
+122253
-13702
No files matched your search
@@ -77,15 +77,7 @@ func (a *Adaptor) ConvertClaudeRequest(c *gin.Context, info *relaycommon.RelayIn
|
||||
case relayconvert.ConverterNone:
|
||||
return a.claudeAdaptor.ConvertClaudeRequest(c, info, request)
|
||||
case relayconvert.ConverterClaudeMessagesToOpenAIChat:
|
||||
result, err := service.ConvertRequestByID(c, info, converter, request)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
chatRequest, ok := result.Value.(*dto.GeneralOpenAIRequest)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("expected OpenAI chat completions request, got %T", result.Value)
|
||||
}
|
||||
return a.convertOpenAICompatibleRequest(c, info, chatRequest)
|
||||
return a.convertCrossProtocolChatRequest(c, info, converter, request)
|
||||
default:
|
||||
return nil, fmt.Errorf("converter %q does not support Anthropic Messages requests", converter)
|
||||
}
|
||||
@@ -101,15 +93,7 @@ func (a *Adaptor) ConvertGeminiRequest(c *gin.Context, info *relaycommon.RelayIn
|
||||
case relayconvert.ConverterNone:
|
||||
return a.geminiAdaptor.ConvertGeminiRequest(c, info, request)
|
||||
case relayconvert.ConverterGeminiContentToOpenAIChat:
|
||||
result, err := service.ConvertRequestByID(c, info, converter, request)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
chatRequest, ok := result.Value.(*dto.GeneralOpenAIRequest)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("expected OpenAI chat completions request, got %T", result.Value)
|
||||
}
|
||||
return a.convertOpenAICompatibleRequest(c, info, chatRequest)
|
||||
return a.convertCrossProtocolChatRequest(c, info, converter, request)
|
||||
default:
|
||||
return nil, fmt.Errorf("converter %q does not support Gemini generateContent requests", converter)
|
||||
}
|
||||
@@ -124,15 +108,7 @@ func (a *Adaptor) ConvertOpenAIResponsesRequest(c *gin.Context, info *relaycommo
|
||||
case relayconvert.ConverterNone:
|
||||
return a.convertOpenAICompatibleResponsesRequest(c, info, request)
|
||||
case relayconvert.ConverterOpenAIResponsesToOpenAIChat:
|
||||
result, err := service.ConvertRequestByID(c, info, converter, request)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
chatRequest, ok := result.Value.(*dto.GeneralOpenAIRequest)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("expected OpenAI chat completions request, got %T", result.Value)
|
||||
}
|
||||
return a.convertOpenAICompatibleRequest(c, info, chatRequest)
|
||||
return a.convertCrossProtocolChatRequest(c, info, converter, request)
|
||||
case relayconvert.ConverterOpenAIResponsesToGemini:
|
||||
result, err := service.ConvertRequestByID(c, info, converter, request)
|
||||
if err != nil {
|
||||
@@ -286,7 +262,7 @@ func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, request
|
||||
if err := a.resolve(c, info); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if !a.converted && a.converter != relayconvert.ConverterNone {
|
||||
if !a.converted && !a.route.SupportsPassThroughBody() {
|
||||
return nil, errors.New("advanced custom converter routes cannot be used with pass-through request body")
|
||||
}
|
||||
|
||||
@@ -307,6 +283,8 @@ func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycom
|
||||
}
|
||||
|
||||
switch a.converter {
|
||||
case dto.AdvancedCustomConverterSGLangRerank:
|
||||
return a.doSGLangRerankResponse(c, resp, info)
|
||||
case relayconvert.ConverterNone:
|
||||
return a.doNativeResponse(c, resp, info)
|
||||
case relayconvert.ConverterClaudeMessagesToOpenAIChat,
|
||||
@@ -534,6 +512,25 @@ func isJSONRequest(c *gin.Context) bool {
|
||||
return strings.Contains(strings.ToLower(c.Request.Header.Get("Content-Type")), "application/json")
|
||||
}
|
||||
|
||||
// convertCrossProtocolChatRequest converts a Claude, Gemini, or Responses request into an
|
||||
// OpenAI chat completions request for an OpenAI-compatible upstream. Streaming requests
|
||||
// always ask the upstream for usage, because the downstream protocol reports usage in its
|
||||
// own format and never carries stream_options itself.
|
||||
func (a *Adaptor) convertCrossProtocolChatRequest(c *gin.Context, info *relaycommon.RelayInfo, converter string, request any) (any, error) {
|
||||
result, err := service.ConvertRequestByID(c, info, converter, request)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
chatRequest, ok := result.Value.(*dto.GeneralOpenAIRequest)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("expected OpenAI chat completions request, got %T", result.Value)
|
||||
}
|
||||
if info.SupportStreamOptions && info.IsStream {
|
||||
chatRequest.StreamOptions = &dto.StreamOptions{IncludeUsage: true}
|
||||
}
|
||||
return a.convertOpenAICompatibleRequest(c, info, chatRequest)
|
||||
}
|
||||
|
||||
func (a *Adaptor) convertOpenAICompatibleRequest(c *gin.Context, info *relaycommon.RelayInfo, request *dto.GeneralOpenAIRequest) (any, error) {
|
||||
old := info.ChannelType
|
||||
info.ChannelType = constant.ChannelTypeOpenAI
|
||||
|
||||
@@ -829,6 +829,143 @@ func TestAdaptorConvertsGeminiRequestToOpenAIChatUpstream(t *testing.T) {
|
||||
assert.Equal(t, "user", chatReq.Messages[0].Role)
|
||||
}
|
||||
|
||||
func TestAdaptorCrossProtocolChatUpstreamRequestsStreamUsage(t *testing.T) {
|
||||
claudeReq := &dto.ClaudeRequest{
|
||||
Model: "gpt-test",
|
||||
Messages: []dto.ClaudeMessage{{Role: "user", Content: "hello"}},
|
||||
}
|
||||
geminiReq := &dto.GeminiChatRequest{
|
||||
Contents: []dto.GeminiChatContent{{Role: "user", Parts: []dto.GeminiPart{{Text: "hello"}}}},
|
||||
}
|
||||
responsesReq := dto.OpenAIResponsesRequest{
|
||||
Model: "gpt-test",
|
||||
Input: mustAdvancedCustomRawMessage(t, "hello"),
|
||||
}
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
route dto.AdvancedCustomRoute
|
||||
relayFormat types.RelayFormat
|
||||
relayMode int
|
||||
requestPath string
|
||||
convert func(*Adaptor, *gin.Context, *relaycommon.RelayInfo) (any, error)
|
||||
isStream bool
|
||||
supportUsage bool
|
||||
wantUsage bool
|
||||
}{
|
||||
{
|
||||
name: "claude stream requests usage",
|
||||
route: dto.AdvancedCustomRoute{
|
||||
IncomingPath: "/v1/messages",
|
||||
UpstreamPath: "/v1/chat/completions",
|
||||
Converter: relayconvert.ConverterClaudeMessagesToOpenAIChat,
|
||||
},
|
||||
relayFormat: types.RelayFormatClaude,
|
||||
relayMode: relayconstant.RelayModeChatCompletions,
|
||||
requestPath: "/v1/messages",
|
||||
convert: func(a *Adaptor, c *gin.Context, info *relaycommon.RelayInfo) (any, error) {
|
||||
return a.ConvertClaudeRequest(c, info, claudeReq)
|
||||
},
|
||||
isStream: true,
|
||||
supportUsage: true,
|
||||
wantUsage: true,
|
||||
},
|
||||
{
|
||||
name: "gemini stream requests usage",
|
||||
route: dto.AdvancedCustomRoute{
|
||||
IncomingPath: "/v1beta/models/{model}:generateContent",
|
||||
UpstreamPath: "/v1/chat/completions",
|
||||
Converter: relayconvert.ConverterGeminiContentToOpenAIChat,
|
||||
},
|
||||
relayFormat: types.RelayFormatGemini,
|
||||
relayMode: relayconstant.RelayModeGemini,
|
||||
requestPath: "/v1beta/models/gpt-test:generateContent",
|
||||
convert: func(a *Adaptor, c *gin.Context, info *relaycommon.RelayInfo) (any, error) {
|
||||
return a.ConvertGeminiRequest(c, info, geminiReq)
|
||||
},
|
||||
isStream: true,
|
||||
supportUsage: true,
|
||||
wantUsage: true,
|
||||
},
|
||||
{
|
||||
name: "responses stream requests usage",
|
||||
route: dto.AdvancedCustomRoute{
|
||||
IncomingPath: "/v1/responses",
|
||||
UpstreamPath: "/v1/chat/completions",
|
||||
Converter: relayconvert.ConverterOpenAIResponsesToOpenAIChat,
|
||||
},
|
||||
relayFormat: types.RelayFormatOpenAIResponses,
|
||||
relayMode: relayconstant.RelayModeResponses,
|
||||
requestPath: "/v1/responses",
|
||||
convert: func(a *Adaptor, c *gin.Context, info *relaycommon.RelayInfo) (any, error) {
|
||||
return a.ConvertOpenAIResponsesRequest(c, info, responsesReq)
|
||||
},
|
||||
isStream: true,
|
||||
supportUsage: true,
|
||||
wantUsage: true,
|
||||
},
|
||||
{
|
||||
name: "claude non-stream leaves stream options unset",
|
||||
route: dto.AdvancedCustomRoute{
|
||||
IncomingPath: "/v1/messages",
|
||||
UpstreamPath: "/v1/chat/completions",
|
||||
Converter: relayconvert.ConverterClaudeMessagesToOpenAIChat,
|
||||
},
|
||||
relayFormat: types.RelayFormatClaude,
|
||||
relayMode: relayconstant.RelayModeChatCompletions,
|
||||
requestPath: "/v1/messages",
|
||||
convert: func(a *Adaptor, c *gin.Context, info *relaycommon.RelayInfo) (any, error) {
|
||||
return a.ConvertClaudeRequest(c, info, claudeReq)
|
||||
},
|
||||
isStream: false,
|
||||
supportUsage: true,
|
||||
wantUsage: false,
|
||||
},
|
||||
{
|
||||
name: "claude stream without stream options support",
|
||||
route: dto.AdvancedCustomRoute{
|
||||
IncomingPath: "/v1/messages",
|
||||
UpstreamPath: "/v1/chat/completions",
|
||||
Converter: relayconvert.ConverterClaudeMessagesToOpenAIChat,
|
||||
},
|
||||
relayFormat: types.RelayFormatClaude,
|
||||
relayMode: relayconstant.RelayModeChatCompletions,
|
||||
requestPath: "/v1/messages",
|
||||
convert: func(a *Adaptor, c *gin.Context, info *relaycommon.RelayInfo) (any, error) {
|
||||
return a.ConvertClaudeRequest(c, info, claudeReq)
|
||||
},
|
||||
isStream: true,
|
||||
supportUsage: false,
|
||||
wantUsage: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
adaptor := &Adaptor{}
|
||||
info := advancedCustomRelayInfo(&dto.AdvancedCustomConfig{Routes: []dto.AdvancedCustomRoute{tt.route}})
|
||||
info.RelayFormat = tt.relayFormat
|
||||
info.RelayMode = tt.relayMode
|
||||
info.RequestURLPath = tt.requestPath
|
||||
info.IsStream = tt.isStream
|
||||
info.SupportStreamOptions = tt.supportUsage
|
||||
c := advancedCustomGinContext(tt.requestPath)
|
||||
|
||||
converted, err := tt.convert(adaptor, c, info)
|
||||
require.NoError(t, err)
|
||||
chatReq, ok := converted.(*dto.GeneralOpenAIRequest)
|
||||
require.True(t, ok)
|
||||
|
||||
if !tt.wantUsage {
|
||||
assert.Nil(t, chatReq.StreamOptions)
|
||||
return
|
||||
}
|
||||
require.NotNil(t, chatReq.StreamOptions)
|
||||
assert.True(t, chatReq.StreamOptions.IncludeUsage)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func advancedCustomRelayInfo(config *dto.AdvancedCustomConfig) *relaycommon.RelayInfo {
|
||||
return &relaycommon.RelayInfo{
|
||||
RelayFormat: types.RelayFormatOpenAI,
|
||||
|
||||
@@ -0,0 +1,134 @@
|
||||
package advancedcustom_test
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/relay"
|
||||
"github.com/QuantumNous/new-api/relay/channel/advancedcustom"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
relayconstant "github.com/QuantumNous/new-api/relay/constant"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
"github.com/QuantumNous/new-api/relaykit/types"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestSGLangChannelProtocols(t *testing.T) {
|
||||
apiType, ok := common.ChannelType2APIType(constant.ChannelTypeSGLang)
|
||||
require.True(t, ok)
|
||||
adaptor := relay.GetAdaptor(apiType)
|
||||
require.NotNil(t, adaptor)
|
||||
assert.Equal(t, "advanced_custom", adaptor.GetChannelName())
|
||||
assert.Contains(t, common.GetEndpointTypesByChannelType(constant.ChannelTypeSGLang, "served-model"), constant.EndpointTypeAnthropic)
|
||||
assert.Contains(t, common.GetEndpointTypesByChannelType(constant.ChannelTypeSGLang, "served-model"), constant.EndpointTypeOpenAIResponse)
|
||||
info := &relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{ChannelType: constant.ChannelTypeSGLang, ChannelBaseUrl: "https://inference.example/prefix", ApiKey: "test-key"}}
|
||||
info.ChannelOtherSettings.AdvancedCustom = common.GetAdvancedCustomPreset(constant.ChannelTypeSGLang)
|
||||
adaptor.Init(info)
|
||||
for _, path := range []string{"/v1/chat/completions", "/v1/completions", "/v1/embeddings", "/v1/responses", "/v1/messages"} {
|
||||
info.RequestURLPath = path
|
||||
adaptor = &advancedcustom.Adaptor{}
|
||||
adaptor.Init(info)
|
||||
url, err := adaptor.GetRequestURL(info)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "https://inference.example/prefix"+path, url)
|
||||
}
|
||||
for _, format := range []types.RelayFormat{types.RelayFormatOpenAI, types.RelayFormatClaude} {
|
||||
info.RelayFormat = format
|
||||
info.RequestURLPath = "/v1/messages"
|
||||
adaptor = &advancedcustom.Adaptor{}
|
||||
adaptor.Init(info)
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil)
|
||||
headers := http.Header{}
|
||||
require.NoError(t, adaptor.SetupRequestHeader(c, &headers, info))
|
||||
assert.Equal(t, "Bearer test-key", headers.Get("Authorization"))
|
||||
}
|
||||
info.RequestURLPath = "/v1/chat/completions"
|
||||
adaptor = &advancedcustom.Adaptor{}
|
||||
adaptor.Init(info)
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
|
||||
request := &dto.GeneralOpenAIRequest{Model: "served-model"}
|
||||
converted, err := adaptor.ConvertOpenAIRequest(c, info, request)
|
||||
require.NoError(t, err)
|
||||
assert.Same(t, request, converted)
|
||||
info.RequestURLPath = "/v1/messages"
|
||||
adaptor = &advancedcustom.Adaptor{}
|
||||
adaptor.Init(info)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil)
|
||||
claudeRequest := &dto.ClaudeRequest{Model: "served-model", Thinking: &dto.Thinking{Type: "adaptive"}}
|
||||
converted, err = adaptor.ConvertClaudeRequest(c, info, claudeRequest)
|
||||
require.NoError(t, err)
|
||||
assert.Same(t, claudeRequest, converted)
|
||||
info.RequestURLPath = "/v1beta/models/test:generateContent"
|
||||
adaptor = &advancedcustom.Adaptor{}
|
||||
adaptor.Init(info)
|
||||
_, err = adaptor.ConvertGeminiRequest(nil, info, &dto.GeminiChatRequest{})
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
func TestSGLangRerankRequestAndResponse(t *testing.T) {
|
||||
service.InitHttpClient()
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
assert.Equal(t, "/prefix/v1/rerank", r.URL.Path)
|
||||
assert.Equal(t, "Bearer test-key", r.Header.Get("Authorization"))
|
||||
var request dto.RerankRequest
|
||||
require.NoError(t, common.DecodeJson(r.Body, &request))
|
||||
assert.Equal(t, "served-reranker", request.Model)
|
||||
assert.Equal(t, []any{"first", "second"}, request.Documents)
|
||||
_, _ = io.WriteString(w, `[{"index":1,"score":0.9,"document":"second","meta_info":{"prompt_tokens":5}}]`)
|
||||
}))
|
||||
defer upstream.Close()
|
||||
info := &relaycommon.RelayInfo{
|
||||
ChannelMeta: &relaycommon.ChannelMeta{ChannelType: constant.ChannelTypeSGLang, ChannelBaseUrl: upstream.URL + "/prefix", ApiKey: "test-key"},
|
||||
RequestURLPath: "/rerank", RelayMode: relayconstant.RelayModeRerank,
|
||||
RelayFormat: types.RelayFormatOpenAI,
|
||||
RerankerInfo: &relaycommon.RerankerInfo{Documents: []any{"first", "second"}, ReturnDocuments: true},
|
||||
}
|
||||
info.SetEstimatePromptTokens(12)
|
||||
adaptor := &advancedcustom.Adaptor{}
|
||||
info.ChannelOtherSettings.AdvancedCustom = common.GetAdvancedCustomPreset(constant.ChannelTypeSGLang)
|
||||
adaptor.Init(info)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/rerank", nil).WithContext(context.Background())
|
||||
converted, err := adaptor.ConvertRerankRequest(c, info.RelayMode, dto.RerankRequest{Model: "served-reranker", Query: "query", Documents: info.Documents})
|
||||
require.NoError(t, err)
|
||||
body, err := common.Marshal(converted)
|
||||
require.NoError(t, err)
|
||||
response, err := adaptor.DoRequest(c, info, bytes.NewReader(body))
|
||||
require.NoError(t, err)
|
||||
usage, apiErr := adaptor.DoResponse(c, response.(*http.Response), info)
|
||||
require.Nil(t, apiErr)
|
||||
assert.Equal(t, 12, usage.(*dto.Usage).TotalTokens)
|
||||
var decoded dto.RerankResponse
|
||||
require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &decoded))
|
||||
assert.Equal(t, []dto.RerankResponseResult{{Index: 1, RelevanceScore: 0.9, Document: map[string]any{"text": "second"}}}, decoded.Results)
|
||||
assert.Equal(t, 12, decoded.Usage.PromptTokens)
|
||||
assert.Zero(t, decoded.Usage.CompletionTokens)
|
||||
}
|
||||
|
||||
func TestSGLangRerankRejectsMalformedResults(t *testing.T) {
|
||||
for _, body := range []string{`null`, `{}`, `[{"index":99,"score":1}]`, `[{"index":-1,"score":0}]`, `[{"index":0}]`, `[{"score":1}]`} {
|
||||
t.Run(body, func(t *testing.T) {
|
||||
adaptor := &advancedcustom.Adaptor{}
|
||||
info := &relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{ChannelOtherSettings: dto.ChannelOtherSettings{AdvancedCustom: common.GetAdvancedCustomPreset(constant.ChannelTypeSGLang)}}, RequestURLPath: "/v1/rerank", RelayMode: relayconstant.RelayModeRerank, RerankerInfo: &relaycommon.RerankerInfo{Documents: []any{"only document"}}}
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
usage, err := adaptor.DoResponse(c, &http.Response{Body: io.NopCloser(strings.NewReader(body))}, info)
|
||||
require.NotNil(t, err)
|
||||
assert.Nil(t, usage)
|
||||
assert.Empty(t, recorder.Body.String())
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,46 @@
|
||||
package advancedcustom
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"net/http"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
"github.com/QuantumNous/new-api/relaykit/types"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func (a *Adaptor) doSGLangRerankResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (any, *types.NewAPIError) {
|
||||
defer service.CloseResponseBodyGracefully(resp)
|
||||
// SGLang v0.5.19 returns a sorted array with score, not a Jina results envelope.
|
||||
// https://github.com/sgl-project/sglang/blob/v0.5.19/python/sglang/srt/entrypoints/openai/serving_rerank.py
|
||||
var results []struct {
|
||||
Index *int `json:"index"`
|
||||
Score *float64 `json:"score"`
|
||||
}
|
||||
if err := common.DecodeJson(resp.Body, &results); err != nil {
|
||||
return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusBadGateway)
|
||||
}
|
||||
if results == nil {
|
||||
return nil, types.NewOpenAIError(errors.New("invalid SGLang rerank response"), types.ErrorCodeBadResponseBody, http.StatusBadGateway)
|
||||
}
|
||||
response := dto.RerankResponse{Results: make([]dto.RerankResponseResult, 0, len(results))}
|
||||
for _, result := range results {
|
||||
if result.Index == nil || result.Score == nil || *result.Index < 0 || *result.Index >= len(info.Documents) {
|
||||
return nil, types.NewOpenAIError(errors.New("invalid SGLang rerank result"), types.ErrorCodeBadResponseBody, http.StatusBadGateway)
|
||||
}
|
||||
item := dto.RerankResponseResult{Index: *result.Index, RelevanceScore: *result.Score}
|
||||
if info.ReturnDocuments {
|
||||
item.Document = dto.RerankDocument{Text: info.Documents[*result.Index]}
|
||||
}
|
||||
response.Results = append(response.Results, item)
|
||||
}
|
||||
// top_n can omit scored documents, and decoder rerankers can omit meta_info.
|
||||
// Per-result counts therefore cannot account for the full request; use the
|
||||
// existing prompt estimate, as with other rerank providers without usage.
|
||||
response.Usage = dto.Usage{PromptTokens: info.GetEstimatePromptTokens(), TotalTokens: info.GetEstimatePromptTokens()}
|
||||
c.JSON(http.StatusOK, response)
|
||||
return &response.Usage, nil
|
||||
}
|
||||
@@ -1,10 +1,12 @@
|
||||
package ali
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"mime/multipart"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
@@ -19,6 +21,7 @@ import (
|
||||
relayhelper "github.com/QuantumNous/new-api/relay/helper"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/QuantumNous/new-api/setting/ratio_setting"
|
||||
"github.com/QuantumNous/new-api/setting/system_setting"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/assert"
|
||||
@@ -26,6 +29,43 @@ import (
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
func TestAliMultipartEditsUseValidatedProviderQuantity(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
convert func(*gin.Context, *relaycommon.RelayInfo, dto.ImageRequest) (*AliImageRequest, error)
|
||||
}{
|
||||
{"multimodal edit", oaiFormEdit2AliImageEdit},
|
||||
{"legacy Wan edit", oaiFormEdit2WanxImageEdit},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
var body bytes.Buffer
|
||||
writer := multipart.NewWriter(&body)
|
||||
require.NoError(t, writer.WriteField("model", "fixture-image"))
|
||||
require.NoError(t, writer.WriteField("n", "2"))
|
||||
require.NoError(t, writer.WriteField("parameters", `{"n":3,"prompt_extend":false}`))
|
||||
part, err := writer.CreateFormFile("image", "input.png")
|
||||
require.NoError(t, err)
|
||||
_, err = part.Write([]byte("fixture image"))
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, writer.Close())
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/edits", &body)
|
||||
c.Request.Header.Set("Content-Type", writer.FormDataContentType())
|
||||
common.SetContextKey(c, rootconstant.ContextKeyChannelType, rootconstant.ChannelTypeAli)
|
||||
request, err := relayhelper.GetAndValidOpenAIImageRequest(c, constant.RelayModeImagesEdits)
|
||||
require.NoError(t, err)
|
||||
info := &relaycommon.RelayInfo{Request: request}
|
||||
converted, err := tc.convert(c, info, *request)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, converted.Parameters.N)
|
||||
assert.Equal(t, uint(3), *converted.Parameters.N)
|
||||
require.NotNil(t, converted.Parameters.PromptExtend)
|
||||
assert.False(t, *converted.Parameters.PromptExtend)
|
||||
assert.Equal(t, uint(2), *request.N, "conversion must preserve the incoming request")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestConvertOpenAIRequestFiltersThinkingBudgetByUpstreamModel(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
@@ -169,6 +209,42 @@ func TestMappedAliImageModelUsesUpstreamProtocol(t *testing.T) {
|
||||
assert.IsType(t, &AliImageRequest{}, converted)
|
||||
}
|
||||
|
||||
func TestAliImageCountMatchesLegacyReservation(t *testing.T) {
|
||||
saved := ratio_setting.ModelPrice2JSONString()
|
||||
require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(`{"z-image":0.04}`))
|
||||
t.Cleanup(func() { require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(saved)) })
|
||||
for _, tc := range []struct {
|
||||
name, body string
|
||||
count int
|
||||
multiplier float64
|
||||
}{
|
||||
{"provider override", `{"model":"z-image","n":2,"parameters":{"n":4}}`, 4, 1},
|
||||
{"empty parameters inherit top-level", `{"model":"z-image","n":2,"parameters":{}}`, 2, 1},
|
||||
{"null count inherits top-level", `{"model":"z-image","n":2,"parameters":{"n":null}}`, 2, 1},
|
||||
{"default count", `{"model":"z-image"}`, 1, 1},
|
||||
{"prompt extension", `{"model":"z-image","parameters":{"n":3,"prompt_extend":true}}`, 3, 2},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", strings.NewReader(tc.body))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
common.SetContextKey(c, rootconstant.ContextKeyChannelType, rootconstant.ChannelTypeAli)
|
||||
request, err := relayhelper.GetAndValidOpenAIImageRequest(c, constant.RelayModeImagesGenerations)
|
||||
require.NoError(t, err)
|
||||
info := &relaycommon.RelayInfo{Request: request, OriginModelName: request.Model, UserGroup: "default", UsingGroup: "default"}
|
||||
price, err := relayhelper.ModelPriceHelper(c, info, 0, request.GetTokenCountMeta())
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, common.QuotaFromFloat(0.04*float64(tc.count)*tc.multiplier*common.QuotaPerUnit), price.QuotaToPreConsume)
|
||||
info.ChannelMeta = &relaycommon.ChannelMeta{ChannelType: rootconstant.ChannelTypeAli, UpstreamModelName: request.Model}
|
||||
converted, err := oaiImage2AliImageRequest(info, *request, true)
|
||||
require.NoError(t, err)
|
||||
encoded, err := common.Marshal(converted)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(tc.count), gjson.GetBytes(encoded, "parameters.n").Int())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestAliImageHandlerHonorsRequestResponseFormat(t *testing.T) {
|
||||
imageBytes := []byte("ali-image")
|
||||
var downloads atomic.Int32
|
||||
|
||||
@@ -172,7 +172,7 @@ type AliImageRequest struct {
|
||||
|
||||
type AliImageParameters struct {
|
||||
Size string `json:"size,omitempty"`
|
||||
N int `json:"n,omitempty"`
|
||||
N *uint `json:"n,omitempty"`
|
||||
Steps string `json:"steps,omitempty"`
|
||||
Scale string `json:"scale,omitempty"`
|
||||
Watermark *bool `json:"watermark,omitempty"`
|
||||
|
||||
+43
-26
@@ -18,25 +18,29 @@ import (
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/samber/lo"
|
||||
)
|
||||
|
||||
func oaiImage2AliImageRequest(info *relaycommon.RelayInfo, request dto.ImageRequest, isSync bool) (*AliImageRequest, error) {
|
||||
var imageRequest AliImageRequest
|
||||
imageRequest.Model = request.Model
|
||||
imageRequest.ResponseFormat = request.ResponseFormat
|
||||
imageRequest.Parameters = AliImageParameters{
|
||||
Size: strings.ReplaceAll(request.Size, "x", "*"),
|
||||
Watermark: request.Watermark,
|
||||
}
|
||||
if request.Extra != nil {
|
||||
if val, ok := request.Extra["parameters"]; ok {
|
||||
imageRequest.Parameters = AliImageParameters{}
|
||||
err := common.Unmarshal(val, &imageRequest.Parameters)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("invalid parameters field: %w", err)
|
||||
}
|
||||
} else {
|
||||
// 兼容没有parameters字段的情况,从openai标准字段中提取参数
|
||||
imageRequest.Parameters = AliImageParameters{
|
||||
Size: strings.Replace(request.Size, "x", "*", -1),
|
||||
N: int(lo.FromPtrOr(request.N, uint(1))),
|
||||
Watermark: request.Watermark,
|
||||
// Direct adaptor callers can bypass ingress validation. Reuse the
|
||||
// decoded provider scalars and the same quantity validation below.
|
||||
if request.BillingParameters == nil {
|
||||
request.BillingParameters = &dto.ImageBillingParameters{
|
||||
N: imageRequest.Parameters.N, PromptExtend: imageRequest.Parameters.PromptExtend,
|
||||
}
|
||||
}
|
||||
}
|
||||
if val, ok := request.Extra["input"]; ok {
|
||||
@@ -47,23 +51,21 @@ func oaiImage2AliImageRequest(info *relaycommon.RelayInfo, request dto.ImageRequ
|
||||
}
|
||||
}
|
||||
|
||||
if strings.Contains(request.Model, "z-image") {
|
||||
// z-image 开启prompt_extend后,按2倍计费
|
||||
if imageRequest.Parameters.PromptExtendValue() {
|
||||
info.PriceData.AddOtherRatio("prompt_extend", 2)
|
||||
count, err := request.ImageCount(true)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
imageRequest.Parameters.N = common.GetPointer(uint(count))
|
||||
if request.BillingParameters != nil {
|
||||
imageRequest.Parameters.PromptExtend = request.BillingParameters.PromptExtend
|
||||
}
|
||||
if info.TieredBillingSnapshot == nil {
|
||||
info.PriceData.AddOtherRatio("n", float64(count))
|
||||
if strings.Contains(request.Model, "z-image") && imageRequest.Parameters.PromptExtendValue() {
|
||||
info.PriceData.AddOtherRatio("prompt_extend", common.ZImagePromptExtendMultiplier)
|
||||
}
|
||||
}
|
||||
|
||||
// Parameters may come from Extra["parameters"], bypassing the standard
|
||||
// top-level n validation; enforce the same bound before it becomes a
|
||||
// billing multiplier.
|
||||
if imageRequest.Parameters.N < 0 || imageRequest.Parameters.N > dto.MaxImageN {
|
||||
return nil, fmt.Errorf("parameters.n must be an integer between 1 and %d", dto.MaxImageN)
|
||||
}
|
||||
if imageRequest.Parameters.N != 0 {
|
||||
info.PriceData.AddOtherRatio("n", float64(imageRequest.Parameters.N))
|
||||
}
|
||||
|
||||
// 同步图片模型和异步图片模型请求格式不一样
|
||||
if isSync {
|
||||
if imageRequest.Input == nil {
|
||||
@@ -159,6 +161,10 @@ func getImageBase64sFromForm(c *gin.Context, fieldName string) ([]string, error)
|
||||
}
|
||||
|
||||
func oaiFormEdit2AliImageEdit(c *gin.Context, info *relaycommon.RelayInfo, request dto.ImageRequest) (*AliImageRequest, error) {
|
||||
count, err := request.ImageCount(true)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var imageRequest AliImageRequest
|
||||
imageRequest.Model = request.Model
|
||||
imageRequest.ResponseFormat = request.ResponseFormat
|
||||
@@ -186,9 +192,12 @@ func oaiFormEdit2AliImageEdit(c *gin.Context, info *relaycommon.RelayInfo, reque
|
||||
},
|
||||
}
|
||||
imageRequest.Parameters = AliImageParameters{
|
||||
N: int(lo.FromPtrOr(request.N, uint(1))),
|
||||
N: common.GetPointer(uint(count)),
|
||||
Watermark: request.Watermark,
|
||||
}
|
||||
if request.BillingParameters != nil {
|
||||
imageRequest.Parameters.PromptExtend = request.BillingParameters.PromptExtend
|
||||
}
|
||||
return &imageRequest, nil
|
||||
}
|
||||
|
||||
@@ -336,10 +345,18 @@ func aliImageHandler(a *Adaptor, c *gin.Context, resp *http.Response, info *rela
|
||||
}
|
||||
|
||||
imageResponses := responseAli2OpenAIImage(c, aliResponse, originRespBody, info, responseFormat)
|
||||
if aliResponse.Usage.ImageCount != 0 {
|
||||
info.PriceData.AddOtherRatio("n", float64(aliResponse.Usage.ImageCount))
|
||||
} else if len(imageResponses.Data) != 0 {
|
||||
info.PriceData.AddOtherRatio("n", float64(len(imageResponses.Data)))
|
||||
count := int64(aliResponse.Usage.ImageCount)
|
||||
if count < 0 || count > dto.MaxImageN {
|
||||
logger.LogWarn(c, "invalid Ali image usage count %d; falling back to response images or requested quantity", count)
|
||||
count = 0
|
||||
}
|
||||
if count == 0 {
|
||||
count = int64(len(imageResponses.Data))
|
||||
}
|
||||
if count > dto.MaxImageN {
|
||||
logger.LogWarn(c, "invalid Ali response image count %d; retaining requested quantity", count)
|
||||
} else if count > 0 {
|
||||
info.UpdateImageCount(count)
|
||||
}
|
||||
jsonResponse, err := common.Marshal(imageResponses)
|
||||
if err != nil {
|
||||
|
||||
@@ -13,7 +13,10 @@ import (
|
||||
)
|
||||
|
||||
func oaiFormEdit2WanxImageEdit(c *gin.Context, info *relaycommon.RelayInfo, request dto.ImageRequest) (*AliImageRequest, error) {
|
||||
var err error
|
||||
count, err := request.ImageCount(true)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var imageRequest AliImageRequest
|
||||
imageRequest.Model = request.Model
|
||||
imageRequest.ResponseFormat = request.ResponseFormat
|
||||
@@ -27,14 +30,13 @@ func oaiFormEdit2WanxImageEdit(c *gin.Context, info *relaycommon.RelayInfo, requ
|
||||
if wanInput.Images, err = getImageBase64sFromForm(c, "image"); err != nil {
|
||||
return nil, fmt.Errorf("get image base64s from form failed: %w", err)
|
||||
}
|
||||
//wanParams := WanImageParameters{
|
||||
// N: int(request.N),
|
||||
//}
|
||||
imageRequest.Input = wanInput
|
||||
imageRequest.Parameters = AliImageParameters{
|
||||
N: int(lo.FromPtrOr(request.N, uint(1))),
|
||||
N: common.GetPointer(uint(count)),
|
||||
}
|
||||
if request.BillingParameters != nil {
|
||||
imageRequest.Parameters.PromptExtend = request.BillingParameters.PromptExtend
|
||||
}
|
||||
info.PriceData.AddOtherRatio("n", float64(imageRequest.Parameters.N))
|
||||
|
||||
return &imageRequest, nil
|
||||
}
|
||||
|
||||
@@ -372,11 +372,16 @@ func DoFormRequest(a Adaptor, c *gin.Context, info *common.RelayInfo, requestBod
|
||||
return resp, nil
|
||||
}
|
||||
|
||||
// DoWssRequest dials the adaptor's upstream over WebSocket for the realtime
|
||||
// and Responses WebSocket relays. It honors the channel proxy, is bound to the
|
||||
// request context, and reports a rejected handshake as a *types.NewAPIError
|
||||
// carrying the upstream status code (types.NewError preserves it).
|
||||
func DoWssRequest(a Adaptor, c *gin.Context, info *common.RelayInfo, requestBody io.Reader) (*websocket.Conn, error) {
|
||||
fullRequestURL, err := a.GetRequestURL(info)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("get request url failed: %w", err)
|
||||
}
|
||||
fullRequestURL = toWebSocketURL(fullRequestURL)
|
||||
targetHeader := http.Header{}
|
||||
err = a.SetupRequestHeader(c, &targetHeader, info)
|
||||
if err != nil {
|
||||
@@ -392,16 +397,41 @@ func DoWssRequest(a Adaptor, c *gin.Context, info *common.RelayInfo, requestBody
|
||||
targetHeader.Set(key, value)
|
||||
}
|
||||
targetHeader.Set("Content-Type", c.Request.Header.Get("Content-Type"))
|
||||
targetConn, _, err := websocket.DefaultDialer.Dial(fullRequestURL, targetHeader)
|
||||
dialer := *websocket.DefaultDialer
|
||||
if info.ChannelSetting.Proxy != "" {
|
||||
proxyURL, _, proxyErr := common2.ParseProxyURLRuntime(info.ChannelSetting.Proxy)
|
||||
if proxyErr != nil {
|
||||
return nil, proxyErr
|
||||
}
|
||||
dialer.Proxy = http.ProxyURL(proxyURL)
|
||||
}
|
||||
targetConn, resp, err := dialer.DialContext(c.Request.Context(), fullRequestURL, targetHeader)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("dial failed to %s: %w", common.SanitizeURLForLog(fullRequestURL), err)
|
||||
statusCode := http.StatusInternalServerError
|
||||
if resp != nil {
|
||||
statusCode = resp.StatusCode
|
||||
if resp.Body != nil {
|
||||
_ = resp.Body.Close()
|
||||
}
|
||||
}
|
||||
return nil, types.NewErrorWithStatusCode(fmt.Errorf("dial failed to %s: %w", common.SanitizeURLForLog(fullRequestURL), err), types.ErrorCodeDoRequestFailed, statusCode)
|
||||
}
|
||||
// send request body
|
||||
//all, err := io.ReadAll(requestBody)
|
||||
//err = service.WssString(c, targetConn, string(all))
|
||||
return targetConn, nil
|
||||
}
|
||||
|
||||
// toWebSocketURL maps an http(s) endpoint to ws(s). Realtime adaptors already
|
||||
// return ws(s) URLs, which pass through unchanged.
|
||||
func toWebSocketURL(raw string) string {
|
||||
switch {
|
||||
case strings.HasPrefix(raw, "https://"):
|
||||
return "wss://" + strings.TrimPrefix(raw, "https://")
|
||||
case strings.HasPrefix(raw, "http://"):
|
||||
return "ws://" + strings.TrimPrefix(raw, "http://")
|
||||
default:
|
||||
return raw
|
||||
}
|
||||
}
|
||||
|
||||
func startPingKeepAlive(c *gin.Context, pingInterval time.Duration) (context.CancelFunc, <-chan struct{}) {
|
||||
pingerCtx, stopPinger := context.WithCancel(context.Background())
|
||||
done := make(chan struct{})
|
||||
@@ -488,6 +518,13 @@ func keepUpstreamRedirectResponse(_ *http.Request, _ []*http.Request) error {
|
||||
}
|
||||
|
||||
func doRequest(c *gin.Context, req *http.Request, info *common.RelayInfo) (*http.Response, error) {
|
||||
guardEnabled := info.ResponseGuard != nil && info.AstraAutoSwitch
|
||||
if guardEnabled && info.UpstreamRequests >= common.ResponseGuardMaxAttempts {
|
||||
return nil, types.NewErrorWithStatusCode(errors.New("response model attempt budget exhausted"), common.ResponseGuardErrorCode, http.StatusBadGateway, types.ErrOptionWithSkipRetry())
|
||||
}
|
||||
if guardEnabled && c.Request.Context().Err() != nil {
|
||||
return nil, types.NewErrorWithStatusCode(c.Request.Context().Err(), common.ResponseGuardErrorCode, http.StatusBadGateway, types.ErrOptionWithSkipRetry())
|
||||
}
|
||||
client, err := service.GetHttpClientWithProxySettings(info.ChannelSetting.Proxy, info.ChannelSetting)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("new proxy http client failed: %w", err)
|
||||
@@ -511,7 +548,7 @@ func doRequest(c *gin.Context, req *http.Request, info *common.RelayInfo) (*http
|
||||
|
||||
var stopPinger context.CancelFunc
|
||||
var pingerDone <-chan struct{}
|
||||
if info.IsStream {
|
||||
if info.IsStream && !guardEnabled {
|
||||
helper.SetEventStreamHeaders(c)
|
||||
// 处理流式请求的 ping 保活
|
||||
generalSettings := operation_setting.GetGeneralSetting()
|
||||
@@ -529,14 +566,55 @@ func doRequest(c *gin.Context, req *http.Request, info *common.RelayInfo) (*http
|
||||
}
|
||||
}
|
||||
|
||||
info.UpstreamRequests++
|
||||
var guard *common.ResponseGuardAttempt
|
||||
var guardUsage *service.GuardUsageCollector
|
||||
if guardEnabled {
|
||||
var upstream context.Context
|
||||
upstream, guard = info.ResponseGuard.Begin(c.Request.Context(), info.UpstreamRequests, info.ChannelId)
|
||||
if info.DurableAttemptID != 0 {
|
||||
guardUsage = service.NewGuardUsageCollector(c, info)
|
||||
guard.OnEvent = guardUsage.Observe
|
||||
guard.OnClose = func() { guardUsage.CheckpointFinal(guard.Reason) }
|
||||
}
|
||||
req = req.WithContext(upstream)
|
||||
}
|
||||
resp, err := relayClient.Do(req)
|
||||
if err != nil {
|
||||
if guard != nil {
|
||||
guard.Close()
|
||||
guard.Finish(false)
|
||||
if c.Request.Context().Err() != nil {
|
||||
guard.Reject("client_cancelled")
|
||||
return nil, guard.APIError()
|
||||
}
|
||||
if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
|
||||
guard.Reject("confirmation_timeout")
|
||||
return nil, guard.APIError()
|
||||
}
|
||||
guard.Reason = "transport_error"
|
||||
}
|
||||
logger.LogError(c, "do request failed: "+err.Error())
|
||||
return nil, types.NewError(err, types.ErrorCodeDoRequestFailed, types.ErrOptionWithHideErrMsg("upstream error: do request failed"))
|
||||
}
|
||||
if resp == nil {
|
||||
if guard != nil {
|
||||
guard.Close()
|
||||
}
|
||||
return nil, errors.New("resp is nil")
|
||||
}
|
||||
if guard != nil {
|
||||
if err := guard.Confirm(resp, c.Request.Context()); err != nil {
|
||||
if guardUsage != nil && guardUsage.FinishRejected(info, guard, resp) {
|
||||
guard.Finish(false)
|
||||
return nil, guard.APIError()
|
||||
}
|
||||
guard.Close()
|
||||
_ = resp.Body.Close()
|
||||
guard.Finish(false)
|
||||
return nil, guard.APIError()
|
||||
}
|
||||
}
|
||||
if common2.DebugEnabled {
|
||||
policy := service.NormalizeHTTPTransportPolicy(info.ChannelSetting)
|
||||
logger.LogDebug(c, fmt.Sprintf(
|
||||
|
||||
@@ -8,6 +8,7 @@ import (
|
||||
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
@@ -205,3 +206,14 @@ func TestProcessHeaderOverride_PassHeadersTemplateSetsRuntimeHeaders(t *testing.
|
||||
require.Equal(t, "sess-123", upstreamReq.Header.Get("Session_id"))
|
||||
require.Empty(t, upstreamReq.Header.Get("X-Codex-Beta-Features"))
|
||||
}
|
||||
|
||||
func TestToWebSocketURL(t *testing.T) {
|
||||
for input, want := range map[string]string{
|
||||
"https://api.openai.com/v1/responses": "wss://api.openai.com/v1/responses",
|
||||
"http://127.0.0.1:3000/v1/responses": "ws://127.0.0.1:3000/v1/responses",
|
||||
"wss://chatgpt.com/backend-api/codex/responses": "wss://chatgpt.com/backend-api/codex/responses",
|
||||
"ws://127.0.0.1:3000/backend-api/codex/responses": "ws://127.0.0.1:3000/backend-api/codex/responses",
|
||||
} {
|
||||
assert.Equal(t, want, toWebSocketURL(input), input)
|
||||
}
|
||||
}
|
||||
@@ -24,11 +24,12 @@ func stopReasonClaude2OpenAI(reason string) string {
|
||||
return relayconvert.StopReasonClaudeToOpenAI(reason)
|
||||
}
|
||||
|
||||
func maybeMarkClaudeRefusal(c *gin.Context, stopReason string) {
|
||||
func maybeMarkClaudeRefusal(c *gin.Context, info *relaycommon.RelayInfo, stopReason string) {
|
||||
if c == nil {
|
||||
return
|
||||
}
|
||||
if strings.EqualFold(stopReason, "refusal") {
|
||||
info.PerformanceBusinessRejection = true
|
||||
common.SetContextKey(c, constant.ContextKeyAdminRejectReason, "claude_stop_reason=refusal")
|
||||
}
|
||||
}
|
||||
@@ -94,11 +95,17 @@ func HandleStreamResponseData(c *gin.Context, info *relaycommon.RelayInfo, claud
|
||||
if claudeError := claudeResponse.GetClaudeError(); claudeError != nil && claudeError.Type != "" {
|
||||
return types.WithClaudeError(*claudeError, http.StatusInternalServerError)
|
||||
}
|
||||
if claudeResponse.Type == "message_start" && claudeResponse.Message != nil {
|
||||
info.ObserveResponseModel(claudeResponse.Message.Model)
|
||||
}
|
||||
if claudeResponse.StopReason != "" {
|
||||
maybeMarkClaudeRefusal(c, claudeResponse.StopReason)
|
||||
maybeMarkClaudeRefusal(c, info, claudeResponse.StopReason)
|
||||
}
|
||||
if claudeResponse.Delta != nil && claudeResponse.Delta.StopReason != nil {
|
||||
maybeMarkClaudeRefusal(c, *claudeResponse.Delta.StopReason)
|
||||
maybeMarkClaudeRefusal(c, info, *claudeResponse.Delta.StopReason)
|
||||
}
|
||||
if claudeResponse.Type == "message_stop" {
|
||||
info.StreamStatus.MarkCompleted()
|
||||
}
|
||||
if info.RelayFormat == types.RelayFormatClaude {
|
||||
FormatClaudeResponseInfo(&claudeResponse, nil, claudeInfo)
|
||||
@@ -299,6 +306,7 @@ func ClaudeStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon.
|
||||
sr.Stop(err)
|
||||
}
|
||||
})
|
||||
info.StreamStatus.RequireTerminal()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -316,7 +324,8 @@ func HandleClaudeResponseData(c *gin.Context, info *relaycommon.RelayInfo, claud
|
||||
if claudeError := claudeResponse.GetClaudeError(); claudeError != nil && claudeError.Type != "" {
|
||||
return types.WithClaudeError(*claudeError, http.StatusInternalServerError)
|
||||
}
|
||||
maybeMarkClaudeRefusal(c, claudeResponse.StopReason)
|
||||
info.ObserveResponseModel(claudeResponse.Model)
|
||||
maybeMarkClaudeRefusal(c, info, claudeResponse.StopReason)
|
||||
if claudeInfo.Usage == nil {
|
||||
claudeInfo.Usage = &dto.Usage{}
|
||||
}
|
||||
|
||||
@@ -107,12 +107,16 @@ func ClaudeResponsesStreamHandler(c *gin.Context, resp *http.Response, info *rel
|
||||
}
|
||||
|
||||
if claudeResponse.StopReason != "" {
|
||||
maybeMarkClaudeRefusal(c, claudeResponse.StopReason)
|
||||
maybeMarkClaudeRefusal(c, info, claudeResponse.StopReason)
|
||||
}
|
||||
if claudeResponse.Delta != nil && claudeResponse.Delta.StopReason != nil {
|
||||
maybeMarkClaudeRefusal(c, *claudeResponse.Delta.StopReason)
|
||||
maybeMarkClaudeRefusal(c, info, *claudeResponse.Delta.StopReason)
|
||||
}
|
||||
if claudeResponse.Type == "message_stop" {
|
||||
info.StreamStatus.MarkCompleted()
|
||||
}
|
||||
if claudeResponse.Type == "message_start" && claudeResponse.Message != nil {
|
||||
info.ObserveResponseModel(claudeResponse.Message.Model)
|
||||
info.UpstreamModelName = claudeResponse.Message.Model
|
||||
}
|
||||
FormatClaudeResponseInfo(&claudeResponse, nil, claudeInfo)
|
||||
|
||||
@@ -15,6 +15,7 @@ import (
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
func GeminiTextGenerationHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) {
|
||||
@@ -34,9 +35,11 @@ func GeminiTextGenerationHandler(c *gin.Context, info *relaycommon.RelayInfo, re
|
||||
if err != nil {
|
||||
return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
|
||||
}
|
||||
info.ObserveResponseModel(gjson.GetBytes(responseBody, "modelVersion").Str)
|
||||
countGeminiBillableFunctionCalls(info, &geminiResponse)
|
||||
|
||||
if len(geminiResponse.Candidates) == 0 && geminiResponse.PromptFeedback != nil && geminiResponse.PromptFeedback.BlockReason != nil {
|
||||
info.PerformanceBusinessRejection = true
|
||||
common.SetContextKey(c, constant.ContextKeyAdminRejectReason, fmt.Sprintf("gemini_block_reason=%s", *geminiResponse.PromptFeedback.BlockReason))
|
||||
}
|
||||
|
||||
|
||||
@@ -20,6 +20,7 @@ import (
|
||||
"github.com/QuantumNous/new-api/relaykit/types"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
func buildUsageFromGeminiMetadata(metadata *dto.GeminiUsageMetadata, fallbackPromptTokens int) dto.Usage {
|
||||
@@ -181,8 +182,20 @@ func geminiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http
|
||||
}
|
||||
|
||||
if len(geminiResponse.Candidates) == 0 && geminiResponse.PromptFeedback != nil && geminiResponse.PromptFeedback.BlockReason != nil {
|
||||
info.PerformanceBusinessRejection = true
|
||||
common.SetContextKey(c, constant.ContextKeyAdminRejectReason, fmt.Sprintf("gemini_block_reason=%s", *geminiResponse.PromptFeedback.BlockReason))
|
||||
}
|
||||
info.ObserveResponseModel(gjson.Get(data, "modelVersion").Str)
|
||||
for _, candidate := range geminiResponse.Candidates {
|
||||
if candidate.FinishReason == nil || *candidate.FinishReason == "" || *candidate.FinishReason == "FINISH_REASON_UNSPECIFIED" {
|
||||
continue
|
||||
}
|
||||
switch *candidate.FinishReason {
|
||||
case "SAFETY", "RECITATION", "BLOCKLIST", "PROHIBITED_CONTENT", "SPII", "IMAGE_SAFETY", "IMAGE_PROHIBITED_CONTENT", "IMAGE_RECITATION":
|
||||
info.PerformanceBusinessRejection = true
|
||||
}
|
||||
info.StreamStatus.MarkCompleted()
|
||||
}
|
||||
|
||||
markGeminiGoogleSearchCall(c, &geminiResponse)
|
||||
countGeminiBillableFunctionCalls(info, &geminiResponse)
|
||||
@@ -216,6 +229,7 @@ func geminiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http
|
||||
sr.Stop(streamErr)
|
||||
}
|
||||
})
|
||||
info.StreamStatus.RequireTerminal()
|
||||
|
||||
if !hasBillableUsageMetadata {
|
||||
if info.ReceivedResponseCount > 0 {
|
||||
@@ -370,6 +384,7 @@ func GeminiChatHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.R
|
||||
if err != nil {
|
||||
return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
|
||||
}
|
||||
info.ObserveResponseModel(gjson.GetBytes(responseBody, "modelVersion").Str)
|
||||
markGeminiGoogleSearchCall(c, &geminiResponse)
|
||||
countGeminiBillableFunctionCalls(info, &geminiResponse)
|
||||
if len(geminiResponse.Candidates) == 0 {
|
||||
@@ -377,6 +392,7 @@ func GeminiChatHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.R
|
||||
|
||||
var newAPIError *types.NewAPIError
|
||||
if geminiResponse.PromptFeedback != nil && geminiResponse.PromptFeedback.BlockReason != nil {
|
||||
info.PerformanceBusinessRejection = true
|
||||
common.SetContextKey(c, constant.ContextKeyAdminRejectReason, fmt.Sprintf("gemini_block_reason=%s", *geminiResponse.PromptFeedback.BlockReason))
|
||||
newAPIError = types.NewOpenAIError(
|
||||
errors.New("request blocked by Gemini API: "+*geminiResponse.PromptFeedback.BlockReason),
|
||||
|
||||
@@ -16,6 +16,7 @@ import (
|
||||
"github.com/QuantumNous/new-api/relaykit/types"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
func GeminiResponsesHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) {
|
||||
@@ -31,11 +32,13 @@ func GeminiResponsesHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *h
|
||||
if err := common.Unmarshal(responseBody, &geminiResponse); err != nil {
|
||||
return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
|
||||
}
|
||||
info.ObserveResponseModel(gjson.GetBytes(responseBody, "modelVersion").Str)
|
||||
markGeminiGoogleSearchCall(c, &geminiResponse)
|
||||
countGeminiBillableFunctionCalls(info, &geminiResponse)
|
||||
if len(geminiResponse.Candidates) == 0 {
|
||||
usage := buildUsageFromGeminiResponse(c, info, &geminiResponse)
|
||||
if geminiResponse.PromptFeedback != nil && geminiResponse.PromptFeedback.BlockReason != nil {
|
||||
info.PerformanceBusinessRejection = true
|
||||
common.SetContextKey(c, constant.ContextKeyAdminRejectReason, fmt.Sprintf("gemini_block_reason=%s", *geminiResponse.PromptFeedback.BlockReason))
|
||||
return &usage, types.NewOpenAIError(
|
||||
errors.New("request blocked by Gemini API: "+*geminiResponse.PromptFeedback.BlockReason),
|
||||
|
||||
@@ -38,6 +38,9 @@ func (a *Adaptor) ConvertImageRequest(c *gin.Context, info *relaycommon.RelayInf
|
||||
}
|
||||
|
||||
func (a *Adaptor) Init(info *relaycommon.RelayInfo) {
|
||||
// OpenAI-compatible chat responses are handled by the OpenAI adaptor, which
|
||||
// relies on the thinking-to-content state initialized here.
|
||||
(&openai.Adaptor{}).Init(info)
|
||||
}
|
||||
|
||||
func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
|
||||
@@ -55,6 +58,9 @@ func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
|
||||
case relayconstant.RelayModeCompletions:
|
||||
return fmt.Sprintf("%s/api/generate", info.ChannelBaseUrl), nil
|
||||
default:
|
||||
if info.ChannelOtherSettings.OllamaOpenAIChat {
|
||||
return fmt.Sprintf("%s/v1/chat/completions", info.ChannelBaseUrl), nil
|
||||
}
|
||||
return fmt.Sprintf("%s/api/chat", info.ChannelBaseUrl), nil
|
||||
}
|
||||
}
|
||||
@@ -83,6 +89,9 @@ func (a *Adaptor) ConvertOpenAIRequest(c *gin.Context, info *relaycommon.RelayIn
|
||||
case relayconstant.RelayModeCompletions:
|
||||
return openAIToGenerate(c, request)
|
||||
default:
|
||||
if info.ChannelOtherSettings.OllamaOpenAIChat {
|
||||
return (&openai.Adaptor{}).ConvertOpenAIRequest(c, info, request)
|
||||
}
|
||||
return openAIChatToOllamaChat(c, request)
|
||||
}
|
||||
}
|
||||
@@ -117,6 +126,11 @@ func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycom
|
||||
adaptor := openai.Adaptor{}
|
||||
return adaptor.DoResponse(c, resp, info)
|
||||
default:
|
||||
// /api/generate always returns native NDJSON; chat follows the channel setting.
|
||||
if info.RelayMode != relayconstant.RelayModeCompletions && info.ChannelOtherSettings.OllamaOpenAIChat {
|
||||
adaptor := openai.Adaptor{}
|
||||
return adaptor.DoResponse(c, resp, info)
|
||||
}
|
||||
if info.IsStream {
|
||||
return ollamaStreamHandler(c, info, resp)
|
||||
}
|
||||
|
||||
@@ -79,6 +79,51 @@ func ollamaToolCallsToOpenAI(toolCalls []OllamaToolCall, startIndex int, include
|
||||
return result, startIndex
|
||||
}
|
||||
|
||||
func buildOllamaStreamDelta(chunk *ollamaChatStreamChunk, responseID string, created int64, model string, toolCallIndex *int) (dto.ChatCompletionsStreamResponse, bool) {
|
||||
delta := dto.ChatCompletionsStreamResponse{
|
||||
Id: responseID,
|
||||
Object: "chat.completion.chunk",
|
||||
Created: created,
|
||||
Model: model,
|
||||
Choices: []dto.ChatCompletionsStreamResponseChoice{{
|
||||
Index: 0,
|
||||
Delta: dto.ChatCompletionsStreamResponseChoiceDelta{Role: "assistant"},
|
||||
}},
|
||||
}
|
||||
|
||||
var content string
|
||||
if chunk.Message != nil {
|
||||
content = chunk.Message.Content
|
||||
} else {
|
||||
content = chunk.Response
|
||||
}
|
||||
if content != "" {
|
||||
delta.Choices[0].Delta.SetContentString(content)
|
||||
}
|
||||
|
||||
hasPayload := content != ""
|
||||
if chunk.Message != nil && len(chunk.Message.Thinking) > 0 {
|
||||
raw := strings.TrimSpace(string(chunk.Message.Thinking))
|
||||
if raw != "" && raw != "null" {
|
||||
// Unmarshal the JSON string to get the actual content without quotes
|
||||
var thinkingContent string
|
||||
if err := common.Unmarshal(chunk.Message.Thinking, &thinkingContent); err == nil {
|
||||
delta.Choices[0].Delta.SetReasoningContent(thinkingContent)
|
||||
} else {
|
||||
// Fallback to raw string if it's not a JSON string
|
||||
delta.Choices[0].Delta.SetReasoningContent(raw)
|
||||
}
|
||||
hasPayload = true
|
||||
}
|
||||
}
|
||||
if chunk.Message != nil && len(chunk.Message.ToolCalls) > 0 {
|
||||
delta.Choices[0].Delta.ToolCalls, *toolCallIndex = ollamaToolCallsToOpenAI(chunk.Message.ToolCalls, *toolCallIndex, true)
|
||||
hasPayload = true
|
||||
}
|
||||
|
||||
return delta, hasPayload
|
||||
}
|
||||
|
||||
func toUnix(ts string) int64 {
|
||||
if ts == "" {
|
||||
return time.Now().Unix()
|
||||
@@ -130,49 +175,18 @@ func ollamaStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http
|
||||
created = toUnix(chunk.CreatedAt)
|
||||
|
||||
if !chunk.Done {
|
||||
// delta content
|
||||
var content string
|
||||
if chunk.Message != nil {
|
||||
content = chunk.Message.Content
|
||||
} else {
|
||||
content = chunk.Response
|
||||
}
|
||||
delta := dto.ChatCompletionsStreamResponse{
|
||||
Id: responseId,
|
||||
Object: "chat.completion.chunk",
|
||||
Created: created,
|
||||
Model: model,
|
||||
Choices: []dto.ChatCompletionsStreamResponseChoice{{
|
||||
Index: 0,
|
||||
Delta: dto.ChatCompletionsStreamResponseChoiceDelta{Role: "assistant"},
|
||||
}},
|
||||
}
|
||||
if content != "" {
|
||||
delta.Choices[0].Delta.SetContentString(content)
|
||||
}
|
||||
if chunk.Message != nil && len(chunk.Message.Thinking) > 0 {
|
||||
raw := strings.TrimSpace(string(chunk.Message.Thinking))
|
||||
if raw != "" && raw != "null" {
|
||||
// Unmarshal the JSON string to get the actual content without quotes
|
||||
var thinkingContent string
|
||||
if err := common.Unmarshal(chunk.Message.Thinking, &thinkingContent); err == nil {
|
||||
delta.Choices[0].Delta.SetReasoningContent(thinkingContent)
|
||||
} else {
|
||||
// Fallback to raw string if it's not a JSON string
|
||||
delta.Choices[0].Delta.SetReasoningContent(raw)
|
||||
}
|
||||
}
|
||||
}
|
||||
// tool calls
|
||||
if chunk.Message != nil && len(chunk.Message.ToolCalls) > 0 {
|
||||
delta.Choices[0].Delta.ToolCalls, toolCallIndex = ollamaToolCallsToOpenAI(chunk.Message.ToolCalls, toolCallIndex, true)
|
||||
}
|
||||
delta, _ := buildOllamaStreamDelta(&chunk, responseId, created, model, &toolCallIndex)
|
||||
if data, err := common.Marshal(delta); err == nil {
|
||||
_ = helper.StringData(c, string(data))
|
||||
}
|
||||
continue
|
||||
}
|
||||
// done frame
|
||||
if delta, hasPayload := buildOllamaStreamDelta(&chunk, responseId, created, model, &toolCallIndex); hasPayload {
|
||||
if data, err := common.Marshal(delta); err == nil {
|
||||
_ = helper.StringData(c, string(data))
|
||||
}
|
||||
}
|
||||
// finalize once and break loop
|
||||
usage.PromptTokens = chunk.PromptEvalCount
|
||||
usage.CompletionTokens = chunk.EvalCount
|
||||
|
||||
@@ -97,3 +97,148 @@ func TestOllamaChatHandlerNonStreamToolCalls(t *testing.T) {
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildOllamaStreamDelta(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
raw string
|
||||
startIndex int
|
||||
wantContent string
|
||||
wantReasoning string
|
||||
wantToolCalls []struct {
|
||||
id string
|
||||
name string
|
||||
index int
|
||||
arguments string
|
||||
}
|
||||
wantNextIndex int
|
||||
wantPayload bool
|
||||
}{
|
||||
{
|
||||
name: "chat content",
|
||||
raw: `{"message":{"content":"hello"}}`,
|
||||
wantContent: "hello",
|
||||
wantPayload: true,
|
||||
},
|
||||
{
|
||||
name: "generate content",
|
||||
raw: `{"response":"hello"}`,
|
||||
wantContent: "hello",
|
||||
wantPayload: true,
|
||||
},
|
||||
{
|
||||
name: "thinking json string",
|
||||
raw: `{"message":{"thinking":"consider this"}}`,
|
||||
wantReasoning: "consider this",
|
||||
wantPayload: true,
|
||||
},
|
||||
{
|
||||
name: "thinking raw fallback",
|
||||
raw: `{"message":{"thinking":{"step":"consider this"}}}`,
|
||||
wantReasoning: `{"step":"consider this"}`,
|
||||
wantPayload: true,
|
||||
},
|
||||
{
|
||||
name: "multiple tool calls",
|
||||
raw: `{"message":{"tool_calls":[{"id":"call_upstream","function":{"name":"get_weather","arguments":{"city":"Paris"}}},{"function":{"name":"get_time","arguments":{"timezone":"UTC"}}}]}}`,
|
||||
startIndex: 2,
|
||||
wantToolCalls: []struct {
|
||||
id string
|
||||
name string
|
||||
index int
|
||||
arguments string
|
||||
}{
|
||||
{id: "call_upstream", name: "get_weather", index: 2, arguments: `{"city":"Paris"}`},
|
||||
{id: "call_3", name: "get_time", index: 3, arguments: `{"timezone":"UTC"}`},
|
||||
},
|
||||
wantNextIndex: 4,
|
||||
wantPayload: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
var chunk ollamaChatStreamChunk
|
||||
require.NoError(t, common.Unmarshal([]byte(tt.raw), &chunk))
|
||||
|
||||
toolCallIndex := tt.startIndex
|
||||
delta, hasPayload := buildOllamaStreamDelta(&chunk, "response-id", 123, "model", &toolCallIndex)
|
||||
assert.Equal(t, tt.wantPayload, hasPayload)
|
||||
assert.Equal(t, tt.wantContent, delta.Choices[0].Delta.GetContentString())
|
||||
assert.Equal(t, tt.wantReasoning, delta.Choices[0].Delta.GetReasoningContent())
|
||||
assert.Equal(t, tt.wantNextIndex, toolCallIndex)
|
||||
|
||||
if tt.wantToolCalls == nil {
|
||||
assert.Empty(t, delta.Choices[0].Delta.ToolCalls)
|
||||
return
|
||||
}
|
||||
require.Len(t, delta.Choices[0].Delta.ToolCalls, len(tt.wantToolCalls))
|
||||
for i, want := range tt.wantToolCalls {
|
||||
got := delta.Choices[0].Delta.ToolCalls[i]
|
||||
assert.Equal(t, want.id, got.ID)
|
||||
assert.Equal(t, "function", got.Type)
|
||||
assert.Equal(t, want.name, got.Function.Name)
|
||||
assert.Equal(t, want.arguments, got.Function.Arguments)
|
||||
require.NotNil(t, got.Index)
|
||||
assert.Equal(t, want.index, *got.Index)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestOllamaStreamHandlerDoneToolCalls(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
|
||||
resp := &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: make(http.Header),
|
||||
Body: io.NopCloser(strings.NewReader(`{"model":"qwen3-coder","created_at":"2026-05-27T12:00:00Z","message":{"role":"assistant","content":"","tool_calls":[{"function":{"name":"get_weather","arguments":{"city":"北京"}}},{"function":{"name":"get_time","arguments":{"timezone":"Asia/Shanghai"}}}]},"done":true,"done_reason":"stop","prompt_eval_count":5,"eval_count":7}`)),
|
||||
}
|
||||
|
||||
usage, apiErr := ollamaStreamHandler(c, &relaycommon.RelayInfo{
|
||||
ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "fallback-model"},
|
||||
}, resp)
|
||||
require.Nil(t, apiErr)
|
||||
require.NotNil(t, usage)
|
||||
assert.Equal(t, 12, usage.TotalTokens)
|
||||
|
||||
var chunks []dto.ChatCompletionsStreamResponse
|
||||
for _, line := range strings.Split(w.Body.String(), "\n") {
|
||||
data, ok := strings.CutPrefix(line, "data: ")
|
||||
if !ok || data == "[DONE]" {
|
||||
continue
|
||||
}
|
||||
var chunk dto.ChatCompletionsStreamResponse
|
||||
require.NoError(t, common.Unmarshal([]byte(data), &chunk))
|
||||
chunks = append(chunks, chunk)
|
||||
}
|
||||
|
||||
var toolChunkIndex, stopChunkIndex int = -1, -1
|
||||
for i := range chunks {
|
||||
if len(chunks[i].Choices) == 0 {
|
||||
continue
|
||||
}
|
||||
if len(chunks[i].Choices[0].Delta.ToolCalls) > 0 {
|
||||
toolChunkIndex = i
|
||||
}
|
||||
if chunks[i].Choices[0].FinishReason != nil {
|
||||
stopChunkIndex = i
|
||||
}
|
||||
}
|
||||
require.NotEqual(t, -1, toolChunkIndex)
|
||||
require.NotEqual(t, -1, stopChunkIndex)
|
||||
assert.Less(t, toolChunkIndex, stopChunkIndex)
|
||||
|
||||
toolCalls := chunks[toolChunkIndex].Choices[0].Delta.ToolCalls
|
||||
require.Len(t, toolCalls, 2)
|
||||
assert.Equal(t, "call_0", toolCalls[0].ID)
|
||||
assert.Equal(t, "call_1", toolCalls[1].ID)
|
||||
require.NotNil(t, toolCalls[0].Index)
|
||||
require.NotNil(t, toolCalls[1].Index)
|
||||
assert.Equal(t, 0, *toolCalls[0].Index)
|
||||
assert.Equal(t, 1, *toolCalls[1].Index)
|
||||
assert.Equal(t, constant.FinishReasonToolCalls, *chunks[stopChunkIndex].Choices[0].FinishReason)
|
||||
assert.Contains(t, w.Body.String(), "data: [DONE]")
|
||||
}
|
||||
@@ -51,6 +51,11 @@ func (a *Adaptor) ConvertGeminiRequest(c *gin.Context, info *relaycommon.RelayIn
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("expected OpenAI chat completions request, got %T", result.Value)
|
||||
}
|
||||
if info.SupportStreamOptions && info.IsStream {
|
||||
openaiRequest.StreamOptions = &dto.StreamOptions{
|
||||
IncludeUsage: true,
|
||||
}
|
||||
}
|
||||
return a.ConvertOpenAIRequest(c, info, openaiRequest)
|
||||
}
|
||||
|
||||
@@ -262,23 +267,26 @@ func (a *Adaptor) ConvertOpenAIRequest(c *gin.Context, info *relaycommon.RelayIn
|
||||
info.SetReasoningEffort(request.ReasoningEffort)
|
||||
}
|
||||
if info.ChannelType == constant.ChannelTypeOpenRouter {
|
||||
initialIntent, err := kitreasoning.FromOpenAIChat(request)
|
||||
initialIntent, diagnostics, err := kitreasoning.FromOpenAIChat(request)
|
||||
if err != nil {
|
||||
return nil, kitreasoning.AsClientError(err)
|
||||
}
|
||||
recordReasoningDiagnostics(c, info, diagnostics)
|
||||
if request.THINKING != nil && strings.HasPrefix(info.UpstreamModelName, "anthropic") {
|
||||
var thinking dto.Thinking
|
||||
if err := common.Unmarshal(request.THINKING, &thinking); err != nil {
|
||||
return nil, fmt.Errorf("error Unmarshal thinking: %w", err)
|
||||
}
|
||||
legacyIntent, err := kitreasoning.FromClaude(&dto.ClaudeRequest{Thinking: &thinking})
|
||||
legacyIntent, diagnostics, err := kitreasoning.FromClaude(&dto.ClaudeRequest{Thinking: &thinking})
|
||||
if err != nil {
|
||||
return nil, kitreasoning.AsClientError(err)
|
||||
}
|
||||
initialIntent, err = kitreasoning.MergeExplicit(initialIntent, legacyIntent, request.Model)
|
||||
recordReasoningDiagnostics(c, info, diagnostics)
|
||||
initialIntent, diagnostics, err = kitreasoning.MergeExplicit(initialIntent, legacyIntent, request.Model)
|
||||
if err != nil {
|
||||
return nil, kitreasoning.AsClientError(err)
|
||||
}
|
||||
recordReasoningDiagnostics(c, info, diagnostics)
|
||||
request.THINKING = nil
|
||||
}
|
||||
if len(request.Usage) == 0 {
|
||||
@@ -298,7 +306,9 @@ func (a *Adaptor) ConvertOpenAIRequest(c *gin.Context, info *relaycommon.RelayIn
|
||||
if effort == kitreasoning.EffortNone {
|
||||
mode = kitreasoning.ModeDisabled
|
||||
}
|
||||
initialIntent, err = kitreasoning.MergeExplicitAndSuffix(initialIntent, kitreasoning.Intent{Mode: mode, Effort: effort, Source: kitreasoning.SourceSuffix}, modelName)
|
||||
var diagnostics []types.ConversionDiagnostic
|
||||
initialIntent, diagnostics, err = kitreasoning.MergeExplicitAndSuffix(initialIntent, kitreasoning.Intent{Mode: mode, Effort: effort, Source: kitreasoning.SourceSuffix}, modelName)
|
||||
recordReasoningDiagnostics(c, info, diagnostics)
|
||||
return err
|
||||
}
|
||||
if !preserveSuffix {
|
||||
@@ -362,10 +372,11 @@ func (a *Adaptor) ConvertOpenAIRequest(c *gin.Context, info *relaycommon.RelayIn
|
||||
if preserveSuffix {
|
||||
effort = ""
|
||||
}
|
||||
currentIntent, err := kitreasoning.FromOpenAIChat(request)
|
||||
currentIntent, diagnostics, err := kitreasoning.FromOpenAIChat(request)
|
||||
if err != nil {
|
||||
return nil, kitreasoning.AsClientError(err)
|
||||
}
|
||||
recordReasoningDiagnostics(c, info, diagnostics)
|
||||
mergeSuffix := func(modelName, rawEffort string) error {
|
||||
if rawEffort == "" {
|
||||
return nil
|
||||
@@ -378,7 +389,9 @@ func (a *Adaptor) ConvertOpenAIRequest(c *gin.Context, info *relaycommon.RelayIn
|
||||
if suffixEffort == kitreasoning.EffortNone {
|
||||
mode = kitreasoning.ModeDisabled
|
||||
}
|
||||
currentIntent, err = kitreasoning.MergeExplicitAndSuffix(currentIntent, kitreasoning.Intent{Mode: mode, Effort: suffixEffort, Source: kitreasoning.SourceSuffix}, modelName)
|
||||
var diagnostics []types.ConversionDiagnostic
|
||||
currentIntent, diagnostics, err = kitreasoning.MergeExplicitAndSuffix(currentIntent, kitreasoning.Intent{Mode: mode, Effort: suffixEffort, Source: kitreasoning.SourceSuffix}, modelName)
|
||||
recordReasoningDiagnostics(c, info, diagnostics)
|
||||
return err
|
||||
}
|
||||
if err := mergeSuffix(info.UpstreamModelName, effort); err != nil {
|
||||
@@ -684,10 +697,11 @@ func (a *Adaptor) ConvertOpenAIResponsesRequest(c *gin.Context, info *relaycommo
|
||||
}
|
||||
return request, nil
|
||||
}
|
||||
currentIntent, err := kitreasoning.FromOpenAIResponses(&request)
|
||||
currentIntent, diagnostics, err := kitreasoning.FromOpenAIResponses(&request)
|
||||
if err != nil {
|
||||
return nil, kitreasoning.AsClientError(err)
|
||||
}
|
||||
recordReasoningDiagnostics(c, info, diagnostics)
|
||||
mergeSuffix := func(modelName, rawEffort string) error {
|
||||
if rawEffort == "" {
|
||||
return nil
|
||||
@@ -700,7 +714,9 @@ func (a *Adaptor) ConvertOpenAIResponsesRequest(c *gin.Context, info *relaycommo
|
||||
if suffixEffort == kitreasoning.EffortNone {
|
||||
mode = kitreasoning.ModeDisabled
|
||||
}
|
||||
currentIntent, err = kitreasoning.MergeExplicitAndSuffix(currentIntent, kitreasoning.Intent{Mode: mode, Effort: suffixEffort, Source: kitreasoning.SourceSuffix}, modelName)
|
||||
var diagnostics []types.ConversionDiagnostic
|
||||
currentIntent, diagnostics, err = kitreasoning.MergeExplicitAndSuffix(currentIntent, kitreasoning.Intent{Mode: mode, Effort: suffixEffort, Source: kitreasoning.SourceSuffix}, modelName)
|
||||
recordReasoningDiagnostics(c, info, diagnostics)
|
||||
return err
|
||||
}
|
||||
if err := mergeSuffix(request.Model, effort); err != nil {
|
||||
@@ -811,3 +827,15 @@ func (a *Adaptor) GetChannelName() string {
|
||||
return ChannelName
|
||||
}
|
||||
}
|
||||
|
||||
// recordReasoningDiagnostics attaches best-effort reasoning resolutions made
|
||||
// while merging request fields with model-name effort tails to the request log.
|
||||
func recordReasoningDiagnostics(c *gin.Context, info *relaycommon.RelayInfo, diagnostics []types.ConversionDiagnostic) {
|
||||
if info == nil || len(diagnostics) == 0 {
|
||||
return
|
||||
}
|
||||
for i := range diagnostics {
|
||||
diagnostics[i].From = info.RelayFormat
|
||||
}
|
||||
info.RecordConversionDiagnostics(c, diagnostics)
|
||||
}
|
||||
@@ -41,6 +41,7 @@ func OaiResponsesToChatHandler(c *gin.Context, info *relaycommon.RelayInfo, resp
|
||||
return nil, types.WithOpenAIError(*oaiError, resp.StatusCode)
|
||||
}
|
||||
|
||||
info.ObserveResponseModel(responsesResp.Model)
|
||||
responseValue, usage, err := convertResponsesResponseForClient(c, info, &responsesResp)
|
||||
if err != nil {
|
||||
return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
|
||||
@@ -60,6 +61,8 @@ func OaiResponsesToChatBufferedStreamHandler(c *gin.Context, info *relaycommon.R
|
||||
}
|
||||
defer service.CloseResponseBodyGracefully(resp)
|
||||
|
||||
info.StreamStatus = relaycommon.NewStreamStatus()
|
||||
info.StreamStatus.RequireTerminal()
|
||||
accumulator := relayconvert.NewResponsesBufferedAccumulator()
|
||||
var finalResponse *dto.OpenAIResponsesResponse
|
||||
var streamErr *types.NewAPIError
|
||||
@@ -86,6 +89,10 @@ func OaiResponsesToChatBufferedStreamHandler(c *gin.Context, info *relaycommon.R
|
||||
streamErr = types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
|
||||
break
|
||||
}
|
||||
if streamResp.Response != nil {
|
||||
info.ObserveResponseModel(streamResp.Response.Model)
|
||||
}
|
||||
service.ObserveResponsesOutcome(info, &streamResp)
|
||||
accumulator.ProcessEvent(&streamResp)
|
||||
switch streamResp.Type {
|
||||
case "response.completed", "response.done", "response.incomplete":
|
||||
@@ -120,7 +127,7 @@ func OaiResponsesToChatBufferedStreamHandler(c *gin.Context, info *relaycommon.R
|
||||
if finalResponse == nil {
|
||||
finalResponse = &dto.OpenAIResponsesResponse{
|
||||
ID: helper.GetResponseID(c),
|
||||
CreatedAt: int(time.Now().Unix()),
|
||||
CreatedAt: dto.IntValue(time.Now().Unix()),
|
||||
Model: info.UpstreamModelName,
|
||||
Status: []byte(`"completed"`),
|
||||
}
|
||||
@@ -257,6 +264,9 @@ func OaiResponsesToChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo
|
||||
return
|
||||
}
|
||||
|
||||
if streamResp.Response != nil {
|
||||
info.ObserveResponseModel(streamResp.Response.Model)
|
||||
}
|
||||
if streamResp.Type == "response.error" || streamResp.Type == "response.failed" {
|
||||
if streamResp.Response != nil {
|
||||
if oaiErr := streamResp.Response.GetOpenAIError(); oaiErr != nil && oaiErr.Type != "" {
|
||||
|
||||
@@ -137,13 +137,14 @@ func ProcessStreamResponse(streamResponse dto.ChatCompletionsStreamResponse, res
|
||||
return nil
|
||||
}
|
||||
|
||||
func processTokenData(relayMode int, data string, responseTextBuilder *strings.Builder, toolCount *int) error {
|
||||
switch relayMode {
|
||||
func processTokenData(info *relaycommon.RelayInfo, data string, responseTextBuilder *strings.Builder, toolCount *int) error {
|
||||
switch info.RelayMode {
|
||||
case relayconstant.RelayModeChatCompletions:
|
||||
var streamResponse dto.ChatCompletionsStreamResponse
|
||||
if err := common.UnmarshalJsonStr(data, &streamResponse); err != nil {
|
||||
return err
|
||||
}
|
||||
info.ObserveResponseModel(streamResponse.Model)
|
||||
return ProcessStreamResponse(streamResponse, responseTextBuilder, toolCount)
|
||||
case relayconstant.RelayModeCompletions:
|
||||
var streamResponse dto.CompletionsStreamResponse
|
||||
|
||||
@@ -2,6 +2,7 @@ package openai
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
@@ -10,9 +11,12 @@ import (
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/pkg/billingexpr"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
relayconstant "github.com/QuantumNous/new-api/relay/constant"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
@@ -35,6 +39,112 @@ func newImageTestContext(t *testing.T, body, contentType string, isStream bool)
|
||||
return c, recorder, resp, info
|
||||
}
|
||||
|
||||
func TestImageExpressionUsesCompletedCountAndProtectsAbortedStreams(t *testing.T) {
|
||||
oldTimeout := constant.StreamingTimeout
|
||||
constant.StreamingTimeout = 30
|
||||
t.Cleanup(func() { constant.StreamingTimeout = oldTimeout })
|
||||
for _, tc := range []struct {
|
||||
name, body string
|
||||
stream, abort bool
|
||||
requested, want int
|
||||
}{
|
||||
{"JSON uses actual count", `{"data":[{"b64_json":"first"},{"b64_json":"second"}]}`, false, false, 3, 2},
|
||||
{"JSON wrapped as SSE uses actual count", `{"data":[{"b64_json":"first"}]}`, true, false, 3, 1},
|
||||
{"empty response retains request", `{"data":[]}`, false, false, 3, 3},
|
||||
{"JSON object data counts one image", `{"data":{"url":"https://example.com/a.png","b64_json":"first"}}`, false, false, 3, 1},
|
||||
{"JSON wrapped as SSE counts object data once", `{"data":{"b64_json":"first"}}`, true, false, 3, 1},
|
||||
{"completed stream refunds missing images", "data: {\"type\":\"image_generation.completed\",\"b64_json\":\"first\"}\n\ndata: [DONE]\n\n", true, false, 3, 1},
|
||||
{"client abort cannot reduce count", "data: {\"type\":\"image_generation.completed\",\"b64_json\":\"first\"}\n\n", true, true, 3, 3},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
contentType := "application/json"
|
||||
if strings.HasPrefix(tc.body, "data:") {
|
||||
contentType = "text/event-stream"
|
||||
}
|
||||
c, _, resp, info := newImageTestContext(t, tc.body, contentType, tc.stream)
|
||||
if tc.abort {
|
||||
c, _, resp, info = newDisconnectingImageStream(t, tc.body, "first")
|
||||
}
|
||||
info.TieredBillingSnapshot = &billingexpr.BillingSnapshot{EstimatedImageCount: &tc.requested}
|
||||
if tc.stream {
|
||||
_, err := OpenaiImageStreamHandler(c, info, resp)
|
||||
require.Nil(t, err)
|
||||
} else {
|
||||
_, err := OpenaiImageHandler(c, info, resp)
|
||||
require.Nil(t, err)
|
||||
}
|
||||
count := info.RequestedImageCount()
|
||||
if info.BillingImageCount != nil {
|
||||
count = *info.BillingImageCount
|
||||
}
|
||||
assert.Equal(t, tc.want, count)
|
||||
assert.Empty(t, info.PriceData.OtherRatios(), "expression quantities must not add a legacy multiplier")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestImageCacheUsageAcrossResponseFormats(t *testing.T) {
|
||||
previousTimeout := constant.StreamingTimeout
|
||||
constant.StreamingTimeout = 30
|
||||
t.Cleanup(func() { constant.StreamingTimeout = previousTimeout })
|
||||
// Contract fixture for compatible upstreams. Live OpenAI cache response
|
||||
// verification is a separate release check; this is not a captured response.
|
||||
const usageJSON = `{"input_tokens":1000,"output_tokens":100,"total_tokens":1100,"input_tokens_details":{"text_tokens":400,"image_tokens":600,"cached_tokens":300,"cached_tokens_details":{"text_tokens":100,"image_tokens":200}}}`
|
||||
for _, mode := range []string{"image_generation", "image_edit"} {
|
||||
for _, stream := range []bool{false, true} {
|
||||
t.Run(fmt.Sprintf("%s/stream=%t", mode, stream), func(t *testing.T) {
|
||||
body := `{"data":[{"b64_json":"image"}],"usage":` + usageJSON + `}`
|
||||
contentType := "application/json"
|
||||
if stream {
|
||||
body = "data: {\"type\":\"" + mode + ".completed\",\"usage\":" + usageJSON + "}\n\ndata: [DONE]\n\n"
|
||||
contentType = "text/event-stream"
|
||||
}
|
||||
ctx, recorder, response, info := newImageTestContext(t, body, contentType, stream)
|
||||
info.RelayMode = relayconstant.RelayModeImagesGenerations
|
||||
if mode == "image_edit" {
|
||||
info.RelayMode = relayconstant.RelayModeImagesEdits
|
||||
ctx.Request.URL.Path = "/v1/images/edits"
|
||||
}
|
||||
result, apiErr := (&Adaptor{}).DoResponse(ctx, response, info)
|
||||
require.Nil(t, apiErr)
|
||||
usage := result.(*dto.Usage)
|
||||
assert.Equal(t, 1000, usage.PromptTokens)
|
||||
assert.Equal(t, 300, usage.PromptTokensDetails.CachedTokens)
|
||||
require.NotNil(t, usage.PromptTokensDetails.CachedTokensDetails)
|
||||
assert.Equal(t, 200, *usage.PromptTokensDetails.CachedTokensDetails.ImageTokens)
|
||||
assert.Contains(t, recorder.Body.String(), `"cached_tokens_details":{"text_tokens":100,"image_tokens":200}`)
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestNormalizeOpenAIUsageMapsOutputImageTokens(t *testing.T) {
|
||||
previousTimeout := constant.StreamingTimeout
|
||||
constant.StreamingTimeout = 30
|
||||
t.Cleanup(func() { constant.StreamingTimeout = previousTimeout })
|
||||
|
||||
const usageJSON = `{"input_tokens":15,"output_tokens":1352,"total_tokens":1367,"input_tokens_details":{"text_tokens":15,"image_tokens":0},"output_tokens_details":{"image_tokens":1120,"text_tokens":232}}`
|
||||
for _, stream := range []bool{false, true} {
|
||||
t.Run(fmt.Sprintf("stream=%t", stream), func(t *testing.T) {
|
||||
body := `{"data":[{"b64_json":"image"}],"usage":` + usageJSON + `}`
|
||||
contentType := "application/json"
|
||||
if stream {
|
||||
body = "data: {\"type\":\"image_generation.completed\",\"usage\":" + usageJSON + "}\n\ndata: [DONE]\n\n"
|
||||
contentType = "text/event-stream"
|
||||
}
|
||||
ctx, _, response, info := newImageTestContext(t, body, contentType, stream)
|
||||
info.RelayMode = relayconstant.RelayModeImagesGenerations
|
||||
result, apiErr := (&Adaptor{}).DoResponse(ctx, response, info)
|
||||
require.Nil(t, apiErr)
|
||||
usage := result.(*dto.Usage)
|
||||
assert.Equal(t, 15, usage.PromptTokens)
|
||||
assert.Equal(t, 1352, usage.CompletionTokens)
|
||||
assert.Equal(t, 1120, usage.CompletionTokenDetails.ImageTokens)
|
||||
assert.Equal(t, 232, usage.CompletionTokenDetails.TextTokens)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenaiImageDoResponseUsesInfoIsStream(t *testing.T) {
|
||||
oldMode := gin.Mode()
|
||||
gin.SetMode(gin.TestMode)
|
||||
@@ -344,6 +454,24 @@ func TestOpenaiImageHandlerUsesPositiveActualCountForFixedPrice(t *testing.T) {
|
||||
usePrice: false,
|
||||
wantCount: 3,
|
||||
},
|
||||
{
|
||||
name: "object data with url and b64_json counts one image",
|
||||
body: `{"data":{"url":"https://example.com/a.png","b64_json":"` + longImage + `"}}`,
|
||||
usePrice: true,
|
||||
wantCount: 1,
|
||||
},
|
||||
{
|
||||
name: "url and b64_json split across entries count one image",
|
||||
body: `{"data":[{"url":"https://example.com/a.png"},{"b64_json":"` + longImage + `"}]}`,
|
||||
usePrice: true,
|
||||
wantCount: 1,
|
||||
},
|
||||
{
|
||||
name: "entries without image payload keep requested count",
|
||||
body: `{"data":[{"revised_prompt":"draw a cat"}]}`,
|
||||
usePrice: true,
|
||||
wantCount: 3,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
@@ -361,6 +489,62 @@ func TestOpenaiImageHandlerUsesPositiveActualCountForFixedPrice(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// TestOpenaiImageStreamHandlerWrapsNonStandardDataShapes covers the JSON-to-SSE
|
||||
// fallback when upstream returns data as an object, splits one image across a
|
||||
// url entry and a b64_json entry, or returns entries without any image payload.
|
||||
// Forwarded events and the billed quantity must both follow the real images.
|
||||
func TestOpenaiImageStreamHandlerWrapsNonStandardDataShapes(t *testing.T) {
|
||||
oldMode := gin.Mode()
|
||||
gin.SetMode(gin.TestMode)
|
||||
t.Cleanup(func() { gin.SetMode(oldMode) })
|
||||
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
body string
|
||||
wantEvents int
|
||||
wantCount float64
|
||||
wantBody []string
|
||||
}{
|
||||
{
|
||||
name: "object data is forwarded as one completed event",
|
||||
body: `{"data":{"url":"https://example.com/a.png","b64_json":"first"}}`,
|
||||
wantEvents: 1,
|
||||
wantCount: 1,
|
||||
wantBody: []string{`"url":"https://example.com/a.png"`, `"b64_json":"first"`},
|
||||
},
|
||||
{
|
||||
name: "split url and b64_json entries bill one image",
|
||||
body: `{"data":[{"url":"https://example.com/a.png"},{"b64_json":"first"}]}`,
|
||||
wantEvents: 2,
|
||||
wantCount: 1,
|
||||
wantBody: []string{`"url":"https://example.com/a.png"`, `"b64_json":"first"`},
|
||||
},
|
||||
{
|
||||
name: "entries without image payload are dropped and keep requested count",
|
||||
body: `{"data":[{"revised_prompt":"draw a cat"}]}`,
|
||||
wantEvents: 0,
|
||||
wantCount: 3,
|
||||
},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
c, recorder, resp, info := newImageTestContext(t, tc.body, "application/json", true)
|
||||
info.PriceData.UsePrice = true
|
||||
info.PriceData.AddOtherRatio("n", 3)
|
||||
|
||||
_, err := OpenaiImageStreamHandler(c, info, resp)
|
||||
require.Nil(t, err)
|
||||
out := recorder.Body.String()
|
||||
assert.Equal(t, tc.wantEvents, strings.Count(out, "event: image_generation.completed"))
|
||||
assert.Equal(t, tc.wantEvents, info.ReceivedResponseCount)
|
||||
for _, want := range tc.wantBody {
|
||||
assert.Contains(t, out, want)
|
||||
}
|
||||
assert.True(t, strings.HasSuffix(out, "data: [DONE]\n\n"))
|
||||
assert.Equal(t, tc.wantCount, info.PriceData.OtherRatios()["n"])
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestOpenaiImageHandlersReturnJSONError covers JSON error responses for both
|
||||
// entry points: the non-streaming handler and the stream handler's non-SSE
|
||||
// fallback. Neither must leak the error body to the client.
|
||||
|
||||
@@ -134,14 +134,16 @@ func OaiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Re
|
||||
}
|
||||
|
||||
lastStreamData = data
|
||||
collectStreamFunctionCallNames(data, seenStreamToolCalls, &streamFunctionCallNames)
|
||||
if err := processTokenData(info.RelayMode, data, &responseTextBuilder, &toolCount); err != nil {
|
||||
observeStreamChoices(info, data, seenStreamToolCalls, &streamFunctionCallNames)
|
||||
if err := processTokenData(info, data, &responseTextBuilder, &toolCount); err != nil {
|
||||
logger.LogError(c, "error processing stream token data: "+err.Error())
|
||||
sr.Error(err)
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
info.StreamStatus.RequireTerminal()
|
||||
|
||||
// 处理最后的响应
|
||||
shouldSendLastResp := true
|
||||
if err := handleLastResponse(lastStreamData, &responseId, &createAt, &systemFingerprint, &model, &usage,
|
||||
@@ -194,12 +196,20 @@ func OaiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Re
|
||||
return usage, nil
|
||||
}
|
||||
|
||||
func collectStreamFunctionCallNames(data string, seen map[string]struct{}, names *[]string) {
|
||||
// observeStreamChoices collects billable function call names and records the
|
||||
// finish reason facts used by health sampling from one parsed chunk.
|
||||
func observeStreamChoices(info *relaycommon.RelayInfo, data string, seen map[string]struct{}, names *[]string) {
|
||||
var streamResponse dto.ChatCompletionsStreamResponse
|
||||
if err := common.UnmarshalJsonStr(data, &streamResponse); err != nil {
|
||||
return
|
||||
}
|
||||
for _, choice := range streamResponse.Choices {
|
||||
if choice.FinishReason != nil && *choice.FinishReason != "" {
|
||||
if *choice.FinishReason == constant.FinishReasonContentFilter {
|
||||
info.PerformanceBusinessRejection = true
|
||||
}
|
||||
info.StreamStatus.MarkCompleted()
|
||||
}
|
||||
for i, tc := range choice.Delta.ToolCalls {
|
||||
name := strings.TrimSpace(tc.Function.Name)
|
||||
if name == "" {
|
||||
@@ -272,8 +282,10 @@ func OpenaiHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Respo
|
||||
return nil, types.WithOpenAIError(*oaiError, resp.StatusCode)
|
||||
}
|
||||
|
||||
info.ObserveResponseModel(simpleResponse.Model)
|
||||
for _, choice := range simpleResponse.Choices {
|
||||
if choice.FinishReason == constant.FinishReasonContentFilter {
|
||||
info.PerformanceBusinessRejection = true
|
||||
common.SetContextKey(c, constant.ContextKeyAdminRejectReason, "openai_finish_reason=content_filter")
|
||||
break
|
||||
}
|
||||
|
||||
@@ -5,7 +5,6 @@ import (
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
@@ -22,13 +21,6 @@ import (
|
||||
"github.com/tidwall/sjson"
|
||||
)
|
||||
|
||||
func updateOpenAIImageCount(info *relaycommon.RelayInfo, count int64) {
|
||||
if info == nil || !info.PriceData.UsePrice || count <= 0 || count > int64(dto.MaxImageN) {
|
||||
return
|
||||
}
|
||||
info.PriceData.AddOtherRatio("n", float64(count))
|
||||
}
|
||||
|
||||
// OpenaiImageHandler handles non-streaming OpenAI image responses
|
||||
// (generations/edits), returning the parsed usage for billing.
|
||||
func OpenaiImageHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) {
|
||||
@@ -49,7 +41,7 @@ func OpenaiImageHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.
|
||||
return nil, types.WithOpenAIError(*oaiError, resp.StatusCode)
|
||||
}
|
||||
|
||||
updateOpenAIImageCount(info, gjson.GetBytes(responseBody, "data.#").Int())
|
||||
info.UpdateImageCount(openaiImageResponseCount(responseBody))
|
||||
|
||||
// 写入新的 response body
|
||||
service.IOCopyBytesGracefully(c, resp, responseBody)
|
||||
@@ -59,14 +51,54 @@ func OpenaiImageHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.
|
||||
return &usageResp.Usage, nil
|
||||
}
|
||||
|
||||
// openaiImageResponseCount counts billable images in an OpenAI-format image
|
||||
// response body. An object-shaped data is one image when it carries url or
|
||||
// b64_json. For arrays the count is the larger of the url-bearing and the
|
||||
// b64_json-bearing entry counts: a standard response uses one response_format
|
||||
// so this equals the entry count, an upstream that splits one image into a url
|
||||
// entry and a b64_json entry bills once, and entries without any image payload
|
||||
// bill nothing. A zero result leaves the requested quantity in place because
|
||||
// UpdateImageCount ignores non-positive counts.
|
||||
func openaiImageResponseCount(responseBody []byte) int64 {
|
||||
data := gjson.GetBytes(responseBody, "data")
|
||||
if data.IsObject() {
|
||||
if openaiImageDataHasField(data, "url") || openaiImageDataHasField(data, "b64_json") {
|
||||
return 1
|
||||
}
|
||||
return 0
|
||||
}
|
||||
if !data.IsArray() {
|
||||
return 0
|
||||
}
|
||||
var urls, b64s int64
|
||||
data.ForEach(func(_, item gjson.Result) bool {
|
||||
if openaiImageDataHasField(item, "url") {
|
||||
urls++
|
||||
}
|
||||
if openaiImageDataHasField(item, "b64_json") {
|
||||
b64s++
|
||||
}
|
||||
return true
|
||||
})
|
||||
return max(urls, b64s)
|
||||
}
|
||||
|
||||
// openaiImageDataHasField reports whether an image data entry carries a
|
||||
// non-empty string value for field.
|
||||
func openaiImageDataHasField(item gjson.Result, field string) bool {
|
||||
value := item.Get(field)
|
||||
return value.Type == gjson.String && value.Raw != `""`
|
||||
}
|
||||
|
||||
// normalizeOpenAIUsage maps the OpenAI Images usage shape (input_tokens /
|
||||
// output_tokens / input_tokens_details) onto the canonical prompt/completion
|
||||
// fields. It is used only on the OpenAI image relay paths (generations/edits,
|
||||
// streaming and non-streaming): the image API never returns prompt_tokens /
|
||||
// completion_tokens, so the overwrite (=) semantics here are equivalent to the
|
||||
// previous additive (+=) behavior while avoiding any future double-counting if
|
||||
// both field sets are ever populated. Do not reuse this on chat/embedding paths
|
||||
// without revisiting the overwrite semantics.
|
||||
// output_tokens / input_tokens_details / output_tokens_details) onto the
|
||||
// canonical prompt/completion fields. It is used only on the OpenAI image
|
||||
// relay paths (generations/edits, streaming and non-streaming): the image
|
||||
// API never returns prompt_tokens / completion_tokens, so the overwrite (=)
|
||||
// semantics here are equivalent to the previous additive (+=) behavior while
|
||||
// avoiding any future double-counting if both field sets are ever populated.
|
||||
// Do not reuse this on chat/embedding paths without revisiting the overwrite
|
||||
// semantics.
|
||||
func normalizeOpenAIUsage(usage *dto.Usage) {
|
||||
if usage == nil {
|
||||
return
|
||||
@@ -78,12 +110,10 @@ func normalizeOpenAIUsage(usage *dto.Usage) {
|
||||
usage.CompletionTokens = usage.OutputTokens
|
||||
}
|
||||
if usage.InputTokensDetails != nil {
|
||||
usage.PromptTokensDetails.CachedTokens = usage.InputTokensDetails.CachedTokens
|
||||
usage.PromptTokensDetails.CachedCreationTokens = usage.InputTokensDetails.CachedCreationTokens
|
||||
usage.PromptTokensDetails.CacheWriteTokens = usage.InputTokensDetails.CacheWriteTokens
|
||||
usage.PromptTokensDetails.ImageTokens = usage.InputTokensDetails.ImageTokens
|
||||
usage.PromptTokensDetails.TextTokens = usage.InputTokensDetails.TextTokens
|
||||
usage.PromptTokensDetails.AudioTokens = usage.InputTokensDetails.AudioTokens
|
||||
usage.PromptTokensDetails = usage.InputTokensDetails.Clone()
|
||||
}
|
||||
if usage.OutputTokensDetails != nil {
|
||||
usage.CompletionTokenDetails = *usage.OutputTokensDetails
|
||||
}
|
||||
if usage.TotalTokens == 0 {
|
||||
usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens
|
||||
@@ -156,12 +186,8 @@ func OpenaiImageStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp
|
||||
if info.StreamStatus != nil {
|
||||
upstreamFinished := info.StreamStatus.EndReason == relaycommon.StreamEndReasonDone ||
|
||||
info.StreamStatus.EndReason == relaycommon.StreamEndReasonEOF
|
||||
requestedN := 1.0
|
||||
if n, ok := info.PriceData.OtherRatios()["n"]; ok {
|
||||
requestedN = n
|
||||
}
|
||||
if upstreamFinished || float64(completedImages) > requestedN {
|
||||
updateOpenAIImageCount(info, completedImages)
|
||||
if upstreamFinished || completedImages > int64(info.RequestedImageCount()) {
|
||||
info.UpdateImageCount(completedImages)
|
||||
}
|
||||
}
|
||||
return usage, nil
|
||||
@@ -252,8 +278,7 @@ func openaiImageJSONAsStreamHandler(c *gin.Context, info *relaycommon.RelayInfo,
|
||||
normalizeOpenAIUsage(&usageResp.Usage)
|
||||
applyUsagePostProcessing(info, &usageResp.Usage, responseBody)
|
||||
|
||||
imageCount := gjson.GetBytes(responseBody, "data.#").Int()
|
||||
updateOpenAIImageCount(info, imageCount)
|
||||
info.UpdateImageCount(openaiImageResponseCount(responseBody))
|
||||
|
||||
helper.SetEventStreamHeaders(c)
|
||||
c.Status(http.StatusOK)
|
||||
@@ -275,8 +300,15 @@ func openaiImageJSONAsStreamHandler(c *gin.Context, info *relaycommon.RelayInfo,
|
||||
}
|
||||
}
|
||||
|
||||
for i := range imageCount {
|
||||
image := gjson.GetBytes(responseBody, "data."+strconv.FormatInt(i, 10))
|
||||
// gjson.Result.Array returns the element list for an array and a single
|
||||
// element for an object-shaped data, keeping document-relative indexes so
|
||||
// the zero-copy field forwarding below stays valid for both shapes.
|
||||
// Entries without url or b64_json carry no image and are not forwarded.
|
||||
emitted := 0
|
||||
for _, image := range gjson.GetBytes(responseBody, "data").Array() {
|
||||
if !openaiImageDataHasField(image, "url") && !openaiImageDataHasField(image, "b64_json") {
|
||||
continue
|
||||
}
|
||||
payload := []byte(`{"type":"image_generation.completed"}`)
|
||||
payload, err = sjson.SetBytes(payload, "created_at", created)
|
||||
if err != nil {
|
||||
@@ -311,6 +343,7 @@ func openaiImageJSONAsStreamHandler(c *gin.Context, info *relaycommon.RelayInfo,
|
||||
}
|
||||
return &usageResp.Usage, nil
|
||||
}
|
||||
emitted++
|
||||
}
|
||||
if err := writeOpenaiImageStreamDone(c); err != nil {
|
||||
if info != nil && info.StreamStatus != nil {
|
||||
@@ -319,7 +352,7 @@ func openaiImageJSONAsStreamHandler(c *gin.Context, info *relaycommon.RelayInfo,
|
||||
return &usageResp.Usage, nil
|
||||
}
|
||||
if info != nil {
|
||||
info.ReceivedResponseCount += int(imageCount)
|
||||
info.ReceivedResponseCount += emitted
|
||||
if info.StreamStatus == nil {
|
||||
info.StreamStatus = relaycommon.NewStreamStatus()
|
||||
}
|
||||
|
||||
@@ -1,12 +1,15 @@
|
||||
package openai
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/logger"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/relay/helper"
|
||||
@@ -16,10 +19,25 @@ import (
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/tidwall/gjson"
|
||||
"github.com/tidwall/sjson"
|
||||
)
|
||||
|
||||
func OaiResponsesHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) {
|
||||
defer service.CloseResponseBodyGracefully(resp)
|
||||
if strings.HasPrefix(strings.ToLower(resp.Header.Get("Content-Type")), "text/event-stream") {
|
||||
body, err := readBufferedResponsesSSE(resp.Body)
|
||||
if err != nil {
|
||||
return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusBadGateway)
|
||||
}
|
||||
converted := *resp
|
||||
converted.Header = resp.Header.Clone()
|
||||
converted.Header.Set("Content-Type", "application/json")
|
||||
converted.Header.Del("Content-Length")
|
||||
converted.Body = io.NopCloser(bytes.NewReader(body))
|
||||
converted.ContentLength = int64(len(body))
|
||||
return OaiResponsesHandler(c, info, &converted)
|
||||
}
|
||||
|
||||
// read response body
|
||||
var responsesResponse dto.OpenAIResponsesResponse
|
||||
@@ -35,11 +53,15 @@ func OaiResponsesHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http
|
||||
return nil, types.WithOpenAIError(*oaiError, resp.StatusCode)
|
||||
}
|
||||
|
||||
info.ObserveResponseModel(responsesResponse.Model)
|
||||
responseBody = rewriteSGLangResponsesCreatedAt(info, responseBody, "created_at", responsesResponse.CreatedAt)
|
||||
|
||||
// 写入新的 response body
|
||||
service.IOCopyBytesGracefully(c, resp, responseBody)
|
||||
|
||||
// compute usage
|
||||
usage := relayconvert.NormalizeResponsesUsage(responsesResponse.Usage)
|
||||
usage := &dto.Usage{}
|
||||
service.ApplyResponsesUsage(usage, responsesResponse.Usage)
|
||||
// Count actual tool invocations from Output (not tool declarations).
|
||||
for _, output := range responsesResponse.Output {
|
||||
switch output.Type {
|
||||
@@ -64,6 +86,58 @@ func OaiResponsesHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http
|
||||
return usage, nil
|
||||
}
|
||||
|
||||
// Some upstreams return SSE even for stream=false. Require a terminal response
|
||||
// and reuse JSON delivery and settlement so upstream usage is charged once.
|
||||
func readBufferedResponsesSSE(body io.Reader) ([]byte, error) {
|
||||
scanner := helper.NewStreamScanner(body)
|
||||
scanner.Split(bufio.ScanLines)
|
||||
accumulator := relayconvert.NewResponsesBufferedAccumulator()
|
||||
var data []string
|
||||
for scanner.Scan() {
|
||||
line := strings.TrimSuffix(scanner.Text(), "\r")
|
||||
if payload, ok := strings.CutPrefix(line, "data:"); ok {
|
||||
data = append(data, strings.TrimPrefix(payload, " "))
|
||||
}
|
||||
if line != "" || len(data) == 0 {
|
||||
continue
|
||||
}
|
||||
payload := strings.Join(data, "\n")
|
||||
data = nil
|
||||
if payload == "[DONE]" {
|
||||
break
|
||||
}
|
||||
var event dto.ResponsesStreamResponse
|
||||
if err := common.UnmarshalJsonStr(payload, &event); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
accumulator.ProcessEvent(&event)
|
||||
switch event.Type {
|
||||
case "error", "response.error", "response.failed", "response.cancelled", "response.canceled":
|
||||
return nil, fmt.Errorf("responses stream ended with %s", event.Type)
|
||||
case "response.completed", "response.done", "response.incomplete":
|
||||
if event.Response == nil {
|
||||
return nil, fmt.Errorf("responses terminal event has no response")
|
||||
}
|
||||
result := []byte(gjson.Get(payload, "response").Raw)
|
||||
if len(event.Response.Output) == 0 {
|
||||
accumulator.SupplementResponseOutput(event.Response)
|
||||
if len(event.Response.Output) != 0 {
|
||||
output, err := common.Marshal(event.Response.Output)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return sjson.SetRawBytes(result, "output", output)
|
||||
}
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
}
|
||||
if err := scanner.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return nil, fmt.Errorf("responses stream ended without a terminal response")
|
||||
}
|
||||
|
||||
func OaiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) {
|
||||
if resp == nil || resp.Body == nil {
|
||||
logger.LogError(c, "invalid response or response body")
|
||||
@@ -72,10 +146,7 @@ func OaiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp
|
||||
|
||||
defer service.CloseResponseBodyGracefully(resp)
|
||||
|
||||
var usage = &dto.Usage{}
|
||||
var responseTextBuilder strings.Builder
|
||||
imageCounter := &relaycommon.ImageGenerationCallCounter{}
|
||||
imageCommitted := false
|
||||
accumulator := service.NewResponsesUsageAccumulator(info)
|
||||
|
||||
helper.StreamScannerHandler(c, resp, info, func(data string, sr *helper.StreamResult) {
|
||||
|
||||
@@ -86,77 +157,28 @@ func OaiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp
|
||||
sr.Error(err)
|
||||
return
|
||||
}
|
||||
sendResponsesStreamData(c, streamResponse, data)
|
||||
switch streamResponse.Type {
|
||||
case "response.completed", "response.done":
|
||||
if streamResponse.Response != nil {
|
||||
if streamResponse.Response.Usage != nil {
|
||||
incomingUsage := relayconvert.NormalizeResponsesUsage(streamResponse.Response.Usage)
|
||||
usage = dto.MergeUsageNonZero(usage, incomingUsage)
|
||||
}
|
||||
if !imageCommitted {
|
||||
if relaycommon.IsNonBillableResponsesStatus(streamResponse.Response.Status) {
|
||||
imageCounter.Reset()
|
||||
imageCounter.Commit(info)
|
||||
imageCommitted = true
|
||||
} else {
|
||||
for i := range streamResponse.Response.Output {
|
||||
idx := i
|
||||
imageCounter.Observe(&streamResponse.Response.Output[i], &idx)
|
||||
}
|
||||
imageCounter.Commit(info)
|
||||
imageCommitted = true
|
||||
}
|
||||
}
|
||||
} else if !imageCommitted {
|
||||
imageCounter.Commit(info)
|
||||
imageCommitted = true
|
||||
}
|
||||
case "response.failed", "response.incomplete", "response.cancelled", "response.canceled":
|
||||
if !imageCommitted {
|
||||
imageCounter.Reset()
|
||||
imageCounter.Commit(info)
|
||||
imageCommitted = true
|
||||
}
|
||||
case "response.output_text.delta":
|
||||
// 处理输出文本
|
||||
responseTextBuilder.WriteString(streamResponse.Delta)
|
||||
case dto.ResponsesOutputTypeItemDone:
|
||||
if streamResponse.Item != nil {
|
||||
switch streamResponse.Item.Type {
|
||||
case dto.BuildInCallWebSearchCall:
|
||||
info.CountBillableToolCall(dto.BuildInCallWebSearchCall, "")
|
||||
case dto.BuildInCallFileSearchCall:
|
||||
info.CountBillableToolCall(dto.BuildInCallFileSearchCall, "")
|
||||
case dto.BuildInCallFunctionCall:
|
||||
info.CountBillableToolCall(dto.BuildInCallFunctionCall, streamResponse.Item.Name)
|
||||
case dto.ResponsesOutputTypeImageGenerationCall:
|
||||
if !imageCommitted {
|
||||
imageCounter.Observe(streamResponse.Item, streamResponse.OutputIndex)
|
||||
}
|
||||
}
|
||||
}
|
||||
if streamResponse.Response != nil {
|
||||
data = string(rewriteSGLangResponsesCreatedAt(info, []byte(data), "response.created_at", streamResponse.Response.CreatedAt))
|
||||
}
|
||||
sendResponsesStreamData(c, streamResponse, data)
|
||||
accumulator.Observe(&streamResponse)
|
||||
})
|
||||
|
||||
if usage.CompletionTokens == 0 {
|
||||
// 计算输出文本的 token 数量
|
||||
tempStr := responseTextBuilder.String()
|
||||
if len(tempStr) > 0 {
|
||||
// 非正常结束,使用输出文本的 token 数量
|
||||
completionTokens := service.CountTextToken(tempStr, info.UpstreamModelName)
|
||||
usage.CompletionTokens = completionTokens
|
||||
}
|
||||
common.SetContextKey(c, constant.ContextKeyResponseStreamStatus, info.StreamStatus)
|
||||
info.StreamStatus.RequireTerminal()
|
||||
return accumulator.Finish(), nil
|
||||
}
|
||||
|
||||
func rewriteSGLangResponsesCreatedAt(info *relaycommon.RelayInfo, payload []byte, path string, createdAt dto.IntValue) []byte {
|
||||
if info.GetChannelType() != constant.ChannelTypeSGLang {
|
||||
return payload
|
||||
}
|
||||
|
||||
if usage.PromptTokens == 0 && usage.CompletionTokens != 0 {
|
||||
usage.PromptTokens = info.GetEstimatePromptTokens()
|
||||
if !gjson.GetBytes(payload, path).Exists() {
|
||||
return payload
|
||||
}
|
||||
|
||||
usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens
|
||||
if usage.BillingUsage != nil {
|
||||
usage.BillingUsage = dto.CloneBillingUsageWithEstimatedCompletion(usage.BillingUsage, usage.CompletionTokens)
|
||||
patched, err := sjson.SetBytes(payload, path, int(createdAt))
|
||||
if err != nil {
|
||||
return payload
|
||||
}
|
||||
|
||||
return usage, nil
|
||||
return patched
|
||||
}
|
||||
@@ -18,6 +18,44 @@ import (
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestOaiResponsesHandlerBuffersUnexpectedSSEAndPreservesUsage(t *testing.T) {
|
||||
terminal := `{"id":"resp_test","object":"response","model":"deepseek-v4.1-flash","status":"completed","extra_provider_field":true,"output":[{"type":"message","role":"assistant","content":[{"type":"output_text","text":"OK"}]}],"usage":{"input_tokens":86926,"output_tokens":163,"total_tokens":87089,"input_tokens_details":{"cached_tokens":86784}}}`
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
body string
|
||||
wantError bool
|
||||
}{
|
||||
{"completed", ": ping\n\nevent: response.completed\ndata: {\"type\":\"response.completed\",\"response\":" + terminal + "}\n\ndata: [DONE]\n\n", false},
|
||||
{"incomplete", "data: {\"type\":\"response.incomplete\",\"response\":" + strings.Replace(terminal, `"completed"`, `"incomplete"`, 1) + "}\n\n", false},
|
||||
{"truncated", "data: {\"type\":\"response.created\",\"response\":{\"model\":\"deepseek-v4.1-flash\"}}\n\ndata: [DONE]\n\n", true},
|
||||
{"failed", "data: {\"type\":\"response.failed\",\"response\":{\"status\":\"failed\"}}\n\n", true},
|
||||
{"invalid", "data: not-json\n\n", true},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||||
info := &relaycommon.RelayInfo{OriginModelName: "deepseek-v4.1-flash", ChannelMeta: &relaycommon.ChannelMeta{}}
|
||||
resp := &http.Response{StatusCode: 200, Header: http.Header{"Content-Type": {"text/event-stream; charset=utf-8"}}, Body: io.NopCloser(strings.NewReader(tc.body))}
|
||||
usage, apiErr := OaiResponsesHandler(c, info, resp)
|
||||
if tc.wantError {
|
||||
require.NotNil(t, apiErr)
|
||||
assert.Nil(t, usage)
|
||||
assert.False(t, c.Writer.Written())
|
||||
return
|
||||
}
|
||||
require.Nil(t, apiErr)
|
||||
require.NotNil(t, usage)
|
||||
assert.Equal(t, 86926, usage.PromptTokens)
|
||||
assert.Equal(t, 163, usage.CompletionTokens)
|
||||
assert.Equal(t, 86784, usage.PromptTokensDetails.CachedTokens)
|
||||
assert.Equal(t, "application/json", w.Header().Get("Content-Type"))
|
||||
assert.Contains(t, w.Body.String(), `"extra_provider_field":true`)
|
||||
assert.Contains(t, w.Body.String(), `"text":"OK"`)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestOaiResponsesHandlerCountsOutputCallsNotDeclarations(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
operation_setting.SetToolPriceForTest("priced_fn", 5.0)
|
||||
@@ -265,3 +303,92 @@ func TestOaiResponsesStreamHandlerDoesNotCountPartialImageEvent(t *testing.T) {
|
||||
|
||||
assert.Equal(t, 0, info.ResponsesUsageInfo.BuiltInTools[dto.BuildInToolImageGeneration].CallCount)
|
||||
}
|
||||
|
||||
func TestOaiResponsesHandlerRewritesSGLangCreatedAtToInt(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||||
|
||||
info := &relaycommon.RelayInfo{
|
||||
ChannelMeta: &relaycommon.ChannelMeta{ChannelType: constant.ChannelTypeSGLang},
|
||||
}
|
||||
resp := &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Body: io.NopCloser(strings.NewReader(`{"id":"resp_1","object":"response","created_at":1786588600.0,"status":"completed","output":[]}`)),
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
}
|
||||
|
||||
_, apiErr := OaiResponsesHandler(c, info, resp)
|
||||
require.Nil(t, apiErr)
|
||||
assert.Contains(t, w.Body.String(), `"created_at":1786588600`)
|
||||
assert.NotContains(t, w.Body.String(), `1786588600.0`)
|
||||
}
|
||||
|
||||
func TestOaiResponsesStreamHandlerRewritesSGLangCreatedAtToInt(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
oldTimeout := constant.StreamingTimeout
|
||||
constant.StreamingTimeout = 30
|
||||
t.Cleanup(func() {
|
||||
constant.StreamingTimeout = oldTimeout
|
||||
})
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||||
c.Set(common.RequestIdKey, "sglang-created-at-test")
|
||||
|
||||
info := &relaycommon.RelayInfo{
|
||||
DisablePing: true,
|
||||
ChannelMeta: &relaycommon.ChannelMeta{
|
||||
ChannelType: constant.ChannelTypeSGLang,
|
||||
UpstreamModelName: "served-model",
|
||||
},
|
||||
}
|
||||
resp := &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Body: io.NopCloser(strings.NewReader(
|
||||
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_1\",\"created_at\":1.7865886E9,\"status\":\"completed\",\"output\":[]}}\n\n" +
|
||||
"data: [DONE]\n\n",
|
||||
)),
|
||||
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
|
||||
}
|
||||
|
||||
_, apiErr := OaiResponsesStreamHandler(c, info, resp)
|
||||
require.Nil(t, apiErr)
|
||||
assert.Contains(t, w.Body.String(), `"created_at":1786588600`)
|
||||
assert.NotContains(t, w.Body.String(), `1.7865886E9`)
|
||||
assert.NotContains(t, w.Body.String(), `1786588600.0`)
|
||||
}
|
||||
|
||||
func TestOaiResponsesStreamHandlerKeepsNonSGLangCreatedAt(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
oldTimeout := constant.StreamingTimeout
|
||||
constant.StreamingTimeout = 30
|
||||
t.Cleanup(func() {
|
||||
constant.StreamingTimeout = oldTimeout
|
||||
})
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||||
c.Set(common.RequestIdKey, "openai-created-at-test")
|
||||
|
||||
info := &relaycommon.RelayInfo{
|
||||
DisablePing: true,
|
||||
ChannelMeta: &relaycommon.ChannelMeta{ChannelType: constant.ChannelTypeOpenAI, UpstreamModelName: "gpt-test"},
|
||||
}
|
||||
resp := &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Body: io.NopCloser(strings.NewReader(
|
||||
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_1\",\"created_at\":1786588600.0,\"status\":\"completed\",\"output\":[]}}\n\n" +
|
||||
"data: [DONE]\n\n",
|
||||
)),
|
||||
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
|
||||
}
|
||||
|
||||
_, apiErr := OaiResponsesStreamHandler(c, info, resp)
|
||||
require.Nil(t, apiErr)
|
||||
assert.Contains(t, w.Body.String(), `1786588600.0`)
|
||||
}
|
||||
@@ -35,6 +35,7 @@ func OaiChatToResponsesHandler(c *gin.Context, info *relaycommon.RelayInfo, resp
|
||||
return nil, types.WithOpenAIError(*oaiError, resp.StatusCode)
|
||||
}
|
||||
|
||||
info.ObserveResponseModel(chatResp.Model)
|
||||
if responseID := helper.GetResponseID(c); responseID != "" {
|
||||
chatResp.Id = responseID
|
||||
}
|
||||
@@ -140,6 +141,7 @@ func OaiChatToResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo
|
||||
return
|
||||
}
|
||||
|
||||
info.ObserveResponseModel(chunk.Model)
|
||||
results, err := service.ConvertStreamResponseChunk(c, info, state, &chunk)
|
||||
if err != nil {
|
||||
if failResponsesStream(err) {
|
||||
|
||||
@@ -3,11 +3,13 @@ package openai
|
||||
import (
|
||||
"testing"
|
||||
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestCollectStreamFunctionCallNamesDedupesSameIndex(t *testing.T) {
|
||||
func TestObserveStreamChoicesDedupesFunctionCallNames(t *testing.T) {
|
||||
info := &relaycommon.RelayInfo{StreamStatus: relaycommon.NewStreamStatus()}
|
||||
seen := make(map[string]struct{})
|
||||
var names []string
|
||||
|
||||
@@ -19,9 +21,18 @@ func TestCollectStreamFunctionCallNamesDedupesSameIndex(t *testing.T) {
|
||||
`{"choices":[{"index":0,"delta":{"tool_calls":[{"index":1,"function":{"arguments":"{}"}}]}}]}`,
|
||||
}
|
||||
for _, chunk := range chunks {
|
||||
collectStreamFunctionCallNames(chunk, seen, &names)
|
||||
observeStreamChoices(info, chunk, seen, &names)
|
||||
}
|
||||
|
||||
require.Len(t, names, 2)
|
||||
assert.Equal(t, []string{"get_weather", "get_time"}, names)
|
||||
assert.Empty(t, info.StreamStatus.ResponseOutcome(), "tool call deltas carry no finish reason")
|
||||
|
||||
observeStreamChoices(info, `{"choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}]}`, seen, &names)
|
||||
assert.Equal(t, "completed", info.StreamStatus.ResponseOutcome())
|
||||
assert.False(t, info.PerformanceBusinessRejection)
|
||||
|
||||
filtered := &relaycommon.RelayInfo{StreamStatus: relaycommon.NewStreamStatus()}
|
||||
observeStreamChoices(filtered, `{"choices":[{"index":0,"delta":{},"finish_reason":"content_filter"}]}`, map[string]struct{}{}, &names)
|
||||
assert.True(t, filtered.PerformanceBusinessRejection)
|
||||
}
|
||||
@@ -2,6 +2,7 @@ package jsplugin
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"cmp"
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"fmt"
|
||||
@@ -18,6 +19,7 @@ import (
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
"unicode/utf8"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
@@ -33,6 +35,7 @@ import (
|
||||
)
|
||||
|
||||
type requestDescriptor struct {
|
||||
ResponseType string `json:"responseType"`
|
||||
URL string `json:"url"`
|
||||
Method string `json:"method"`
|
||||
Headers map[string]string `json:"headers"`
|
||||
@@ -120,14 +123,24 @@ func (a *TaskAdaptor) ValidateRequestAndSetAction(c *gin.Context, info *relaycom
|
||||
return taskErr
|
||||
}
|
||||
}
|
||||
if request, exists := c.Get("task_request"); exists {
|
||||
if err := a.validateResolvedUsageRequest(request); err != nil {
|
||||
request, hasRequest := c.Get("task_request")
|
||||
hasUsageProfiles := len(a.plugin.Meta.UsageProfiles) > 0
|
||||
if hasRequest && !hasUsageProfiles {
|
||||
if err := a.validateResolvedUsageRequest(request, ""); err != nil {
|
||||
return service.TaskErrorWrapperLocal(err, "plugin_usage_invalid", http.StatusBadRequest)
|
||||
}
|
||||
}
|
||||
if _, err := a.buildSubmit(c, info); err != nil {
|
||||
return service.TaskErrorWrapperLocal(err, "plugin_request_invalid", http.StatusBadRequest)
|
||||
}
|
||||
// The descriptor may rewrite the model. Validate profiled requests against
|
||||
// that final model before any quota calculation or upstream submission.
|
||||
if hasRequest && hasUsageProfiles {
|
||||
usageModel := cmp.Or(info.UpstreamModelName, info.OriginModelName)
|
||||
if err := a.validateResolvedUsageRequest(request, usageModel); err != nil {
|
||||
return service.TaskErrorWrapperLocal(err, "plugin_usage_invalid", http.StatusBadRequest)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -143,7 +156,7 @@ func (a *TaskAdaptor) EstimateBilling(c *gin.Context, info *relaycommon.RelayInf
|
||||
func (a *TaskAdaptor) EstimateBillingValidated(c *gin.Context, info *relaycommon.RelayInfo) (map[string]float64, error) {
|
||||
usageContext := a.submitContext(c, info)
|
||||
usageContext["usagePurpose"] = "billing_ratios"
|
||||
return a.usageRatios(c.Request.Context(), "extractUsage", usageContext)
|
||||
return a.usageRatios(c.Request.Context(), cmp.Or(info.UpstreamModelName, info.OriginModelName), "extractUsage", usageContext)
|
||||
}
|
||||
|
||||
func (a *TaskAdaptor) ExtractUsageFacts(c *gin.Context, info *relaycommon.RelayInfo) map[string]any {
|
||||
@@ -172,7 +185,7 @@ func (a *TaskAdaptor) ExtractUsageFactsValidated(c *gin.Context, info *relaycomm
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("plugin usage hook must return an object")
|
||||
}
|
||||
if _, err = a.validatedUsageRatios(facts); err != nil {
|
||||
if _, err = a.validatedUsageRatios(facts, cmp.Or(info.UpstreamModelName, info.OriginModelName)); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return facts, nil
|
||||
@@ -183,7 +196,7 @@ func (a *TaskAdaptor) AdjustBillingOnSubmit(info *relaycommon.RelayInfo, taskDat
|
||||
if err := common.Unmarshal(taskData, &data); err != nil {
|
||||
data = string(taskData)
|
||||
}
|
||||
ratios, err := a.usageRatios(context.Background(), "extractUsageOnSubmit", a.submitContext(nil, info), data)
|
||||
ratios, err := a.usageRatios(context.Background(), cmp.Or(info.UpstreamModelName, info.OriginModelName), "extractUsageOnSubmit", a.submitContext(nil, info), data)
|
||||
if err != nil {
|
||||
a.logRejectedUsage("extractUsageOnSubmit", err)
|
||||
return nil
|
||||
@@ -199,7 +212,11 @@ func (a *TaskAdaptor) AdjustBillingOnComplete(task *model.Task, result *relaycom
|
||||
if err != nil {
|
||||
return 0
|
||||
}
|
||||
a.applyCompletionUsageFacts(result, value)
|
||||
usageModel := ""
|
||||
if task != nil {
|
||||
usageModel = cmp.Or(task.Properties.UpstreamModelName, task.Properties.OriginModelName)
|
||||
}
|
||||
a.applyCompletionUsageFacts(result, value, usageModel)
|
||||
return 0
|
||||
}
|
||||
|
||||
@@ -423,6 +440,9 @@ func encodeFilePlaceholder(placeholder map[string]any, form *multipart.Form, lim
|
||||
}
|
||||
|
||||
func (a *TaskAdaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, body io.Reader) (*http.Response, error) {
|
||||
wasStream := info.IsStream
|
||||
info.IsStream = false
|
||||
defer func() { info.IsStream = wasStream }()
|
||||
if a.submit != nil && strings.TrimSpace(a.submit.Method) != "" {
|
||||
originalMethod := c.Request.Method
|
||||
c.Request.Method = strings.ToUpper(a.submit.Method)
|
||||
@@ -431,24 +451,42 @@ func (a *TaskAdaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, bod
|
||||
return channel.DoTaskApiRequest(a, c, info, body)
|
||||
}
|
||||
|
||||
func (a *TaskAdaptor) ParseResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (*channel.TaskSubmitResponse, *dto.TaskError) {
|
||||
func (a *TaskAdaptor) ParseResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (_ *channel.TaskSubmitResponse, taskErr *dto.TaskError) {
|
||||
started := time.Now()
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
logger.LogDebug(c, "task_plugin subsystem=adaptor event=parse_submit_failed plugin=%q stage=read_response reason=read_failed status=%d", a.plugin.Meta.Key, resp.StatusCode)
|
||||
return nil, service.TaskErrorWrapper(err, "read_response_body_failed", http.StatusInternalServerError)
|
||||
var responseBody any
|
||||
var err error
|
||||
streaming := a.submit != nil && a.submit.ResponseType == "sse"
|
||||
mediaType, _, _ := mime.ParseMediaType(resp.Header.Get("Content-Type"))
|
||||
acceptedStream := streaming || mediaType == "text/event-stream"
|
||||
defer func() {
|
||||
if acceptedStream && taskErr != nil {
|
||||
taskErr.NoRetry = true
|
||||
}
|
||||
}()
|
||||
if !streaming && acceptedStream {
|
||||
return nil, service.TaskErrorWrapperLocal(fmt.Errorf("unexpected SSE response for a JSON submission"), "plugin_submit_response_invalid", http.StatusBadGateway)
|
||||
}
|
||||
logger.LogDebug(
|
||||
c,
|
||||
"task_plugin subsystem=adaptor event=submit_response_received plugin=%q status=%d body_bytes=%d",
|
||||
a.plugin.Meta.Key,
|
||||
resp.StatusCode,
|
||||
len(body),
|
||||
)
|
||||
responseBody := any(string(body))
|
||||
var decoded any
|
||||
if common.Unmarshal(body, &decoded) == nil {
|
||||
responseBody = decoded
|
||||
if streaming {
|
||||
// Every failure after accepting this stream is non-retryable: the
|
||||
// upstream may already have performed billable work.
|
||||
defer resp.Body.Close()
|
||||
responseBody, err = a.readSubmitEvents(c.Request.Context(), resp, a.submitContext(c, info))
|
||||
} else {
|
||||
var body []byte
|
||||
body, err = io.ReadAll(io.LimitReader(resp.Body, maxTaskPluginPersistedJSONBytes+1))
|
||||
if err == nil && len(body) > maxTaskPluginPersistedJSONBytes {
|
||||
err = fmt.Errorf("task submit response exceeds size limit")
|
||||
}
|
||||
responseBody = string(body)
|
||||
var decoded any
|
||||
if err == nil && common.Unmarshal(body, &decoded) == nil {
|
||||
responseBody = decoded
|
||||
}
|
||||
}
|
||||
if err != nil {
|
||||
failure := service.TaskErrorWrapper(err, "read_response_body_failed", http.StatusBadGateway)
|
||||
failure.NoRetry = streaming
|
||||
return nil, failure
|
||||
}
|
||||
headers := make(map[string][]string, len(resp.Header))
|
||||
maps.Copy(headers, resp.Header)
|
||||
@@ -485,34 +523,57 @@ func (a *TaskAdaptor) ParseResponse(c *gin.Context, resp *http.Response, info *r
|
||||
var taskData []byte
|
||||
if parsed.TaskData != nil {
|
||||
taskData, err = common.Marshal(parsed.TaskData)
|
||||
if err != nil {
|
||||
if err != nil || len(taskData) > maxTaskPluginPersistedJSONBytes {
|
||||
if err == nil {
|
||||
err = fmt.Errorf("task data exceeds size limit")
|
||||
}
|
||||
return nil, service.TaskErrorWrapper(err, "plugin_submit_response_invalid", http.StatusBadGateway)
|
||||
}
|
||||
}
|
||||
var immediate *relaycommon.TaskInfo
|
||||
if parsed.Immediate != nil {
|
||||
immediate = &relaycommon.TaskInfo{Code: parsed.Immediate.Code, TaskID: parsed.Immediate.TaskID, Status: parsed.Immediate.Status, Progress: parsed.Immediate.Progress, Reason: parsed.Immediate.Reason, Url: parsed.Immediate.URL, RemoteUrl: parsed.Immediate.RemoteURL}
|
||||
}
|
||||
logger.LogDebug(
|
||||
c,
|
||||
"task_plugin subsystem=adaptor event=parse_submit_complete plugin=%q status=%d task_data_bytes=%d immediate=%t elapsed_ms=%d",
|
||||
a.plugin.Meta.Key,
|
||||
resp.StatusCode,
|
||||
len(taskData),
|
||||
immediate != nil,
|
||||
time.Since(started).Milliseconds(),
|
||||
)
|
||||
pluginState, _ := encodeReturnedPluginState(value)
|
||||
if len(pluginState) > maxTaskPluginPersistedJSONBytes {
|
||||
logger.LogWarn(c, fmt.Sprintf("task plugin %s rejected oversized submit state (%d bytes)", a.plugin.Meta.Key, len(pluginState)))
|
||||
pluginState = nil
|
||||
}
|
||||
return &channel.TaskSubmitResponse{
|
||||
response := &channel.TaskSubmitResponse{
|
||||
UpstreamTaskID: parsed.TaskID,
|
||||
TaskData: taskData,
|
||||
Immediate: immediate,
|
||||
PluginState: pluginState,
|
||||
}, nil
|
||||
}
|
||||
logger.LogDebug(c, "task_plugin subsystem=adaptor event=parse_submit_complete plugin=%q status=%d task_data_bytes=%d immediate=%t elapsed_ms=%d",
|
||||
a.plugin.Meta.Key, resp.StatusCode, len(taskData), parsed.Immediate != nil, time.Since(started).Milliseconds())
|
||||
if parsed.Immediate == nil {
|
||||
return response, nil
|
||||
}
|
||||
immediate := &relaycommon.TaskInfo{
|
||||
Code: parsed.Immediate.Code, TaskID: parsed.Immediate.TaskID,
|
||||
Status: parsed.Immediate.Status, Progress: parsed.Immediate.Progress,
|
||||
Reason: parsed.Immediate.Reason, Url: parsed.Immediate.URL, RemoteUrl: parsed.Immediate.RemoteURL,
|
||||
CompletionTokens: positiveInt(parsed.Immediate.CompletionTokens), TotalTokens: positiveInt(parsed.Immediate.TotalTokens),
|
||||
}
|
||||
response.Immediate = immediate
|
||||
if immediate.Status != model.TaskStatusSuccess || !a.hasHook(c.Request.Context(), "extractUsageOnComplete") {
|
||||
return response, nil
|
||||
}
|
||||
// Match the polling hook context without persisting a temporary task.
|
||||
task := &model.Task{TaskID: info.PublicTaskID, Action: info.Action, Data: taskData,
|
||||
Properties: model.Properties{OriginModelName: info.OriginModelName, UpstreamModelName: info.UpstreamModelName}}
|
||||
task.PrivateData.UpstreamTaskID = parsed.TaskID
|
||||
task.PrivateData.PluginState = pluginState
|
||||
ctx, err := a.queryContext(task, info.ApiKey, info.ChannelBaseUrl, info.ChannelSetting.Proxy)
|
||||
if err != nil {
|
||||
logger.LogWarn(c, fmt.Sprintf("task plugin %s completion context failed; retaining reserved quota: %v", a.plugin.Meta.Key, err))
|
||||
return response, nil
|
||||
}
|
||||
facts, err := a.plugin.Engine.Call(c.Request.Context(), "extractUsageOnComplete", ctx, jsonValue(immediate), parsed.TaskData)
|
||||
if err != nil {
|
||||
logger.LogWarn(c, fmt.Sprintf("task plugin %s completion usage failed; retaining reserved quota: %v", a.plugin.Meta.Key, err))
|
||||
return response, nil
|
||||
}
|
||||
if err := a.applyCompletionUsageFacts(immediate, facts, cmp.Or(info.UpstreamModelName, info.OriginModelName)); err != nil {
|
||||
logger.LogWarn(c, fmt.Sprintf("task plugin %s completion usage rejected; retaining reserved quota: %v", a.plugin.Meta.Key, err))
|
||||
}
|
||||
return response, nil
|
||||
}
|
||||
|
||||
func (a *TaskAdaptor) GetModelList() []string { return append([]string(nil), a.plugin.Meta.Models...) }
|
||||
@@ -680,7 +741,8 @@ func (a *TaskAdaptor) ParseBatchResult(tasks []*model.Task, resp *http.Response,
|
||||
}
|
||||
facts, hookErr := a.plugin.Engine.Call(context.Background(), "extractUsageOnComplete", itemCtx, jsonValue(&info), usageBody)
|
||||
if hookErr == nil {
|
||||
a.applyCompletionUsageFacts(&info, facts)
|
||||
usageModel, _ := itemCtx["upstreamModel"].(string)
|
||||
a.applyCompletionUsageFacts(&info, facts, usageModel)
|
||||
}
|
||||
}
|
||||
results[item.TaskID] = &service.BatchTaskResult{TaskInfo: info, Action: item.Action, SubmitTime: item.SubmitTime, StartTime: item.StartTime, FinishTime: item.FinishTime, Data: item.Data}
|
||||
@@ -745,7 +807,8 @@ func (a *TaskAdaptor) ParseTaskResult(task *model.Task, resp *http.Response, bod
|
||||
if a.hasHook(context.Background(), "extractUsageOnComplete") {
|
||||
facts, hookErr := a.plugin.Engine.Call(context.Background(), "extractUsageOnComplete", ctx, jsonValue(result), input)
|
||||
if hookErr == nil {
|
||||
a.applyCompletionUsageFacts(result, facts)
|
||||
usageModel, _ := ctx["upstreamModel"].(string)
|
||||
a.applyCompletionUsageFacts(result, facts, usageModel)
|
||||
}
|
||||
}
|
||||
taskStatus := model.TaskStatus(result.Status)
|
||||
@@ -760,20 +823,20 @@ func (a *TaskAdaptor) ParseTaskResult(task *model.Task, resp *http.Response, bod
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (a *TaskAdaptor) applyCompletionUsageFacts(result *relaycommon.TaskInfo, facts any) {
|
||||
values, err := a.validatedCompletionUsageFacts(facts)
|
||||
func (a *TaskAdaptor) applyCompletionUsageFacts(result *relaycommon.TaskInfo, facts any, modelName string) error {
|
||||
values, err := a.validatedCompletionUsageFacts(facts, modelName)
|
||||
if err != nil {
|
||||
a.logRejectedUsage("extractUsageOnComplete", err)
|
||||
return
|
||||
return err
|
||||
}
|
||||
if len(values) == 0 {
|
||||
return
|
||||
return nil
|
||||
}
|
||||
result.UsageFacts = values
|
||||
if units := positiveInt(values["upstreamUnits"]); units > 0 {
|
||||
result.CompletionTokens = units
|
||||
result.TotalTokens = units
|
||||
return
|
||||
return nil
|
||||
}
|
||||
if completionTokens, exists := values["completionTokens"]; exists {
|
||||
result.CompletionTokens = positiveInt(completionTokens)
|
||||
@@ -781,6 +844,7 @@ func (a *TaskAdaptor) applyCompletionUsageFacts(result *relaycommon.TaskInfo, fa
|
||||
if totalTokens, exists := values["totalTokens"]; exists {
|
||||
result.TotalTokens = positiveInt(totalTokens)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (a *TaskAdaptor) ConvertToOpenAIVideo(task *model.Task) ([]byte, error) {
|
||||
@@ -805,26 +869,24 @@ func (a *TaskAdaptor) ConvertToOpenAIVideo(task *model.Task) ([]byte, error) {
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
rendered := kitdto.NewOpenAIVideo()
|
||||
if err = common.Unmarshal(encoded, rendered); err != nil {
|
||||
var rendered map[string]any
|
||||
if err = common.Unmarshal(encoded, &rendered); err != nil || rendered == nil {
|
||||
return nil, fmt.Errorf("plugin returned an invalid OpenAI video object")
|
||||
}
|
||||
// Keep provider extensions intact while the host owns the public task's
|
||||
// identity and lifecycle, including completion timestamps after settlement.
|
||||
host := task.ToOpenAIVideo()
|
||||
rendered.ID = host.ID
|
||||
rendered.Object = host.Object
|
||||
rendered.TaskID = ""
|
||||
rendered.Status = host.Status
|
||||
rendered.Progress = host.Progress
|
||||
rendered.CreatedAt = host.CreatedAt
|
||||
rendered.Model = host.Model
|
||||
rendered.CompletedAt = host.CompletedAt
|
||||
for key := range rendered.Metadata {
|
||||
if strings.EqualFold(key, "url") {
|
||||
delete(rendered.Metadata, key)
|
||||
}
|
||||
}
|
||||
if len(rendered.Metadata) == 0 {
|
||||
rendered.Metadata = nil
|
||||
rendered["id"] = host.ID
|
||||
rendered["object"] = host.Object
|
||||
delete(rendered, "task_id")
|
||||
rendered["status"] = host.Status
|
||||
rendered["progress"] = host.Progress
|
||||
rendered["created_at"] = host.CreatedAt
|
||||
rendered["model"] = host.Model
|
||||
if host.CompletedAt != 0 {
|
||||
rendered["completed_at"] = host.CompletedAt
|
||||
} else {
|
||||
delete(rendered, "completed_at")
|
||||
}
|
||||
return common.Marshal(rendered)
|
||||
}
|
||||
@@ -1139,6 +1201,16 @@ func (a *TaskAdaptor) buildSubmit(c *gin.Context, info *relaycommon.RelayInfo) (
|
||||
)
|
||||
return nil, err
|
||||
}
|
||||
if descriptor.ResponseType == "" {
|
||||
descriptor.ResponseType = "json"
|
||||
}
|
||||
allowed := a.plugin.Meta.SubmitResponseTypes
|
||||
if len(allowed) == 0 {
|
||||
allowed = []string{"json"}
|
||||
}
|
||||
if !slices.Contains(allowed, descriptor.ResponseType) {
|
||||
return nil, fmt.Errorf("plugin does not support submit response type %q", descriptor.ResponseType)
|
||||
}
|
||||
if strings.TrimSpace(descriptor.URL) == "" {
|
||||
logger.LogDebug(c, "task_plugin subsystem=adaptor event=build_submit_failed plugin=%q stage=validate_url reason=empty_url", a.plugin.Meta.Key)
|
||||
return nil, fmt.Errorf("plugin returned an empty submit URL")
|
||||
@@ -1294,7 +1366,7 @@ func (a *TaskAdaptor) submitContext(c *gin.Context, info *relaycommon.RelayInfo)
|
||||
return ctx
|
||||
}
|
||||
|
||||
func (a *TaskAdaptor) usageRatios(ctx context.Context, hook string, args ...any) (map[string]float64, error) {
|
||||
func (a *TaskAdaptor) usageRatios(ctx context.Context, modelName, hook string, args ...any) (map[string]float64, error) {
|
||||
if !a.hasHook(ctx, hook) {
|
||||
return nil, nil
|
||||
}
|
||||
@@ -1312,7 +1384,7 @@ func (a *TaskAdaptor) usageRatios(ctx context.Context, hook string, args ...any)
|
||||
logger.LogDebug(ctx, "task_plugin subsystem=adaptor event=usage_hook_failed plugin=%q hook=%q reason=result_not_object elapsed_ms=%d", a.plugin.Meta.Key, hook, time.Since(started).Milliseconds())
|
||||
return nil, fmt.Errorf("plugin usage hook must return an object")
|
||||
}
|
||||
ratios, err := a.validatedUsageRatios(facts)
|
||||
ratios, err := a.validatedUsageRatios(facts, modelName)
|
||||
if err != nil {
|
||||
logger.LogDebug(ctx, "task_plugin subsystem=adaptor event=usage_hook_failed plugin=%q hook=%q reason=invalid_usage elapsed_ms=%d", a.plugin.Meta.Key, hook, time.Since(started).Milliseconds())
|
||||
return nil, err
|
||||
@@ -1329,15 +1401,16 @@ func (a *TaskAdaptor) usageRatios(ctx context.Context, hook string, args ...any)
|
||||
return ratios, nil
|
||||
}
|
||||
|
||||
func (a *TaskAdaptor) validateResolvedUsageRequest(request any) error {
|
||||
return a.validateResolvedUsageValue(jsonValue(request))
|
||||
func (a *TaskAdaptor) validateResolvedUsageRequest(request any, modelName string) error {
|
||||
schema, _ := a.plugin.Meta.UsageForModel(modelName)
|
||||
return a.validateResolvedUsageValue(jsonValue(request), schema)
|
||||
}
|
||||
|
||||
func (a *TaskAdaptor) validateResolvedUsageValue(value any) error {
|
||||
func (a *TaskAdaptor) validateResolvedUsageValue(value any, usageSchema map[string]pluginruntime.UsageFieldSchema) error {
|
||||
switch typed := value.(type) {
|
||||
case map[string]any:
|
||||
for key, item := range typed {
|
||||
if schema, declared := a.plugin.Meta.UsageSchema[key]; declared {
|
||||
if schema, declared := usageSchema[key]; declared {
|
||||
if _, err := validateUsageValue(item, schema, true); err != nil {
|
||||
return err
|
||||
}
|
||||
@@ -1346,13 +1419,13 @@ func (a *TaskAdaptor) validateResolvedUsageValue(value any) error {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if err := a.validateResolvedUsageValue(item); err != nil {
|
||||
if err := a.validateResolvedUsageValue(item, usageSchema); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
case []any:
|
||||
for _, item := range typed {
|
||||
if err := a.validateResolvedUsageValue(item); err != nil {
|
||||
if err := a.validateResolvedUsageValue(item, usageSchema); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
@@ -1360,10 +1433,11 @@ func (a *TaskAdaptor) validateResolvedUsageValue(value any) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (a *TaskAdaptor) validatedUsageRatios(facts map[string]any) (map[string]float64, error) {
|
||||
func (a *TaskAdaptor) validatedUsageRatios(facts map[string]any, modelName string) (map[string]float64, error) {
|
||||
usageSchema, _ := a.plugin.Meta.UsageForModel(modelName)
|
||||
ratios := make(map[string]float64)
|
||||
for key, value := range facts {
|
||||
if schema, declared := a.plugin.Meta.UsageSchema[key]; declared {
|
||||
if schema, declared := usageSchema[key]; declared {
|
||||
number, err := validateUsageValue(value, schema, false)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
@@ -1390,6 +1464,9 @@ func (a *TaskAdaptor) validatedUsageRatios(facts map[string]any) (map[string]flo
|
||||
if err := validateUsageNumberLimit(number, limit); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
// Keep numeric facts usable by existing expressions even when the
|
||||
// selected profile no longer declares that legacy extension field.
|
||||
facts[key] = number
|
||||
if number > 0 {
|
||||
ratios[key] = number
|
||||
}
|
||||
@@ -1397,7 +1474,8 @@ func (a *TaskAdaptor) validatedUsageRatios(facts map[string]any) (map[string]flo
|
||||
return ratios, nil
|
||||
}
|
||||
|
||||
func (a *TaskAdaptor) validatedCompletionUsageFacts(facts any) (map[string]any, error) {
|
||||
func (a *TaskAdaptor) validatedCompletionUsageFacts(facts any, modelName string) (map[string]any, error) {
|
||||
usageSchema, _ := a.plugin.Meta.UsageForModel(modelName)
|
||||
if facts == nil {
|
||||
return nil, nil
|
||||
}
|
||||
@@ -1408,7 +1486,7 @@ func (a *TaskAdaptor) validatedCompletionUsageFacts(facts any) (map[string]any,
|
||||
validated := make(map[string]any, len(values))
|
||||
for key, value := range values {
|
||||
validated[key] = value
|
||||
if schema, declared := a.plugin.Meta.UsageSchema[key]; declared {
|
||||
if schema, declared := usageSchema[key]; declared {
|
||||
number, err := validateUsageValue(value, schema, false)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
@@ -1436,6 +1514,10 @@ func (a *TaskAdaptor) validatedCompletionUsageFacts(facts any) (map[string]any,
|
||||
return nil, fmt.Errorf("plugin usage value must be a finite non-negative number")
|
||||
}
|
||||
validated[key] = float64(common.QuotaFromFloat(number))
|
||||
default:
|
||||
if number, numeric := usageNumber(value, false); numeric {
|
||||
validated[key] = number
|
||||
}
|
||||
}
|
||||
}
|
||||
return validated, nil
|
||||
@@ -1549,6 +1631,9 @@ func convert(value any, target any) error {
|
||||
}
|
||||
|
||||
func jsonValue(value any) any {
|
||||
if normalized, ok := cloneJSONValue(value, 0); ok {
|
||||
return normalized
|
||||
}
|
||||
data, err := common.Marshal(value)
|
||||
if err != nil {
|
||||
return value
|
||||
@@ -1559,6 +1644,58 @@ func jsonValue(value any) any {
|
||||
}
|
||||
return normalized
|
||||
}
|
||||
|
||||
// Already-parsed JSON and Sobek's plain exported values do not need another
|
||||
// encode/decode pass. Copy containers to isolate hooks, preserving the codec's
|
||||
// float64 numbers and nulls. Other Go types still use the configured codec.
|
||||
func cloneJSONValue(value any, depth int) (any, bool) {
|
||||
if depth > 64 {
|
||||
return nil, false
|
||||
}
|
||||
switch typed := value.(type) {
|
||||
case nil, bool:
|
||||
return typed, true
|
||||
case string:
|
||||
return typed, utf8.ValidString(typed)
|
||||
case float64:
|
||||
return typed, !math.IsNaN(typed) && !math.IsInf(typed, 0)
|
||||
case int64:
|
||||
return float64(typed), true
|
||||
case int:
|
||||
return float64(typed), true
|
||||
case map[string]any:
|
||||
if typed == nil {
|
||||
return nil, true
|
||||
}
|
||||
cloned := make(map[string]any, len(typed))
|
||||
for key, item := range typed {
|
||||
if !utf8.ValidString(key) {
|
||||
return nil, false
|
||||
}
|
||||
child, ok := cloneJSONValue(item, depth+1)
|
||||
if !ok {
|
||||
return nil, false
|
||||
}
|
||||
cloned[key] = child
|
||||
}
|
||||
return cloned, true
|
||||
case []any:
|
||||
if typed == nil {
|
||||
return nil, true
|
||||
}
|
||||
cloned := make([]any, len(typed))
|
||||
for index, item := range typed {
|
||||
child, ok := cloneJSONValue(item, depth+1)
|
||||
if !ok {
|
||||
return nil, false
|
||||
}
|
||||
cloned[index] = child
|
||||
}
|
||||
return cloned, true
|
||||
default:
|
||||
return nil, false
|
||||
}
|
||||
}
|
||||
func positiveInt(value any) int {
|
||||
switch number := value.(type) {
|
||||
case int64:
|
||||
|
||||
@@ -4,6 +4,7 @@ import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"math"
|
||||
"mime/multipart"
|
||||
@@ -16,7 +17,9 @@ import (
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/pkg/billingexpr"
|
||||
pluginruntime "github.com/QuantumNous/new-api/pkg/jsplugin"
|
||||
"github.com/QuantumNous/new-api/plugins"
|
||||
"github.com/QuantumNous/new-api/relay/channel"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/relay/helper"
|
||||
@@ -532,69 +535,90 @@ func TestTaskAdaptorMapsJSContract(t *testing.T) {
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
func TestTaskAdaptorSanitizesOpenAIVideoRendererOutput(t *testing.T) {
|
||||
source := `
|
||||
export const meta = {
|
||||
apiVersion: 1, key: "safe-video", name: "Safe Video", version: "1.0.0",
|
||||
author: {name: "Test"}, models: ["model"], fetchMode: "per_task", protocols: ["openai_video"],
|
||||
};
|
||||
export function buildSubmitRequest(ctx) { return {url: ctx.baseUrl + "/submit"}; }
|
||||
export function parseSubmitResponse() { return {taskId: "upstream"}; }
|
||||
export function buildQueryRequest(ctx) { return {url: ctx.baseUrl + "/query"}; }
|
||||
export function parseTaskResult() { return {status: "SUCCESS"}; }
|
||||
export function listArtifacts() { return []; }
|
||||
export function buildContentRequest() { throw new Error("artifact_not_found"); }
|
||||
export const protocols = {openai_video: {
|
||||
decodeRequest: function(ctx) { return {kind:"submit", model:ctx.model, requestBody:ctx.body.value}; },
|
||||
render: function() {
|
||||
return {
|
||||
id: "upstream-id",
|
||||
task_id: "upstream-task-id",
|
||||
object: "provider-video",
|
||||
model: "model",
|
||||
status: "completed",
|
||||
progress: 100,
|
||||
created_at: 10,
|
||||
completed_at: 20,
|
||||
metadata: {
|
||||
url: "https://upstream.example/video.mp4",
|
||||
URL: "https://upstream.example/uppercase.mp4",
|
||||
label: "safe",
|
||||
},
|
||||
url: "https://upstream.example/top-level.mp4",
|
||||
upstream_url: "https://upstream.example/unknown.mp4",
|
||||
provider_payload: {task_id: "upstream-task-id"},
|
||||
};
|
||||
}}};
|
||||
`
|
||||
func TestTaskAdaptorPreservesSoraVideoResponseFields(t *testing.T) {
|
||||
source, err := plugins.Source("sora")
|
||||
require.NoError(t, err)
|
||||
plugin, err := pluginruntime.NewRegistry().Register(source, pluginruntime.Options{})
|
||||
require.NoError(t, err)
|
||||
adaptor := New(plugin)
|
||||
|
||||
rendered, err := adaptor.ConvertToOpenAIVideo(&model.Task{
|
||||
TaskID: "task_public",
|
||||
Status: model.TaskStatusInProgress,
|
||||
Properties: model.Properties{OriginModelName: "origin-model"},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
for _, tc := range []struct {
|
||||
status model.TaskStatus
|
||||
want string
|
||||
}{
|
||||
{model.TaskStatusInProgress, "in_progress"},
|
||||
{model.TaskStatusSuccess, "completed"},
|
||||
{model.TaskStatusFailure, "failed"},
|
||||
} {
|
||||
t.Run(string(tc.status), func(t *testing.T) {
|
||||
task := &model.Task{
|
||||
TaskID: "task_public",
|
||||
Status: tc.status,
|
||||
Progress: "42%",
|
||||
CreatedAt: 100,
|
||||
FinishTime: 200,
|
||||
Properties: model.Properties{OriginModelName: "origin-model"},
|
||||
PrivateData: model.TaskPrivateData{UpstreamTaskID: "upstream-task-id"},
|
||||
Data: []byte(`{
|
||||
"id":"upstream-task-id","task_id":"upstream-task-id",
|
||||
"object":"provider-video","model":"provider-model",
|
||||
"status":"completed","progress":100,"created_at":10,"completed_at":20,
|
||||
"url":"https://cdn.example/video.mp4",
|
||||
"metadata":{"url":"https://cdn.example/video.mp4","URL":"https://cdn.example/uppercase.mp4"},
|
||||
"provider_payload":{"task_id":"upstream-task-id","items":[{"enabled":false,"count":0,"value":null}]},
|
||||
"seconds":8,"resolution":"720p","aspect_ratio":"16:9",
|
||||
"reference_images":["https://cdn.example/reference.png"],
|
||||
"error":{"code":"provider_error","message":"provider rejected request","detail":{"retryable":false}}
|
||||
}`),
|
||||
}
|
||||
rendered, err := adaptor.ConvertToOpenAIVideo(task)
|
||||
require.NoError(t, err)
|
||||
|
||||
var video dto.OpenAIVideo
|
||||
require.NoError(t, common.Unmarshal(rendered, &video))
|
||||
assert.Equal(t, "task_public", video.ID)
|
||||
assert.Equal(t, "video", video.Object)
|
||||
assert.Empty(t, video.TaskID)
|
||||
assert.Equal(t, "origin-model", video.Model)
|
||||
assert.Zero(t, video.CompletedAt)
|
||||
assert.Equal(t, map[string]any{"label": "safe"}, video.Metadata)
|
||||
var fields map[string]any
|
||||
require.NoError(t, common.Unmarshal(rendered, &fields))
|
||||
assert.Equal(t, "task_public", fields["id"])
|
||||
assert.NotContains(t, fields, "task_id")
|
||||
assert.Equal(t, "video", fields["object"])
|
||||
assert.Equal(t, "origin-model", fields["model"])
|
||||
assert.Equal(t, tc.want, fields["status"])
|
||||
assert.Equal(t, float64(42), fields["progress"])
|
||||
assert.Equal(t, float64(100), fields["created_at"])
|
||||
if tc.status == model.TaskStatusSuccess {
|
||||
assert.Equal(t, float64(200), fields["completed_at"])
|
||||
} else {
|
||||
assert.NotContains(t, fields, "completed_at")
|
||||
}
|
||||
assert.Equal(t, "https://cdn.example/video.mp4", fields["url"])
|
||||
assert.Equal(t, map[string]any{
|
||||
"url": "https://cdn.example/video.mp4",
|
||||
"URL": "https://cdn.example/uppercase.mp4",
|
||||
}, fields["metadata"])
|
||||
assert.Equal(t, map[string]any{
|
||||
"task_id": "task_public",
|
||||
"items": []any{map[string]any{"enabled": false, "count": float64(0), "value": nil}},
|
||||
}, fields["provider_payload"])
|
||||
assert.Equal(t, float64(8), fields["seconds"])
|
||||
assert.Equal(t, "720p", fields["resolution"])
|
||||
assert.Equal(t, "16:9", fields["aspect_ratio"])
|
||||
assert.Equal(t, []any{"https://cdn.example/reference.png"}, fields["reference_images"])
|
||||
assert.Equal(t, map[string]any{
|
||||
"code": "provider_error", "message": "provider rejected request",
|
||||
"detail": map[string]any{"retryable": false},
|
||||
}, fields["error"])
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
var fields map[string]any
|
||||
require.NoError(t, common.Unmarshal(rendered, &fields))
|
||||
assert.NotContains(t, fields, "url")
|
||||
assert.NotContains(t, fields, "upstream_url")
|
||||
assert.NotContains(t, fields, "provider_payload")
|
||||
assert.NotContains(t, fields, "completed_at")
|
||||
assert.NotContains(t, string(rendered), "upstream.example")
|
||||
assert.NotContains(t, string(rendered), "upstream-task-id")
|
||||
func TestTaskAdaptorRejectsNonObjectOpenAIVideoRendererOutput(t *testing.T) {
|
||||
for _, value := range []string{"null", "[]", `"video"`, "42", "false"} {
|
||||
t.Run(value, func(t *testing.T) {
|
||||
source := strings.Replace(mockPlugin, `return {id: task.task_id, status: "completed"};`, "return "+value+";", 1)
|
||||
plugin, err := pluginruntime.NewRegistry().Register(source, pluginruntime.Options{})
|
||||
require.NoError(t, err)
|
||||
_, err = New(plugin).ConvertToOpenAIVideo(&model.Task{TaskID: "task_public"})
|
||||
require.ErrorContains(t, err, "invalid OpenAI video object")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestTaskAdaptorPreservesOpenAIVideoFailureSlotsAndOwnsLifecycle(t *testing.T) {
|
||||
@@ -1413,3 +1437,407 @@ export function parseBatchResult(){return [];}
|
||||
assert.Equal(t, []any{"model-a", "model-b"}, captured["models"])
|
||||
assert.Equal(t, false, captured["hasRequestBody"])
|
||||
}
|
||||
|
||||
func TestTaskAdaptorUsageProfilesFollowExecutionModel(t *testing.T) {
|
||||
const source = `
|
||||
export const meta = {
|
||||
apiVersion:1, key:"profile-usage", name:"Profile Usage", version:"1.0.0", author:{name:"Test"},
|
||||
models:["image", "video"], fetchMode:"batch",
|
||||
usageSchema:{units:{type:"number",unit:"second"}},
|
||||
usageProfiles:[
|
||||
{models:["image"],schema:{units:{type:"number",unit:"count"},mode:{enum:["image"]}}},
|
||||
{models:["video"],schema:{units:{type:"number",unit:"token"},mode:{enum:["video"]}},
|
||||
examples:[{label:"video",facts:{units:1,mode:"video"}}]}
|
||||
]
|
||||
};
|
||||
export function buildSubmitRequest(ctx) {
|
||||
return {url:ctx.baseUrl+"/submit",rewriteModel:ctx.requestBody.rewriteTo || ""};
|
||||
}
|
||||
export function parseSubmitResponse(){return {taskId:"task"};}
|
||||
export function buildQueryRequest(){return {url:"https://provider.example/task"};}
|
||||
export function buildBatchQueryRequest(){return {url:"https://provider.example/tasks"};}
|
||||
export function parseTaskResult(){return {status:"SUCCESS"};}
|
||||
export function parseBatchResult(ctx,body){return body.items;}
|
||||
export function extractUsage(ctx){return {units:ctx.requestBody.hookUnits,mode:ctx.requestBody.hookMode,legacyRatio:2};}
|
||||
export function extractUsageOnSubmit(ctx,body){return body.usage;}
|
||||
export function extractUsageOnComplete(ctx,result,body){return body.usage;}
|
||||
`
|
||||
plugin, err := pluginruntime.NewRegistry().Register(source, pluginruntime.Options{})
|
||||
require.NoError(t, err)
|
||||
newRequest := func(t *testing.T, upstream string, body map[string]any) (*TaskAdaptor, *gin.Context, *relaycommon.RelayInfo) {
|
||||
t.Helper()
|
||||
info := &relaycommon.RelayInfo{
|
||||
OriginModelName: "public-alias",
|
||||
ChannelMeta: &relaycommon.ChannelMeta{
|
||||
UpstreamModelName: upstream, ChannelBaseUrl: "https://provider.example",
|
||||
},
|
||||
TaskRelayInfo: &relaycommon.TaskRelayInfo{},
|
||||
}
|
||||
adaptor := New(plugin)
|
||||
adaptor.Init(info)
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/native/submit", nil)
|
||||
c.Set("task_request", body)
|
||||
return adaptor, c, info
|
||||
}
|
||||
|
||||
for _, tc := range []struct {
|
||||
name, upstream, rewrite, mode string
|
||||
units float64
|
||||
wantError bool
|
||||
}{
|
||||
{"mapped image", "image", "", "image", 2, false},
|
||||
{"mapped video token units", "video", "", "video", 500000, false},
|
||||
{"rewritten model", "image", "video", "video", 500000, false},
|
||||
{"image count ceiling", "image", "", "image", float64(dto.MaxImageN + 1), true},
|
||||
{"profile enum", "image", "", "video", 2, true},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
adaptor, c, info := newRequest(t, tc.upstream, map[string]any{
|
||||
"hookUnits": tc.units, "hookMode": tc.mode, "rewriteTo": tc.rewrite,
|
||||
})
|
||||
require.Nil(t, adaptor.ValidateRequestAndSetAction(c, info))
|
||||
facts, err := adaptor.ExtractUsageFactsValidated(c, info)
|
||||
if tc.wantError {
|
||||
require.Error(t, err)
|
||||
return
|
||||
}
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, map[string]any{"units": tc.units, "mode": tc.mode, "legacyRatio": 2.0}, facts)
|
||||
ratios, err := adaptor.EstimateBillingValidated(c, info)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, map[string]float64{"units": tc.units, "legacyRatio": 2}, ratios)
|
||||
body, err := common.Marshal(map[string]any{"usage": facts})
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, ratios, adaptor.AdjustBillingOnSubmit(info, body))
|
||||
})
|
||||
}
|
||||
|
||||
t.Run("request revalidated after model rewrite", func(t *testing.T) {
|
||||
adaptor, c, info := newRequest(t, "video", map[string]any{
|
||||
"rewriteTo": "image", "metadata": map[string]any{"units": dto.MaxImageN + 1},
|
||||
})
|
||||
taskErr := adaptor.ValidateRequestAndSetAction(c, info)
|
||||
require.NotNil(t, taskErr)
|
||||
assert.Equal(t, "plugin_usage_invalid", taskErr.Code)
|
||||
})
|
||||
|
||||
t.Run("request enum belongs to the rewritten model", func(t *testing.T) {
|
||||
adaptor, c, info := newRequest(t, "video", map[string]any{
|
||||
"rewriteTo": "image", "mode": "image", "units": 2,
|
||||
})
|
||||
require.Nil(t, adaptor.ValidateRequestAndSetAction(c, info))
|
||||
assert.Equal(t, "image", info.UpstreamModelName)
|
||||
})
|
||||
|
||||
t.Run("completion selects each persisted model including legacy origin fallback", func(t *testing.T) {
|
||||
adaptor, _, _ := newRequest(t, "image", map[string]any{})
|
||||
for _, tc := range []struct {
|
||||
upstream, origin, mode string
|
||||
units float64
|
||||
accepted bool
|
||||
}{
|
||||
{"video", "public-alias", "video", 500000, true},
|
||||
{"image", "public-alias", "image", 2, true},
|
||||
{"", "video", "video", 500000, true},
|
||||
{"image", "public-alias", "image", float64(dto.MaxImageN + 1), false},
|
||||
{"video", "public-alias", "image", 1, false},
|
||||
} {
|
||||
task := &model.Task{Properties: model.Properties{OriginModelName: tc.origin, UpstreamModelName: tc.upstream}}
|
||||
body, err := common.Marshal(map[string]any{"usage": map[string]any{"units": tc.units, "mode": tc.mode}})
|
||||
require.NoError(t, err)
|
||||
result, err := adaptor.ParseTaskResult(task, &http.Response{StatusCode: http.StatusOK}, body)
|
||||
require.NoError(t, err)
|
||||
if tc.accepted {
|
||||
assert.Equal(t, map[string]any{"units": tc.units, "mode": tc.mode}, result.UsageFacts)
|
||||
} else {
|
||||
assert.Nil(t, result.UsageFacts)
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("legacy completion facts remain usable by stored expressions", func(t *testing.T) {
|
||||
updated, err := pluginruntime.NewRegistry().Register(strings.Replace(source,
|
||||
`export function extractUsageOnComplete(ctx,result,body){return body.usage;}`,
|
||||
`export function extractUsageOnComplete(){return {legacyRatio:3};}`, 1), pluginruntime.Options{})
|
||||
require.NoError(t, err)
|
||||
_, _, info := newRequest(t, "image", map[string]any{})
|
||||
adaptor := New(updated)
|
||||
adaptor.Init(info)
|
||||
result, err := adaptor.ParseTaskResult(&model.Task{Properties: model.Properties{UpstreamModelName: "image"}},
|
||||
&http.Response{StatusCode: http.StatusOK}, []byte(`{}`))
|
||||
require.NoError(t, err)
|
||||
cost, _, err := billingexpr.RunExprWithRequest(`tier("legacy", u("legacyRatio") * 2)`,
|
||||
billingexpr.TokenParams{}, billingexpr.RequestInput{Usage: result.UsageFacts})
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 6.0, cost)
|
||||
})
|
||||
|
||||
t.Run("mixed batch selects schema by task rather than adaptor", func(t *testing.T) {
|
||||
adaptor, _, _ := newRequest(t, "image", map[string]any{})
|
||||
tasks := []*model.Task{
|
||||
{TaskID: "video-task", Properties: model.Properties{UpstreamModelName: "video"}},
|
||||
{TaskID: "image-task", Properties: model.Properties{UpstreamModelName: "image"}},
|
||||
{TaskID: "invalid-image", Properties: model.Properties{UpstreamModelName: "image"}},
|
||||
}
|
||||
body := []byte(`{"items":[
|
||||
{"taskId":"image-task","status":"SUCCESS","data":{"usage":{"units":2,"mode":"image"}}},
|
||||
{"taskId":"video-task","status":"SUCCESS","data":{"usage":{"units":500000,"mode":"video"}}},
|
||||
{"taskId":"invalid-image","status":"SUCCESS","data":{"usage":{"units":129,"mode":"image"}}}
|
||||
]}`)
|
||||
results, err := adaptor.ParseBatchResult(tasks, &http.Response{StatusCode: http.StatusOK}, body)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, results, 3)
|
||||
assert.Equal(t, map[string]any{"units": 500000.0, "mode": "video"}, results["video-task"].TaskInfo.UsageFacts)
|
||||
assert.Equal(t, map[string]any{"units": 2.0, "mode": "image"}, results["image-task"].TaskInfo.UsageFacts)
|
||||
assert.Nil(t, results["invalid-image"].TaskInfo.UsageFacts)
|
||||
})
|
||||
}
|
||||
|
||||
const documentStreamPlugin = `
|
||||
export const meta = {apiVersion:1,key:"document-stream",name:"Document stream",version:"1.0.0",author:{name:"Test"},models:["document"],fetchMode:"per_task",submitResponseTypes:["json","sse"],usageSchema:{units:{type:"number",unit:"count"}}};
|
||||
export function buildSubmitRequest(ctx){return {url:ctx.baseUrl+"/compile",responseType:ctx.requestBody.responseType || "sse"};}
|
||||
export function parseSubmitEvent(ctx,event,previous) {
|
||||
const chunk = JSON.parse(event.data);
|
||||
if (chunk.error) throw new Error("provider stream failure");
|
||||
if (chunk.badState) return {state:null};
|
||||
if (chunk.largeState) return {state:{document:"x".repeat(1048577)},done:chunk.complete === true};
|
||||
if (chunk.escapedState) return {state:{document:"<".repeat(200000)},done:true};
|
||||
if (chunk.resetState) return {state:{document:"",units:0},done:true};
|
||||
const state = Object.assign({}, previous || {document:"",units:0});
|
||||
state.document += chunk.part || "";
|
||||
if (chunk.units !== undefined) state.units = chunk.units;
|
||||
state.event = event.event; state.id = event.id;
|
||||
return {state:state,done:chunk.complete === true};
|
||||
}
|
||||
export function parseSubmitResponse(ctx,response){return {taskId:"vendor-document",taskData:response.body,immediate:{status:"SUCCESS"},state:{revision:1}};}
|
||||
export function parseTaskResult(){return {status:"SUCCESS"};}
|
||||
export function buildQueryRequest(ctx){return {url:ctx.baseUrl+"/query"};}
|
||||
export function extractUsageOnComplete(ctx,result,body){
|
||||
if(ctx.model!=="alias" || ctx.upstreamModel!=="document" || ctx.taskId!=="vendor-document" || ctx.publicTaskId!=="public-document" || ctx.state.revision!==1 || ctx.data.document!==body.document) throw new Error("invalid completion context");
|
||||
return {units:body.units};
|
||||
}
|
||||
`
|
||||
|
||||
func TestTaskSubmitStreamContract(t *testing.T) {
|
||||
plugin, err := pluginruntime.CompilePlugin(documentStreamPlugin, pluginruntime.Options{})
|
||||
require.NoError(t, err)
|
||||
for _, tc := range []struct {
|
||||
name, body, contentType string
|
||||
valid bool
|
||||
}{
|
||||
{"multiline and CRLF", ": heartbeat\r\nid: doc-1\r\nevent: update\r\ndata: {\"part\":\r\ndata: \"hello\",\"units\":2}\r\n\r\ndata: {\"part\":\"world\",\"units\":3,\"complete\":true}\n\n", "text/event-stream; charset=utf-8", true},
|
||||
{"zero actual units", "data: {\"part\":\"free\",\"units\":0,\"complete\":true}\n\n", "text/event-stream", true},
|
||||
{"premature EOF", "data: {\"part\":\"partial\"}\n\n", "text/event-stream", false},
|
||||
{"unterminated event", "data: {\"complete\":true}", "text/event-stream", false},
|
||||
{"provider error", "data: {\"error\":true}\n\n", "text/event-stream", false},
|
||||
{"invalid hook result", "data: {\"badState\":true}\n\n", "text/event-stream", false},
|
||||
{"state limit", "data: {\"largeState\":true}\n\n", "text/event-stream", false},
|
||||
{"intermediate state limit", "data: {\"largeState\":true}\n\ndata: {\"resetState\":true}\n\n", "text/event-stream", false},
|
||||
{"encoded state limit", "data: {\"escapedState\":true}\n\n", "text/event-stream", false},
|
||||
{"event limit", "data: " + strings.Repeat("x", maxTaskPluginPersistedJSONBytes) + "\n\n", "text/event-stream", false},
|
||||
{"wrong response type", "{}", "application/json", false},
|
||||
{"undeclared stream", "data: {}\n\n", "text/event-stream", false},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
info := &relaycommon.RelayInfo{OriginModelName: "alias", ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "document", ChannelBaseUrl: "https://provider.example"}, TaskRelayInfo: &relaycommon.TaskRelayInfo{PublicTaskID: "public-document"}}
|
||||
adaptor := New(plugin)
|
||||
adaptor.Init(info)
|
||||
c, recorder := gin.CreateTestContext(httptest.NewRecorder())
|
||||
_ = recorder
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/compile", nil)
|
||||
requestBody := map[string]any{}
|
||||
if tc.name == "undeclared stream" {
|
||||
requestBody["responseType"] = "json"
|
||||
}
|
||||
c.Set("task_request", requestBody)
|
||||
require.Nil(t, adaptor.ValidateRequestAndSetAction(c, info))
|
||||
response := &http.Response{StatusCode: 200, Header: http.Header{"Content-Type": {tc.contentType}}, Body: io.NopCloser(strings.NewReader(tc.body))}
|
||||
parsed, taskErr := adaptor.ParseResponse(c, response, info)
|
||||
assert.False(t, c.Writer.Written())
|
||||
if !tc.valid {
|
||||
require.NotNil(t, taskErr)
|
||||
assert.True(t, taskErr.NoRetry)
|
||||
return
|
||||
}
|
||||
require.Nil(t, taskErr)
|
||||
require.NotNil(t, parsed.Immediate)
|
||||
if tc.name == "zero actual units" {
|
||||
assert.Equal(t, map[string]any{"units": float64(0)}, parsed.Immediate.UsageFacts)
|
||||
return
|
||||
}
|
||||
assert.Equal(t, map[string]any{"units": float64(3)}, parsed.Immediate.UsageFacts)
|
||||
assert.JSONEq(t, `{"document":"helloworld","units":3,"event":"message","id":"doc-1"}`, string(parsed.TaskData))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestTaskSubmitStreamCancellationClosesReader(t *testing.T) {
|
||||
plugin, err := pluginruntime.CompilePlugin(documentStreamPlugin, pluginruntime.Options{})
|
||||
require.NoError(t, err)
|
||||
reader, writer := io.Pipe()
|
||||
defer writer.Close()
|
||||
ctx, cancel := context.WithCancel(t.Context())
|
||||
cancel()
|
||||
_, err = New(plugin).readSubmitEvents(ctx, &http.Response{Header: http.Header{"Content-Type": {"text/event-stream"}}, Body: reader}, nil)
|
||||
require.ErrorIs(t, err, context.Canceled)
|
||||
_, err = writer.Write([]byte("data: {}\n\n"))
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
func TestTaskSubmitStreamIdleTimeout(t *testing.T) {
|
||||
previous := constant.StreamingTimeout
|
||||
constant.StreamingTimeout = 1
|
||||
t.Cleanup(func() { constant.StreamingTimeout = previous })
|
||||
plugin, err := pluginruntime.CompilePlugin(documentStreamPlugin, pluginruntime.Options{})
|
||||
require.NoError(t, err)
|
||||
reader, writer := io.Pipe()
|
||||
defer writer.Close()
|
||||
_, err = New(plugin).readSubmitEvents(t.Context(), &http.Response{Header: http.Header{"Content-Type": {"text/event-stream"}}, Body: reader}, nil)
|
||||
require.ErrorContains(t, err, "idle timeout")
|
||||
}
|
||||
|
||||
func TestPluginJSONValuesPreserveCodecNormalizationAndIsolation(t *testing.T) {
|
||||
for _, value := range []any{
|
||||
map[string]any{"units": int64(3), "nested": []any{true, "<image> / 图像", math.Copysign(0, -1), nil}},
|
||||
map[string]any{"empty": []any{}, "null": []any(nil), "object": map[string]any(nil)},
|
||||
map[string]any{"large": int64(math.MaxInt64), "invalid UTF-8": string([]byte{0xff, 0xfe})},
|
||||
map[string]any{string([]byte{0xff}): "invalid key"},
|
||||
map[string]any{"bytes": []byte{0, 1, 255}, "number": json.Number("9007199254740993")},
|
||||
json.RawMessage(`{"units":2,"enabled":false}`),
|
||||
struct {
|
||||
Units int `json:"units"`
|
||||
}{Units: 0},
|
||||
} {
|
||||
encoded, err := common.Marshal(value)
|
||||
require.NoError(t, err)
|
||||
var expected any
|
||||
require.NoError(t, common.Unmarshal(encoded, &expected))
|
||||
assert.Equal(t, expected, jsonValue(value))
|
||||
}
|
||||
source := map[string]any{"items": []any{map[string]any{"label": "original"}}}
|
||||
copy := jsonValue(source).(map[string]any)
|
||||
copy["items"].([]any)[0].(map[string]any)["label"] = "changed"
|
||||
assert.Equal(t, "original", source["items"].([]any)[0].(map[string]any)["label"])
|
||||
}
|
||||
|
||||
func TestTaskSubmitHooksReceiveIndependentRequestSnapshots(t *testing.T) {
|
||||
source := `
|
||||
export const meta={apiVersion:1,key:"request-copy",name:"Request copy",version:"1.0.0",author:{name:"Test"},models:["copy"],fetchMode:"per_task",usageSchema:{units:{type:"number",unit:"count"}}};
|
||||
export function buildSubmitRequest(ctx){ctx.requestBody.units=3;ctx.requestBody.nested.label="built";return {url:ctx.baseUrl+"/submit",body:ctx.requestBody};}
|
||||
export function extractUsage(ctx){if(ctx.requestBody.nested.label!=="original")throw new Error("request changed");return {units:ctx.requestBody.units};}
|
||||
export function parseSubmitResponse(){return {taskId:"copy"};}
|
||||
export function buildQueryRequest(ctx){return {url:ctx.baseUrl+"/query"};}
|
||||
export function parseTaskResult(){return {status:"SUCCESS"};}
|
||||
`
|
||||
plugin, err := pluginruntime.CompilePlugin(source, pluginruntime.Options{})
|
||||
require.NoError(t, err)
|
||||
request := map[string]any{"units": int64(2), "nested": map[string]any{"label": "original"}}
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/submit", nil)
|
||||
c.Set("task_request", request)
|
||||
info := &relaycommon.RelayInfo{OriginModelName: "copy", ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "copy", ChannelBaseUrl: "https://provider.example"}, TaskRelayInfo: &relaycommon.TaskRelayInfo{}}
|
||||
adaptor := New(plugin)
|
||||
adaptor.Init(info)
|
||||
require.Nil(t, adaptor.ValidateRequestAndSetAction(c, info))
|
||||
facts, err := adaptor.ExtractUsageFactsValidated(c, info)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, map[string]any{"units": float64(2)}, facts)
|
||||
body, err := adaptor.BuildRequestBody(c, info)
|
||||
require.NoError(t, err)
|
||||
encoded, err := io.ReadAll(body)
|
||||
require.NoError(t, err)
|
||||
assert.JSONEq(t, `{"units":3,"nested":{"label":"built"}}`, string(encoded))
|
||||
assert.Equal(t, int64(2), request["units"])
|
||||
assert.Equal(t, "original", request["nested"].(map[string]any)["label"])
|
||||
}
|
||||
|
||||
func TestTaskSubmitDeltaStreamContract(t *testing.T) {
|
||||
source := strings.Replace(documentStreamPlugin, `submitResponseTypes:["json","sse"],`, `submitResponseTypes:["json","sse"],requiredCapabilities:["submit-sse-delta@1"],`, 1) + `
|
||||
export function parseSubmitEventDelta(ctx,event,previous) {
|
||||
if(previous && previous.document !== undefined) throw new Error("full result leaked into control state");
|
||||
const chunk=JSON.parse(event.data);
|
||||
let changes=chunk.changes;
|
||||
if(chunk.largeResult) changes=[{op:"set",path:[],value:{document:"<".repeat(200000)}}];
|
||||
if(chunk.resetAfterLarge) changes.push({op:"set",path:[],value:{document:"",units:0}});
|
||||
const result={changes:changes,state:{events:(previous ? previous.events : 0)+1},done:chunk.complete===true};
|
||||
if(chunk.largeControl) result.state={text:"x".repeat(65537)};
|
||||
if(chunk.extra) result.extra=true;
|
||||
return result;
|
||||
}`
|
||||
plugin, err := pluginruntime.CompilePlugin(source, pluginruntime.Options{})
|
||||
require.NoError(t, err)
|
||||
const first = `{"changes":[{"op":"set","path":[],"value":{"document":"hello","units":2}}]}`
|
||||
const last = `{"changes":[{"op":"appendText","path":["document"],"value":"world"},{"op":"set","path":["units"],"value":0}],"complete":true}`
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
frames []string
|
||||
valid bool
|
||||
}{
|
||||
{"control state and zero usage", []string{first, last}, true},
|
||||
{"oversized intermediate result", []string{`{"largeResult":true}`, last}, false},
|
||||
{"oversized operation before reset", []string{`{"largeResult":true,"resetAfterLarge":true,"complete":true}`}, false},
|
||||
{"control state limit", []string{first, `{"changes":[],"largeControl":true,"complete":true}`}, false},
|
||||
{"missing changes", []string{`{"complete":true}`}, false},
|
||||
{"extra result fields", []string{first, `{"changes":[],"extra":true,"complete":true}`}, false},
|
||||
{"unfinished delta stream", []string{first}, false},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
var stream strings.Builder
|
||||
for _, frame := range tc.frames {
|
||||
stream.WriteString("data: " + frame + "\n\n")
|
||||
}
|
||||
info := &relaycommon.RelayInfo{OriginModelName: "alias", ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "document", ChannelBaseUrl: "https://provider.example"}, TaskRelayInfo: &relaycommon.TaskRelayInfo{PublicTaskID: "public-document"}}
|
||||
adaptor := New(plugin)
|
||||
adaptor.Init(info)
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/compile", nil)
|
||||
c.Set("task_request", map[string]any{})
|
||||
require.Nil(t, adaptor.ValidateRequestAndSetAction(c, info))
|
||||
response := &http.Response{StatusCode: http.StatusOK, Header: http.Header{"Content-Type": {"text/event-stream"}}, Body: io.NopCloser(strings.NewReader(stream.String()))}
|
||||
parsed, taskErr := adaptor.ParseResponse(c, response, info)
|
||||
assert.False(t, c.Writer.Written())
|
||||
if !tc.valid {
|
||||
require.NotNil(t, taskErr)
|
||||
assert.True(t, taskErr.NoRetry)
|
||||
return
|
||||
}
|
||||
require.Nil(t, taskErr)
|
||||
require.NotNil(t, parsed.Immediate)
|
||||
assert.JSONEq(t, `{"document":"helloworld","units":0}`, string(parsed.TaskData))
|
||||
assert.Equal(t, map[string]any{"units": float64(0)}, parsed.Immediate.UsageFacts)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestAlibabaSubmitDeltaDoesNotMutateControlState(t *testing.T) {
|
||||
source, err := plugins.Source("alibaba")
|
||||
require.NoError(t, err)
|
||||
plugin, err := pluginruntime.CompilePlugin(source, pluginruntime.Options{})
|
||||
require.NoError(t, err)
|
||||
const before = `{"request_id":"old","usage":{"image_count":1},"output":{"finished":false,"choices":[{"message":{"role":"assistant","content":[{"text":"hello"}]}}]}}`
|
||||
const control = `{"choices":[{"count":1,"lastText":true,"finishReason":""}],"hasUsage":true}`
|
||||
var previous, initial any
|
||||
require.NoError(t, common.UnmarshalJsonStr(control, &previous))
|
||||
require.NoError(t, common.UnmarshalJsonStr(before, &initial))
|
||||
accumulated := pluginruntime.NewJSONState(maxTaskPluginPersistedJSONBytes)
|
||||
require.NoError(t, accumulated.Apply(t.Context(), []any{map[string]any{"op": "set", "path": []any{}, "value": initial}}))
|
||||
value, err := plugin.Engine.Call(t.Context(), "parseSubmitEventDelta", map[string]any{}, map[string]any{
|
||||
"event": "message", "data": `{"request_id":"new","usage":{"image_count":2},"output":{"choices":[{"finish_reason":"stop","message":{"content":[{"text":" world"},{"image":"https://cdn.example/image.png"}]}}]}}`,
|
||||
}, previous)
|
||||
require.NoError(t, err)
|
||||
unchanged, err := common.Marshal(previous)
|
||||
require.NoError(t, err)
|
||||
assert.JSONEq(t, control, string(unchanged))
|
||||
result := value.(map[string]any)
|
||||
assert.Equal(t, true, result["done"])
|
||||
require.NoError(t, accumulated.Apply(t.Context(), result["changes"]))
|
||||
final, err := accumulated.Value()
|
||||
require.NoError(t, err)
|
||||
encoded, err := common.Marshal(final)
|
||||
require.NoError(t, err)
|
||||
assert.JSONEq(t, `{"request_id":"new","usage":{"image_count":2},"output":{"finished":true,"choices":[{"finish_reason":"stop","message":{"role":"assistant","content":[{"text":"hello world"},{"image":"https://cdn.example/image.png"}]}}]}}`, string(encoded))
|
||||
encoded, err = common.Marshal(result["state"])
|
||||
require.NoError(t, err)
|
||||
assert.JSONEq(t, `{"choices":[{"count":2,"lastText":false,"finishReason":"stop"}],"hasUsage":true}`, string(encoded))
|
||||
}
|
||||
@@ -0,0 +1,190 @@
|
||||
package jsplugin
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"runtime"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
pluginruntime "github.com/QuantumNous/new-api/pkg/jsplugin"
|
||||
"github.com/QuantumNous/new-api/plugins"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func benchmarkAlibaba(b *testing.B) (*pluginruntime.LoadedPlugin, *atomic.Int64) {
|
||||
b.Helper()
|
||||
source, err := plugins.Source("alibaba")
|
||||
require.NoError(b, err)
|
||||
initializations := &atomic.Int64{}
|
||||
plugin, err := pluginruntime.CompilePlugin(source+"\n"+`console.log("benchmark_runtime_initialized");`, pluginruntime.Options{
|
||||
Log: func(message string) {
|
||||
if strings.HasSuffix(message, "benchmark_runtime_initialized") {
|
||||
initializations.Add(1)
|
||||
}
|
||||
},
|
||||
})
|
||||
require.NoError(b, err)
|
||||
return plugin, initializations
|
||||
}
|
||||
|
||||
func BenchmarkTaskPluginRuntime(b *testing.B) {
|
||||
for _, name := range []string{"HasExport", "MissingExport", "HookCall", "ParallelHookCall"} {
|
||||
b.Run(name, func(b *testing.B) {
|
||||
plugin, initializations := benchmarkAlibaba(b)
|
||||
ctx := context.Background()
|
||||
input := map[string]any{"model": "wan2.5-i2v-preview", "upstreamModel": "wan2.5-i2v-preview", "requestBody": map[string]any{
|
||||
"prompt": "animate the landscape", "image": "https://cdn.example/image.png", "duration": 5, "size": "720P",
|
||||
}}
|
||||
before := initializations.Load()
|
||||
b.ReportAllocs()
|
||||
if name == "ParallelHookCall" {
|
||||
b.ResetTimer()
|
||||
b.RunParallel(func(pb *testing.PB) {
|
||||
for pb.Next() {
|
||||
if _, err := plugin.Engine.Call(ctx, "extractUsage", input); err != nil {
|
||||
b.Error(err)
|
||||
return
|
||||
}
|
||||
}
|
||||
})
|
||||
} else {
|
||||
for b.Loop() {
|
||||
var err error
|
||||
switch name {
|
||||
case "HasExport":
|
||||
_, err = plugin.Engine.HasExport(ctx, "extractUsage")
|
||||
case "MissingExport":
|
||||
_, err = plugin.Engine.HasExport(ctx, "absentHook")
|
||||
default:
|
||||
_, err = plugin.Engine.Call(ctx, "extractUsage", input)
|
||||
}
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
b.ReportMetric(float64(initializations.Load()-before)/float64(b.N), "runtime_inits/op")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// Run with -benchtime=50x: GC is excluded so this measures the first hook after
|
||||
// reclamation, not the runtime.GC operation or an artificial throughput workload.
|
||||
func BenchmarkTaskPluginAfterGC(b *testing.B) {
|
||||
plugin, initializations := benchmarkAlibaba(b)
|
||||
input := map[string]any{"model": "wan2.5-i2v-preview", "upstreamModel": "wan2.5-i2v-preview", "requestBody": map[string]any{
|
||||
"prompt": "animate the landscape", "image": "https://cdn.example/image.png", "duration": 5, "size": "720P",
|
||||
}}
|
||||
before := initializations.Load()
|
||||
b.ReportAllocs()
|
||||
b.ResetTimer()
|
||||
for range b.N {
|
||||
b.StopTimer()
|
||||
runtime.GC()
|
||||
runtime.GC()
|
||||
b.StartTimer()
|
||||
if _, err := plugin.Engine.Call(b.Context(), "extractUsage", input); err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
b.ReportMetric(float64(initializations.Load()-before)/float64(b.N), "runtime_inits/op")
|
||||
}
|
||||
|
||||
func BenchmarkTaskPluginSubmit(b *testing.B) {
|
||||
for _, size := range []int{4 << 10, 1 << 20, 8 << 20} {
|
||||
b.Run(fmt.Sprintf("%dKiB", size>>10), func(b *testing.B) {
|
||||
plugin, initializations := benchmarkAlibaba(b)
|
||||
const modelName = "wan2.5-i2v-preview"
|
||||
wire, err := common.Marshal(map[string]any{
|
||||
"model": modelName, "duration": 5, "size": "720P",
|
||||
"input": []any{map[string]any{"type": "input_text", "text": "animate the landscape"}, map[string]any{"type": "input_image", "image_url": "data:image/png;base64," + strings.Repeat("A", size)}},
|
||||
})
|
||||
require.NoError(b, err)
|
||||
before := initializations.Load()
|
||||
b.SetBytes(int64(len(wire)))
|
||||
b.ReportAllocs()
|
||||
for b.Loop() {
|
||||
var body any
|
||||
if err := common.Unmarshal(wire, &body); err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
protocol := pluginruntime.ProtocolRequestContext{
|
||||
RouteRequestContext: pluginruntime.RouteRequestContext{Path: "/v1/responses", Method: http.MethodPost, Body: map[string]any{"kind": "json", "value": body}},
|
||||
Protocol: "openai_responses", Operation: "create", Model: modelName,
|
||||
}
|
||||
// The middleware decoder runs before channel selection; the adaptor
|
||||
// decodes again after the executing provider has been pinned.
|
||||
decoded, err := plugin.Engine.CallPath(b.Context(), "protocols", []string{"openai_responses", "decodeRequest"}, protocol.JSValue())
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||||
c.Set("task_request", decoded.(map[string]any)["requestBody"])
|
||||
c.Set(pluginruntime.ContextKeyRouteRequest, protocol.RouteRequestContext)
|
||||
c.Set(pluginruntime.ContextKeyProtocolRequest, protocol)
|
||||
c.Set(pluginruntime.ContextKeyPinnedEndpoint, pluginruntime.PinnedEndpoint{Plugin: plugin, Protocol: protocol.Protocol, Model: modelName})
|
||||
info := &relaycommon.RelayInfo{OriginModelName: modelName, ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: modelName, ChannelBaseUrl: "https://provider.example"}, TaskRelayInfo: &relaycommon.TaskRelayInfo{PublicTaskID: "benchmark-task"}}
|
||||
adaptor := New(plugin)
|
||||
adaptor.Init(info)
|
||||
if taskErr := adaptor.ValidateRequestAndSetAction(c, info); taskErr != nil {
|
||||
b.Fatal(taskErr)
|
||||
}
|
||||
if _, err := adaptor.ExtractUsageFactsValidated(c, info); err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
upstream, err := adaptor.BuildRequestBody(c, info)
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
if _, err := io.Copy(io.Discard, upstream); err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
b.ReportMetric(float64(initializations.Load()-before)/float64(b.N), "runtime_inits/op")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkTaskPluginSSE(b *testing.B) {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
events, textBytes int
|
||||
}{{"small_state", 32, 48}, {"growing_state", 64, 8 << 10}} {
|
||||
b.Run(tc.name, func(b *testing.B) {
|
||||
plugin, initializations := benchmarkAlibaba(b)
|
||||
var stream strings.Builder
|
||||
for index := range tc.events {
|
||||
choice := map[string]any{"message": map[string]any{"role": "assistant", "content": []any{map[string]any{"text": strings.Repeat("a", tc.textBytes)}}}}
|
||||
if index == tc.events-1 {
|
||||
choice["finish_reason"] = "stop"
|
||||
}
|
||||
frame, err := common.Marshal(map[string]any{"request_id": "benchmark-sse", "output": map[string]any{"choices": []any{choice}}, "usage": map[string]any{"image_count": 1}})
|
||||
require.NoError(b, err)
|
||||
stream.WriteString("data: ")
|
||||
stream.Write(frame)
|
||||
stream.WriteString("\n\n")
|
||||
}
|
||||
wire := stream.String()
|
||||
input := map[string]any{"model": "wan2.6-image"}
|
||||
before := initializations.Load()
|
||||
b.SetBytes(int64(len(wire)))
|
||||
b.ReportAllocs()
|
||||
for b.Loop() {
|
||||
response := &http.Response{Header: http.Header{"Content-Type": {"text/event-stream"}}, Body: io.NopCloser(strings.NewReader(wire))}
|
||||
if _, err := New(plugin).readSubmitEvents(b.Context(), response, input); err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
b.ReportMetric(float64(initializations.Load()-before)/float64(b.N), "runtime_inits/op")
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,146 @@
|
||||
package jsplugin
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"mime"
|
||||
"net/http"
|
||||
"slices"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
pluginruntime "github.com/QuantumNous/new-api/pkg/jsplugin"
|
||||
"github.com/QuantumNous/new-api/relay/helper"
|
||||
)
|
||||
|
||||
// readSubmitEvents owns only SSE framing and bounded JSON state. The plugin
|
||||
// interprets data, accumulates provider output and decides when it is complete.
|
||||
// No response bytes are written to the client at this boundary.
|
||||
func (a *TaskAdaptor) readSubmitEvents(parent context.Context, resp *http.Response, driverContext map[string]any) (any, error) {
|
||||
mediaType, _, err := mime.ParseMediaType(resp.Header.Get("Content-Type"))
|
||||
if err != nil || mediaType != "text/event-stream" {
|
||||
return nil, fmt.Errorf("expected a text/event-stream submit response")
|
||||
}
|
||||
ctx, cancel := context.WithCancelCause(parent)
|
||||
defer cancel(nil)
|
||||
stopClose := context.AfterFunc(ctx, func() { _ = resp.Body.Close() })
|
||||
defer stopClose()
|
||||
idleTimeout := time.Duration(constant.StreamingTimeout) * time.Second
|
||||
if idleTimeout <= 0 {
|
||||
idleTimeout = 30 * time.Second
|
||||
}
|
||||
timer := time.AfterFunc(idleTimeout, func() { cancel(fmt.Errorf("task submit stream idle timeout")) })
|
||||
defer timer.Stop()
|
||||
|
||||
scanner := helper.NewStreamScanner(resp.Body, maxTaskPluginPersistedJSONBytes+1)
|
||||
hook := "parseSubmitEvent"
|
||||
var accumulated *pluginruntime.JSONState
|
||||
if slices.Contains(a.plugin.Meta.RequiredCapabilities, pluginruntime.CapabilitySubmitSSEDelta) {
|
||||
hook = "parseSubmitEventDelta"
|
||||
accumulated = pluginruntime.NewJSONState(maxTaskPluginPersistedJSONBytes)
|
||||
}
|
||||
var state any
|
||||
var data []string
|
||||
var eventName, lastID string
|
||||
frameBytes := 0
|
||||
firstLine := true
|
||||
for scanner.Scan() {
|
||||
if err := context.Cause(ctx); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
timer.Reset(idleTimeout)
|
||||
line := scanner.Text()
|
||||
if firstLine {
|
||||
line = strings.TrimPrefix(line, "\ufeff")
|
||||
firstLine = false
|
||||
}
|
||||
frameBytes += len(line) + 1
|
||||
if frameBytes > maxTaskPluginPersistedJSONBytes {
|
||||
return nil, fmt.Errorf("task submit SSE event exceeds size limit")
|
||||
}
|
||||
if line != "" {
|
||||
field, value, _ := strings.Cut(line, ":")
|
||||
value, _ = strings.CutPrefix(value, " ")
|
||||
switch field {
|
||||
case "data":
|
||||
data = append(data, value)
|
||||
case "event":
|
||||
eventName = value
|
||||
case "id":
|
||||
if !strings.ContainsRune(value, '\x00') {
|
||||
lastID = value
|
||||
}
|
||||
}
|
||||
continue
|
||||
}
|
||||
frameBytes = 0
|
||||
if len(data) == 0 {
|
||||
eventName = ""
|
||||
continue
|
||||
}
|
||||
if eventName == "" {
|
||||
eventName = "message"
|
||||
}
|
||||
event := map[string]any{"event": eventName, "id": lastID, "data": strings.Join(data, "\n")}
|
||||
value, err := a.plugin.Engine.Call(ctx, hook, driverContext, event, state)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
result, ok := value.(map[string]any)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("%s must return an object", hook)
|
||||
}
|
||||
nextState, hasState := result["state"]
|
||||
done, hasDone := result["done"].(bool)
|
||||
if !hasState || !hasDone {
|
||||
return nil, fmt.Errorf("%s must return state and a boolean done", hook)
|
||||
}
|
||||
encoded, err := common.Marshal(nextState)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("invalid submit stream state: %w", err)
|
||||
}
|
||||
if len(encoded) > maxTaskPluginPersistedJSONBytes {
|
||||
return nil, fmt.Errorf("task submit stream state exceeds size limit")
|
||||
}
|
||||
if accumulated != nil {
|
||||
if len(result) != 3 {
|
||||
return nil, fmt.Errorf("parseSubmitEventDelta must return only changes, state and done")
|
||||
}
|
||||
// Control state is separate from the full result and is the only
|
||||
// snapshot passed back to JS on the next event.
|
||||
if len(encoded) > 64<<10 {
|
||||
return nil, fmt.Errorf("task submit stream control state exceeds size limit")
|
||||
}
|
||||
if err = accumulated.Apply(ctx, result["changes"]); err != nil {
|
||||
return nil, fmt.Errorf("invalid submit stream changes: %w", err)
|
||||
}
|
||||
}
|
||||
// Keep the exact encoded-byte limit above. Plain JSON state can be
|
||||
// isolated without parsing large strings again; codec-specific values
|
||||
// (for example exported typed arrays) retain the old normalization.
|
||||
var plainJSON bool
|
||||
state, plainJSON = cloneJSONValue(nextState, 0)
|
||||
if !plainJSON {
|
||||
if err = common.Unmarshal(encoded, &state); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
if done {
|
||||
if accumulated != nil {
|
||||
return accumulated.Value()
|
||||
}
|
||||
return state, nil
|
||||
}
|
||||
data = nil
|
||||
eventName = ""
|
||||
}
|
||||
if err := context.Cause(ctx); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := scanner.Err(); err != nil {
|
||||
return nil, fmt.Errorf("read task submit stream: %w", err)
|
||||
}
|
||||
return nil, fmt.Errorf("task submit stream ended before the plugin reported completion")
|
||||
}
|
||||
@@ -14,6 +14,7 @@ import (
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
relayconstant "github.com/QuantumNous/new-api/relay/constant"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
"github.com/QuantumNous/new-api/relaykit/relayconvert"
|
||||
relaytypes "github.com/QuantumNous/new-api/relaykit/types"
|
||||
hosttypes "github.com/QuantumNous/new-api/types"
|
||||
"github.com/gin-gonic/gin"
|
||||
@@ -21,6 +22,59 @@ import (
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestPrepareResponsesRequestRetainsConvertedAdaptor(t *testing.T) {
|
||||
type capturedRequest struct {
|
||||
path string
|
||||
body []byte
|
||||
}
|
||||
captured := make(chan capturedRequest, 1)
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
body, err := io.ReadAll(r.Body)
|
||||
if err != nil {
|
||||
http.Error(w, err.Error(), http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
captured <- capturedRequest{path: r.URL.Path, body: body}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(`{"id":"chatcmpl_1","object":"chat.completion","model":"gpt-4o","choices":[{"index":0,"message":{"role":"assistant","content":"ok"},"finish_reason":"stop"}],"usage":{"prompt_tokens":3,"completion_tokens":2,"total_tokens":5}}`))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
common.SetContextKey(c, constant.ContextKeyOriginalModel, "gpt-4o")
|
||||
common.SetContextKey(c, constant.ContextKeyChannelType, constant.ChannelTypeAdvancedCustom)
|
||||
common.SetContextKey(c, constant.ContextKeyChannelBaseUrl, server.URL)
|
||||
common.SetContextKey(c, constant.ContextKeyChannelKey, "test-key")
|
||||
common.SetContextKey(c, constant.ContextKeyChannelOtherSetting, dto.ChannelOtherSettings{
|
||||
AdvancedCustom: &dto.AdvancedCustomConfig{Routes: []dto.AdvancedCustomRoute{{
|
||||
IncomingPath: "/v1/responses", UpstreamPath: "/v1/chat/completions", Converter: relayconvert.ConverterOpenAIResponsesToOpenAIChat,
|
||||
}}},
|
||||
})
|
||||
request := &dto.OpenAIResponsesRequest{Model: "gpt-4o", Input: []byte(`"hello"`)}
|
||||
info := relaycommon.GenRelayInfoResponses(c, request)
|
||||
adaptor, body, closer, apiErr := PrepareResponsesRequest(c, info, request)
|
||||
require.Nil(t, apiErr)
|
||||
defer closer.Close()
|
||||
response, err := adaptor.DoRequest(c, info, body)
|
||||
require.NoError(t, err)
|
||||
httpResponse, ok := response.(*http.Response)
|
||||
require.True(t, ok)
|
||||
usage, apiErr := adaptor.DoResponse(c, httpResponse, info)
|
||||
require.Nil(t, apiErr)
|
||||
require.IsType(t, &dto.Usage{}, usage)
|
||||
assert.Equal(t, 5, usage.(*dto.Usage).TotalTokens)
|
||||
|
||||
upstream := <-captured
|
||||
assert.Equal(t, "/v1/chat/completions", upstream.path)
|
||||
var upstreamRequest dto.GeneralOpenAIRequest
|
||||
require.NoError(t, common.Unmarshal(upstream.body, &upstreamRequest))
|
||||
require.Len(t, upstreamRequest.Messages, 1)
|
||||
assert.Equal(t, "hello", upstreamRequest.Messages[0].StringContent())
|
||||
assert.Equal(t, []relaytypes.RelayFormat{relaytypes.RelayFormatOpenAIResponses, relaytypes.RelayFormatOpenAI}, info.RequestConversionChain)
|
||||
}
|
||||
|
||||
func TestIsResponsesEventStreamContentType(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
|
||||
@@ -99,6 +99,11 @@ type RelayInfo struct {
|
||||
UsePrice bool
|
||||
RelayMode int
|
||||
OriginModelName string
|
||||
ResponseModel *ResponseModel
|
||||
ResponseGuard *ResponseModelGuard
|
||||
AstraDetectionDisabled bool
|
||||
AstraAutoSwitch bool
|
||||
UpstreamRequests int
|
||||
|
||||
// BillingModelName is the pricing identity for this request. It is kept
|
||||
// separate from OriginModelName and UpstreamModelName so virtual pricing
|
||||
@@ -138,7 +143,11 @@ type RelayInfo struct {
|
||||
ForcePreConsume bool
|
||||
// Billing 是计费会话,封装了预扣费/结算/退款的统一生命周期。
|
||||
// 初始免费组可为 nil;若 auto 重试切换到付费组,会在发送前创建。
|
||||
Billing BillingSettler
|
||||
Billing BillingSettler
|
||||
DurableAttemptBilling bool
|
||||
DurableAttemptID int64
|
||||
AttemptDetached bool
|
||||
GuardDeadline time.Time
|
||||
// BillingSource indicates whether this request is billed from wallet quota or subscription.
|
||||
// "" or "wallet" => wallet; "subscription" => subscription
|
||||
BillingSource string
|
||||
@@ -176,6 +185,11 @@ type RelayInfo struct {
|
||||
// and again before settlement. Non-nil only when billing mode is "tiered_expr".
|
||||
TieredBillingSnapshot *billingexpr.BillingSnapshot
|
||||
BillingRequestInput *billingexpr.RequestInput
|
||||
BillingImageCount *int
|
||||
// ImageRequestCount is the effective quantity sent on the current attempt;
|
||||
// ImageQuotaBeforeGroup is the frozen legacy estimate before request ratios.
|
||||
ImageRequestCount int
|
||||
ImageQuotaBeforeGroup float64
|
||||
|
||||
Request dto.Request
|
||||
|
||||
@@ -187,6 +201,10 @@ type RelayInfo struct {
|
||||
FinalRequestRelayFormat types.RelayFormat
|
||||
|
||||
StreamStatus *StreamStatus
|
||||
// PerformanceOutputTokens is captured by settlement and sampled once at
|
||||
// the request boundary, independently of billing success or failure.
|
||||
PerformanceOutputTokens int64
|
||||
PerformanceBusinessRejection bool
|
||||
|
||||
// convOptions caches the converter settings snapshot (see ConvOptions).
|
||||
convOptions *convmeta.Options
|
||||
@@ -204,7 +222,36 @@ type RelayInfo struct {
|
||||
*TaskRelayInfo
|
||||
}
|
||||
|
||||
// UpdateImageCount replaces the billable quantity without changing the frozen
|
||||
// request parameters or multiplying the legacy and expression prices together.
|
||||
func (info *RelayInfo) UpdateImageCount(count int64) {
|
||||
if info == nil || count <= 0 || count > int64(dto.MaxImageN) {
|
||||
return
|
||||
}
|
||||
if info.PriceData.UsePrice {
|
||||
info.PriceData.AddOtherRatio("n", float64(count))
|
||||
}
|
||||
if info.TieredBillingSnapshot != nil && info.TieredBillingSnapshot.EstimatedImageCount != nil {
|
||||
n := int(count)
|
||||
info.BillingImageCount = &n
|
||||
}
|
||||
}
|
||||
|
||||
func (info *RelayInfo) RequestedImageCount() int {
|
||||
if info.ImageRequestCount > 0 {
|
||||
return info.ImageRequestCount
|
||||
}
|
||||
if info.TieredBillingSnapshot != nil && info.TieredBillingSnapshot.EstimatedImageCount != nil {
|
||||
return *info.TieredBillingSnapshot.EstimatedImageCount
|
||||
}
|
||||
if count, ok := info.PriceData.OtherRatios()["n"]; ok && count >= 1 && count <= dto.MaxImageN {
|
||||
return int(count)
|
||||
}
|
||||
return 1
|
||||
}
|
||||
|
||||
func (info *RelayInfo) InitChannelMeta(c *gin.Context) {
|
||||
info.ResponseModel = nil
|
||||
info.FinalRequestRelayFormat = ""
|
||||
info.RequestConversionChain = nil
|
||||
info.InitRequestConversionChain()
|
||||
@@ -252,6 +299,15 @@ func (info *RelayInfo) InitChannelMeta(c *gin.Context) {
|
||||
channelMeta.ChannelOtherSettings = channelOtherSettings
|
||||
}
|
||||
|
||||
if channelType == constant.ChannelTypeAdvancedCustom &&
|
||||
!channelMeta.ChannelSetting.PassThroughBodyEnabled &&
|
||||
c.Request != nil && c.Request.URL != nil {
|
||||
route, matched := channelMeta.ChannelOtherSettings.AdvancedCustom.MatchPathForModel(c.Request.URL.Path, info.OriginModelName)
|
||||
if matched && route.PassThroughBodyEnabled {
|
||||
channelMeta.ChannelSetting.PassThroughBodyEnabled = true
|
||||
}
|
||||
}
|
||||
|
||||
if streamSupportedChannels[channelMeta.ChannelType] {
|
||||
channelMeta.SupportStreamOptions = true
|
||||
}
|
||||
@@ -377,6 +433,8 @@ var streamSupportedChannels = map[int]bool{
|
||||
constant.ChannelTypeAdvancedCustom: true,
|
||||
constant.ChannelTypeSub2API: true,
|
||||
constant.ChannelTypeNewAPI: true,
|
||||
constant.ChannelTypeVLLM: true,
|
||||
constant.ChannelTypeSGLang: true,
|
||||
constant.ChannelTypeTencent: true,
|
||||
}
|
||||
|
||||
@@ -498,6 +556,9 @@ func reasoningEffortFromRequest(request dto.Request) string {
|
||||
if req != nil && req.GenerationConfig.ThinkingConfig != nil {
|
||||
config := req.GenerationConfig.ThinkingConfig
|
||||
effort = config.ThinkingLevel
|
||||
if canonical, err := kitreasoning.ParseEffort(effort); err == nil {
|
||||
effort = string(canonical)
|
||||
}
|
||||
if effort == "" && config.ThinkingBudget != nil {
|
||||
effort = string(kitreasoning.EffortFromBudget(*config.ThinkingBudget))
|
||||
}
|
||||
@@ -551,10 +612,12 @@ func genBaseRelayInfo(c *gin.Context, request dto.Request) *RelayInfo {
|
||||
|
||||
OriginModelName: originModelName,
|
||||
|
||||
TokenId: common.GetContextKeyInt(c, constant.ContextKeyTokenId),
|
||||
TokenKey: common.GetContextKeyString(c, constant.ContextKeyTokenKey),
|
||||
TokenUnlimited: common.GetContextKeyBool(c, constant.ContextKeyTokenUnlimited),
|
||||
TokenGroup: tokenGroup,
|
||||
TokenId: common.GetContextKeyInt(c, constant.ContextKeyTokenId),
|
||||
TokenKey: common.GetContextKeyString(c, constant.ContextKeyTokenKey),
|
||||
TokenUnlimited: common.GetContextKeyBool(c, constant.ContextKeyTokenUnlimited),
|
||||
TokenGroup: tokenGroup,
|
||||
AstraAutoSwitch: common.GetContextKeyBool(c, constant.ContextKeyTokenAstraAutoSwitch),
|
||||
AstraDetectionDisabled: common.GetContextKeyBool(c, constant.ContextKeyTokenAstraDetectionDisabled),
|
||||
|
||||
isFirstResponse: true,
|
||||
RelayMode: relayconstant.Path2RelayMode(c.Request.URL.Path),
|
||||
|
||||
@@ -6,6 +6,8 @@ import (
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
"github.com/QuantumNous/new-api/relaykit/relayconvert"
|
||||
"github.com/QuantumNous/new-api/relaykit/relayconvert/convmeta"
|
||||
@@ -147,6 +149,24 @@ func TestGenRelayInfoCapturesRequestReasoningEffort(t *testing.T) {
|
||||
}},
|
||||
expected: "low",
|
||||
},
|
||||
{
|
||||
name: "Gemini uppercase enum thinking level is canonicalized",
|
||||
path: "/v1beta/models/gemini-3.7-flash:generateContent",
|
||||
relayFormat: types.RelayFormatGemini,
|
||||
request: &dto.GeminiChatRequest{GenerationConfig: dto.GeminiChatGenerationConfig{
|
||||
ThinkingConfig: &dto.GeminiThinkingConfig{ThinkingLevel: " MEDIUM "},
|
||||
}},
|
||||
expected: "medium",
|
||||
},
|
||||
{
|
||||
name: "Gemini unknown thinking level is recorded as sent",
|
||||
path: "/v1beta/models/gemini-3.7-flash:generateContent",
|
||||
relayFormat: types.RelayFormatGemini,
|
||||
request: &dto.GeminiChatRequest{GenerationConfig: dto.GeminiChatGenerationConfig{
|
||||
ThinkingConfig: &dto.GeminiThinkingConfig{ThinkingLevel: "ULTRA"},
|
||||
}},
|
||||
expected: "ULTRA",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
@@ -195,6 +215,84 @@ func TestInitChannelMetaRestoresRequestReasoningEffortForRetry(t *testing.T) {
|
||||
assert.Equal(t, "max", info.ReasoningEffort)
|
||||
}
|
||||
|
||||
func TestInitChannelMetaAppliesAdvancedCustomRoutePassThrough(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
advancedCustom := &dto.AdvancedCustomConfig{
|
||||
Routes: []dto.AdvancedCustomRoute{
|
||||
{
|
||||
IncomingPath: "/v1/chat/completions",
|
||||
UpstreamPath: "/v1/chat/completions",
|
||||
Models: []string{"native-model"},
|
||||
PassThroughBodyEnabled: true,
|
||||
},
|
||||
{
|
||||
IncomingPath: "/v1/chat/completions",
|
||||
UpstreamPath: "/v1/messages",
|
||||
Converter: relayconvert.ConverterOpenAIChatToClaudeMessages,
|
||||
},
|
||||
},
|
||||
}
|
||||
require.NoError(t, advancedCustom.Validate())
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
channelType int
|
||||
channelSetting dto.ChannelSettings
|
||||
model string
|
||||
wantPassThrough bool
|
||||
wantEffort string
|
||||
}{
|
||||
{
|
||||
name: "matched route enables pass-through and drops reasoning effort",
|
||||
channelType: constant.ChannelTypeAdvancedCustom,
|
||||
model: "native-model",
|
||||
wantPassThrough: true,
|
||||
wantEffort: "",
|
||||
},
|
||||
{
|
||||
name: "fallback converter route keeps conversion",
|
||||
channelType: constant.ChannelTypeAdvancedCustom,
|
||||
model: "other-model",
|
||||
wantPassThrough: false,
|
||||
wantEffort: "high",
|
||||
},
|
||||
{
|
||||
name: "channel-level pass-through still applies to converter route",
|
||||
channelType: constant.ChannelTypeAdvancedCustom,
|
||||
channelSetting: dto.ChannelSettings{PassThroughBodyEnabled: true},
|
||||
model: "other-model",
|
||||
wantPassThrough: true,
|
||||
wantEffort: "",
|
||||
},
|
||||
{
|
||||
name: "route flag is ignored for other channel types",
|
||||
channelType: constant.ChannelTypeOpenAI,
|
||||
model: "native-model",
|
||||
wantPassThrough: false,
|
||||
wantEffort: "high",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
ctx.Request = httptest.NewRequest("POST", "/v1/chat/completions", nil)
|
||||
ctx.Set("original_model", tt.model)
|
||||
common.SetContextKey(ctx, constant.ContextKeyChannelType, tt.channelType)
|
||||
common.SetContextKey(ctx, constant.ContextKeyChannelSetting, tt.channelSetting)
|
||||
common.SetContextKey(ctx, constant.ContextKeyChannelOtherSetting, dto.ChannelOtherSettings{AdvancedCustom: advancedCustom})
|
||||
|
||||
info, err := GenRelayInfo(ctx, types.RelayFormatOpenAI, &dto.GeneralOpenAIRequest{Model: tt.model, ReasoningEffort: "high"}, nil)
|
||||
require.NoError(t, err)
|
||||
|
||||
info.InitChannelMeta(ctx)
|
||||
|
||||
assert.Equal(t, tt.wantPassThrough, info.ChannelSetting.PassThroughBodyEnabled)
|
||||
assert.Equal(t, tt.wantEffort, info.ReasoningEffort)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestInitChannelMetaResetsPerAttemptStreamStateAndPreservesRequestState(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
|
||||
@@ -0,0 +1,48 @@
|
||||
package common
|
||||
|
||||
import "strings"
|
||||
|
||||
// ResponseModel records upstream declarations before response conversion. It is
|
||||
// diagnostic only: it must never change routing, pricing, or downstream output.
|
||||
type ResponseModel struct {
|
||||
RequestedModel string `json:"requested_model"`
|
||||
UpstreamModel string `json:"upstream_model"`
|
||||
ReturnedModel string `json:"returned_model"`
|
||||
Mismatch bool `json:"mismatch"`
|
||||
}
|
||||
|
||||
// ObserveResponseModel retains the first differing model for inspection, with
|
||||
// mismatches taking priority over prefix or case-only differences. A later
|
||||
// matching or empty event cannot erase it. Only observe upstream declarations,
|
||||
// never models synthesized by a response converter.
|
||||
func (info *RelayInfo) ObserveResponseModel(model string) {
|
||||
if info == nil || strings.TrimSpace(model) == "" || (info.OriginModelName == "gpt-6-astra" && info.AstraDetectionDisabled) {
|
||||
return
|
||||
}
|
||||
if info.ResponseModel == nil {
|
||||
info.ResponseModel = &ResponseModel{
|
||||
RequestedModel: info.OriginModelName,
|
||||
UpstreamModel: info.GetUpstreamModelName(),
|
||||
}
|
||||
}
|
||||
observation := info.ResponseModel
|
||||
if observation.Mismatch {
|
||||
return
|
||||
}
|
||||
mismatch := true
|
||||
for _, expected := range []string{observation.RequestedModel, observation.UpstreamModel} {
|
||||
if expected != "" && (strings.HasPrefix(model, expected) || strings.EqualFold(model, expected)) {
|
||||
mismatch = false
|
||||
break
|
||||
}
|
||||
}
|
||||
if observation.RequestedModel == "gpt-6-astra" {
|
||||
mismatch = strings.TrimSpace(model) != "gpt-6-astra"
|
||||
}
|
||||
if !mismatch && observation.ReturnedModel != "" &&
|
||||
observation.ReturnedModel != observation.RequestedModel && observation.ReturnedModel != observation.UpstreamModel {
|
||||
return
|
||||
}
|
||||
observation.ReturnedModel = model
|
||||
observation.Mismatch = mismatch
|
||||
}
|
||||
@@ -0,0 +1,358 @@
|
||||
package common
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
appcommon "github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/relaykit/types"
|
||||
)
|
||||
|
||||
const ResponseGuardMaxAttempts = 3
|
||||
const ResponseGuardErrorCode types.ErrorCode = "response_model_rejected"
|
||||
|
||||
// The guard is request-scoped, unlike channel metadata which is reset on retry.
|
||||
type ResponseModelGuard struct {
|
||||
Attempts []*ResponseGuardAttempt `json:"attempts"`
|
||||
Deadline time.Time `json:"-"`
|
||||
}
|
||||
|
||||
type ResponseGuardUsage struct {
|
||||
Input int `json:"input_tokens"`
|
||||
Output int `json:"output_tokens"`
|
||||
Total int `json:"total_tokens"`
|
||||
}
|
||||
|
||||
type ResponseGuardAttempt struct {
|
||||
OnEvent func([]byte, bool) `json:"-"`
|
||||
OnClose func() `json:"-"`
|
||||
closeOnce sync.Once
|
||||
Drain func() `json:"-"`
|
||||
Number int `json:"number"`
|
||||
ChannelID int `json:"channel_id"`
|
||||
Expected string `json:"expected_model"`
|
||||
Returned string `json:"returned_model,omitempty"`
|
||||
Reason string `json:"reason"`
|
||||
ElapsedMS int64 `json:"elapsed_ms"`
|
||||
Usage *ResponseGuardUsage `json:"usage,omitempty"`
|
||||
Delivered bool `json:"delivered"`
|
||||
Status int `json:"status,omitempty"`
|
||||
ResponseID string `json:"-"`
|
||||
Sequence int `json:"-"`
|
||||
started time.Time
|
||||
timeout time.Duration
|
||||
confirmed bool
|
||||
err error
|
||||
cancel context.CancelFunc
|
||||
timer *time.Timer
|
||||
stopClient func() bool
|
||||
}
|
||||
|
||||
func (g *ResponseModelGuard) Last() *ResponseGuardAttempt {
|
||||
if g == nil || len(g.Attempts) == 0 {
|
||||
return nil
|
||||
}
|
||||
return g.Attempts[len(g.Attempts)-1]
|
||||
}
|
||||
|
||||
func (g *ResponseModelGuard) Begin(ctx context.Context, number, channel int) (context.Context, *ResponseGuardAttempt) {
|
||||
upstream, cancel := context.WithCancel(context.WithoutCancel(ctx))
|
||||
a := &ResponseGuardAttempt{Number: number, ChannelID: channel, Expected: "gpt-6-astra", Reason: "pending", started: time.Now(), cancel: cancel}
|
||||
// Model confirmation must allow the same initial wait as an ordinary stream.
|
||||
// A separate 10-second cutoff rejects valid slow upstream responses.
|
||||
a.timeout = time.Duration(constant.StreamingTimeout) * time.Second
|
||||
if a.timeout <= 0 {
|
||||
a.timeout = 300 * time.Second
|
||||
}
|
||||
if !g.Deadline.IsZero() {
|
||||
a.timeout = max(time.Until(g.Deadline), time.Nanosecond)
|
||||
}
|
||||
a.timer = time.AfterFunc(a.timeout, cancel)
|
||||
a.stopClient = context.AfterFunc(ctx, cancel)
|
||||
g.Attempts = append(g.Attempts, a)
|
||||
return upstream, a
|
||||
}
|
||||
|
||||
func (a *ResponseGuardAttempt) Finish(delivered bool) {
|
||||
a.ElapsedMS = time.Since(a.started).Milliseconds()
|
||||
a.Delivered = delivered
|
||||
}
|
||||
|
||||
func (a *ResponseGuardAttempt) Close() {
|
||||
a.closeOnce.Do(func() {
|
||||
if a.OnClose != nil {
|
||||
a.OnClose()
|
||||
}
|
||||
a.timer.Stop()
|
||||
a.stopClient()
|
||||
a.cancel()
|
||||
})
|
||||
}
|
||||
|
||||
func (a *ResponseGuardAttempt) Reject(reason string) error {
|
||||
if a.err == nil {
|
||||
a.Reason = reason
|
||||
a.err = errors.New("response model guard: " + reason)
|
||||
}
|
||||
return a.err
|
||||
}
|
||||
|
||||
func (a *ResponseGuardAttempt) APIError() *types.NewAPIError {
|
||||
if a == nil || a.err == nil {
|
||||
return nil
|
||||
}
|
||||
return types.NewErrorWithStatusCode(errors.New(a.PublicErrorMessage()), ResponseGuardErrorCode, http.StatusConflict, types.ErrOptionWithSkipRetry())
|
||||
}
|
||||
|
||||
func (a *ResponseGuardAttempt) PublicErrorMessage() string {
|
||||
if a != nil && a.Reason == "model_mismatch" && a.Returned != "" {
|
||||
return "Requested model " + a.Expected + ", but the upstream response declared " + a.Returned + ". No verified result could be delivered. Usage from each upstream attempt is charged independently, including rejected results."
|
||||
}
|
||||
if a != nil && a.Expected != "" {
|
||||
return "Could not verify that the upstream response used " + a.Expected + ". Usage from each upstream attempt is charged independently, including rejected results."
|
||||
}
|
||||
return "Could not verify the upstream response model."
|
||||
}
|
||||
|
||||
func (g *ResponseModelGuard) APIError() *types.NewAPIError {
|
||||
if g == nil {
|
||||
return nil
|
||||
}
|
||||
for i := len(g.Attempts) - 1; i >= 0; i-- {
|
||||
attempt := g.Attempts[i]
|
||||
if attempt != nil && attempt.Reason == "model_mismatch" && attempt.Returned != "" {
|
||||
return attempt.APIError()
|
||||
}
|
||||
}
|
||||
last := g.Last()
|
||||
if last == nil {
|
||||
return nil
|
||||
}
|
||||
return last.APIError()
|
||||
}
|
||||
|
||||
// Only read the provider's declaration and numeric usage, never generated text.
|
||||
func (a *ResponseGuardAttempt) observe(payload []byte, stream bool) error {
|
||||
if a.OnEvent != nil {
|
||||
a.OnEvent(payload, stream)
|
||||
}
|
||||
var event struct {
|
||||
Model string `json:"model"`
|
||||
ID string `json:"id"`
|
||||
Sequence int `json:"sequence_number"`
|
||||
Usage *ResponseGuardUsage `json:"usage"`
|
||||
Response *struct {
|
||||
Model string `json:"model"`
|
||||
ID string `json:"id"`
|
||||
Usage *ResponseGuardUsage `json:"usage"`
|
||||
} `json:"response"`
|
||||
}
|
||||
if err := appcommon.Unmarshal(payload, &event); err != nil {
|
||||
return a.Reject("invalid_response")
|
||||
}
|
||||
model, id, usage := event.Model, event.ID, event.Usage
|
||||
if stream {
|
||||
model, id, usage = "", "", nil
|
||||
if event.Response != nil {
|
||||
model, id, usage = event.Response.Model, event.Response.ID, event.Response.Usage
|
||||
}
|
||||
}
|
||||
if usage != nil {
|
||||
a.Usage = usage
|
||||
}
|
||||
a.Sequence = max(a.Sequence, event.Sequence)
|
||||
model = strings.TrimSpace(model)
|
||||
if model == "" {
|
||||
return nil
|
||||
}
|
||||
// Bound diagnostics from an untrusted upstream.
|
||||
a.Returned = string([]rune(model)[:min(len([]rune(model)), 128)])
|
||||
if model != a.Expected {
|
||||
return a.Reject("model_mismatch")
|
||||
}
|
||||
if id != "" {
|
||||
a.ResponseID = string([]rune(id)[:min(len([]rune(id)), 256)])
|
||||
}
|
||||
if !a.confirmed {
|
||||
if !a.timer.Stop() {
|
||||
return a.Reject("confirmation_timeout")
|
||||
}
|
||||
if !a.stopClient() {
|
||||
return a.Reject("client_cancelled")
|
||||
}
|
||||
a.confirmed = true
|
||||
a.Reason = "accepted"
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Confirm buffers only the prefix until a provider model is known. No caller
|
||||
// may commit headers or start a downstream pinger before it returns success.
|
||||
func (a *ResponseGuardAttempt) Confirm(resp *http.Response, ctx context.Context) error {
|
||||
a.Status = resp.StatusCode
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
a.timer.Stop()
|
||||
a.stopClient()
|
||||
a.Reason = "http_error"
|
||||
resp.Body = &guardCloseBody{Reader: resp.Body, body: resp.Body, close: a.Close}
|
||||
return nil
|
||||
}
|
||||
if !strings.HasPrefix(strings.ToLower(resp.Header.Get("Content-Type")), "text/event-stream") {
|
||||
body, err := io.ReadAll(io.LimitReader(resp.Body, (1<<20)+1))
|
||||
if err != nil {
|
||||
return a.readError(ctx, err)
|
||||
}
|
||||
if len(body) > 1<<20 {
|
||||
return a.Reject("confirmation_buffer_limit")
|
||||
}
|
||||
if err := a.observe(body, false); err != nil {
|
||||
return err
|
||||
}
|
||||
if !a.confirmed {
|
||||
return a.Reject("missing_model")
|
||||
}
|
||||
_ = resp.Body.Close()
|
||||
resp.Body = &guardCloseBody{Reader: bytes.NewReader(body), body: io.NopCloser(bytes.NewReader(nil)), close: a.Close}
|
||||
return nil
|
||||
}
|
||||
b := &guardStreamBody{reader: bufio.NewReader(resp.Body), body: resp.Body, attempt: a}
|
||||
a.Drain = func() {
|
||||
b.draining = true
|
||||
for {
|
||||
if _, err := b.readFrame(8 << 20); err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
var prefix bytes.Buffer
|
||||
for !a.confirmed {
|
||||
frame, err := b.readFrame((1 << 20) - prefix.Len())
|
||||
if err != nil {
|
||||
return a.readError(ctx, err)
|
||||
}
|
||||
prefix.Write(frame)
|
||||
}
|
||||
b.pending = prefix.Bytes()
|
||||
resp.Body = b
|
||||
return nil
|
||||
}
|
||||
|
||||
func (a *ResponseGuardAttempt) readError(ctx context.Context, err error) error {
|
||||
if a.err != nil {
|
||||
return a.err
|
||||
}
|
||||
if ctx.Err() != nil {
|
||||
return a.Reject("client_cancelled")
|
||||
}
|
||||
if time.Since(a.started) >= a.timeout {
|
||||
return a.Reject("confirmation_timeout")
|
||||
}
|
||||
if errors.Is(err, io.EOF) {
|
||||
return a.Reject("missing_model")
|
||||
}
|
||||
return a.Reject("response_read_error")
|
||||
}
|
||||
|
||||
type guardCloseBody struct {
|
||||
io.Reader
|
||||
body io.Closer
|
||||
close func()
|
||||
once sync.Once
|
||||
}
|
||||
|
||||
func (b *guardCloseBody) Close() error { b.once.Do(b.close); return b.body.Close() }
|
||||
|
||||
// Validate complete SSE records before exposing their bytes to the existing
|
||||
// scanner, so a late wrong-model event cannot race past the check.
|
||||
type guardStreamBody struct {
|
||||
draining bool
|
||||
reader *bufio.Reader
|
||||
body io.ReadCloser
|
||||
attempt *ResponseGuardAttempt
|
||||
pending []byte
|
||||
once sync.Once
|
||||
}
|
||||
|
||||
func (b *guardStreamBody) Close() error { b.once.Do(b.attempt.Close); return b.body.Close() }
|
||||
func (b *guardStreamBody) Read(p []byte) (int, error) {
|
||||
if len(p) == 0 {
|
||||
return 0, nil
|
||||
}
|
||||
if len(b.pending) == 0 {
|
||||
frame, err := b.readFrame(128 << 20)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
b.pending = frame
|
||||
}
|
||||
n := copy(p, b.pending)
|
||||
b.pending = b.pending[n:]
|
||||
return n, nil
|
||||
}
|
||||
|
||||
func (b *guardStreamBody) readFrame(limit int) ([]byte, error) {
|
||||
var frame, line, data bytes.Buffer
|
||||
for {
|
||||
part, err := b.reader.ReadSlice('\n')
|
||||
if frame.Len()+len(part) > limit {
|
||||
if b.draining {
|
||||
return nil, io.ErrShortBuffer
|
||||
}
|
||||
return nil, b.attempt.Reject("confirmation_buffer_limit")
|
||||
}
|
||||
frame.Write(part)
|
||||
line.Write(part)
|
||||
if errors.Is(err, bufio.ErrBufferFull) {
|
||||
continue
|
||||
}
|
||||
if err != nil && !errors.Is(err, io.EOF) {
|
||||
return nil, err
|
||||
}
|
||||
text := strings.TrimSuffix(strings.TrimSuffix(line.String(), "\n"), "\r")
|
||||
line.Reset()
|
||||
if value, ok := strings.CutPrefix(text, "data:"); ok {
|
||||
if data.Len() > 0 {
|
||||
data.WriteByte('\n')
|
||||
}
|
||||
data.WriteString(strings.TrimPrefix(value, " "))
|
||||
}
|
||||
if text == "" || err != nil {
|
||||
if frame.Len() == 0 {
|
||||
return nil, io.EOF
|
||||
}
|
||||
payload := bytes.TrimSpace(data.Bytes())
|
||||
if len(payload) > 0 && !bytes.Equal(payload, []byte("[DONE]")) {
|
||||
if b.draining {
|
||||
if b.attempt.OnEvent != nil {
|
||||
b.attempt.OnEvent(payload, true)
|
||||
}
|
||||
} else if e := b.attempt.observe(payload, true); e != nil {
|
||||
return nil, e
|
||||
}
|
||||
}
|
||||
if b.draining && bytes.Equal(payload, []byte("[DONE]")) {
|
||||
return nil, io.EOF
|
||||
}
|
||||
if !b.attempt.confirmed && bytes.Equal(payload, []byte("[DONE]")) {
|
||||
return nil, b.attempt.Reject("missing_model")
|
||||
}
|
||||
return frame.Bytes(), nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Detach keeps only this upstream attempt alive while a different channel is
|
||||
// tried; its numeric usage is collected for at most one minute.
|
||||
func (a *ResponseGuardAttempt) Detach() {
|
||||
a.timer.Stop()
|
||||
a.stopClient()
|
||||
a.timer = time.AfterFunc(60*time.Second, a.cancel)
|
||||
}
|
||||
+159
-10
@@ -21,6 +21,18 @@ const (
|
||||
StreamEndReasonPingFail StreamEndReason = "ping_fail"
|
||||
)
|
||||
|
||||
// ResponseOutcome is the protocol-level result of one response, independent of
|
||||
// how the transport ended. Adaptors mark it from the events they already parse.
|
||||
type ResponseOutcome string
|
||||
|
||||
const (
|
||||
ResponseOutcomeUnknown ResponseOutcome = ""
|
||||
ResponseOutcomeCompleted ResponseOutcome = "completed"
|
||||
ResponseOutcomeFailed ResponseOutcome = "failed"
|
||||
ResponseOutcomeIncomplete ResponseOutcome = "incomplete"
|
||||
ResponseOutcomeCancelled ResponseOutcome = "cancelled"
|
||||
)
|
||||
|
||||
const maxStreamErrorEntries = 20
|
||||
|
||||
type StreamErrorEntry struct {
|
||||
@@ -31,11 +43,32 @@ type StreamErrorEntry struct {
|
||||
type StreamStatus struct {
|
||||
EndReason StreamEndReason
|
||||
EndError error
|
||||
endOnce sync.Once
|
||||
endMu sync.Mutex
|
||||
endSet bool
|
||||
|
||||
mu sync.Mutex
|
||||
Errors []StreamErrorEntry
|
||||
ErrorCount int
|
||||
|
||||
response ResponseOutcome
|
||||
errorCode string
|
||||
errorType string
|
||||
errorStatus int
|
||||
incompleteReason string
|
||||
expectsTerminal bool
|
||||
}
|
||||
|
||||
// StreamOutcome holds classification facts only; upstream messages never
|
||||
// enter it because they may contain credentials or request content.
|
||||
type StreamOutcome struct {
|
||||
EndReason StreamEndReason
|
||||
HasErrors bool
|
||||
ExpectsTerminal bool
|
||||
Response ResponseOutcome
|
||||
ErrorCode string
|
||||
ErrorType string
|
||||
ErrorStatus int
|
||||
IncompleteReason string
|
||||
}
|
||||
|
||||
func NewStreamStatus() *StreamStatus {
|
||||
@@ -46,10 +79,30 @@ func (s *StreamStatus) SetEndReason(reason StreamEndReason, err error) {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
s.endOnce.Do(func() {
|
||||
s.EndReason = reason
|
||||
s.endMu.Lock()
|
||||
defer s.endMu.Unlock()
|
||||
if s.endSet {
|
||||
return
|
||||
}
|
||||
s.EndReason = reason
|
||||
s.EndError = err
|
||||
s.endSet = true
|
||||
}
|
||||
|
||||
// SetClientGone preserves the client disconnect as the diagnostic reason when
|
||||
// it races with a normal upstream terminal event. Fatal and handler-stop reasons
|
||||
// remain authoritative.
|
||||
func (s *StreamStatus) SetClientGone(err error) {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
s.endMu.Lock()
|
||||
defer s.endMu.Unlock()
|
||||
if !s.endSet || s.EndReason == StreamEndReasonDone || s.EndReason == StreamEndReasonEOF {
|
||||
s.EndReason = StreamEndReasonClientGone
|
||||
s.EndError = err
|
||||
})
|
||||
s.endSet = true
|
||||
}
|
||||
}
|
||||
|
||||
func (s *StreamStatus) RecordError(msg string) {
|
||||
@@ -67,6 +120,95 @@ func (s *StreamStatus) RecordError(msg string) {
|
||||
}
|
||||
}
|
||||
|
||||
// RequireTerminal declares that the protocol always ends with an explicit
|
||||
// terminal event, so a stream that ends without one was cut short.
|
||||
func (s *StreamStatus) RequireTerminal() {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.expectsTerminal = true
|
||||
}
|
||||
|
||||
// MarkCompleted, MarkIncomplete and MarkCancelled keep the first terminal seen;
|
||||
// MarkFailed always wins because an error after completion is still a failure.
|
||||
func (s *StreamStatus) MarkCompleted() {
|
||||
s.markTerminal(ResponseOutcomeCompleted, "")
|
||||
}
|
||||
|
||||
func (s *StreamStatus) MarkIncomplete(reason string) {
|
||||
s.markTerminal(ResponseOutcomeIncomplete, reason)
|
||||
}
|
||||
|
||||
func (s *StreamStatus) MarkCancelled() {
|
||||
s.markTerminal(ResponseOutcomeCancelled, "")
|
||||
}
|
||||
|
||||
func (s *StreamStatus) markTerminal(outcome ResponseOutcome, incompleteReason string) {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if s.response != ResponseOutcomeUnknown {
|
||||
return
|
||||
}
|
||||
s.response = outcome
|
||||
s.incompleteReason = incompleteReason
|
||||
}
|
||||
|
||||
// MarkFailed records a protocol failure. Empty details never erase details
|
||||
// recorded earlier, so a bare error envelope keeps the structured error.
|
||||
func (s *StreamStatus) MarkFailed(code, errorType string, status int) {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.response = ResponseOutcomeFailed
|
||||
if code != "" {
|
||||
s.errorCode = code
|
||||
}
|
||||
if errorType != "" {
|
||||
s.errorType = errorType
|
||||
}
|
||||
if status != 0 {
|
||||
s.errorStatus = status
|
||||
}
|
||||
}
|
||||
|
||||
func (s *StreamStatus) ResponseOutcome() string {
|
||||
if s == nil {
|
||||
return ""
|
||||
}
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
return string(s.response)
|
||||
}
|
||||
|
||||
func (s *StreamStatus) ResponseFailed() bool {
|
||||
return s.ResponseOutcome() == string(ResponseOutcomeFailed)
|
||||
}
|
||||
|
||||
func (s *StreamStatus) OutcomeSnapshot() StreamOutcome {
|
||||
if s == nil {
|
||||
return StreamOutcome{}
|
||||
}
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
return StreamOutcome{
|
||||
EndReason: s.EndReason,
|
||||
HasErrors: s.ErrorCount > 0,
|
||||
ExpectsTerminal: s.expectsTerminal,
|
||||
Response: s.response,
|
||||
ErrorCode: s.errorCode,
|
||||
ErrorType: s.errorType,
|
||||
ErrorStatus: s.errorStatus,
|
||||
IncompleteReason: s.incompleteReason,
|
||||
}
|
||||
}
|
||||
|
||||
func (s *StreamStatus) HasErrors() bool {
|
||||
if s == nil {
|
||||
return false
|
||||
@@ -89,9 +231,12 @@ func (s *StreamStatus) IsNormalEnd() bool {
|
||||
if s == nil {
|
||||
return true
|
||||
}
|
||||
return s.EndReason == StreamEndReasonDone ||
|
||||
s.EndReason == StreamEndReasonEOF ||
|
||||
s.EndReason == StreamEndReasonHandlerStop
|
||||
s.endMu.Lock()
|
||||
reason := s.EndReason
|
||||
s.endMu.Unlock()
|
||||
return reason == StreamEndReasonDone ||
|
||||
reason == StreamEndReasonEOF ||
|
||||
reason == StreamEndReasonHandlerStop
|
||||
}
|
||||
|
||||
func (s *StreamStatus) Summary() string {
|
||||
@@ -99,9 +244,13 @@ func (s *StreamStatus) Summary() string {
|
||||
return "StreamStatus<nil>"
|
||||
}
|
||||
b := &strings.Builder{}
|
||||
fmt.Fprintf(b, "reason=%s", s.EndReason)
|
||||
if s.EndError != nil {
|
||||
fmt.Fprintf(b, " end_error=%q", s.EndError.Error())
|
||||
s.endMu.Lock()
|
||||
reason := s.EndReason
|
||||
endError := s.EndError
|
||||
s.endMu.Unlock()
|
||||
fmt.Fprintf(b, "reason=%s", reason)
|
||||
if endError != nil {
|
||||
fmt.Fprintf(b, " end_error=%q", endError.Error())
|
||||
}
|
||||
s.mu.Lock()
|
||||
if s.ErrorCount > 0 {
|
||||
|
||||
@@ -20,6 +20,17 @@ func TestStreamStatus_SetEndReason_FirstWins(t *testing.T) {
|
||||
assert.Nil(t, s.EndError)
|
||||
}
|
||||
|
||||
func TestStreamStatus_SetClientGoneOverridesNormalTerminal(t *testing.T) {
|
||||
s := NewStreamStatus()
|
||||
s.SetEndReason(StreamEndReasonDone, nil)
|
||||
s.SetClientGone(fmt.Errorf("context canceled"))
|
||||
assert.Equal(t, StreamEndReasonClientGone, s.EndReason)
|
||||
assert.EqualError(t, s.EndError, "context canceled")
|
||||
|
||||
s.SetEndReason(StreamEndReasonHandlerStop, fmt.Errorf("fatal"))
|
||||
assert.Equal(t, StreamEndReasonClientGone, s.EndReason)
|
||||
}
|
||||
|
||||
func TestStreamStatus_SetEndReason_WithError(t *testing.T) {
|
||||
t.Parallel()
|
||||
s := NewStreamStatus()
|
||||
|
||||
@@ -1,22 +1,138 @@
|
||||
package relay
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/pkg/billingexpr"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
relayconstant "github.com/QuantumNous/new-api/relay/constant"
|
||||
"github.com/QuantumNous/new-api/relay/helper"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
kitreasoning "github.com/QuantumNous/new-api/relaykit/relayconvert/reasoning"
|
||||
"github.com/QuantumNous/new-api/relaykit/types"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/QuantumNous/new-api/setting/model_setting"
|
||||
hosttypes "github.com/QuantumNous/new-api/types"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
type imageReservation struct {
|
||||
held, limit int
|
||||
}
|
||||
|
||||
func (s *imageReservation) Reserve(quota int) error {
|
||||
if quota > s.limit {
|
||||
return errors.New("insufficient image quota")
|
||||
}
|
||||
s.held = max(s.held, quota)
|
||||
return nil
|
||||
}
|
||||
func (s *imageReservation) GetPreConsumedQuota() int { return s.held }
|
||||
func (*imageReservation) Settle(int) error { return nil }
|
||||
func (*imageReservation) Refund(*gin.Context) {}
|
||||
func (*imageReservation) NeedsRefund() bool { return false }
|
||||
|
||||
func TestImageRequestReservesFinalQuantityBeforeUpstream(t *testing.T) {
|
||||
service.InitHttpClient()
|
||||
for _, tc := range []struct {
|
||||
name, body string
|
||||
tiered, passThrough, insufficient, retryToOpenAI bool
|
||||
override any
|
||||
count, status int
|
||||
}{
|
||||
{name: "Ali nested count", body: `{"model":"z-image","n":1,"parameters":{"n":4}}`, count: 4},
|
||||
{name: "Ali empty parameters", body: `{"model":"z-image","n":2,"parameters":{}}`, count: 2},
|
||||
{name: "legacy channel quantity override", body: `{"model":"z-image","n":1}`, override: 4, count: 4},
|
||||
{name: "expression channel quantity override", body: `{"model":"z-image","n":1}`, override: 4, count: 4, tiered: true},
|
||||
{name: "pass-through quantity", body: `{"model":"z-image","n":2,"parameters":{}}`, count: 2, passThrough: true},
|
||||
{name: "zero override rejected", body: `{"model":"z-image","n":1}`, override: 0, status: http.StatusBadRequest},
|
||||
{name: "oversized override rejected", body: `{"model":"z-image","n":1}`, override: 129, status: http.StatusBadRequest},
|
||||
{name: "insufficient reservation blocks upstream", body: `{"model":"z-image","n":1}`, override: 4, count: 4, insufficient: true, status: http.StatusForbidden},
|
||||
{name: "retry drops Ali quantity and surcharge", body: `{"model":"z-image","n":1,"parameters":{"n":4,"prompt_extend":true}}`, count: 1, retryToOpenAI: true},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
received := make(chan []byte, 1)
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
body, err := io.ReadAll(r.Body)
|
||||
if err == nil {
|
||||
received <- body
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusBadGateway)
|
||||
_, _ = io.WriteString(w, `{"error":{"message":"fixture upstream failure","type":"upstream_error"}}`)
|
||||
}))
|
||||
t.Cleanup(upstream.Close)
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", strings.NewReader(tc.body))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
channel := constant.ChannelTypeAli
|
||||
if tc.retryToOpenAI {
|
||||
channel = constant.ChannelTypeOpenAI
|
||||
}
|
||||
common.SetContextKey(c, constant.ContextKeyChannelType, channel)
|
||||
common.SetContextKey(c, constant.ContextKeyChannelBaseUrl, upstream.URL)
|
||||
common.SetContextKey(c, constant.ContextKeyChannelSetting, dto.ChannelSettings{PassThroughBodyEnabled: tc.passThrough})
|
||||
if tc.override != nil {
|
||||
common.SetContextKey(c, constant.ContextKeyChannelParamOverride, map[string]any{"operations": []any{map[string]any{"path": "parameters.n", "mode": "set", "value": tc.override}}})
|
||||
}
|
||||
request, err := helper.GetAndValidOpenAIImageRequest(c, relayconstant.RelayModeImagesGenerations)
|
||||
require.NoError(t, err)
|
||||
reservation := &imageReservation{held: 20000, limit: 500000}
|
||||
if tc.insufficient {
|
||||
reservation.limit = reservation.held
|
||||
}
|
||||
info := &relaycommon.RelayInfo{Request: request, OriginModelName: "z-image", RelayMode: relayconstant.RelayModeImagesGenerations,
|
||||
RequestURLPath: c.Request.URL.Path, Billing: reservation,
|
||||
PriceData: hosttypes.PriceData{UsePrice: true, ModelPrice: 0.04, GroupRatioInfo: hosttypes.GroupRatioInfo{GroupRatio: 1}},
|
||||
}
|
||||
if tc.tiered {
|
||||
expr := `tier("image", fixed(0.04)) * image_count`
|
||||
info.TieredBillingSnapshot = &billingexpr.BillingSnapshot{BillingMode: "tiered_expr", ExprString: expr, ExprHash: billingexpr.ExprHashString(expr), GroupRatio: 1, QuotaPerUnit: common.QuotaPerUnit, EstimatedImageCount: common.GetPointer(1)}
|
||||
info.BillingRequestInput = &billingexpr.RequestInput{Body: []byte(tc.body), ImageCount: common.GetPointer(1)}
|
||||
info.PriceData.UsePrice = false
|
||||
}
|
||||
if tc.retryToOpenAI {
|
||||
reservation.held = 160000
|
||||
info.PriceData.AddOtherRatio("n", 4)
|
||||
info.PriceData.AddOtherRatio("prompt_extend", 2)
|
||||
info.BillingImageCount = common.GetPointer(3)
|
||||
}
|
||||
apiErr := ImageHelper(c, info)
|
||||
require.NotNil(t, apiErr)
|
||||
if tc.status != 0 {
|
||||
assert.Equal(t, tc.status, apiErr.StatusCode)
|
||||
assert.Empty(t, received, "no upstream submission before quantity validation and reservation")
|
||||
return
|
||||
}
|
||||
assert.Equal(t, http.StatusBadGateway, apiErr.StatusCode)
|
||||
require.Len(t, received, 1)
|
||||
body := <-received
|
||||
path := "parameters.n"
|
||||
if tc.retryToOpenAI {
|
||||
path = "n"
|
||||
}
|
||||
assert.Equal(t, int64(tc.count), gjson.GetBytes(body, path).Int())
|
||||
assert.Equal(t, tc.count*20000, info.PriceData.QuotaToPreConsume)
|
||||
assert.GreaterOrEqual(t, reservation.held, info.PriceData.QuotaToPreConsume)
|
||||
assert.Nil(t, info.BillingImageCount)
|
||||
if tc.tiered {
|
||||
assert.Equal(t, tc.body, string(info.BillingRequestInput.Body))
|
||||
assert.Equal(t, 1, *info.BillingRequestInput.ImageCount)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestOptInSafeToolLossRejectedAsBadRequestWithAdminDiagnostics(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
@@ -96,3 +212,90 @@ func hasHostDiagnosticCode(diagnostics []types.ConversionDiagnostic, code string
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func TestGeminiThinkingControlsConvertBestEffort(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
settings := model_setting.GetGeminiSettings()
|
||||
originalAdapter := settings.ThinkingAdapterEnabled
|
||||
t.Cleanup(func() { settings.ThinkingAdapterEnabled = originalAdapter })
|
||||
|
||||
t.Run("openai thinking_budget on gemini 3 becomes thinkingLevel", func(t *testing.T) {
|
||||
settings.ThinkingAdapterEnabled = false
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
|
||||
req := &dto.GeneralOpenAIRequest{
|
||||
Model: "gemini-3.1-pro-preview-thinking",
|
||||
Messages: []dto.Message{{Role: "user", Content: "hello"}},
|
||||
ExtraBody: []byte(`{"google":{"thinking_config":{"thinking_budget":8192}}}`),
|
||||
}
|
||||
info := &relaycommon.RelayInfo{
|
||||
RelayFormat: types.RelayFormatOpenAI,
|
||||
OriginModelName: "gemini-3.1-pro-preview-thinking",
|
||||
Request: req,
|
||||
ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "gemini-3.1-pro-preview"},
|
||||
}
|
||||
require.NoError(t, helper.ApplyReasoningModelSuffix(c, info, req))
|
||||
|
||||
result, err := service.ConvertRequest(c, info, types.RelayFormatGemini, req)
|
||||
require.NoError(t, err)
|
||||
converted, ok := result.Value.(*dto.GeminiChatRequest)
|
||||
require.True(t, ok)
|
||||
require.NotNil(t, converted.GenerationConfig.ThinkingConfig)
|
||||
assert.Equal(t, "medium", converted.GenerationConfig.ThinkingConfig.ThinkingLevel)
|
||||
assert.Nil(t, converted.GenerationConfig.ThinkingConfig.ThinkingBudget)
|
||||
assert.True(t, hasHostDiagnosticCode(info.ConversionDiagnostics(), "gemini_budget_to_level"))
|
||||
})
|
||||
|
||||
t.Run("thinking alias with native thinking_level keeps the level", func(t *testing.T) {
|
||||
settings.ThinkingAdapterEnabled = true
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
|
||||
req := &dto.GeneralOpenAIRequest{
|
||||
Model: "gemini-3.1-pro-preview-thinking",
|
||||
Messages: []dto.Message{{Role: "user", Content: "hello"}},
|
||||
ExtraBody: []byte(`{"google":{"thinking_config":{"thinking_level":"low"}}}`),
|
||||
}
|
||||
info := &relaycommon.RelayInfo{
|
||||
RelayFormat: types.RelayFormatOpenAI,
|
||||
OriginModelName: "gemini-3.1-pro-preview-thinking",
|
||||
Request: req,
|
||||
ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "gemini-3.1-pro-preview-thinking"},
|
||||
}
|
||||
require.NoError(t, helper.ApplyReasoningModelSuffix(c, info, req))
|
||||
assert.Equal(t, "gemini-3.1-pro-preview", info.UpstreamModelName)
|
||||
|
||||
result, err := service.ConvertRequest(c, info, types.RelayFormatGemini, req)
|
||||
require.NoError(t, err)
|
||||
converted, ok := result.Value.(*dto.GeminiChatRequest)
|
||||
require.True(t, ok)
|
||||
require.NotNil(t, converted.GenerationConfig.ThinkingConfig)
|
||||
assert.Equal(t, "low", converted.GenerationConfig.ThinkingConfig.ThinkingLevel)
|
||||
assert.Nil(t, converted.GenerationConfig.ThinkingConfig.ThinkingBudget)
|
||||
})
|
||||
|
||||
t.Run("gemini thinkingBudget converts to openai reasoning_effort", func(t *testing.T) {
|
||||
settings.ThinkingAdapterEnabled = false
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1beta/models/gemini-3.1-pro-preview:generateContent", nil)
|
||||
budget := 8192
|
||||
req := &dto.GeminiChatRequest{
|
||||
Contents: []dto.GeminiChatContent{{Role: "user", Parts: []dto.GeminiPart{{Text: "hello"}}}},
|
||||
GenerationConfig: dto.GeminiChatGenerationConfig{
|
||||
ThinkingConfig: &dto.GeminiThinkingConfig{ThinkingBudget: &budget},
|
||||
},
|
||||
}
|
||||
info := &relaycommon.RelayInfo{
|
||||
RelayFormat: types.RelayFormatGemini,
|
||||
OriginModelName: "gemini-3.1-pro-preview",
|
||||
Request: req,
|
||||
ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "gpt-5"},
|
||||
}
|
||||
|
||||
result, err := service.ConvertRequest(c, info, types.RelayFormatOpenAI, req)
|
||||
require.NoError(t, err)
|
||||
converted, ok := result.Value.(*dto.GeneralOpenAIRequest)
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, "medium", converted.ReasoningEffort)
|
||||
assert.True(t, hasHostDiagnosticCode(info.ConversionDiagnostics(), "gemini_budget_to_level"))
|
||||
})
|
||||
}
|
||||
@@ -5,6 +5,7 @@ import (
|
||||
"strings"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/pkg/billingexpr"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
@@ -33,6 +34,39 @@ func ResolveIncomingBillingExprRequestInput(c *gin.Context, info *relaycommon.Re
|
||||
return input, nil
|
||||
}
|
||||
|
||||
// ResolveImageBillingRequestInput freezes only the validated scalar image
|
||||
// parameters needed by pricing. Image files, prompts and base64 payloads are
|
||||
// deliberately excluded, including for multipart edits.
|
||||
func ResolveImageBillingRequestInput(c *gin.Context, info *relaycommon.RelayInfo, input billingexpr.RequestInput) (billingexpr.RequestInput, error) {
|
||||
request, ok := info.Request.(*dto.ImageRequest)
|
||||
if !ok {
|
||||
return input, nil
|
||||
}
|
||||
channelType := common.GetContextKeyInt(c, constant.ContextKeyChannelType)
|
||||
if info.ChannelMeta != nil {
|
||||
channelType = info.ChannelType
|
||||
}
|
||||
count, err := request.ImageCount(channelType == constant.ChannelTypeAli)
|
||||
if err != nil {
|
||||
return input, err
|
||||
}
|
||||
topLevelCount, err := request.ImageCount(false)
|
||||
if err != nil {
|
||||
return input, err
|
||||
}
|
||||
body := map[string]any{"model": request.Model, "n": topLevelCount, "size": request.Size, "quality": request.Quality}
|
||||
if request.BillingParameters != nil {
|
||||
body["parameters"] = request.BillingParameters
|
||||
}
|
||||
encoded, err := common.Marshal(body)
|
||||
if err != nil {
|
||||
return input, err
|
||||
}
|
||||
input.Body = encoded
|
||||
input.ImageCount = &count
|
||||
return input, nil
|
||||
}
|
||||
|
||||
func BuildBillingExprRequestInputFromRequest(request dto.Request, headers map[string]string) (billingexpr.RequestInput, error) {
|
||||
input := billingexpr.RequestInput{
|
||||
Headers: cloneStringMap(headers),
|
||||
@@ -64,6 +98,10 @@ func cloneRequestInput(src billingexpr.RequestInput) billingexpr.RequestInput {
|
||||
input := billingexpr.RequestInput{
|
||||
Headers: cloneStringMap(src.Headers),
|
||||
}
|
||||
if src.ImageCount != nil {
|
||||
count := *src.ImageCount
|
||||
input.ImageCount = &count
|
||||
}
|
||||
if len(src.Body) > 0 {
|
||||
input.Body = append([]byte(nil), src.Body...)
|
||||
}
|
||||
|
||||
@@ -70,3 +70,25 @@ func TestMaxTokensBounds(t *testing.T) {
|
||||
require.Contains(t, err.Error(), "max_output_tokens is invalid")
|
||||
})
|
||||
}
|
||||
|
||||
func TestSGLangMinTokensBounds(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
value string
|
||||
valid bool
|
||||
}{
|
||||
{"0", true}, {"8192", true}, {"-1", false}, {"18446744073686646784", false},
|
||||
} {
|
||||
t.Run(tc.value, func(t *testing.T) {
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewBufferString(`{"model":"served-model","messages":[{"role":"user","content":"hi"}],"min_tokens":`+tc.value+`}`))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
request, err := GetAndValidateTextRequest(c, relayconstant.RelayModeChatCompletions)
|
||||
if !tc.valid {
|
||||
require.Error(t, err)
|
||||
return
|
||||
}
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, request.MinTokens)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -11,12 +11,50 @@ import (
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/pkg/billingexpr"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
relayconstant "github.com/QuantumNous/new-api/relay/constant"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestLegacyDalleValidationAndPricesRemainCompatible(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
body, size, quality string
|
||||
ratio float64
|
||||
invalid bool
|
||||
}{
|
||||
{body: `{"model":"dall-e-2"}`, size: "1024x1024", ratio: 1},
|
||||
{body: `{"model":"dall-e"}`, size: "1024x1024", ratio: 1},
|
||||
{body: `{"model":"dall-e-3"}`, size: "1024x1024", quality: "standard", ratio: 1},
|
||||
{body: `{"model":"dall-e-2","size":"256x256"}`, size: "256x256", ratio: 0.4},
|
||||
{body: `{"model":"dall-e-2","size":"512x512","quality":"hd"}`, size: "512x512", quality: "hd", ratio: 0.45},
|
||||
{body: `{"model":"dall-e-3","quality":"hd"}`, size: "1024x1024", quality: "hd", ratio: 2},
|
||||
{body: `{"model":"dall-e-3","size":"1024x1792","quality":"hd"}`, size: "1024x1792", quality: "hd", ratio: 3},
|
||||
{body: `{"model":"dall-e-3","size":"1792x1024"}`, size: "1792x1024", quality: "standard", ratio: 2},
|
||||
{body: `{"model":"dall-e-3","size":"256x256"}`, invalid: true},
|
||||
{body: `{"model":"dall-e-2","size":"1024x1792"}`, invalid: true},
|
||||
} {
|
||||
t.Run(tc.body, func(t *testing.T) {
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewBufferString(tc.body))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
request, err := GetAndValidOpenAIImageRequest(c, relayconstant.RelayModeImagesGenerations)
|
||||
if tc.invalid {
|
||||
require.Error(t, err)
|
||||
return
|
||||
}
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, tc.size, request.Size)
|
||||
assert.Equal(t, tc.quality, request.Quality)
|
||||
assert.Equal(t, tc.ratio, request.GetTokenCountMeta().ImagePriceRatio)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestGetAndValidOpenAIImageRequestMultipartStream verifies multipart image
|
||||
// edit parsing: the stream field is parsed and validated, and the request body
|
||||
// stays replayable for the upstream request.
|
||||
@@ -61,6 +99,11 @@ func TestGetAndValidOpenAIImageRequestMultipartStream(t *testing.T) {
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "true", url.Values(form.Value).Get("stream"))
|
||||
require.Len(t, form.File["image"], 1)
|
||||
billing, err := ResolveImageBillingRequestInput(c, &relaycommon.RelayInfo{Request: req}, billingexpr.RequestInput{})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 1, *billing.ImageCount)
|
||||
require.NotContains(t, string(billing.Body), "fake image")
|
||||
require.NotContains(t, string(billing.Body), "edit this image")
|
||||
})
|
||||
|
||||
t.Run("invalid stream value is rejected", func(t *testing.T) {
|
||||
@@ -72,6 +115,44 @@ func TestGetAndValidOpenAIImageRequestMultipartStream(t *testing.T) {
|
||||
})
|
||||
}
|
||||
|
||||
func TestImageBillingRequestUsesValidatedProviderCount(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
body string
|
||||
channel, count int
|
||||
invalid bool
|
||||
}{
|
||||
{`{"model":"z-image","n":2,"parameters":{"n":3,"prompt_extend":true}}`, constant.ChannelTypeAli, 3, false},
|
||||
{`{"model":"gpt-image-2","n":2,"parameters":{"n":3}}`, constant.ChannelTypeOpenAI, 2, false},
|
||||
{`{"model":"z-image","n":2,"parameters":{"n":129}}`, constant.ChannelTypeAli, 0, true},
|
||||
{`{"model":"z-image","parameters":{"n":-1}}`, constant.ChannelTypeAli, 0, true},
|
||||
{`{"model":"z-image","n":2,"parameters":{}}`, constant.ChannelTypeAli, 2, false},
|
||||
{`{"model":"z-image","n":2,"parameters":{"n":null}}`, constant.ChannelTypeAli, 2, false},
|
||||
{`{"model":"z-image","n":2,"parameters":{"n":0}}`, constant.ChannelTypeAli, 0, true},
|
||||
{`{"model":"z-image","parameters":{"n":1.5}}`, constant.ChannelTypeAli, 0, true},
|
||||
{`{"model":"z-image","parameters":{"n":18446744073686646784}}`, constant.ChannelTypeAli, 0, true},
|
||||
{`{"model":"z-image","parameters":{"n":128}}`, constant.ChannelTypeAli, 128, false},
|
||||
{`{"model":"z-image","n":0}`, constant.ChannelTypeAli, 1, false},
|
||||
{`{"model":"gpt-image-2","n":2,"parameters":{"n":0}}`, constant.ChannelTypeOpenAI, 2, false},
|
||||
} {
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewBufferString(tc.body))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
common.SetContextKey(c, constant.ContextKeyChannelType, tc.channel)
|
||||
request, err := GetAndValidOpenAIImageRequest(c, relayconstant.RelayModeImagesGenerations)
|
||||
if tc.invalid {
|
||||
require.Error(t, err)
|
||||
continue
|
||||
}
|
||||
require.NoError(t, err)
|
||||
input, err := ResolveImageBillingRequestInput(c, &relaycommon.RelayInfo{Request: request}, billingexpr.RequestInput{})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, tc.count, *input.ImageCount)
|
||||
cost, _, err := billingexpr.RunExprWithRequest(`tier("image", fixed(0.04)) * image_count`, billingexpr.TokenParams{}, input)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, float64(tc.count)*40000, cost)
|
||||
}
|
||||
}
|
||||
|
||||
// TestGetAndValidOpenAIImageRequestNBounds guards the billing invariant that
|
||||
// the image generation count can never reach quota calculation with a value
|
||||
// large enough to overflow int64 into a negative charge.
|
||||
|
||||
+52
-2
@@ -5,10 +5,12 @@ import (
|
||||
"strings"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/logger"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/pkg/billingexpr"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
"github.com/QuantumNous/new-api/relaykit/relayconvert/reasoning"
|
||||
"github.com/QuantumNous/new-api/relaykit/types"
|
||||
"github.com/QuantumNous/new-api/setting/billing_setting"
|
||||
@@ -132,10 +134,16 @@ func ModelPriceHelper(c *gin.Context, info *relaycommon.RelayInfo, promptTokens
|
||||
return hosttypes.PriceData{}, err
|
||||
}
|
||||
preConsumedQuota = quota
|
||||
if _, image := info.Request.(*dto.ImageRequest); image {
|
||||
info.ImageQuotaBeforeGroup = float64(preConsumedTokens) * modelRatio
|
||||
}
|
||||
} else {
|
||||
if meta.ImagePriceRatio != 0 {
|
||||
modelPrice = modelPrice * meta.ImagePriceRatio
|
||||
}
|
||||
if _, image := info.Request.(*dto.ImageRequest); image {
|
||||
info.ImageQuotaBeforeGroup = modelPrice * common.QuotaPerUnit
|
||||
}
|
||||
}
|
||||
|
||||
// check if free model pre-consume is disabled
|
||||
@@ -177,6 +185,36 @@ func ModelPriceHelper(c *gin.Context, info *relaycommon.RelayInfo, promptTokens
|
||||
for name, ratio := range meta.BillingRatios {
|
||||
priceData.AddOtherRatio(name, ratio)
|
||||
}
|
||||
}
|
||||
if request, image := info.Request.(*dto.ImageRequest); image {
|
||||
channelType := common.GetContextKeyInt(c, constant.ContextKeyChannelType)
|
||||
count, err := request.ImageCount(channelType == constant.ChannelTypeAli)
|
||||
if err != nil {
|
||||
return hosttypes.PriceData{}, err
|
||||
}
|
||||
if usePrice || channelType == constant.ChannelTypeAli {
|
||||
priceData.AddOtherRatio("n", float64(count))
|
||||
}
|
||||
if channelType == constant.ChannelTypeAli && request.BillingParameters != nil && request.BillingParameters.PromptExtend != nil && *request.BillingParameters.PromptExtend {
|
||||
// Resolve only routing identity; do not initialize ChannelMeta on the
|
||||
// real request, which also distinguishes the first channel attempt.
|
||||
mapped := &relaycommon.RelayInfo{OriginModelName: info.OriginModelName, ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: info.OriginModelName}}
|
||||
if err := ModelMappedHelper(c, mapped, nil); err != nil {
|
||||
return hosttypes.PriceData{}, err
|
||||
}
|
||||
if strings.Contains(mapped.UpstreamModelName, "z-image") {
|
||||
priceData.AddOtherRatio("prompt_extend", common.ZImagePromptExtendMultiplier)
|
||||
}
|
||||
}
|
||||
if !usePrice {
|
||||
quota, err := common.QuotaFromFloatStrict(priceData.ApplyOtherRatiosToFloat(info.ImageQuotaBeforeGroup * groupRatioInfo.GroupRatio))
|
||||
if err != nil {
|
||||
return hosttypes.PriceData{}, err
|
||||
}
|
||||
priceData.QuotaToPreConsume = quota
|
||||
}
|
||||
}
|
||||
if usePrice {
|
||||
quotaToPreConsume := priceData.ApplyOtherRatiosToFloat(modelPrice * common.QuotaPerUnit * groupRatioInfo.GroupRatio)
|
||||
quota, err := common.QuotaFromFloatStrict(quotaToPreConsume)
|
||||
if err != nil {
|
||||
@@ -324,6 +362,10 @@ func modelPriceHelperTiered(c *gin.Context, info *relaycommon.RelayInfo, billing
|
||||
if !ok {
|
||||
return hosttypes.PriceData{}, fmt.Errorf("model %s is configured as tiered_expr but has no billing expression", billingModelName)
|
||||
}
|
||||
exprHash := billingexpr.ExprHashString(exprStr)
|
||||
if info.RelayFormat == types.RelayFormatOpenAIRealtime && billingexpr.UsesFixedPricingByHash(exprStr, exprHash) {
|
||||
return hosttypes.PriceData{}, fmt.Errorf("fixed pricing is not supported for Realtime requests")
|
||||
}
|
||||
|
||||
estimatedCompletionTokens := meta.MaxTokens
|
||||
if estimatedCompletionTokens == 0 && groupRatioInfo.GroupRatio != 0 {
|
||||
@@ -334,8 +376,14 @@ func modelPriceHelperTiered(c *gin.Context, info *relaycommon.RelayInfo, billing
|
||||
if err != nil {
|
||||
return hosttypes.PriceData{}, err
|
||||
}
|
||||
if billingexpr.UsedVarsByHash(exprStr, exprHash)["image_count"] {
|
||||
requestInput, err = ResolveImageBillingRequestInput(c, info, requestInput)
|
||||
if err != nil {
|
||||
return hosttypes.PriceData{}, err
|
||||
}
|
||||
}
|
||||
|
||||
rawCost, trace, err := billingexpr.RunExprWithRequest(exprStr, billingexpr.TokenParams{
|
||||
rawCost, trace, err := billingexpr.RunExprByHashWithRequest(exprStr, exprHash, billingexpr.TokenParams{
|
||||
P: float64(promptTokens),
|
||||
C: float64(estimatedCompletionTokens),
|
||||
Len: float64(promptTokens),
|
||||
@@ -359,8 +407,8 @@ func modelPriceHelperTiered(c *gin.Context, info *relaycommon.RelayInfo, billing
|
||||
}
|
||||
}
|
||||
|
||||
exprHash := billingexpr.ExprHashString(exprStr)
|
||||
snapshot := &billingexpr.BillingSnapshot{
|
||||
EstimatedImageCount: trace.ImageCount,
|
||||
BillingMode: billing_setting.BillingModeTieredExpr,
|
||||
ModelName: billingModelName,
|
||||
ExprString: exprStr,
|
||||
@@ -371,6 +419,8 @@ func modelPriceHelperTiered(c *gin.Context, info *relaycommon.RelayInfo, billing
|
||||
EstimatedQuotaBeforeGroup: quotaBeforeGroup,
|
||||
EstimatedQuotaAfterGroup: preConsumedQuota,
|
||||
EstimatedTier: trace.MatchedTier,
|
||||
EstimatedBillingUnit: trace.BillingUnit,
|
||||
EstimatedFixedPrice: trace.FixedPrice,
|
||||
QuotaPerUnit: common.QuotaPerUnit,
|
||||
ExprVersion: billingexpr.ExprVersion(exprStr),
|
||||
}
|
||||
|
||||
@@ -68,6 +68,47 @@ func TestModelPriceHelperTieredUsesPreloadedRequestInput(t *testing.T) {
|
||||
require.Equal(t, common.QuotaPerUnit, info.TieredBillingSnapshot.QuotaPerUnit)
|
||||
}
|
||||
|
||||
func TestFixedPricePreConsumeAndRealtimeRejection(t *testing.T) {
|
||||
saved := map[string]string{}
|
||||
require.NoError(t, config.GlobalConfig.SaveToDB(func(key, value string) error {
|
||||
saved[key] = value
|
||||
return nil
|
||||
}))
|
||||
t.Cleanup(func() { require.NoError(t, config.GlobalConfig.LoadFromDB(saved)) })
|
||||
require.NoError(t, config.GlobalConfig.LoadFromDB(map[string]string{
|
||||
"billing_setting.billing_mode": `{"fixed-test":"tiered_expr"}`,
|
||||
"billing_setting.billing_expr": `{"fixed-test":"len <= 32000 ? tier(\"short\", fixed(0.01)) : tier(\"long\", p * 2)"}`,
|
||||
"group_ratio_setting.group_ratio": `{"default":1.5}`,
|
||||
}))
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
format types.RelayFormat
|
||||
prompt int
|
||||
wantError bool
|
||||
}{
|
||||
{"HTTP charges once", types.RelayFormatOpenAI, 0, false},
|
||||
{"Realtime rejects even unselected fixed branch", types.RelayFormatOpenAIRealtime, 50000, true},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
|
||||
info := &relaycommon.RelayInfo{OriginModelName: "fixed-test", UserGroup: "default", UsingGroup: "default", RelayFormat: tc.format, BillingRequestInput: &billingexpr.RequestInput{}}
|
||||
price, err := ModelPriceHelper(ctx, info, tc.prompt, &types.TokenCountMeta{})
|
||||
if tc.wantError {
|
||||
require.ErrorContains(t, err, "Realtime")
|
||||
assert.Nil(t, info.TieredBillingSnapshot)
|
||||
return
|
||||
}
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 7500, price.QuotaToPreConsume)
|
||||
require.NotNil(t, info.TieredBillingSnapshot)
|
||||
assert.Equal(t, billingexpr.BillingUnitRequest, info.TieredBillingSnapshot.EstimatedBillingUnit)
|
||||
require.NotNil(t, info.TieredBillingSnapshot.EstimatedFixedPrice)
|
||||
assert.Equal(t, 0.01, *info.TieredBillingSnapshot.EstimatedFixedPrice)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestModelPriceHelperTieredPreConsumeMaxTokensFallback(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
|
||||
@@ -59,10 +59,11 @@ func ApplyReasoningModelSuffix(c *gin.Context, info *relaycommon.RelayInfo, outb
|
||||
}
|
||||
|
||||
if selected.hasThinking {
|
||||
explicit, err := explicitIntentFromRequest(info.Request)
|
||||
explicit, explicitDiagnostics, err := explicitIntentFromRequest(info.Request)
|
||||
if err != nil {
|
||||
return reasoning.AsClientError(err)
|
||||
}
|
||||
diagnostics = append(diagnostics, explicitDiagnostics...)
|
||||
diagnostics = append(diagnostics, modifierRequestOverrideDiagnostics(explicit, selected.intent)...)
|
||||
if selected.intent.IncludeThoughts == nil {
|
||||
selected.intent.IncludeThoughts = explicit.IncludeThoughts
|
||||
@@ -215,7 +216,7 @@ func parseHostModelSuffix(name string, opts *convmeta.Options) (string, reasonin
|
||||
)
|
||||
}
|
||||
|
||||
func explicitIntentFromRequest(req dto.Request) (reasoning.Intent, error) {
|
||||
func explicitIntentFromRequest(req dto.Request) (reasoning.Intent, []types.ConversionDiagnostic, error) {
|
||||
switch r := req.(type) {
|
||||
case *dto.ClaudeRequest:
|
||||
return reasoning.FromClaude(r)
|
||||
@@ -226,6 +227,6 @@ func explicitIntentFromRequest(req dto.Request) (reasoning.Intent, error) {
|
||||
case *dto.OpenAIResponsesRequest:
|
||||
return reasoning.FromOpenAIResponses(r)
|
||||
default:
|
||||
return reasoning.Intent{}, nil
|
||||
return reasoning.Intent{}, nil, nil
|
||||
}
|
||||
}
|
||||
@@ -1,7 +1,10 @@
|
||||
package helper
|
||||
|
||||
import (
|
||||
"errors"
|
||||
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/relaykit/types"
|
||||
)
|
||||
|
||||
// StreamResult is passed to each dataHandler invocation, providing methods
|
||||
@@ -23,13 +26,15 @@ func (r *StreamResult) Error(err error) {
|
||||
return
|
||||
}
|
||||
r.status.RecordError(err.Error())
|
||||
var apiErr *types.NewAPIError
|
||||
if errors.As(err, &apiErr) {
|
||||
r.status.MarkFailed(string(apiErr.GetErrorCode()), apiErr.ToOpenAIError().Type, apiErr.StatusCode)
|
||||
}
|
||||
}
|
||||
|
||||
// Stop records a fatal error and marks the stream to stop after this chunk.
|
||||
func (r *StreamResult) Stop(err error) {
|
||||
if err != nil {
|
||||
r.status.RecordError(err.Error())
|
||||
}
|
||||
r.Error(err)
|
||||
r.status.SetEndReason(relaycommon.StreamEndReasonHandlerStop, err)
|
||||
r.stopped = true
|
||||
}
|
||||
|
||||
@@ -30,7 +30,8 @@ const (
|
||||
// unconditional wg.Wait() in cleanup can always finish. Without it, a slow
|
||||
// but connected client (full TCP buffer, no server WriteTimeout) could hang
|
||||
// the handler forever.
|
||||
streamWriteTimeout = 30 * time.Second
|
||||
streamWriteTimeout = 30 * time.Second
|
||||
clientGoneDrainTimeout = 2 * time.Minute
|
||||
)
|
||||
|
||||
func getScannerBufferSize() int {
|
||||
@@ -40,9 +41,15 @@ func getScannerBufferSize() int {
|
||||
return DefaultMaxScannerBufferSize
|
||||
}
|
||||
|
||||
func NewStreamScanner(reader io.Reader) *bufio.Scanner {
|
||||
// NewStreamScanner shares relay scanner configuration. Callers buffering bounded
|
||||
// task state may additionally cap a line without increasing the configured limit.
|
||||
func NewStreamScanner(reader io.Reader, maxBytes ...int) *bufio.Scanner {
|
||||
limit := getScannerBufferSize()
|
||||
if len(maxBytes) > 0 && maxBytes[0] > 0 {
|
||||
limit = min(limit, maxBytes[0])
|
||||
}
|
||||
scanner := bufio.NewScanner(reader)
|
||||
scanner.Buffer(make([]byte, InitialScannerBufferSize), getScannerBufferSize())
|
||||
scanner.Buffer(make([]byte, min(InitialScannerBufferSize, limit)), limit)
|
||||
return scanner
|
||||
}
|
||||
|
||||
@@ -121,11 +128,11 @@ func StreamScannerHandler(c *gin.Context, resp *http.Response, info *relaycommon
|
||||
logger.LogDebug(c, "streaming timeout seconds: %d", int64(streamingTimeout.Seconds()))
|
||||
logger.LogDebug(c, "ping interval seconds: %d", int64(pingInterval.Seconds()))
|
||||
|
||||
cleanup := func() {
|
||||
cleanup := func(closeUpstream bool) {
|
||||
cleanupOnce.Do(func() {
|
||||
cancel()
|
||||
stop()
|
||||
if resp.Body != nil {
|
||||
if closeUpstream && resp.Body != nil {
|
||||
_ = resp.Body.Close()
|
||||
}
|
||||
|
||||
@@ -138,7 +145,9 @@ func StreamScannerHandler(c *gin.Context, resp *http.Response, info *relaycommon
|
||||
})
|
||||
}
|
||||
// Ensure gin.Context is not returned to Gin's pool while any stream goroutine can still use it.
|
||||
defer cleanup()
|
||||
defer func() {
|
||||
cleanup(true)
|
||||
}()
|
||||
|
||||
scanner.Split(bufio.ScanLines)
|
||||
copyCodexSSEHeaders(c, resp)
|
||||
@@ -289,19 +298,44 @@ func StreamScannerHandler(c *gin.Context, resp *http.Response, info *relaycommon
|
||||
info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonEOF, nil)
|
||||
})
|
||||
|
||||
// 客户端断开后仍读取上游终止事件中的 usage,避免上游已计费而
|
||||
// NewAPI 无法结算。各写入函数会自行跳过已断开的客户端。
|
||||
drainClientGone := func(err error) {
|
||||
info.StreamStatus.SetClientGone(err)
|
||||
logger.LogInfo(c, "client disconnected; draining upstream for billing usage")
|
||||
drainTimer := time.NewTimer(clientGoneDrainTimeout)
|
||||
select {
|
||||
case <-stopChan:
|
||||
case <-drainTimer.C:
|
||||
logger.LogError(c, "client_gone upstream drain timeout; usage may be unavailable")
|
||||
}
|
||||
if !drainTimer.Stop() {
|
||||
select {
|
||||
case <-drainTimer.C:
|
||||
default:
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 主循环等待完成或超时
|
||||
select {
|
||||
case <-ticker.C:
|
||||
info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonTimeout, nil)
|
||||
if err := c.Request.Context().Err(); err != nil {
|
||||
drainClientGone(err)
|
||||
} else {
|
||||
info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonTimeout, nil)
|
||||
}
|
||||
case <-stopChan:
|
||||
// EndReason already set by the goroutine that triggered stopChan
|
||||
// A client disconnect can race with upstream [DONE]. Preserve it as the
|
||||
// diagnostic reason while keeping fatal handler/scanner reasons intact.
|
||||
if err := c.Request.Context().Err(); err != nil {
|
||||
info.StreamStatus.SetClientGone(err)
|
||||
}
|
||||
case <-c.Request.Context().Done():
|
||||
// 客户端断开:立即 cleanup 关闭上游 resp.Body,解除 scanner 阻塞并让上游停止生成,
|
||||
// 避免为已放弃的请求继续消费上游 token。
|
||||
info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonClientGone, c.Request.Context().Err())
|
||||
drainClientGone(c.Request.Context().Err())
|
||||
}
|
||||
|
||||
cleanup()
|
||||
cleanup(true)
|
||||
if info.StreamStatus.IsNormalEnd() && !info.StreamStatus.HasErrors() {
|
||||
logger.LogInfo(c, fmt.Sprintf("stream ended: %s", info.StreamStatus.Summary()))
|
||||
} else {
|
||||
|
||||
@@ -211,12 +211,9 @@ func TestStreamScannerHandler_DataWithExtraSpaces(t *testing.T) {
|
||||
assert.Equal(t, "{\"trimmed\":true}", got)
|
||||
}
|
||||
|
||||
// TestStreamScannerHandler_ClientCancelAbortsUpstreamAndReturns pins the
|
||||
// disconnect contract: when the client goes away, the handler must return
|
||||
// promptly (all goroutines joined, so the gin.Context can never leak into a
|
||||
// pooled reuse), the upstream body must be closed to stop token generation,
|
||||
// and no data received after the disconnect may be processed or written.
|
||||
func TestStreamScannerHandler_ClientCancelAbortsUpstreamAndReturns(t *testing.T) {
|
||||
// TestStreamScannerHandler_ClientCancelDrainsUpstreamUsage verifies that a client
|
||||
// disconnect stops writes but leaves the upstream readable through its terminal event.
|
||||
func TestStreamScannerHandler_ClientCancelDrainsUpstreamUsage(t *testing.T) {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
|
||||
@@ -252,7 +249,6 @@ func TestStreamScannerHandler_ClientCancelAbortsUpstreamAndReturns(t *testing.T)
|
||||
|
||||
_, err := fmt.Fprint(pw, "data: first\n")
|
||||
require.NoError(t, err)
|
||||
|
||||
select {
|
||||
case <-firstHandled:
|
||||
case <-time.After(2 * time.Second):
|
||||
@@ -260,31 +256,21 @@ func TestStreamScannerHandler_ClientCancelAbortsUpstreamAndReturns(t *testing.T)
|
||||
}
|
||||
|
||||
cancel()
|
||||
_, err = fmt.Fprint(pw, "data: usage\ndata: [DONE]\n")
|
||||
require.NoError(t, err, "client disconnect must not close the upstream body")
|
||||
require.NoError(t, pw.Close())
|
||||
|
||||
// The handler must return without any further upstream input: cleanup
|
||||
// closes resp.Body, which unblocks the scanner goroutine.
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("handler did not return after client disconnect")
|
||||
t.Fatal("handler did not finish draining the upstream stream")
|
||||
}
|
||||
|
||||
// Upstream read side must be closed so the provider stops generating
|
||||
// (and billing) for a request nobody is listening to.
|
||||
_, err = fmt.Fprint(pw, "data: second\n")
|
||||
require.ErrorIs(t, err, io.ErrClosedPipe, "upstream body should be closed after client disconnect")
|
||||
|
||||
assert.Equal(t, int64(1), count.Load(), "no chunk after disconnect should be processed")
|
||||
assert.Equal(t, int64(2), count.Load(), "terminal upstream data should be processed after client disconnect")
|
||||
require.NotNil(t, info.StreamStatus)
|
||||
assert.Equal(t, relaycommon.StreamEndReasonClientGone, info.StreamStatus.EndReason)
|
||||
|
||||
body := recorder.Body.String()
|
||||
assert.Contains(t, body, "first")
|
||||
assert.NotContains(t, body, "second")
|
||||
}
|
||||
|
||||
// ---------- Ping tests ----------
|
||||
|
||||
func TestStreamScannerHandler_PingSentDuringSlowUpstream(t *testing.T) {
|
||||
setting := operation_setting.GetGeneralSetting()
|
||||
oldEnabled := setting.PingIntervalEnabled
|
||||
@@ -571,3 +557,15 @@ func TestStreamScannerHandler_StreamStatus_ReplacesPreInitialized(t *testing.T)
|
||||
assert.Equal(t, relaycommon.StreamEndReasonDone, info.StreamStatus.EndReason)
|
||||
assert.Equal(t, 0, info.StreamStatus.TotalErrorCount())
|
||||
}
|
||||
|
||||
func TestNewStreamScannerCallerLimit(t *testing.T) {
|
||||
// The smaller buffer must actually constrain a line; a preallocated 64 KiB
|
||||
// buffer would otherwise bypass this caller's 1 KiB limit in bufio.Scanner.
|
||||
scanner := NewStreamScanner(strings.NewReader(strings.Repeat("x", 2048)+"\n"), 1024)
|
||||
assert.False(t, scanner.Scan())
|
||||
require.Error(t, scanner.Err())
|
||||
scanner = NewStreamScanner(strings.NewReader("data: ok\n"), 1024)
|
||||
require.True(t, scanner.Scan())
|
||||
assert.Equal(t, "data: ok", scanner.Text())
|
||||
require.NoError(t, scanner.Err())
|
||||
}
|
||||
@@ -1,6 +1,7 @@
|
||||
package helper
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"math"
|
||||
@@ -9,6 +10,7 @@ import (
|
||||
"strings"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/logger"
|
||||
relayconstant "github.com/QuantumNous/new-api/relay/constant"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
@@ -121,7 +123,8 @@ func GetAndValidateEmbeddingRequest(c *gin.Context, relayMode int) (*dto.Embeddi
|
||||
// overflow the conversion and corrupt billing.
|
||||
const maxTokensLimit = math.MaxInt32 / 2
|
||||
|
||||
func exceedsMaxTokensLimit(values ...*uint) bool {
|
||||
// ExceedsMaxTokensLimit checks token limits before they reach billing arithmetic.
|
||||
func ExceedsMaxTokensLimit(values ...*uint) bool {
|
||||
for _, v := range values {
|
||||
if lo.FromPtrOr(v, uint(0)) > maxTokensLimit {
|
||||
return true
|
||||
@@ -142,7 +145,7 @@ func GetAndValidateResponsesRequest(c *gin.Context) (*dto.OpenAIResponsesRequest
|
||||
if request.Input == nil {
|
||||
return nil, errors.New("input is required")
|
||||
}
|
||||
if exceedsMaxTokensLimit(request.MaxOutputTokens) {
|
||||
if ExceedsMaxTokensLimit(request.MaxOutputTokens) {
|
||||
return nil, errors.New("max_output_tokens is invalid")
|
||||
}
|
||||
return request, nil
|
||||
@@ -203,6 +206,9 @@ func GetAndValidOpenAIImageRequest(c *gin.Context, relayMode int) (*dto.ImageReq
|
||||
}
|
||||
imageRequest.Quality = formData.Get("quality")
|
||||
imageRequest.Size = formData.Get("size")
|
||||
if parameters := formData.Get("parameters"); parameters != "" {
|
||||
imageRequest.Extra = map[string]json.RawMessage{"parameters": json.RawMessage(parameters)}
|
||||
}
|
||||
if streamValue := strings.TrimSpace(formData.Get("stream")); streamValue != "" {
|
||||
stream, err := strconv.ParseBool(streamValue)
|
||||
if err != nil {
|
||||
@@ -238,7 +244,6 @@ func GetAndValidOpenAIImageRequest(c *gin.Context, relayMode int) (*dto.ImageReq
|
||||
}
|
||||
|
||||
if imageRequest.Model == "" {
|
||||
//imageRequest.Model = "dall-e-3"
|
||||
return nil, errors.New("model is required")
|
||||
}
|
||||
|
||||
@@ -250,25 +255,10 @@ func GetAndValidOpenAIImageRequest(c *gin.Context, relayMode int) (*dto.ImageReq
|
||||
return nil, fmt.Errorf("n must be an integer between 1 and %d", dto.MaxImageN)
|
||||
}
|
||||
|
||||
// Not "256x256", "512x512", or "1024x1024"
|
||||
if imageRequest.Model == "dall-e-2" || imageRequest.Model == "dall-e" {
|
||||
if imageRequest.Size != "" && imageRequest.Size != "256x256" && imageRequest.Size != "512x512" && imageRequest.Size != "1024x1024" {
|
||||
return nil, errors.New("size must be one of 256x256, 512x512, or 1024x1024 for dall-e-2 or dall-e")
|
||||
}
|
||||
if imageRequest.Size == "" {
|
||||
imageRequest.Size = "1024x1024"
|
||||
}
|
||||
} else if imageRequest.Model == "dall-e-3" {
|
||||
if imageRequest.Size != "" && imageRequest.Size != "1024x1024" && imageRequest.Size != "1024x1792" && imageRequest.Size != "1792x1024" {
|
||||
return nil, errors.New("size must be one of 1024x1024, 1024x1792 or 1792x1024 for dall-e-3")
|
||||
}
|
||||
if imageRequest.Quality == "" {
|
||||
imageRequest.Quality = "standard"
|
||||
}
|
||||
if imageRequest.Size == "" {
|
||||
imageRequest.Size = "1024x1024"
|
||||
}
|
||||
} else if imageRequest.Model == "gpt-image-1" {
|
||||
if err := imageRequest.NormalizeLegacyDalleImageRequest(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if imageRequest.Model == "gpt-image-1" {
|
||||
if imageRequest.Quality == "" {
|
||||
imageRequest.Quality = "auto"
|
||||
}
|
||||
@@ -283,6 +273,19 @@ func GetAndValidOpenAIImageRequest(c *gin.Context, relayMode int) (*dto.ImageReq
|
||||
}
|
||||
}
|
||||
|
||||
// Provider parameters can override the top-level count. Validate before
|
||||
// pricing so malformed multipliers return a client error, not a pricing
|
||||
// failure after reservation has started.
|
||||
if raw, exists := imageRequest.Extra["parameters"]; exists {
|
||||
parameters := &dto.ImageBillingParameters{}
|
||||
if err := common.Unmarshal(raw, parameters); err != nil {
|
||||
return nil, fmt.Errorf("invalid image parameters: %w", err)
|
||||
}
|
||||
imageRequest.BillingParameters = parameters
|
||||
}
|
||||
if _, err := imageRequest.ImageCount(common.GetContextKeyInt(c, constant.ContextKeyChannelType) == constant.ChannelTypeAli); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return imageRequest, nil
|
||||
}
|
||||
|
||||
@@ -298,7 +301,7 @@ func GetAndValidateClaudeRequest(c *gin.Context) (textRequest *dto.ClaudeRequest
|
||||
if textRequest.Model == "" {
|
||||
return nil, errors.New("field model is required")
|
||||
}
|
||||
if exceedsMaxTokensLimit(textRequest.MaxTokens, textRequest.MaxTokensToSample) {
|
||||
if ExceedsMaxTokensLimit(textRequest.MaxTokens, textRequest.MaxTokensToSample) {
|
||||
return nil, errors.New("max_tokens is invalid")
|
||||
}
|
||||
|
||||
@@ -323,9 +326,12 @@ func GetAndValidateTextRequest(c *gin.Context, relayMode int) (*dto.GeneralOpenA
|
||||
textRequest.Model = c.Param("model")
|
||||
}
|
||||
|
||||
if exceedsMaxTokensLimit(textRequest.MaxTokens, textRequest.MaxCompletionTokens) {
|
||||
if ExceedsMaxTokensLimit(textRequest.MaxTokens, textRequest.MaxCompletionTokens) {
|
||||
return nil, errors.New("max_tokens is invalid")
|
||||
}
|
||||
if ExceedsMaxTokensLimit(textRequest.MinTokens) {
|
||||
return nil, errors.New("min_tokens is invalid")
|
||||
}
|
||||
if textRequest.Model == "" {
|
||||
return nil, errors.New("model is required")
|
||||
}
|
||||
@@ -376,7 +382,7 @@ func GetAndValidateGeminiRequest(c *gin.Context) (*dto.GeminiChatRequest, error)
|
||||
if len(request.Contents) == 0 && len(request.Requests) == 0 {
|
||||
return nil, errors.New("contents is required")
|
||||
}
|
||||
if exceedsMaxTokensLimit(request.GenerationConfig.MaxOutputTokens) {
|
||||
if ExceedsMaxTokensLimit(request.GenerationConfig.MaxOutputTokens) {
|
||||
return nil, errors.New("maxOutputTokens is invalid")
|
||||
}
|
||||
|
||||
|
||||
+53
-7
@@ -18,10 +18,13 @@ import (
|
||||
"github.com/QuantumNous/new-api/setting/model_setting"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/tidwall/sjson"
|
||||
)
|
||||
|
||||
func ImageHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *types.NewAPIError) {
|
||||
info.InitChannelMeta(c)
|
||||
info.BillingImageCount = nil
|
||||
info.ImageRequestCount = 0
|
||||
|
||||
imageReq, ok := info.Request.(*dto.ImageRequest)
|
||||
if !ok {
|
||||
@@ -43,15 +46,28 @@ func ImageHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *type
|
||||
return types.NewError(fmt.Errorf("invalid api type: %d", info.ApiType), types.ErrorCodeInvalidApiType, types.ErrOptionWithSkipRetry())
|
||||
}
|
||||
adaptor.Init(info)
|
||||
imageCount, err := request.ImageCount(info.ChannelType == constant.ChannelTypeAli)
|
||||
if err != nil {
|
||||
return types.NewErrorWithStatusCode(err, types.ErrorCodeInvalidRequest, http.StatusBadRequest, types.ErrOptionWithSkipRetry())
|
||||
}
|
||||
promptExtend := request.BillingParameters != nil && request.BillingParameters.PromptExtend != nil && *request.BillingParameters.PromptExtend
|
||||
|
||||
var requestBody io.Reader
|
||||
var jsonData []byte
|
||||
|
||||
if model_setting.GetGlobalSettings().PassThroughRequestEnabled || info.ChannelSetting.PassThroughBodyEnabled {
|
||||
storage, err := common.GetBodyStorage(c)
|
||||
if err != nil {
|
||||
return types.NewErrorWithStatusCode(err, types.ErrorCodeReadRequestBodyFailed, http.StatusBadRequest, types.ErrOptionWithSkipRetry())
|
||||
}
|
||||
requestBody = common.NewReplayableBodyReader(storage)
|
||||
if strings.Contains(c.Request.Header.Get("Content-Type"), "multipart/form-data") {
|
||||
requestBody = common.NewReplayableBodyReader(storage)
|
||||
} else {
|
||||
jsonData, err = storage.Bytes()
|
||||
if err != nil {
|
||||
return types.NewErrorWithStatusCode(err, types.ErrorCodeReadRequestBodyFailed, http.StatusBadRequest, types.ErrOptionWithSkipRetry())
|
||||
}
|
||||
}
|
||||
} else {
|
||||
convertedRequest, err := adaptor.ConvertImageRequest(c, info, *request)
|
||||
if err != nil {
|
||||
@@ -63,7 +79,7 @@ func ImageHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *type
|
||||
case *bytes.Buffer:
|
||||
requestBody = convertedRequest.(io.Reader)
|
||||
default:
|
||||
jsonData, err := common.Marshal(convertedRequest)
|
||||
jsonData, err = common.Marshal(convertedRequest)
|
||||
if err != nil {
|
||||
return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
|
||||
}
|
||||
@@ -76,15 +92,45 @@ func ImageHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *type
|
||||
}
|
||||
}
|
||||
|
||||
logger.LogDebug(c, "image request body: %s", jsonData)
|
||||
body, closer, err := relaycommon.NewOutboundJSONBody(jsonData)
|
||||
}
|
||||
}
|
||||
if jsonData != nil {
|
||||
// This is a different trust boundary from ingress: channel overrides
|
||||
// and pass-through bodies can change the quantity actually submitted.
|
||||
var outbound struct {
|
||||
N *uint `json:"n"`
|
||||
Parameters *dto.ImageBillingParameters `json:"parameters"`
|
||||
}
|
||||
if err := common.Unmarshal(jsonData, &outbound); err != nil {
|
||||
return types.NewErrorWithStatusCode(fmt.Errorf("invalid image billing parameters: %w", err), types.ErrorCodeInvalidRequest, http.StatusBadRequest, types.ErrOptionWithSkipRetry())
|
||||
}
|
||||
quantityRequest := dto.ImageRequest{N: outbound.N, BillingParameters: outbound.Parameters}
|
||||
if quantityRequest.N == nil {
|
||||
quantityRequest.N = common.GetPointer(uint(imageCount))
|
||||
}
|
||||
imageCount, err = quantityRequest.ImageCount(info.ChannelType == constant.ChannelTypeAli)
|
||||
if err != nil {
|
||||
return types.NewErrorWithStatusCode(err, types.ErrorCodeInvalidRequest, http.StatusBadRequest, types.ErrOptionWithSkipRetry())
|
||||
}
|
||||
promptExtend = outbound.Parameters != nil && outbound.Parameters.PromptExtend != nil && *outbound.Parameters.PromptExtend
|
||||
if info.ChannelType == constant.ChannelTypeAli {
|
||||
// Always send the same explicit quantity that is reserved, including
|
||||
// when an empty parameters object accompanies a top-level n.
|
||||
jsonData, err = sjson.SetBytes(jsonData, "parameters.n", imageCount)
|
||||
if err != nil {
|
||||
return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
|
||||
}
|
||||
defer closer.Close()
|
||||
jsonData = nil
|
||||
requestBody = body
|
||||
}
|
||||
logger.LogDebug(c, "image request body: %s", jsonData)
|
||||
body, closer, err := relaycommon.NewOutboundJSONBody(jsonData)
|
||||
if err != nil {
|
||||
return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
|
||||
}
|
||||
defer closer.Close()
|
||||
requestBody = body
|
||||
}
|
||||
if billingErr := service.PrepareImageBillingForRequest(c, info, imageCount, promptExtend); billingErr != nil {
|
||||
return billingErr
|
||||
}
|
||||
|
||||
statusCodeMappingStr := c.GetString("status_code_mapping")
|
||||
|
||||
+25
-11
@@ -2,7 +2,6 @@ package relay
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
@@ -125,7 +124,7 @@ func RelayMidjourneyNotify(c *gin.Context) *dto.MidjourneyResponse {
|
||||
midjourneyTask.FinishTime = midjRequest.FinishTime
|
||||
midjourneyTask.ImageUrl = midjRequest.ImageUrl
|
||||
midjourneyTask.VideoUrl = midjRequest.VideoUrl
|
||||
videoUrlsStr, _ := json.Marshal(midjRequest.VideoUrls)
|
||||
videoUrlsStr, _ := common.Marshal(midjRequest.VideoUrls)
|
||||
midjourneyTask.VideoUrls = string(videoUrlsStr)
|
||||
midjourneyTask.Status = midjRequest.Status
|
||||
midjourneyTask.FailReason = midjRequest.FailReason
|
||||
@@ -167,21 +166,21 @@ func coverMidjourneyTaskDto(c *gin.Context, originTask *model.Midjourney) (midjo
|
||||
midjourneyTask.Prompt = originTask.Prompt
|
||||
if originTask.Buttons != "" {
|
||||
var buttons []dto.ActionButton
|
||||
err := json.Unmarshal([]byte(originTask.Buttons), &buttons)
|
||||
err := common.Unmarshal([]byte(originTask.Buttons), &buttons)
|
||||
if err == nil {
|
||||
midjourneyTask.Buttons = buttons
|
||||
}
|
||||
}
|
||||
if originTask.VideoUrls != "" {
|
||||
var videoUrls []dto.ImgUrls
|
||||
err := json.Unmarshal([]byte(originTask.VideoUrls), &videoUrls)
|
||||
err := common.Unmarshal([]byte(originTask.VideoUrls), &videoUrls)
|
||||
if err == nil {
|
||||
midjourneyTask.VideoUrls = videoUrls
|
||||
}
|
||||
}
|
||||
if originTask.Properties != "" {
|
||||
var properties dto.Properties
|
||||
err := json.Unmarshal([]byte(originTask.Properties), &properties)
|
||||
err := common.Unmarshal([]byte(originTask.Properties), &properties)
|
||||
if err == nil {
|
||||
midjourneyTask.Properties = &properties
|
||||
}
|
||||
@@ -228,7 +227,9 @@ func RelaySwapFace(c *gin.Context, info *relaycommon.RelayInfo) *dto.MidjourneyR
|
||||
requestURL := getMjRequestPath(c.Request.URL.String())
|
||||
baseURL := c.GetString("base_url")
|
||||
fullRequestURL := fmt.Sprintf("%s%s", baseURL, requestURL)
|
||||
service.RequestPolicy(c).BeginAttempt(&model.Channel{Id: c.GetInt("channel_id")}, info.UsingGroup)
|
||||
mjResp, _, err := service.DoMidjourneyHttpRequest(c, time.Second*60, fullRequestURL)
|
||||
accepted := service.RecordMidjourneyPolicyResponse(c, mjResp, err)
|
||||
if err != nil {
|
||||
return &mjResp.Response
|
||||
}
|
||||
@@ -268,11 +269,15 @@ func RelaySwapFace(c *gin.Context, info *relaycommon.RelayInfo) *dto.MidjourneyR
|
||||
if billingErr != nil {
|
||||
common.SysLog("error settling Midjourney quota: " + billingErr.Error())
|
||||
}
|
||||
if accepted {
|
||||
service.MarkRequestPolicySuccess(c, nil)
|
||||
}
|
||||
if billingApplied {
|
||||
billingChannelId := midjourneyTask.GetBillingChannelId()
|
||||
tokenName := c.GetString("token_name")
|
||||
logContent := fmt.Sprintf("模型固定价格 %.2f,分组倍率 %.2f,操作 %s", priceData.ModelPrice, priceData.GroupRatioInfo.GroupRatio, constant.MjActionSwapFace)
|
||||
other := service.GenerateMjOtherInfo(info, priceData)
|
||||
service.AppendRelayLogAdminInfo(c, info, other)
|
||||
model.RecordConsumeLog(c, info.UserId, model.RecordConsumeLogParams{
|
||||
ChannelId: billingChannelId,
|
||||
ModelName: modelName,
|
||||
@@ -287,7 +292,7 @@ func RelaySwapFace(c *gin.Context, info *relaycommon.RelayInfo) *dto.MidjourneyR
|
||||
model.UpdateChannelUsedQuota(billingChannelId, midjourneyTask.Quota)
|
||||
}
|
||||
c.Writer.WriteHeader(mjResp.StatusCode)
|
||||
respBody, err := json.Marshal(midjResponse)
|
||||
respBody, err := common.Marshal(midjResponse)
|
||||
if err != nil {
|
||||
return service.MidjourneyErrorWrapper(constant.MjRequestError, "unmarshal_response_body_failed")
|
||||
}
|
||||
@@ -323,7 +328,7 @@ func RelayMidjourneyTaskImageSeed(c *gin.Context) *dto.MidjourneyResponse {
|
||||
}
|
||||
midjResponse := &midjResponseWithStatus.Response
|
||||
c.Writer.WriteHeader(midjResponseWithStatus.StatusCode)
|
||||
respBody, err := json.Marshal(midjResponse)
|
||||
respBody, err := common.Marshal(midjResponse)
|
||||
if err != nil {
|
||||
return service.MidjourneyErrorWrapper(constant.MjRequestError, "unmarshal_response_body_failed")
|
||||
}
|
||||
@@ -346,7 +351,7 @@ func RelayMidjourneyTask(c *gin.Context, relayMode int) *dto.MidjourneyResponse
|
||||
}
|
||||
}
|
||||
midjourneyTask := coverMidjourneyTaskDto(c, originTask)
|
||||
respBody, err = json.Marshal(midjourneyTask)
|
||||
respBody, err = common.Marshal(midjourneyTask)
|
||||
if err != nil {
|
||||
return &dto.MidjourneyResponse{
|
||||
Code: 4,
|
||||
@@ -375,7 +380,7 @@ func RelayMidjourneyTask(c *gin.Context, relayMode int) *dto.MidjourneyResponse
|
||||
if tasks == nil {
|
||||
tasks = make([]dto.MidjourneyDto, 0)
|
||||
}
|
||||
respBody, err = json.Marshal(tasks)
|
||||
respBody, err = common.Marshal(tasks)
|
||||
if err != nil {
|
||||
return &dto.MidjourneyResponse{
|
||||
Code: 4,
|
||||
@@ -539,7 +544,9 @@ func RelayMidjourneySubmit(c *gin.Context, relayInfo *relaycommon.RelayInfo) *dt
|
||||
}
|
||||
}
|
||||
|
||||
service.RequestPolicy(c).BeginAttempt(&model.Channel{Id: c.GetInt("channel_id")}, relayInfo.UsingGroup)
|
||||
midjResponseWithStatus, responseBody, err := service.DoMidjourneyHttpRequest(c, time.Second*60, fullRequestURL)
|
||||
accepted := service.RecordMidjourneyPolicyResponse(c, midjResponseWithStatus, err)
|
||||
if err != nil {
|
||||
return &midjResponseWithStatus.Response
|
||||
}
|
||||
@@ -576,8 +583,11 @@ func RelayMidjourneySubmit(c *gin.Context, relayInfo *relaycommon.RelayInfo) *dt
|
||||
if err != nil {
|
||||
common.SysLog("get_channel_null: " + err.Error())
|
||||
}
|
||||
if channel.GetAutoBan() && common.AutomaticDisableChannelEnabled {
|
||||
model.UpdateChannelStatus(midjourneyTask.ChannelId, "", 2, "No available account instance")
|
||||
if channel != nil && channel.GetAutoBan() && common.AutomaticDisableChannelEnabled {
|
||||
if model.UpdateChannelStatus(midjourneyTask.ChannelId, "", common.ChannelStatusManuallyDisabled, "No available account instance") {
|
||||
service.CloseActiveWebSocketsForChannel(midjourneyTask.ChannelId, service.ChannelDisabledCloseReason)
|
||||
service.RequestPolicy(c).AddEvent(service.PolicyEvent{ChannelID: midjourneyTask.ChannelId, Decision: service.PolicyDecision{Action: "stop", Reason: "task_failed", Source: "upstream"}, Health: "channel_disable_requested"})
|
||||
}
|
||||
}
|
||||
}
|
||||
if midjResponse.Code != 1 && midjResponse.Code != 21 && midjResponse.Code != 22 {
|
||||
@@ -633,11 +643,15 @@ func RelayMidjourneySubmit(c *gin.Context, relayInfo *relaycommon.RelayInfo) *dt
|
||||
if billingErr != nil {
|
||||
common.SysLog("error settling Midjourney quota: " + billingErr.Error())
|
||||
}
|
||||
if accepted {
|
||||
service.MarkRequestPolicySuccess(c, nil)
|
||||
}
|
||||
if billingApplied {
|
||||
billingChannelId := midjourneyTask.GetBillingChannelId()
|
||||
tokenName := c.GetString("token_name")
|
||||
logContent := fmt.Sprintf("模型固定价格 %.2f,分组倍率 %.2f,操作 %s,ID %s", priceData.ModelPrice, priceData.GroupRatioInfo.GroupRatio, midjRequest.Action, midjResponse.Result)
|
||||
other := service.GenerateMjOtherInfo(relayInfo, priceData)
|
||||
service.AppendRelayLogAdminInfo(c, relayInfo, other)
|
||||
model.RecordConsumeLog(c, relayInfo.UserId, model.RecordConsumeLogParams{
|
||||
ChannelId: billingChannelId,
|
||||
ModelName: modelName,
|
||||
|
||||
@@ -0,0 +1,121 @@
|
||||
package relay
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/pkg/billingexpr"
|
||||
"github.com/QuantumNous/new-api/relay/channel"
|
||||
"github.com/QuantumNous/new-api/relay/channel/openai"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
hosttypes "github.com/QuantumNous/new-api/types"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/glebarez/sqlite"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/gorm"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestRelayAttemptMismatchCollectsUsageAndChargesRequestedPrice(t *testing.T) {
|
||||
service.InitTokenEncoders()
|
||||
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
sqlDB, err := db.DB()
|
||||
require.NoError(t, err)
|
||||
sqlDB.SetMaxOpenConns(1)
|
||||
previousDB, previousLog, previousRedis := model.DB, model.LOG_DB, common.RedisEnabled
|
||||
model.DB = db
|
||||
model.LOG_DB = db
|
||||
common.RedisEnabled = false
|
||||
t.Cleanup(func() {
|
||||
model.DB = previousDB
|
||||
model.LOG_DB = previousLog
|
||||
common.RedisEnabled = previousRedis
|
||||
require.NoError(t, sqlDB.Close())
|
||||
})
|
||||
require.NoError(t, db.AutoMigrate(&model.User{}, &model.Token{}, &model.Channel{}, &model.RelayAttempt{}, &model.Log{}, &model.QuotaData{}))
|
||||
user := model.User{Username: "attempt-relay", Quota: 10000}
|
||||
require.NoError(t, db.Create(&user).Error)
|
||||
token := model.Token{UserId: user.Id, Key: "attempt-relay", Status: common.TokenStatusEnabled, RemainQuota: 10000}
|
||||
require.NoError(t, db.Create(&token).Error)
|
||||
expression := `p * 2 + c * 4 + cr * 0.5`
|
||||
for i, ratio := range []float64{1, 3, 1} {
|
||||
selected := model.Channel{Name: fmt.Sprint(i), Status: common.ChannelStatusEnabled}
|
||||
require.NoError(t, db.Create(&selected).Error)
|
||||
// The final usage is buffered behind the rejected prefix: Drain must reuse
|
||||
// the same reader rather than lose bytes already read from the socket.
|
||||
body := "data: {\"type\":\"response.created\",\"response\":{\"model\":\"gpt-5.6-luna\"}}\n\n" +
|
||||
"data: {\"type\":\"response.output_text.delta\",\"delta\":\"hidden generated text\"}\n\n" +
|
||||
"data: {\"type\":\"response.completed\",\"response\":{\"model\":\"gpt-5.6-luna\",\"usage\":{\"input_tokens\":10,\"output_tokens\":2,\"total_tokens\":12,\"input_tokens_details\":{\"cached_tokens\":4}}}}\n\ndata: [DONE]\n\n"
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
_, _ = io.WriteString(w, body)
|
||||
}))
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
c, writer := gin.CreateTestContext(httptest.NewRecorder())
|
||||
_ = writer
|
||||
c.Request = httptest.NewRequest("POST", "/v1/responses", strings.NewReader(`{"model":"gpt-6-astra","input":"secret prompt"}`)).WithContext(ctx)
|
||||
c.Set("channel_id", selected.Id)
|
||||
c.Set("id", user.Id)
|
||||
c.Set("token_id", token.Id)
|
||||
common.SetContextKey(c, constant.ContextKeyUserName, user.Username)
|
||||
info := &relaycommon.RelayInfo{RequestId: fmt.Sprintf("attempt-request-%d", i), UserId: user.Id, TokenId: token.Id, TokenKey: token.Key, OriginModelName: "gpt-6-astra", UsingGroup: fmt.Sprint(ratio), UserSetting: dto.UserSetting{BillingPreference: "wallet_only"}, StartTime: time.Now(), DurableAttemptBilling: true, ResponseGuard: &relaycommon.ResponseModelGuard{}, ChannelMeta: &relaycommon.ChannelMeta{ChannelId: selected.Id, UpstreamModelName: "gpt-6-astra"}, PriceData: hosttypes.PriceData{GroupRatioInfo: hosttypes.GroupRatioInfo{GroupRatio: ratio}}, TieredBillingSnapshot: &billingexpr.BillingSnapshot{BillingMode: "tiered_expr", ExprString: expression, ExprHash: billingexpr.ExprHashString(expression), QuotaPerUnit: common.QuotaPerUnit, GroupRatio: ratio, EstimatedQuotaAfterGroup: 100}}
|
||||
info.SetEstimatePromptTokens(30)
|
||||
info.AstraAutoSwitch = i < 2
|
||||
if !info.AstraAutoSwitch {
|
||||
info.ResponseGuard = nil
|
||||
}
|
||||
require.Nil(t, service.BeginRelayAttempt(c, info))
|
||||
req, err := http.NewRequest("POST", server.URL, strings.NewReader("{}"))
|
||||
require.NoError(t, err)
|
||||
resp, err := channel.DoRequest(c, req, info)
|
||||
if info.AstraAutoSwitch {
|
||||
require.Error(t, err)
|
||||
require.Nil(t, resp)
|
||||
require.False(t, c.Writer.Written())
|
||||
require.True(t, info.AttemptDetached)
|
||||
require.False(t, service.RequestPolicy(c).Successful, "collector cannot mark the client request successful")
|
||||
} else {
|
||||
require.NoError(t, err)
|
||||
usage, apiErr := openai.OaiResponsesHandler(c, info, resp)
|
||||
require.Nil(t, apiErr)
|
||||
require.True(t, c.Writer.Written())
|
||||
require.True(t, info.ResponseModel.Mismatch)
|
||||
ConsumeResponsesQuota(c, info, usage)
|
||||
}
|
||||
var attempt model.RelayAttempt
|
||||
require.Eventually(t, func() bool {
|
||||
return db.First(&attempt, info.DurableAttemptID).Error == nil && attempt.State == "settled"
|
||||
}, time.Second*5, time.Millisecond*10)
|
||||
require.Equal(t, int(11*ratio), attempt.Quota)
|
||||
if info.AstraAutoSwitch {
|
||||
require.Equal(t, "mismatch_rejected", attempt.GuardResult)
|
||||
} else {
|
||||
require.Equal(t, "mismatch_allowed", attempt.GuardResult)
|
||||
require.Contains(t, attempt.LogJSON, "mismatch_allowed")
|
||||
}
|
||||
require.Equal(t, "upstream", attempt.UsageSource)
|
||||
require.NotContains(t, attempt.LogJSON, "secret prompt")
|
||||
require.NotContains(t, attempt.LogJSON, "hidden generated text")
|
||||
require.NoError(t, model.DeliverRelayAttemptLogs())
|
||||
service.EndFailedRelayAttempt(c, info)
|
||||
cancel()
|
||||
server.Close()
|
||||
}
|
||||
require.NoError(t, db.First(&user, user.Id).Error)
|
||||
require.Equal(t, 9945, user.Quota)
|
||||
require.NoError(t, db.First(&token, token.Id).Error)
|
||||
require.Equal(t, 55, token.UsedQuota)
|
||||
var logs int64
|
||||
require.NoError(t, db.Model(&model.Log{}).Count(&logs).Error)
|
||||
require.EqualValues(t, 3, logs)
|
||||
}
|
||||
+36
-14
@@ -12,8 +12,10 @@ import (
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/dto"
|
||||
"github.com/QuantumNous/new-api/logger"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/pkg/billingexpr"
|
||||
"github.com/QuantumNous/new-api/pkg/jsplugin"
|
||||
"github.com/QuantumNous/new-api/relay/channel"
|
||||
"github.com/QuantumNous/new-api/relay/channel/task/taskcommon"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
@@ -248,25 +250,29 @@ func RelayTaskSubmit(c *gin.Context, info *relaycommon.RelayInfo) (*TaskSubmitRe
|
||||
info.OriginModelName = modelName
|
||||
var priceData types.PriceData
|
||||
var err error
|
||||
useTiered := billing_setting.GetBillingMode(modelName) == billing_setting.BillingModeTieredExpr
|
||||
var exprStr string
|
||||
var exists bool
|
||||
if useTiered {
|
||||
exprStr, exists = billing_setting.GetBillingExpr(modelName)
|
||||
} else if info.IsModelMapped {
|
||||
if billing_setting.GetBillingMode(info.UpstreamModelName) == billing_setting.BillingModeTieredExpr {
|
||||
if tailExpr, tailOK := billing_setting.GetBillingExpr(info.UpstreamModelName); tailOK && strings.TrimSpace(tailExpr) != "" {
|
||||
exprStr = tailExpr
|
||||
exists = true
|
||||
useTiered = true
|
||||
}
|
||||
}
|
||||
pluginKey := c.GetString("task_plugin_key")
|
||||
pinnedValue, _ := c.Get(jsplugin.ContextKeyPinnedPlugin)
|
||||
pinnedPlugin, _ := pinnedValue.(jsplugin.PinnedPlugin)
|
||||
if pinnedPlugin.Plugin != nil {
|
||||
pluginKey = pinnedPlugin.Plugin.Meta.Key
|
||||
}
|
||||
exprStr, exists := billing_setting.ResolveTaskBillingExpr(pluginKey, modelName, info.UpstreamModelName)
|
||||
useTiered := exists || billing_setting.GetBillingMode(modelName) == billing_setting.BillingModeTieredExpr
|
||||
if useTiered {
|
||||
provider, supported := adaptor.(channel.TaskUsageFactsProvider)
|
||||
if billingexpr.UsesFixedPricing(exprStr) {
|
||||
return nil, service.TaskErrorWrapper(fmt.Errorf("fixed pricing is not supported for task usage expressions"), "model_price_error", http.StatusBadRequest)
|
||||
}
|
||||
if !exists || !supported {
|
||||
return nil, service.TaskErrorWrapper(fmt.Errorf("task model %s has no usage expression or meter", modelName), "model_price_error", http.StatusBadRequest)
|
||||
}
|
||||
sharedModel := pinnedPlugin.Generation.SharedModel(modelName) || pinnedPlugin.Generation.SharedModel(info.UpstreamModelName)
|
||||
if sharedModel && pinnedPlugin.Plugin != nil {
|
||||
schema, _ := pinnedPlugin.Plugin.Meta.UsageForModel(info.UpstreamModelName)
|
||||
if !billing_setting.TaskExprCompatible(exprStr, schema) {
|
||||
return nil, service.TaskErrorWrapper(fmt.Errorf("task model %s pricing is not configured for plugin %s", modelName, pluginKey), "model_price_error", http.StatusBadRequest)
|
||||
}
|
||||
}
|
||||
var facts map[string]any
|
||||
if validatedProvider, ok := adaptor.(channel.TaskValidatedUsageFactsProvider); ok {
|
||||
facts, err = validatedProvider.ExtractUsageFactsValidated(c, info)
|
||||
@@ -364,7 +370,21 @@ func RelayTaskSubmit(c *gin.Context, info *relaycommon.RelayInfo) (*TaskSubmitRe
|
||||
|
||||
// 11. 提交后计费调整:让适配器根据上游实际返回调整 OtherRatios
|
||||
finalQuota := info.PriceData.Quota
|
||||
if info.TieredBillingSnapshot == nil {
|
||||
if parsed.Immediate != nil && parsed.Immediate.Status == model.TaskStatusFailure {
|
||||
finalQuota = 0
|
||||
} else if snap := info.TieredBillingSnapshot; snap != nil {
|
||||
if parsed.Immediate != nil && parsed.Immediate.Status == model.TaskStatusSuccess && len(parsed.Immediate.UsageFacts) > 0 {
|
||||
settlement, facts, err := service.EvaluateTaskCompletionUsage(snap, parsed.Immediate.UsageFacts)
|
||||
if err != nil {
|
||||
logger.LogWarn(c, fmt.Sprintf("task immediate usage settlement failed; retaining reserved quota: %v", err))
|
||||
} else {
|
||||
finalQuota = settlement.ActualQuotaAfterGroup
|
||||
snap.UsageFacts = facts
|
||||
snap.EstimatedTier = settlement.MatchedTier
|
||||
noteTaskQuotaClamp(info, settlement.Clamp)
|
||||
}
|
||||
}
|
||||
} else {
|
||||
if adjustedRatios := adaptor.AdjustBillingOnSubmit(info, parsed.TaskData); len(adjustedRatios) > 0 {
|
||||
if adjustedQuota, ok := recalcQuotaFromRatios(info, adjustedRatios); ok {
|
||||
// 基于调整后的 ratios 重新计算 quota
|
||||
@@ -375,6 +395,8 @@ func RelayTaskSubmit(c *gin.Context, info *relaycommon.RelayInfo) (*TaskSubmitRe
|
||||
}
|
||||
}
|
||||
|
||||
info.PriceData.Quota = finalQuota
|
||||
|
||||
return &TaskSubmitResult{
|
||||
UpstreamTaskID: parsed.UpstreamTaskID,
|
||||
TaskData: parsed.TaskData,
|
||||
|
||||
+171
-2
@@ -3,10 +3,12 @@ package relay
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/dto"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/pkg/billingexpr"
|
||||
pluginruntime "github.com/QuantumNous/new-api/pkg/jsplugin"
|
||||
@@ -15,10 +17,82 @@ import (
|
||||
"github.com/QuantumNous/new-api/setting/billing_setting"
|
||||
"github.com/QuantumNous/new-api/setting/config"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/glebarez/sqlite"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
func setupRelayChannelDB(t *testing.T) *gorm.DB {
|
||||
t.Helper()
|
||||
previousDB := model.DB
|
||||
previousType := common.MainDatabaseType()
|
||||
previousCache := common.MemoryCacheEnabled
|
||||
database, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
sqlDB, err := database.DB()
|
||||
require.NoError(t, err)
|
||||
sqlDB.SetMaxOpenConns(1)
|
||||
require.NoError(t, database.AutoMigrate(&model.Channel{}))
|
||||
model.DB = database
|
||||
common.SetMainDatabaseType(common.DatabaseTypeSQLite)
|
||||
common.MemoryCacheEnabled = false
|
||||
t.Cleanup(func() {
|
||||
model.DB = previousDB
|
||||
common.SetMainDatabaseType(previousType)
|
||||
common.MemoryCacheEnabled = previousCache
|
||||
require.NoError(t, sqlDB.Close())
|
||||
})
|
||||
return database
|
||||
}
|
||||
|
||||
func TestApplyChannelPinPreservesOriginTasksAndRetryMode(t *testing.T) {
|
||||
database := setupRelayChannelDB(t)
|
||||
channel := &model.Channel{Name: "origin-channel", Key: "sk-test", Status: common.ChannelStatusEnabled, Type: constant.ChannelTypeDoubaoVideo}
|
||||
require.NoError(t, database.Create(channel).Error)
|
||||
originTask := &model.Task{
|
||||
TaskID: "task-lock", ChannelId: channel.Id, Action: "text_to_video", Status: model.TaskStatusSuccess,
|
||||
PrivateData: model.TaskPrivateData{UpstreamTaskID: "upstream-task-lock"},
|
||||
Data: []byte(`{"id":"upstream-task-lock"}`),
|
||||
}
|
||||
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
tokenPin bool
|
||||
apply func(*gin.Context, *relaycommon.RelayInfo) *dto.TaskError
|
||||
}{
|
||||
{name: "origin affinity", apply: ApplyOriginTaskAffinity},
|
||||
{name: "same channel retry", apply: ApplyChannelPin},
|
||||
{name: "token pin suppresses channel lock", tokenPin: true, apply: ApplyChannelPin},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
common.SetContextKey(c, constant.ContextKeyOriginTasks, []*model.Task{originTask})
|
||||
constraints := service.GetChannelConstraints(c)
|
||||
constraints.AddPin(dto.ChannelPin{ChannelId: channel.Id, Source: dto.PinSourceOriginTask, Rank: dto.PinRankOriginTask, RetryMode: dto.PinRetrySameChannel})
|
||||
if tc.tokenPin {
|
||||
constraints.AddPin(dto.ChannelPin{ChannelId: channel.Id, Source: dto.PinSourceToken, Rank: dto.PinRankToken, RetryMode: dto.PinRetrySingleAttempt})
|
||||
}
|
||||
info := &relaycommon.RelayInfo{TaskRelayInfo: &relaycommon.TaskRelayInfo{}}
|
||||
require.Nil(t, tc.apply(c, info))
|
||||
if tc.tokenPin {
|
||||
assert.Nil(t, info.LockedChannel)
|
||||
} else {
|
||||
locked, ok := info.LockedChannel.(*model.Channel)
|
||||
require.True(t, ok)
|
||||
require.NotNil(t, locked)
|
||||
assert.Equal(t, channel.Id, locked.Id)
|
||||
}
|
||||
require.Len(t, info.OriginTasks, 1)
|
||||
assert.Equal(t, "task-lock", info.OriginTasks[0].TaskID)
|
||||
assert.Equal(t, "upstream-task-lock", info.OriginTasks[0].UpstreamTaskID)
|
||||
assert.Equal(t, "text_to_video", info.OriginTasks[0].Action)
|
||||
assert.Equal(t, string(model.TaskStatusSuccess), info.OriginTasks[0].Status)
|
||||
assert.Equal(t, []byte(originTask.Data), info.OriginTasks[0].Data)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestTaskModel2DtoNormalizesLegacyAction(t *testing.T) {
|
||||
task := &model.Task{Action: "firstTailGenerate"}
|
||||
|
||||
@@ -153,6 +227,7 @@ func TestRelayTaskSubmitAliasBillingIdentityAndExprFallback(t *testing.T) {
|
||||
const mapping = `{"alias-model":"declared-model"}`
|
||||
const aliasExpr = `tier("alias", 2)`
|
||||
const tailExpr = `tier("tail", 3)`
|
||||
const legacyExpr = `tier("legacy", u("old_units") * 2)`
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
@@ -160,6 +235,7 @@ func TestRelayTaskSubmitAliasBillingIdentityAndExprFallback(t *testing.T) {
|
||||
exprs map[string]string
|
||||
wantTiered bool
|
||||
wantExpr string
|
||||
source string
|
||||
}{
|
||||
{
|
||||
name: "alias own tiered wins",
|
||||
@@ -175,6 +251,16 @@ func TestRelayTaskSubmitAliasBillingIdentityAndExprFallback(t *testing.T) {
|
||||
wantTiered: true,
|
||||
wantExpr: tailExpr,
|
||||
},
|
||||
{
|
||||
name: "stored expression keeps running after schema narrows",
|
||||
modes: map[string]string{"declared-model": "tiered_expr"},
|
||||
exprs: map[string]string{"declared-model": legacyExpr},
|
||||
wantTiered: true,
|
||||
wantExpr: legacyExpr,
|
||||
source: strings.Replace(billingFallbackPlugin, `fetchMode:"per_task"`, `fetchMode:"per_task", usageSchema:{old_units:{type:"number",unit:"count"}}, usageProfiles:[{models:["declared-model"],schema:{seconds:{type:"number",unit:"second"}}}]`, 1) + `
|
||||
export function extractUsage(){return {old_units:2};}
|
||||
`,
|
||||
},
|
||||
{
|
||||
name: "neither tiered uses ordinary pricing",
|
||||
wantTiered: false,
|
||||
@@ -204,7 +290,11 @@ func TestRelayTaskSubmitAliasBillingIdentityAndExprFallback(t *testing.T) {
|
||||
c.Set("group", "default")
|
||||
info.UserGroup = "default"
|
||||
info.UsingGroup = "default"
|
||||
pinMappingOrderPlugin(t, c, billingFallbackPlugin)
|
||||
source := testCase.source
|
||||
if source == "" {
|
||||
source = billingFallbackPlugin
|
||||
}
|
||||
pinMappingOrderPlugin(t, c, source)
|
||||
info.OriginModelName = "alias-model"
|
||||
|
||||
_, taskErr := RelayTaskSubmit(c, info)
|
||||
@@ -218,11 +308,14 @@ func TestRelayTaskSubmitAliasBillingIdentityAndExprFallback(t *testing.T) {
|
||||
assert.Equal(t, "declared-model", task.Properties.UpstreamModelName)
|
||||
|
||||
if testCase.wantTiered {
|
||||
require.NotNil(t, info.TieredBillingSnapshot)
|
||||
require.NotNil(t, info.TieredBillingSnapshot, "submission error: %+v", taskErr)
|
||||
assert.Equal(t, "alias-model", info.TieredBillingSnapshot.ModelName)
|
||||
assert.Equal(t, testCase.wantExpr, info.TieredBillingSnapshot.ExprString)
|
||||
assert.Equal(t, billingexpr.ExprHashString(testCase.wantExpr), info.TieredBillingSnapshot.ExprHash)
|
||||
assert.NotEqual(t, "model_price_error", taskErr.Code)
|
||||
if testCase.wantExpr == legacyExpr {
|
||||
assert.Equal(t, 4*common.QuotaPerUnit, info.TieredBillingSnapshot.EstimatedQuotaBeforeGroup)
|
||||
}
|
||||
} else {
|
||||
assert.Nil(t, info.TieredBillingSnapshot)
|
||||
assert.Equal(t, "model_price_error", taskErr.Code)
|
||||
@@ -230,3 +323,79 @@ func TestRelayTaskSubmitAliasBillingIdentityAndExprFallback(t *testing.T) {
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSharedTaskBillingExpressionSelectionAndFrozenSettlement(t *testing.T) {
|
||||
const baseExpr = `tier("base", u("seconds") * 2)`
|
||||
const alphaExpr = `tier("alpha", u("seconds") * 3)`
|
||||
const betaExpr = `tier("beta", u("credits") * 5)`
|
||||
const aliasExpr = `tier("alias", u("credits") * 7)`
|
||||
for _, tc := range []struct {
|
||||
name, plugin, model, mapping, modelExpr, mode, wantExpr string
|
||||
variants map[string]string
|
||||
wantPriceError bool
|
||||
}{
|
||||
{name: "executing plugin override", plugin: "billing-beta", model: "declared-model", modelExpr: baseExpr, mode: "tiered_expr", variants: map[string]string{"billing-alpha::declared-model": alphaExpr, "billing-beta::declared-model": betaExpr}, wantExpr: betaExpr},
|
||||
{name: "override ignores model mode", plugin: "billing-beta", model: "declared-model", modelExpr: baseExpr, mode: "ratio", variants: map[string]string{"billing-beta::declared-model": betaExpr}, wantExpr: betaExpr},
|
||||
{name: "model expression fallback", plugin: "billing-alpha", model: "declared-model", modelExpr: baseExpr, mode: "tiered_expr", wantExpr: baseExpr},
|
||||
{name: "alias override precedes mapped override", plugin: "billing-beta", model: "alias-model", mapping: `{"alias-model":"declared-model"}`, modelExpr: baseExpr, mode: "tiered_expr", variants: map[string]string{"billing-beta::declared-model": betaExpr, "billing-beta::alias-model": aliasExpr}, wantExpr: aliasExpr},
|
||||
{name: "mapped override precedes model fallback", plugin: "billing-beta", model: "alias-model", mapping: `{"alias-model":"declared-model"}`, modelExpr: baseExpr, mode: "tiered_expr", variants: map[string]string{"billing-beta::declared-model": betaExpr}, wantExpr: betaExpr},
|
||||
{name: "unconfigured plugin cannot use another schema", plugin: "billing-beta", model: "declared-model", modelExpr: baseExpr, mode: "tiered_expr", wantPriceError: true},
|
||||
{name: "missing usage in skipped branch remains incompatible", plugin: "billing-beta", model: "declared-model", modelExpr: `true ? tier("free", 0) : tier("missing", u("seconds"))`, mode: "tiered_expr", wantPriceError: true},
|
||||
{name: "fixed pricing is still rejected", plugin: "billing-beta", model: "declared-model", variants: map[string]string{"billing-beta::declared-model": `tier("fixed", fixed(1))`}, wantPriceError: true},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
saveBillingConfig(t)
|
||||
registry := pluginruntime.NewRegistry()
|
||||
for _, spec := range []struct{ key, field, unit string }{{"billing-alpha", "seconds", "second"}, {"billing-beta", "credits", "credit"}} {
|
||||
source := strings.ReplaceAll(billingFallbackPlugin, "bill-fallback", spec.key)
|
||||
source = strings.Replace(source, `fetchMode:"per_task"`, `fetchMode:"per_task",usageSchema:{`+spec.field+`:{type:"number",unit:"`+spec.unit+`"}}`, 1)
|
||||
source += `export function extractUsage(){return {` + spec.field + `:2};}`
|
||||
_, err := registry.Register(source, pluginruntime.Options{})
|
||||
require.NoError(t, err)
|
||||
}
|
||||
variants := tc.variants
|
||||
if variants == nil {
|
||||
variants = map[string]string{}
|
||||
}
|
||||
rawVariants, err := common.Marshal(variants)
|
||||
require.NoError(t, err)
|
||||
modes, err := common.Marshal(map[string]string{"declared-model": tc.mode})
|
||||
require.NoError(t, err)
|
||||
expressions, err := common.Marshal(map[string]string{"declared-model": tc.modelExpr})
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, config.GlobalConfig.LoadFromDB(map[string]string{
|
||||
billing_setting.PluginBillingExprOption: string(rawVariants), "billing_setting.billing_mode": string(modes), "billing_setting.billing_expr": string(expressions),
|
||||
}))
|
||||
c, info := newTaskSubmitContext(t, tc.model, tc.mapping)
|
||||
c.Set("group", "default")
|
||||
c.Set("task_plugin_key", tc.plugin)
|
||||
info.UserGroup = "default"
|
||||
info.UsingGroup = "default"
|
||||
info.OriginModelName = tc.model
|
||||
plugin, ok := registry.Generation().Get(tc.plugin)
|
||||
require.True(t, ok)
|
||||
c.Set(pluginruntime.ContextKeyPinnedPlugin, pluginruntime.PinnedPlugin{Generation: registry.Generation(), Plugin: plugin})
|
||||
_, taskErr := RelayTaskSubmit(c, info)
|
||||
require.NotNil(t, taskErr) // This fixture stops at reservation, before upstream submission.
|
||||
if tc.wantPriceError {
|
||||
assert.Equal(t, "model_price_error", taskErr.Code)
|
||||
assert.Nil(t, info.TieredBillingSnapshot)
|
||||
return
|
||||
}
|
||||
require.NotNil(t, info.TieredBillingSnapshot, "submission error: %+v", taskErr)
|
||||
assert.Equal(t, tc.wantExpr, info.TieredBillingSnapshot.ExprString)
|
||||
assert.NotEqual(t, "model_price_error", taskErr.Code)
|
||||
require.NoError(t, config.GlobalConfig.LoadFromDB(map[string]string{billing_setting.PluginBillingExprOption: `{}`, "billing_setting.billing_expr": `{}`}))
|
||||
field := "seconds"
|
||||
if tc.plugin == "billing-beta" {
|
||||
field = "credits"
|
||||
}
|
||||
result, usage, err := service.EvaluateTaskCompletionUsage(info.TieredBillingSnapshot, map[string]any{field: float64(4)})
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, float64(4), usage[field])
|
||||
assert.Equal(t, float64(2), info.TieredBillingSnapshot.UsageFacts[field])
|
||||
assert.Equal(t, 2*info.TieredBillingSnapshot.EstimatedQuotaAfterGroup, result.ActualQuotaAfterGroup)
|
||||
assert.Equal(t, tc.wantExpr, info.TieredBillingSnapshot.ExprString)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,87 @@
|
||||
package relay
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/logger"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/relay/helper"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
"github.com/QuantumNous/new-api/relaykit/types"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/QuantumNous/new-api/setting"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/samber/lo"
|
||||
)
|
||||
|
||||
// PrepareRequestBilling estimates and reserves one request's charge. Transports
|
||||
// provide the current request body through BodyStorage or BillingRequestInput;
|
||||
// channel retries retain the resulting billing session and pricing snapshot.
|
||||
func PrepareRequestBilling(c *gin.Context, info *relaycommon.RelayInfo) *types.NewAPIError {
|
||||
durable := info.OriginModelName == "gpt-6-astra" && c.Request.Method == http.MethodPost && c.Request.URL.Path == "/v1/responses"
|
||||
needSensitiveCheck := setting.ShouldCheckPromptSensitive()
|
||||
meta := &types.TokenCountMeta{TokenType: types.TokenTypeTokenizer}
|
||||
if info.Request != nil && (needSensitiveCheck || constant.CountToken || durable) {
|
||||
meta = info.Request.GetTokenCountMeta()
|
||||
} else {
|
||||
// Avoid building CombineText when only the pricing quantities are needed.
|
||||
switch request := info.Request.(type) {
|
||||
case *dto.GeneralOpenAIRequest:
|
||||
meta.MaxTokens = int(max(lo.FromPtr(request.MaxTokens), lo.FromPtr(request.MaxCompletionTokens)))
|
||||
case *dto.OpenAIResponsesRequest:
|
||||
meta.MaxTokens = int(lo.FromPtr(request.MaxOutputTokens))
|
||||
case *dto.ClaudeRequest:
|
||||
meta.MaxTokens = int(lo.FromPtr(request.MaxTokens))
|
||||
case *dto.ImageRequest:
|
||||
meta = request.GetTokenCountMeta()
|
||||
}
|
||||
}
|
||||
|
||||
if needSensitiveCheck && meta != nil {
|
||||
if contains, words := service.CheckSensitiveText(meta.CombineText); contains {
|
||||
service.RequestPolicy(c).AddEvent(service.PolicyEvent{ErrorCode: string(types.ErrorCodeSensitiveWordsDetected), ErrorSource: "local", Decision: service.PolicyDecision{Action: "stop", Reason: "local_rejection", Source: "global"}, Health: "unchanged"})
|
||||
message := fmt.Sprintf("user sensitive words detected: %s", strings.Join(words, ", "))
|
||||
logger.LogWarn(c, message)
|
||||
return types.NewError(errors.New(message), types.ErrorCodeSensitiveWordsDetected)
|
||||
}
|
||||
}
|
||||
|
||||
tokens, err := service.EstimateRequestToken(c, meta, info)
|
||||
if durable { tokens,err=service.CountRequestToken(c,meta,info) }
|
||||
if err != nil {
|
||||
return types.NewError(err, types.ErrorCodeCountTokenFailed)
|
||||
}
|
||||
info.SetEstimatePromptTokens(tokens)
|
||||
|
||||
priceData, err := helper.ModelPriceHelper(c, info, tokens, meta)
|
||||
if err != nil {
|
||||
return types.NewError(err, types.ErrorCodeModelPriceError, types.ErrOptionWithStatusCode(http.StatusBadRequest))
|
||||
}
|
||||
if durable {
|
||||
info.DurableAttemptBilling = true
|
||||
return nil
|
||||
}
|
||||
if priceData.FreeModel {
|
||||
logger.LogInfo(c, fmt.Sprintf("模型 %s 免费,跳过预扣费", info.OriginModelName))
|
||||
return nil
|
||||
}
|
||||
return service.PreConsumeBilling(c, priceData.QuotaToPreConsume, info)
|
||||
}
|
||||
|
||||
// RefundFailedRequestBilling applies the common final-failure policy after all
|
||||
// eligible attempts have ended. A settled BillingSession never refunds again.
|
||||
func RefundFailedRequestBilling(c *gin.Context, info *relaycommon.RelayInfo, apiErr *types.NewAPIError) *types.NewAPIError {
|
||||
if apiErr == nil {
|
||||
return nil
|
||||
}
|
||||
apiErr = service.NormalizeViolationFeeError(apiErr)
|
||||
if info.Billing != nil {
|
||||
info.Billing.Refund(c)
|
||||
}
|
||||
service.ChargeViolationFeeIfNeeded(c, info, apiErr)
|
||||
return apiErr
|
||||
}
|
||||
@@ -0,0 +1,447 @@
|
||||
package relay
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/relay/channel"
|
||||
"github.com/QuantumNous/new-api/relay/channel/claude"
|
||||
"github.com/QuantumNous/new-api/relay/channel/gemini"
|
||||
"github.com/QuantumNous/new-api/relay/channel/openai"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
relayconstant "github.com/QuantumNous/new-api/relay/constant"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
"github.com/QuantumNous/new-api/relaykit/types"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestResponseGuardBeforeDelivery(t *testing.T) {
|
||||
for _, tc := range []struct{ name, content, body, reason string }{
|
||||
{"astra", "application/json", `{"model":"gpt-6-astra","output":[]}`, "accepted"},
|
||||
{"luna", "application/json", `{"model":"gpt-5.6-luna","usage":{"input_tokens":8,"output_tokens":2}}`, "model_mismatch"},
|
||||
{"missing", "application/json", `{"output":[]}`, "missing_model"},
|
||||
{"not a prefix match", "application/json", `{"model":"gpt-6-astra-other"}`, "model_mismatch"},
|
||||
{"sse astra", "text/event-stream", "event: response.created\ndata: {\"type\":\"response.created\",\"response\":{\"model\":\"gpt-6-astra\"}}\n\ndata: [DONE]\n\n", "accepted"},
|
||||
{"sse luna", "text/event-stream", "data: {\"type\":\"response.created\",\"response\":{\"model\":\"gpt-5.6-luna\"}}\n\n", "model_mismatch"},
|
||||
{"sse no model", "text/event-stream", "data: {\"type\":\"response.output_text.delta\",\"delta\":\"hidden\"}\n\ndata: [DONE]\n\n", "missing_model"},
|
||||
{"sse buffer bound", "text/event-stream", ":" + strings.Repeat("x", 1<<20), "confirmation_buffer_limit"},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", tc.content)
|
||||
w.Header().Set("X-Codex-Turn-State", "must-not-leak")
|
||||
_, _ = io.WriteString(w, tc.body)
|
||||
}))
|
||||
defer server.Close()
|
||||
writer := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(writer)
|
||||
c.Request = httptest.NewRequest("POST", "/v1/responses", strings.NewReader("{}"))
|
||||
info := &relaycommon.RelayInfo{AstraAutoSwitch: true, ResponseGuard: &relaycommon.ResponseModelGuard{}, ChannelMeta: &relaycommon.ChannelMeta{ChannelId: 1}, IsStream: strings.Contains(tc.content, "event-stream")}
|
||||
req, err := http.NewRequest("POST", server.URL, strings.NewReader("{}"))
|
||||
require.NoError(t, err)
|
||||
resp, err := channel.DoRequest(c, req, info)
|
||||
assert.False(t, c.Writer.Written())
|
||||
assert.Empty(t, writer.Header().Get("X-Codex-Turn-State"))
|
||||
assert.Equal(t, tc.reason, info.ResponseGuard.Last().Reason)
|
||||
if tc.reason == "accepted" {
|
||||
require.NoError(t, err)
|
||||
defer resp.Body.Close()
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, tc.body, string(body))
|
||||
} else {
|
||||
require.Error(t, err)
|
||||
assert.Nil(t, resp)
|
||||
assert.Equal(t, http.StatusConflict, info.ResponseGuard.Last().APIError().StatusCode)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestResponseGuardAcceptsModelAfterFormerTenSecondCutoff(t *testing.T) {
|
||||
previous := constant.StreamingTimeout
|
||||
constant.StreamingTimeout = 30
|
||||
t.Cleanup(func() { constant.StreamingTimeout = previous })
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
select {
|
||||
case <-time.After(11 * time.Second):
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = io.WriteString(w, `{"model":"gpt-6-astra","output":[]}`)
|
||||
case <-r.Context().Done():
|
||||
}
|
||||
}))
|
||||
defer server.Close()
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest("POST", "/v1/responses", nil)
|
||||
info := &relaycommon.RelayInfo{AstraAutoSwitch: true, ResponseGuard: &relaycommon.ResponseModelGuard{}, ChannelMeta: &relaycommon.ChannelMeta{ChannelId: 1}}
|
||||
req, err := http.NewRequest("POST", server.URL, strings.NewReader("{}"))
|
||||
require.NoError(t, err)
|
||||
resp, err := channel.DoRequest(c, req, info)
|
||||
require.NoError(t, err)
|
||||
defer resp.Body.Close()
|
||||
assert.Equal(t, "accepted", info.ResponseGuard.Last().Reason)
|
||||
assert.False(t, c.Writer.Written())
|
||||
}
|
||||
|
||||
func TestResponseGuardOnlyEnablesForAstraRequests(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
path string
|
||||
model string
|
||||
enabled bool
|
||||
want bool
|
||||
}{
|
||||
{name: "astra enabled", path: "/v1/responses", model: "gpt-6-astra", enabled: true, want: true},
|
||||
{name: "direct luna unaffected", path: "/v1/responses", model: "gpt-5.6-luna", enabled: true},
|
||||
{name: "token opt out", path: "/v1/responses", model: "gpt-6-astra"},
|
||||
{name: "other endpoint", path: "/v1/chat/completions", model: "gpt-6-astra", enabled: true},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
info := &relaycommon.RelayInfo{
|
||||
OriginModelName: tt.model,
|
||||
AstraAutoSwitch: tt.enabled,
|
||||
ChannelMeta: &relaycommon.ChannelMeta{
|
||||
ChannelSetting: dto.ChannelSettings{ResponsesAstraModelGuard: common.GetPointer(!tt.enabled)},
|
||||
},
|
||||
}
|
||||
assert.Equal(t, tt.want, shouldEnableAstraResponseGuard(tt.path, info))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestResponseGuardTokenOptOutAllowsMismatchDespiteChannelSetting(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = io.WriteString(w, `{"model":"gpt-5.6-luna","output":[]}`)
|
||||
}))
|
||||
defer server.Close()
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest("POST", "/v1/responses", strings.NewReader("{}"))
|
||||
guard := &relaycommon.ResponseModelGuard{}
|
||||
_, previous := guard.Begin(context.Background(), 1, 1)
|
||||
previous.Close()
|
||||
previous.Reason = "http_error"
|
||||
info := &relaycommon.RelayInfo{
|
||||
ResponseGuard: guard,
|
||||
UpstreamRequests: 1,
|
||||
ChannelMeta: &relaycommon.ChannelMeta{
|
||||
ChannelId: 2,
|
||||
ChannelSetting: dto.ChannelSettings{ResponsesAstraModelGuard: common.GetPointer(true)},
|
||||
},
|
||||
}
|
||||
req, err := http.NewRequest("POST", server.URL, strings.NewReader("{}"))
|
||||
require.NoError(t, err)
|
||||
resp, err := channel.DoRequest(c, req, info)
|
||||
require.NoError(t, err)
|
||||
defer resp.Body.Close()
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
require.NoError(t, err)
|
||||
assert.Contains(t, string(body), `"model":"gpt-5.6-luna"`)
|
||||
assert.Len(t, guard.Attempts, 1)
|
||||
assert.Equal(t, "http_error", previous.Reason)
|
||||
}
|
||||
|
||||
func TestResponseGuardPublicErrorPrefersObservedMismatch(t *testing.T) {
|
||||
guard := &relaycommon.ResponseModelGuard{}
|
||||
_, mismatch := guard.Begin(context.Background(), 1, 1)
|
||||
mismatch.Returned = "gpt-5.6-luna"
|
||||
require.Error(t, mismatch.Reject("model_mismatch"))
|
||||
t.Cleanup(mismatch.Close)
|
||||
err := guard.APIError()
|
||||
require.NotNil(t, err)
|
||||
assert.Equal(t, relaycommon.ResponseGuardErrorCode, err.GetErrorCode())
|
||||
assert.Equal(t, http.StatusConflict, err.StatusCode)
|
||||
assert.Contains(t, err.Error(), "gpt-6-astra")
|
||||
assert.Contains(t, err.Error(), "gpt-5.6-luna")
|
||||
}
|
||||
|
||||
func TestResponseGuardLateMismatchAndCancellation(t *testing.T) {
|
||||
for _, cancelAfterAccept := range []bool{false, true} {
|
||||
t.Run(fmt.Sprint(cancelAfterAccept), func(t *testing.T) {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
guard := &relaycommon.ResponseModelGuard{}
|
||||
_, attempt := guard.Begin(ctx, 1, 1)
|
||||
defer attempt.Close()
|
||||
first := "data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_ok\",\"model\":\"gpt-6-astra\"}}\n\n"
|
||||
model := "gpt-5.6-luna"
|
||||
if cancelAfterAccept {
|
||||
model = "gpt-6-astra"
|
||||
}
|
||||
last := fmt.Sprintf("data: {\"type\":\"response.completed\",\"response\":{\"model\":%q,\"usage\":{\"input_tokens\":8,\"output_tokens\":2}}}\n\n", model)
|
||||
resp := &http.Response{StatusCode: 200, Header: http.Header{"Content-Type": []string{"text/event-stream"}}, Body: io.NopCloser(strings.NewReader(first + last))}
|
||||
require.NoError(t, attempt.Confirm(resp, ctx))
|
||||
defer resp.Body.Close()
|
||||
if cancelAfterAccept {
|
||||
cancel()
|
||||
}
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
if cancelAfterAccept {
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, first+last, string(body))
|
||||
} else {
|
||||
require.Error(t, err)
|
||||
assert.Equal(t, first, string(body))
|
||||
assert.Equal(t, "model_mismatch", attempt.Reason)
|
||||
}
|
||||
require.NotNil(t, attempt.Usage)
|
||||
assert.Equal(t, 8, attempt.Usage.Input)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestResponseGuardRejectsCancelledClientAndExhaustedBudget(t *testing.T) {
|
||||
for _, exhausted := range []bool{false, true} {
|
||||
t.Run(fmt.Sprint(exhausted), func(t *testing.T) {
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
c.Request = httptest.NewRequest("POST", "/v1/responses", strings.NewReader("{}")).WithContext(ctx)
|
||||
info := &relaycommon.RelayInfo{AstraAutoSwitch: true, ResponseGuard: &relaycommon.ResponseModelGuard{}, ChannelMeta: &relaycommon.ChannelMeta{}}
|
||||
if exhausted {
|
||||
info.UpstreamRequests = 4
|
||||
} else {
|
||||
cancel()
|
||||
}
|
||||
req, err := http.NewRequest("POST", "http://127.0.0.1:1", strings.NewReader("{}"))
|
||||
require.NoError(t, err)
|
||||
resp, err := channel.DoRequest(c, req, info)
|
||||
require.Error(t, err)
|
||||
assert.Nil(t, resp)
|
||||
assert.Empty(t, info.ResponseGuard.Attempts)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestResponseModelComparisonAndLog(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
models []string
|
||||
returned string
|
||||
mismatch bool
|
||||
}{
|
||||
{name: "absent", models: []string{"", " "}},
|
||||
{name: "requested model", models: []string{"requested"}, returned: "requested"},
|
||||
{name: "mapped model", models: []string{"mapped"}, returned: "mapped"},
|
||||
{name: "different model", models: []string{"other"}, returned: "other", mismatch: true},
|
||||
{name: "case differs", models: []string{"Requested"}, returned: "Requested"},
|
||||
{name: "requested prefix", models: []string{"requested-2026-09-01"}, returned: "requested-2026-09-01"},
|
||||
{name: "mapped case differs", models: []string{"MAPPED"}, returned: "MAPPED"},
|
||||
{name: "mapped prefix", models: []string{"mapped-2026-09-01"}, returned: "mapped-2026-09-01"},
|
||||
{name: "reverse prefix still warns", models: []string{"request"}, returned: "request", mismatch: true},
|
||||
{name: "substring still warns", models: []string{"other-requested"}, returned: "other-requested", mismatch: true},
|
||||
{name: "compatible difference survives matching frames", models: []string{"mapped", "Requested", "", "requested"}, returned: "Requested"},
|
||||
{name: "warning supersedes compatible difference", models: []string{"Requested", "other"}, returned: "other", mismatch: true},
|
||||
{name: "compatible difference cannot erase warning", models: []string{"other", "Requested"}, returned: "other", mismatch: true},
|
||||
{name: "mismatch survives later frames", models: []string{"mapped", "other", "", "requested"}, returned: "other", mismatch: true},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
info := &relaycommon.RelayInfo{
|
||||
OriginModelName: "requested",
|
||||
ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "mapped", IsModelMapped: true},
|
||||
}
|
||||
for _, name := range tc.models {
|
||||
info.ObserveResponseModel(name)
|
||||
}
|
||||
other := service.GenerateTextOtherInfo(c, info, 1, 1, 1, 0, 0, 0, 1)
|
||||
var stored struct {
|
||||
ResponseModel *relaycommon.ResponseModel `json:"response_model"`
|
||||
}
|
||||
require.NoError(t, common.UnmarshalJsonStr(other.JSONString(), &stored))
|
||||
if tc.returned == "" {
|
||||
assert.Nil(t, stored.ResponseModel)
|
||||
return
|
||||
}
|
||||
require.NotNil(t, stored.ResponseModel)
|
||||
assert.Equal(t, &relaycommon.ResponseModel{
|
||||
RequestedModel: "requested", UpstreamModel: "mapped", ReturnedModel: tc.returned, Mismatch: tc.mismatch,
|
||||
}, stored.ResponseModel)
|
||||
assert.Equal(t, "mapped", info.UpstreamModelName)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestResponseModelLogOmitsUnchangedModel(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
upstream string
|
||||
returned string
|
||||
mapped bool
|
||||
record bool
|
||||
}{
|
||||
{name: "same model", upstream: "requested", returned: "requested"},
|
||||
{name: "no upstream name", returned: "requested"},
|
||||
{name: "mapped model", upstream: "mapped", returned: "mapped", mapped: true, record: true},
|
||||
{name: "mapped response echoes request", upstream: "mapped", returned: "requested", mapped: true, record: true},
|
||||
{name: "prefix difference", upstream: "requested", returned: "requested-2026-09-01", record: true},
|
||||
{name: "case difference", upstream: "requested", returned: "REQUESTED", record: true},
|
||||
{name: "mismatch", upstream: "requested", returned: "other", record: true},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
info := &relaycommon.RelayInfo{
|
||||
OriginModelName: "requested",
|
||||
ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: tc.upstream, IsModelMapped: tc.mapped},
|
||||
}
|
||||
info.ObserveResponseModel(tc.returned)
|
||||
other := service.GenerateTextOtherInfo(c, info, 1, 1, 1, 0, 0, 0, 1)
|
||||
require.NotNil(t, info.ResponseModel)
|
||||
assert.Equal(t, float64(1), other.Snapshot()["model_ratio"])
|
||||
if tc.record {
|
||||
assert.Equal(t, *info.ResponseModel, other.Snapshot()["response_model"])
|
||||
} else {
|
||||
assert.NotContains(t, other.Snapshot(), "response_model")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestResponseModelEmptyExpectedNamesDoNotMatchEveryPrefix(t *testing.T) {
|
||||
for _, tc := range []struct{ requested, upstream string }{
|
||||
{}, {upstream: "mapped"}, {requested: "requested"},
|
||||
} {
|
||||
t.Run(tc.requested+"/"+tc.upstream, func(t *testing.T) {
|
||||
info := &relaycommon.RelayInfo{
|
||||
OriginModelName: tc.requested,
|
||||
ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: tc.upstream},
|
||||
}
|
||||
info.ObserveResponseModel("other")
|
||||
require.NotNil(t, info.ResponseModel)
|
||||
assert.True(t, info.ResponseModel.Mismatch)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestResponseModelRetryResetsObservation(t *testing.T) {
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
common.SetContextKey(c, constant.ContextKeyOriginalModel, "requested")
|
||||
info := &relaycommon.RelayInfo{OriginModelName: "requested"}
|
||||
info.InitChannelMeta(c)
|
||||
info.ObserveResponseModel("other")
|
||||
require.True(t, info.ResponseModel.Mismatch)
|
||||
info.InitChannelMeta(c)
|
||||
assert.Nil(t, info.ResponseModel)
|
||||
info.ObserveResponseModel("requested")
|
||||
assert.False(t, info.ResponseModel.Mismatch)
|
||||
}
|
||||
|
||||
func TestResponseModelHandlersCaptureBeforeConversion(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
previousTimeout := constant.StreamingTimeout
|
||||
constant.StreamingTimeout = 30
|
||||
t.Cleanup(func() { constant.StreamingTimeout = previousTimeout })
|
||||
chat := `{"id":"chat_1","model":"returned","choices":[{"index":0,"message":{"role":"assistant","content":"hello"},"finish_reason":"stop"}],"usage":{"prompt_tokens":2,"completion_tokens":3,"total_tokens":5}}`
|
||||
chatStream := "data: " + `{"id":"chat_1","model":"returned","choices":[{"index":0,"delta":{"role":"assistant","content":"hello"}}]}` + "\n\ndata: " + `{"id":"chat_1","choices":[],"usage":{"prompt_tokens":2,"completion_tokens":3,"total_tokens":5}}` + "\n\ndata: [DONE]\n\n"
|
||||
responses := `{"id":"resp_1","model":"returned","status":"completed","output":[{"type":"message","role":"assistant","content":[{"type":"output_text","text":"hello"}]}],"usage":{"input_tokens":2,"output_tokens":3,"total_tokens":5}}`
|
||||
responsesStream := "data: " + `{"type":"response.created","response":{"id":"resp_1","model":"returned"}}` + "\n\ndata: " + `{"type":"response.completed","response":{"id":"resp_1","status":"completed","output":[],"usage":{"input_tokens":2,"output_tokens":3,"total_tokens":5}}}` + "\n\n"
|
||||
claudeBody := `{"id":"msg_1","type":"message","model":"returned","role":"assistant","content":[{"type":"text","text":"hello"}],"stop_reason":"end_turn","usage":{"input_tokens":2,"output_tokens":3}}`
|
||||
claudeStream := "data: " + `{"type":"message_start","message":{"id":"msg_1","type":"message","model":"returned","role":"assistant","content":[],"usage":{"input_tokens":2,"output_tokens":0}}}` + "\n\ndata: " + `{"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":3}}` + "\n\ndata: " + `{"type":"message_stop"}` + "\n\n"
|
||||
geminiBody := `{"modelVersion":"returned","candidates":[{"content":{"role":"model","parts":[{"text":"hello"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":2,"candidatesTokenCount":3,"totalTokenCount":5}}`
|
||||
geminiStream := "data: " + geminiBody + "\n\n"
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
body string
|
||||
format types.RelayFormat
|
||||
stream bool
|
||||
handler func(*gin.Context, *relaycommon.RelayInfo, *http.Response) (*dto.Usage, *types.NewAPIError)
|
||||
}{
|
||||
{"chat", chat, types.RelayFormatOpenAI, false, openai.OpenaiHandler},
|
||||
{"chat stream", chatStream, types.RelayFormatOpenAI, true, openai.OaiStreamHandler},
|
||||
{"responses", responses, types.RelayFormatOpenAIResponses, false, openai.OaiResponsesHandler},
|
||||
{"responses stream", responsesStream, types.RelayFormatOpenAIResponses, true, openai.OaiResponsesStreamHandler},
|
||||
{"responses to chat", responses, types.RelayFormatOpenAI, false, openai.OaiResponsesToChatHandler},
|
||||
{"responses to chat stream", responsesStream, types.RelayFormatOpenAI, true, openai.OaiResponsesToChatStreamHandler},
|
||||
{"responses buffered to chat", responsesStream, types.RelayFormatOpenAI, false, openai.OaiResponsesToChatBufferedStreamHandler},
|
||||
{"chat to responses", chat, types.RelayFormatOpenAIResponses, false, openai.OaiChatToResponsesHandler},
|
||||
{"chat to responses stream", chatStream, types.RelayFormatOpenAIResponses, true, openai.OaiChatToResponsesStreamHandler},
|
||||
{"claude", claudeBody, types.RelayFormatClaude, false, func(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) {
|
||||
return claude.ClaudeHandler(c, resp, info)
|
||||
}},
|
||||
{"claude stream", claudeStream, types.RelayFormatClaude, true, func(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) {
|
||||
return claude.ClaudeStreamHandler(c, resp, info)
|
||||
}},
|
||||
{"claude to responses stream", claudeStream, types.RelayFormatOpenAIResponses, true, func(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) {
|
||||
return claude.ClaudeResponsesStreamHandler(c, resp, info)
|
||||
}},
|
||||
{"gemini", geminiBody, types.RelayFormatGemini, false, gemini.GeminiTextGenerationHandler},
|
||||
{"gemini stream", geminiStream, types.RelayFormatGemini, true, gemini.GeminiTextGenerationStreamHandler},
|
||||
{"gemini to chat", geminiBody, types.RelayFormatOpenAI, false, gemini.GeminiChatHandler},
|
||||
{"gemini to responses", geminiBody, types.RelayFormatOpenAIResponses, false, gemini.GeminiResponsesHandler},
|
||||
{"gemini to responses stream", geminiStream, types.RelayFormatOpenAIResponses, true, gemini.GeminiResponsesStreamHandler},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||||
info := &relaycommon.RelayInfo{
|
||||
OriginModelName: "requested", RelayFormat: tc.format, RelayMode: relayconstant.RelayModeChatCompletions,
|
||||
IsStream: tc.stream, DisablePing: true, StartTime: time.Now(),
|
||||
ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "mapped"},
|
||||
ThinkingContentInfo: relaycommon.ThinkingContentInfo{IsFirstThinkingContent: true},
|
||||
}
|
||||
usage, apiErr := tc.handler(c, info, &http.Response{StatusCode: http.StatusOK, Header: make(http.Header), Body: io.NopCloser(strings.NewReader(tc.body))})
|
||||
require.Nil(t, apiErr)
|
||||
require.NotNil(t, usage)
|
||||
assert.Equal(t, 5, usage.TotalTokens)
|
||||
require.NotNil(t, info.ResponseModel)
|
||||
assert.Equal(t, &relaycommon.ResponseModel{RequestedModel: "requested", UpstreamModel: "mapped", ReturnedModel: "returned", Mismatch: true}, info.ResponseModel)
|
||||
other := service.GenerateTextOtherInfo(c, info, 1, 1, 1, 0, 0, 0, 1)
|
||||
assert.Equal(t, *info.ResponseModel, other.Snapshot()["response_model"])
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestResponseModelSharedResponsesAccumulator(t *testing.T) {
|
||||
// WebSocket turns and HTTP SSE use this same accumulator. A model-free
|
||||
// terminal must not erase the model declared before an interrupted stream.
|
||||
info := &relaycommon.RelayInfo{OriginModelName: "requested", ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "mapped"}}
|
||||
accumulator := service.NewResponsesUsageAccumulator(info)
|
||||
accumulator.Observe(&dto.ResponsesStreamResponse{Type: "response.created", Response: &dto.OpenAIResponsesResponse{Model: "returned"}})
|
||||
accumulator.Observe(&dto.ResponsesStreamResponse{Type: "response.failed", Response: &dto.OpenAIResponsesResponse{Usage: &dto.Usage{InputTokens: 2, OutputTokens: 3}}})
|
||||
assert.Equal(t, 5, accumulator.Finish().TotalTokens)
|
||||
require.NotNil(t, info.ResponseModel)
|
||||
assert.True(t, info.ResponseModel.Mismatch)
|
||||
assert.Equal(t, "returned", info.ResponseModel.ReturnedModel)
|
||||
}
|
||||
|
||||
func TestResponseModelDoesNotRecordSynthesizedModel(t *testing.T) {
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
|
||||
info := &relaycommon.RelayInfo{
|
||||
OriginModelName: "requested", RelayFormat: types.RelayFormatOpenAI,
|
||||
ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "mapped"},
|
||||
}
|
||||
_, apiErr := openai.OaiResponsesToChatBufferedStreamHandler(c, info, &http.Response{
|
||||
StatusCode: http.StatusOK, Header: make(http.Header), Body: io.NopCloser(strings.NewReader("data: [DONE]\n\n")),
|
||||
})
|
||||
require.Nil(t, apiErr)
|
||||
assert.Nil(t, info.ResponseModel)
|
||||
}
|
||||
|
||||
func TestAstraDetectionCanBeDisabledIndependently(t *testing.T) {
|
||||
for _, disabled := range []bool{false, true} {
|
||||
info := &relaycommon.RelayInfo{OriginModelName: "gpt-6-astra", AstraDetectionDisabled: disabled, AstraAutoSwitch: true, ChannelMeta: &relaycommon.ChannelMeta{}}
|
||||
info.ObserveResponseModel("gpt-5.6-luna")
|
||||
if disabled {
|
||||
require.Nil(t, info.ResponseModel)
|
||||
require.False(t, shouldEnableAstraResponseGuard("/v1/responses", info))
|
||||
} else {
|
||||
require.NotNil(t, info.ResponseModel)
|
||||
require.True(t, info.ResponseModel.Mismatch)
|
||||
require.True(t, shouldEnableAstraResponseGuard("/v1/responses", info))
|
||||
}
|
||||
}
|
||||
}
|
||||
+93
-79
@@ -1,36 +1,40 @@
|
||||
package relay
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/logger"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
relayconstant "github.com/QuantumNous/new-api/relay/constant"
|
||||
"github.com/QuantumNous/new-api/relay/helper"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
"github.com/QuantumNous/new-api/relaykit/types"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/QuantumNous/new-api/setting/model_setting"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func ResponsesHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *types.NewAPIError) {
|
||||
info.InitChannelMeta(c)
|
||||
if info.RelayMode == relayconstant.RelayModeResponsesCompact &&
|
||||
!common.SupportsResponsesCompact(info.ChannelType, info.ApiType) {
|
||||
return types.NewErrorWithStatusCode(
|
||||
fmt.Errorf("unsupported endpoint %q for api type %d", "/v1/responses/compact", info.ApiType),
|
||||
types.ErrorCodeInvalidRequest,
|
||||
http.StatusBadRequest,
|
||||
types.ErrOptionWithSkipRetry(),
|
||||
)
|
||||
}
|
||||
// Independent cursors retain the payload metadata used by net/http redirects
|
||||
// and transport retries without closing the shared request storage.
|
||||
type responsesAttemptBody struct {
|
||||
common.ReplayableBody
|
||||
reader io.Reader
|
||||
}
|
||||
|
||||
func (b responsesAttemptBody) Read(p []byte) (int, error) { return b.reader.Read(p) }
|
||||
|
||||
func shouldEnableAstraResponseGuard(path string, info *relaycommon.RelayInfo) bool {
|
||||
return path == "/v1/responses" &&
|
||||
info.OriginModelName == "gpt-6-astra" &&
|
||||
info.AstraAutoSwitch && !info.AstraDetectionDisabled &&
|
||||
info.ResponseGuard == nil
|
||||
}
|
||||
|
||||
func ResponsesHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *types.NewAPIError) {
|
||||
var responsesReq *dto.OpenAIResponsesRequest
|
||||
switch req := info.Request.(type) {
|
||||
case *dto.OpenAIResponsesRequest:
|
||||
@@ -61,77 +65,64 @@ func ResponsesHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *
|
||||
)
|
||||
}
|
||||
|
||||
request, err := common.DeepCopy(responsesReq)
|
||||
if err != nil {
|
||||
return types.NewError(fmt.Errorf("failed to copy request to GeneralOpenAIRequest: %w", err), types.ErrorCodeInvalidRequest, types.ErrOptionWithSkipRetry())
|
||||
adaptor, requestBody, closer, apiErr := PrepareResponsesRequest(c, info, responsesReq)
|
||||
if apiErr != nil {
|
||||
return apiErr
|
||||
}
|
||||
|
||||
err = helper.ModelMappedHelper(c, info, request)
|
||||
if err != nil {
|
||||
return types.NewError(err, types.ErrorCodeChannelModelMappedError, types.ErrOptionWithSkipRetry())
|
||||
}
|
||||
if err := helper.ApplyReasoningModelSuffix(c, info, request); err != nil {
|
||||
return newConvertRequestFailedError(c, info, err)
|
||||
}
|
||||
|
||||
adaptor := GetAdaptor(info.ApiType)
|
||||
if adaptor == nil {
|
||||
return types.NewError(fmt.Errorf("invalid api type: %d", info.ApiType), types.ErrorCodeInvalidApiType, types.ErrOptionWithSkipRetry())
|
||||
}
|
||||
adaptor.Init(info)
|
||||
var requestBody io.Reader
|
||||
if model_setting.GetGlobalSettings().PassThroughRequestEnabled || info.ChannelSetting.PassThroughBodyEnabled {
|
||||
storage, err := common.GetBodyStorage(c)
|
||||
if err != nil {
|
||||
return types.NewError(err, types.ErrorCodeReadRequestBodyFailed, types.ErrOptionWithSkipRetry())
|
||||
}
|
||||
requestBody = common.NewReplayableBodyReader(storage)
|
||||
} else {
|
||||
convertedRequest, err := adaptor.ConvertOpenAIResponsesRequest(c, info, *request)
|
||||
if err != nil {
|
||||
return newConvertRequestFailedError(c, info, err)
|
||||
}
|
||||
relaycommon.AppendRequestConversionFromRequest(info, convertedRequest)
|
||||
jsonData, err := common.Marshal(convertedRequest)
|
||||
if err != nil {
|
||||
return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
|
||||
}
|
||||
|
||||
// remove disabled fields for OpenAI Responses API
|
||||
jsonData, err = relaycommon.RemoveDisabledFields(jsonData, info.ChannelOtherSettings, info.ChannelSetting.PassThroughBodyEnabled)
|
||||
if err != nil {
|
||||
return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
|
||||
}
|
||||
|
||||
// apply param override
|
||||
if len(info.ParamOverride) > 0 {
|
||||
jsonData, err = relaycommon.ApplyParamOverrideWithRelayInfo(jsonData, info)
|
||||
if err != nil {
|
||||
return newAPIErrorFromParamOverride(err)
|
||||
defer closer.Close()
|
||||
// Missing-usage estimation follows the actual prepared payload, including
|
||||
// parameter overrides, rather than the client's pre-transformation request.
|
||||
if info.DurableAttemptBilling {
|
||||
if cursor, err := requestBody.NewReader(); err == nil {
|
||||
var sent dto.OpenAIResponsesRequest
|
||||
decodeErr := json.NewDecoder(cursor).Decode(&sent)
|
||||
_ = cursor.Close()
|
||||
if decodeErr == nil {
|
||||
if tokens, err := service.CountRequestToken(c, sent.GetTokenCountMeta(), info); err == nil {
|
||||
info.SetEstimatePromptTokens(tokens)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
logger.LogDebug(c, "requestBody: %s", jsonData)
|
||||
body, closer, err := relaycommon.NewOutboundJSONBody(jsonData)
|
||||
if err != nil {
|
||||
return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
|
||||
}
|
||||
defer closer.Close()
|
||||
jsonData = nil
|
||||
requestBody = body
|
||||
}
|
||||
if shouldEnableAstraResponseGuard(c.Request.URL.Path, info) {
|
||||
info.ResponseGuard = &relaycommon.ResponseModelGuard{Deadline: info.GuardDeadline}
|
||||
}
|
||||
|
||||
var httpResp *http.Response
|
||||
resp, err := adaptor.DoRequest(c, info, requestBody)
|
||||
guardEnabled := info.ResponseGuard != nil && info.AstraAutoSwitch
|
||||
|
||||
body, err := requestBody.NewReader()
|
||||
if err != nil {
|
||||
return types.NewOpenAIError(err, types.ErrorCodeDoRequestFailed, http.StatusInternalServerError)
|
||||
return types.NewOpenAIError(err, types.ErrorCodeReadRequestBodyFailed, http.StatusInternalServerError)
|
||||
}
|
||||
resp, requestErr := adaptor.DoRequest(c, info, responsesAttemptBody{ReplayableBody: requestBody, reader: body})
|
||||
_ = body.Close()
|
||||
if requestErr != nil {
|
||||
apiErr := types.NewError(requestErr, types.ErrorCodeDoRequestFailed)
|
||||
guard := info.ResponseGuard
|
||||
if guard == nil || apiErr.GetErrorCode() != relaycommon.ResponseGuardErrorCode {
|
||||
return apiErr
|
||||
}
|
||||
attempt := guard.Last()
|
||||
if attempt != nil {
|
||||
service.RequestPolicy(c).AddEvent(service.PolicyEvent{ChannelID: info.ChannelId, ErrorCode: string(relaycommon.ResponseGuardErrorCode), ErrorSource: "model_guard", Health: "unchanged", Decision: service.PolicyDecision{Action: "reject", Reason: attempt.Reason, Source: "model_guard"}})
|
||||
}
|
||||
return guard.APIError()
|
||||
}
|
||||
var httpResp *http.Response
|
||||
if resp != nil {
|
||||
httpResp = resp.(*http.Response)
|
||||
}
|
||||
if guardEnabled {
|
||||
defer func() {
|
||||
if attempt := info.ResponseGuard.Last(); attempt != nil {
|
||||
attempt.Finish(c.Writer.Written())
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
statusCodeMappingStr := c.GetString("status_code_mapping")
|
||||
|
||||
if resp != nil {
|
||||
httpResp = resp.(*http.Response)
|
||||
|
||||
if httpResp != nil {
|
||||
if httpResp.StatusCode != http.StatusOK {
|
||||
newAPIError = service.RelayErrorHandler(c.Request.Context(), httpResp, false)
|
||||
// reset status code 重置状态码
|
||||
@@ -141,6 +132,23 @@ func ResponsesHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *
|
||||
}
|
||||
|
||||
usage, newAPIError := adaptor.DoResponse(c, httpResp, info)
|
||||
if guardEnabled {
|
||||
if attempt := info.ResponseGuard.Last(); attempt != nil {
|
||||
attempt.Finish(c.Writer.Written())
|
||||
if guardErr := attempt.APIError(); guardErr != nil {
|
||||
if info.StreamStatus != nil {
|
||||
info.StreamStatus.MarkFailed(string(relaycommon.ResponseGuardErrorCode), "api_error", http.StatusConflict)
|
||||
}
|
||||
if c.Writer.Written() && c.Request.Context().Err() == nil {
|
||||
message := common.MessageWithRequestId(attempt.PublicErrorMessage(), c.GetString(common.RequestIdKey))
|
||||
c.SSEvent("response.failed", gin.H{"type": "response.failed", "sequence_number": attempt.Sequence + 1, "response": gin.H{"id": attempt.ResponseID, "object": "response", "status": "failed", "error": gin.H{"code": relaycommon.ResponseGuardErrorCode, "message": message}}})
|
||||
c.Writer.Flush()
|
||||
c.Set("response_guard_error_sent", true)
|
||||
}
|
||||
return guardErr
|
||||
}
|
||||
}
|
||||
}
|
||||
if newAPIError != nil {
|
||||
// reset status code 重置状态码
|
||||
service.ResetStatusCode(newAPIError, statusCodeMappingStr)
|
||||
@@ -165,10 +173,16 @@ func ResponsesHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *
|
||||
return nil
|
||||
}
|
||||
|
||||
if strings.HasPrefix(info.OriginModelName, "gpt-4o-audio") {
|
||||
service.PostAudioConsumeQuota(c, info, usageDto, "")
|
||||
} else {
|
||||
service.PostTextConsumeQuota(c, info, usageDto, nil)
|
||||
}
|
||||
ConsumeResponsesQuota(c, info, usageDto)
|
||||
return nil
|
||||
}
|
||||
|
||||
// ConsumeResponsesQuota applies the same settlement dispatch to HTTP and
|
||||
// WebSocket Responses usage. Compact requests keep their separate repricing.
|
||||
func ConsumeResponsesQuota(c *gin.Context, info *relaycommon.RelayInfo, usage *dto.Usage) {
|
||||
if strings.HasPrefix(info.OriginModelName, "gpt-4o-audio") {
|
||||
service.PostAudioConsumeQuota(c, info, usage, "")
|
||||
return
|
||||
}
|
||||
service.PostTextConsumeQuota(c, info, usage, nil)
|
||||
}
|
||||
@@ -0,0 +1,87 @@
|
||||
package relay
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/logger"
|
||||
relaychannel "github.com/QuantumNous/new-api/relay/channel"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
relayconstant "github.com/QuantumNous/new-api/relay/constant"
|
||||
"github.com/QuantumNous/new-api/relay/helper"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
"github.com/QuantumNous/new-api/relaykit/types"
|
||||
"github.com/QuantumNous/new-api/setting/model_setting"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// PrepareResponsesRequest applies the same model, conversion and channel rules
|
||||
// for HTTP and WebSocket requests. The caller closes closer after the attempt;
|
||||
// passthrough bodies remain owned by the incoming request's BodyStorage.
|
||||
// The returned adaptor retains route/conversion state for DoRequest/DoResponse.
|
||||
func PrepareResponsesRequest(c *gin.Context, info *relaycommon.RelayInfo, req *dto.OpenAIResponsesRequest) (relaychannel.Adaptor, common.ReplayableBody, io.Closer, *types.NewAPIError) {
|
||||
info.InitChannelMeta(c)
|
||||
if info.RelayMode == relayconstant.RelayModeResponsesCompact &&
|
||||
!common.SupportsResponsesCompact(info.ChannelType, info.ApiType) {
|
||||
return nil, nil, nil, types.NewErrorWithStatusCode(
|
||||
fmt.Errorf("unsupported endpoint %q for api type %d", "/v1/responses/compact", info.ApiType),
|
||||
types.ErrorCodeInvalidRequest,
|
||||
http.StatusBadRequest,
|
||||
types.ErrOptionWithSkipRetry(),
|
||||
)
|
||||
}
|
||||
|
||||
request, err := common.DeepCopy(req)
|
||||
if err != nil {
|
||||
return nil, nil, nil, types.NewError(fmt.Errorf("failed to copy responses request: %w", err), types.ErrorCodeInvalidRequest, types.ErrOptionWithSkipRetry())
|
||||
}
|
||||
if err := helper.ModelMappedHelper(c, info, request); err != nil {
|
||||
return nil, nil, nil, types.NewError(err, types.ErrorCodeChannelModelMappedError, types.ErrOptionWithSkipRetry())
|
||||
}
|
||||
if err := helper.ApplyReasoningModelSuffix(c, info, request); err != nil {
|
||||
return nil, nil, nil, newConvertRequestFailedError(c, info, err)
|
||||
}
|
||||
|
||||
adaptor := GetAdaptor(info.ApiType)
|
||||
if adaptor == nil {
|
||||
return nil, nil, nil, types.NewError(fmt.Errorf("invalid api type: %d", info.ApiType), types.ErrorCodeInvalidApiType, types.ErrOptionWithSkipRetry())
|
||||
}
|
||||
adaptor.Init(info)
|
||||
if model_setting.GetGlobalSettings().PassThroughRequestEnabled || info.ChannelSetting.PassThroughBodyEnabled {
|
||||
storage, err := common.GetBodyStorage(c)
|
||||
if err != nil {
|
||||
return nil, nil, nil, types.NewError(err, types.ErrorCodeReadRequestBodyFailed, types.ErrOptionWithSkipRetry())
|
||||
}
|
||||
body := common.NewReplayableBodyReader(storage)
|
||||
return adaptor, body, io.NopCloser(body), nil
|
||||
}
|
||||
|
||||
convertedRequest, err := adaptor.ConvertOpenAIResponsesRequest(c, info, *request)
|
||||
if err != nil {
|
||||
return nil, nil, nil, newConvertRequestFailedError(c, info, err)
|
||||
}
|
||||
relaycommon.AppendRequestConversionFromRequest(info, convertedRequest)
|
||||
jsonData, err := common.Marshal(convertedRequest)
|
||||
if err != nil {
|
||||
return nil, nil, nil, types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
|
||||
}
|
||||
jsonData, err = relaycommon.RemoveDisabledFields(jsonData, info.ChannelOtherSettings, info.ChannelSetting.PassThroughBodyEnabled)
|
||||
if err != nil {
|
||||
return nil, nil, nil, types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
|
||||
}
|
||||
if len(info.ParamOverride) > 0 {
|
||||
jsonData, err = relaycommon.ApplyParamOverrideWithRelayInfo(jsonData, info)
|
||||
if err != nil {
|
||||
return nil, nil, nil, newAPIErrorFromParamOverride(err)
|
||||
}
|
||||
}
|
||||
|
||||
logger.LogDebug(c, "requestBody: %s", jsonData)
|
||||
body, closer, err := relaycommon.NewOutboundJSONBody(jsonData)
|
||||
if err != nil {
|
||||
return nil, nil, nil, types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
|
||||
}
|
||||
return adaptor, body, closer, nil
|
||||
}
|
||||
@@ -0,0 +1,939 @@
|
||||
package relay
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"reflect"
|
||||
"slices"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
appconstant "github.com/QuantumNous/new-api/constant"
|
||||
appdto "github.com/QuantumNous/new-api/dto"
|
||||
"github.com/QuantumNous/new-api/i18n"
|
||||
"github.com/QuantumNous/new-api/middleware"
|
||||
appmodel "github.com/QuantumNous/new-api/model"
|
||||
perfmetrics "github.com/QuantumNous/new-api/pkg/perf_metrics"
|
||||
"github.com/QuantumNous/new-api/pkg/wsmanager"
|
||||
relaychannel "github.com/QuantumNous/new-api/relay/channel"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/relay/helper"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
"github.com/QuantumNous/new-api/relaykit/types"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/gorilla/websocket"
|
||||
)
|
||||
|
||||
const responsesWSEventTypeResponseCreate = "response.create"
|
||||
const responsesWSWriteTimeout = 30 * time.Second
|
||||
|
||||
// ResponsesWSRequestRunner executes the existing authentication and rate-limit
|
||||
// middleware around one complete request, without a second HTTP connection.
|
||||
type ResponsesWSRequestRunner func(*http.Request, string, func(*gin.Context) *types.NewAPIError) *types.NewAPIError
|
||||
|
||||
type responsesWSCreateEvent struct {
|
||||
Type string `json:"type"`
|
||||
EventID string `json:"event_id,omitempty"`
|
||||
StreamID common.RawMessage `json:"stream_id,omitempty"`
|
||||
Generate common.RawMessage `json:"generate,omitempty"`
|
||||
Request common.RawMessage `json:"response,omitempty"`
|
||||
}
|
||||
|
||||
type responsesWSCreateRequest struct {
|
||||
Request dto.OpenAIResponsesRequest
|
||||
Body []byte
|
||||
Generate common.RawMessage
|
||||
StreamID string
|
||||
}
|
||||
|
||||
type responsesWSErrorEvent struct {
|
||||
Type string `json:"type"`
|
||||
Status int `json:"status"`
|
||||
EventID string `json:"event_id,omitempty"`
|
||||
StreamID string `json:"stream_id,omitempty"`
|
||||
ResponseID string `json:"response_id,omitempty"`
|
||||
Error *types.OpenAIError `json:"error"`
|
||||
}
|
||||
|
||||
type responsesWSMessage struct {
|
||||
kind int
|
||||
body []byte
|
||||
err error
|
||||
}
|
||||
|
||||
// responsesWSControl is a client control event (response.cancel) whose
|
||||
// envelope the read loop already parsed.
|
||||
type responsesWSControl struct {
|
||||
body []byte
|
||||
eventID, streamID string
|
||||
}
|
||||
|
||||
// Only the request worker reads or changes billing state. Socket readers pass
|
||||
// bounded messages to it; cancellation never performs an independent refund.
|
||||
type responsesWSCallState struct {
|
||||
inbox chan responsesWSMessage
|
||||
controls chan responsesWSControl
|
||||
done chan struct{}
|
||||
terminal *responsesWSMessage
|
||||
closeAfter bool
|
||||
}
|
||||
|
||||
type responsesWSSession struct {
|
||||
ctx context.Context
|
||||
cancel context.CancelFunc
|
||||
client *websocket.Conn
|
||||
runner ResponsesWSRequestRunner
|
||||
request *http.Request
|
||||
requestID string
|
||||
nextEventIndex int
|
||||
workers sync.WaitGroup
|
||||
|
||||
clientWriteMu sync.Mutex
|
||||
targetWriteMu sync.Mutex
|
||||
connectionMu sync.Mutex
|
||||
target *websocket.Conn
|
||||
unregister func()
|
||||
stateMu sync.Mutex
|
||||
current *responsesWSCallState
|
||||
|
||||
// These fields belong to the serial request worker and describe the actual
|
||||
// established connection. Per-request token/user data is never stored here.
|
||||
lastResponseID string
|
||||
lockedModel string
|
||||
lockedChannelID int
|
||||
lockedGroup string
|
||||
lockedKey string
|
||||
lockedKeyIndex int
|
||||
lockedContext map[appconstant.ContextKey]any
|
||||
lockedRoute dto.AdvancedCustomRoute
|
||||
}
|
||||
|
||||
func ResponsesWebSocketHelper(c *gin.Context, client *websocket.Conn, runner ResponsesWSRequestRunner) *types.NewAPIError {
|
||||
ctx, cancel := context.WithCancel(c.Request.Context())
|
||||
s := &responsesWSSession{ctx: ctx, cancel: cancel, client: client, runner: runner,
|
||||
request: c.Request.Clone(ctx), requestID: c.GetString(common.RequestIdKey)}
|
||||
if s.requestID == "" {
|
||||
s.requestID = common.NewRequestId()
|
||||
}
|
||||
maxMB := appconstant.MaxRequestBodyMB
|
||||
if maxMB <= 0 {
|
||||
maxMB = 128
|
||||
}
|
||||
client.SetReadLimit(int64(maxMB) << 20)
|
||||
defer func() {
|
||||
s.shutdown()
|
||||
s.workers.Wait()
|
||||
}()
|
||||
|
||||
for {
|
||||
_, message, err := client.ReadMessage()
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
envelope, streamID, err := parseResponsesWSEnvelope(message)
|
||||
eventType := envelope.Type
|
||||
if err != nil {
|
||||
s.sendError(envelope.EventID, streamID, newResponsesWSInvalidRequestError(err))
|
||||
continue
|
||||
}
|
||||
if eventType != responsesWSEventTypeResponseCreate {
|
||||
// Controls are owned by the active request too. In particular a cancel
|
||||
// arriving during authentication must not precede its upstream create.
|
||||
state := s.getCurrent()
|
||||
if eventType != "response.cancel" || state == nil {
|
||||
s.sendError(envelope.EventID, streamID, newResponsesWSInvalidRequestError(fmt.Errorf("unsupported websocket event %q", eventType)))
|
||||
continue
|
||||
}
|
||||
select {
|
||||
case state.controls <- responsesWSControl{body: message, eventID: envelope.EventID, streamID: streamID}:
|
||||
case <-state.done:
|
||||
case <-s.ctx.Done():
|
||||
return nil
|
||||
default:
|
||||
s.sendError(envelope.EventID, streamID, newResponsesWSInvalidRequestError(errors.New("a response control event is already pending")))
|
||||
}
|
||||
continue
|
||||
}
|
||||
state := &responsesWSCallState{inbox: make(chan responsesWSMessage), controls: make(chan responsesWSControl, 1), done: make(chan struct{})}
|
||||
if !s.tryReserveCurrent(state) {
|
||||
s.sendError(envelope.EventID, streamID, types.NewErrorWithStatusCode(errors.New("another response.create is already in progress on this websocket connection"), types.ErrorCodeInvalidRequest, http.StatusConflict, types.ErrOptionWithSkipRetry()))
|
||||
continue
|
||||
}
|
||||
requestID := fmt.Sprintf("%s-ws-%d", s.requestID, s.nextEventIndex)
|
||||
s.nextEventIndex++
|
||||
s.workers.Go(func() { s.runRequest(state, message, envelope, streamID, requestID) })
|
||||
}
|
||||
}
|
||||
|
||||
func (s *responsesWSSession) runRequest(state *responsesWSCallState, message []byte, envelope responsesWSCreateEvent, streamID string, requestID string) {
|
||||
create, parseErr := normalizeResponsesWSCreateEvent(message, envelope, streamID)
|
||||
var apiErr *types.NewAPIError
|
||||
defer func() {
|
||||
if recovered := recover(); recovered != nil {
|
||||
apiErr = types.NewError(fmt.Errorf("responses websocket request panic: %v", recovered), types.ErrorCodeBadResponse, types.ErrOptionWithSkipRetry())
|
||||
state.closeAfter = true
|
||||
}
|
||||
// Finish the middleware count before publishing the terminal event. Hold
|
||||
// admission while writing it so an immediate next create cannot race
|
||||
// the current request's release; socket close needs neither lock.
|
||||
outgoing := state.terminal
|
||||
if apiErr != nil {
|
||||
if body, err := buildResponsesWSErrorPayload(envelope.EventID, streamID, apiErr); err == nil {
|
||||
outgoing = &responsesWSMessage{kind: websocket.TextMessage, body: body}
|
||||
}
|
||||
}
|
||||
s.clientWriteMu.Lock()
|
||||
s.stateMu.Lock()
|
||||
if outgoing != nil {
|
||||
if err := s.client.SetWriteDeadline(time.Now().Add(responsesWSWriteTimeout)); err != nil {
|
||||
state.closeAfter = true
|
||||
} else if err := s.client.WriteMessage(outgoing.kind, outgoing.body); err != nil {
|
||||
state.closeAfter = true
|
||||
}
|
||||
}
|
||||
s.current = nil
|
||||
close(state.done)
|
||||
s.stateMu.Unlock()
|
||||
s.clientWriteMu.Unlock()
|
||||
if state.closeAfter {
|
||||
s.shutdown()
|
||||
}
|
||||
}()
|
||||
request := s.request.Clone(s.ctx)
|
||||
request.Body = io.NopCloser(bytes.NewReader(message))
|
||||
request.ContentLength = int64(len(message))
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
apiErr = s.runner(request, requestID, func(c *gin.Context) *types.NewAPIError {
|
||||
if parseErr != nil {
|
||||
return newResponsesWSInvalidRequestError(parseErr)
|
||||
}
|
||||
c.Request.Body = io.NopCloser(bytes.NewReader(create.Body))
|
||||
c.Request.ContentLength = int64(len(create.Body))
|
||||
return s.runCall(c, state, create)
|
||||
})
|
||||
if apiErr != nil && (apiErr.StatusCode == http.StatusUnauthorized || apiErr.StatusCode == http.StatusForbidden) {
|
||||
state.closeAfter = true
|
||||
}
|
||||
}
|
||||
|
||||
func (s *responsesWSSession) runCall(c *gin.Context, state *responsesWSCallState, create responsesWSCreateRequest) (apiErr *types.NewAPIError) {
|
||||
policy := service.RequestPolicy(c)
|
||||
modelName := create.Request.Model
|
||||
started := time.Now()
|
||||
var info *relaycommon.RelayInfo
|
||||
billingPrepared := false
|
||||
defer func() {
|
||||
if recovered := recover(); recovered != nil {
|
||||
apiErr = types.NewError(fmt.Errorf("responses websocket call panic: %v", recovered), types.ErrorCodeBadResponse, types.ErrOptionWithSkipRetry())
|
||||
state.closeAfter = true
|
||||
}
|
||||
if info == nil && modelName != "" {
|
||||
info = &relaycommon.RelayInfo{OriginModelName: modelName, UsingGroup: common.GetContextKeyString(c, appconstant.ContextKeyUsingGroup), StartTime: started}
|
||||
}
|
||||
perfmetrics.RecordRelayResult(c.Request.Context(), info, apiErr)
|
||||
// Settlement already marks the request policy successful, and nothing
|
||||
// reads a termination decision after this point on the WebSocket path,
|
||||
// so neither policy record belongs here.
|
||||
if info != nil && billingPrepared {
|
||||
apiErr = RefundFailedRequestBilling(c, info, apiErr)
|
||||
}
|
||||
}()
|
||||
if modelName == "" {
|
||||
return newResponsesWSInvalidRequestError(errors.New("model is required"))
|
||||
}
|
||||
if s.lockedModel != "" && modelName != s.lockedModel {
|
||||
return newResponsesWSInvalidRequestError(fmt.Errorf("responses websocket connection is locked to model %q", s.lockedModel))
|
||||
}
|
||||
if apiErr = checkResponsesWSModelAccess(c, modelName); apiErr != nil {
|
||||
return apiErr
|
||||
}
|
||||
common.SetContextKey(c, appconstant.ContextKeyOriginalModel, modelName)
|
||||
common.SetContextKey(c, appconstant.ContextKeyRequestStartTime, time.Now())
|
||||
service.GetChannelConstraints(c).AddFilter(appdto.ChannelFilter{Kind: appdto.FilterRequestPath, RequestPath: c.Request.URL.Path})
|
||||
|
||||
if s.lockedChannelID != 0 {
|
||||
if apiErr = s.restoreConnectionContext(c, modelName); apiErr != nil {
|
||||
return apiErr
|
||||
}
|
||||
info = relaycommon.GenRelayInfoResponses(c, &create.Request)
|
||||
info.IsStream = true
|
||||
common.SetContextKey(c, appconstant.ContextKeyIsStream, true)
|
||||
if apiErr = PrepareRequestBilling(c, info); apiErr != nil {
|
||||
return apiErr
|
||||
}
|
||||
billingPrepared = true
|
||||
var payload []byte
|
||||
payload, apiErr = buildResponsesWSCreatePayload(c, info, create.Request, create.Generate, create.StreamID)
|
||||
if apiErr != nil {
|
||||
return apiErr
|
||||
}
|
||||
if err := s.writeTarget(websocket.TextMessage, payload); err != nil {
|
||||
state.closeAfter = true
|
||||
return types.NewError(err, types.ErrorCodeBadResponse, types.ErrOptionWithSkipRetry())
|
||||
}
|
||||
} else {
|
||||
retry := &service.RetryParam{Ctx: c, TokenGroup: common.GetContextKeyString(c, appconstant.ContextKeyUsingGroup), ModelName: modelName, RequestPath: c.Request.URL.Path, Retry: common.GetPointer(0)}
|
||||
for ; retry.GetRetry() <= common.RetryTimes; retry.IncreaseRetry() {
|
||||
var channel *appmodel.Channel
|
||||
channel, apiErr = selectResponsesWSChannel(c, modelName, retry)
|
||||
if apiErr != nil {
|
||||
return apiErr
|
||||
}
|
||||
service.AppendUsedChannel(c, channel.Id)
|
||||
if info == nil {
|
||||
info = relaycommon.GenRelayInfoResponses(c, &create.Request)
|
||||
info.IsStream = true
|
||||
common.SetContextKey(c, appconstant.ContextKeyIsStream, true)
|
||||
if apiErr = PrepareRequestBilling(c, info); apiErr != nil {
|
||||
return apiErr
|
||||
}
|
||||
billingPrepared = true
|
||||
} else {
|
||||
info.PriceData.GroupRatioInfo = helper.HandleGroupRatio(c, info)
|
||||
if apiErr = service.PrepareTieredBillingForSelectedGroup(c, info); apiErr != nil {
|
||||
return apiErr
|
||||
}
|
||||
}
|
||||
info.RetryIndex = retry.GetRetry()
|
||||
policy.BeginAttempt(channel, info.UsingGroup)
|
||||
var payload []byte
|
||||
payload, apiErr = buildResponsesWSCreatePayload(c, info, create.Request, create.Generate, create.StreamID)
|
||||
if apiErr != nil {
|
||||
return apiErr
|
||||
}
|
||||
adaptor := GetAdaptor(info.ApiType)
|
||||
adaptor.Init(info)
|
||||
target, dialErr := relaychannel.DoWssRequest(adaptor, c, info, nil)
|
||||
if dialErr != nil {
|
||||
apiErr = service.NormalizeViolationFeeError(types.NewError(dialErr, types.ErrorCodeDoRequestFailed))
|
||||
service.ResetStatusCode(apiErr, c.GetString("status_code_mapping"))
|
||||
info.LastError = apiErr
|
||||
decision := service.DecideRelayRetry(c, apiErr, common.RetryTimes-retry.GetRetry())
|
||||
service.RecordPolicyFailure(c, channel.Id, apiErr, decision)
|
||||
service.ProcessChannelError(c, *types.NewChannelError(channel.Id, channel.Type, channel.Name, channel.ChannelInfo.IsMultiKey, info.ApiKey, channel.GetAutoBan()), apiErr, info)
|
||||
if decision.Action == "retry" {
|
||||
continue
|
||||
}
|
||||
return apiErr
|
||||
}
|
||||
if !s.setTarget(target) {
|
||||
return types.NewError(context.Canceled, types.ErrorCodeBadResponse, types.ErrOptionWithSkipRetry())
|
||||
}
|
||||
if err := s.writeTarget(websocket.TextMessage, payload); err != nil {
|
||||
state.closeAfter = true
|
||||
return types.NewError(err, types.ErrorCodeBadResponse, types.ErrOptionWithSkipRetry())
|
||||
}
|
||||
s.lockedModel, s.lockedChannelID, s.lockedGroup = modelName, channel.Id, info.UsingGroup
|
||||
s.lockedKey, s.lockedKeyIndex = info.ApiKey, info.ChannelMultiKeyIndex
|
||||
s.lockedRoute, _ = channel.GetOtherSettings().AdvancedCustom.MatchPathForModel(c.Request.URL.Path, modelName)
|
||||
s.lockedContext = make(map[appconstant.ContextKey]any)
|
||||
for _, key := range []appconstant.ContextKey{
|
||||
appconstant.ContextKeyChannelId, appconstant.ContextKeyChannelName, appconstant.ContextKeyChannelType,
|
||||
appconstant.ContextKeyChannelCreateTime, appconstant.ContextKeyChannelSetting, appconstant.ContextKeyChannelOtherSetting,
|
||||
appconstant.ContextKeyChannelParamOverride, appconstant.ContextKeyChannelHeaderOverride, appconstant.ContextKeyChannelOrganization,
|
||||
appconstant.ContextKeyChannelAutoBan, appconstant.ContextKeyChannelModelMapping, appconstant.ContextKeyChannelStatusCodeMapping,
|
||||
appconstant.ContextKeyChannelIsMultiKey, appconstant.ContextKeyChannelMultiKeyIndex, appconstant.ContextKeyChannelKey, appconstant.ContextKeyChannelBaseUrl,
|
||||
} {
|
||||
if value, ok := c.Get(string(key)); ok {
|
||||
s.lockedContext[key] = value
|
||||
}
|
||||
}
|
||||
s.registerChannelClose(channel.Id)
|
||||
s.startTargetReader(target)
|
||||
apiErr = nil
|
||||
break
|
||||
}
|
||||
if apiErr != nil {
|
||||
return apiErr
|
||||
}
|
||||
}
|
||||
|
||||
accumulator := service.NewResponsesUsageAccumulator(info)
|
||||
info.StreamStatus = relaycommon.NewStreamStatus()
|
||||
info.StreamStatus.RequireTerminal()
|
||||
common.SetContextKey(c, appconstant.ContextKeyResponseStreamStatus, info.StreamStatus)
|
||||
timeout := time.Duration(appconstant.StreamingTimeout) * time.Second
|
||||
if timeout <= 0 {
|
||||
timeout = 300 * time.Second
|
||||
}
|
||||
idle := time.NewTimer(timeout)
|
||||
defer idle.Stop()
|
||||
accepted := false
|
||||
var pendingControl []byte
|
||||
var sentControl []byte
|
||||
var responseID string
|
||||
for {
|
||||
select {
|
||||
case incoming := <-state.inbox:
|
||||
idle.Reset(timeout)
|
||||
if incoming.err != nil {
|
||||
info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonScannerErr, incoming.err)
|
||||
state.closeAfter = true
|
||||
ConsumeResponsesQuota(c, info, accumulator.Finish())
|
||||
return nil
|
||||
}
|
||||
info.SetFirstResponseTime()
|
||||
var event struct {
|
||||
dto.ResponsesStreamResponse
|
||||
StreamID string `json:"stream_id"`
|
||||
}
|
||||
if err := common.Unmarshal(incoming.body, &event); err != nil {
|
||||
info.StreamStatus.RecordError("invalid upstream websocket event")
|
||||
} else {
|
||||
if event.Type != "error" && event.StreamID != "" && event.StreamID != create.StreamID {
|
||||
if err := s.writeClient(incoming.kind, incoming.body); err != nil {
|
||||
s.shutdown()
|
||||
}
|
||||
continue
|
||||
}
|
||||
// A repeated terminal from the previous response must never finish
|
||||
// a subsequent request on this persistent connection.
|
||||
if event.Response != nil && event.Response.ID != "" && event.Response.ID == s.lastResponseID {
|
||||
continue
|
||||
}
|
||||
if event.Type == "error" {
|
||||
var rejection responsesWSErrorEvent
|
||||
_ = common.Unmarshal(incoming.body, &rejection)
|
||||
if rejection.ResponseID != "" && rejection.ResponseID == s.lastResponseID {
|
||||
continue
|
||||
}
|
||||
terminal, ambiguous, controlError := responsesWSErrorEndsRequest(rejection, create.StreamID, responseID, sentControl)
|
||||
if !terminal {
|
||||
if err := s.writeClient(incoming.kind, incoming.body); err != nil {
|
||||
s.shutdown()
|
||||
}
|
||||
// Only a control error in this stream resolves its pending control.
|
||||
if controlError {
|
||||
sentControl = nil
|
||||
}
|
||||
continue
|
||||
}
|
||||
if accepted {
|
||||
if rejection.Error != nil {
|
||||
code := ""
|
||||
if rejection.Error.Code != nil {
|
||||
code = fmt.Sprint(rejection.Error.Code)
|
||||
}
|
||||
info.StreamStatus.MarkFailed(code, rejection.Error.Type, rejection.Status)
|
||||
}
|
||||
accumulator.Observe(&event.ResponsesStreamResponse)
|
||||
info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonDone, nil)
|
||||
s.lastResponseID = responseID
|
||||
state.terminal, state.closeAfter = &incoming, ambiguous
|
||||
ConsumeResponsesQuota(c, info, accumulator.Finish())
|
||||
return nil
|
||||
}
|
||||
if rejection.Error == nil {
|
||||
// An error frame without an error object still describes the
|
||||
// rejected request; keep the client-facing type stable.
|
||||
rejection.Error = &types.OpenAIError{Type: "invalid_request_error", Message: event.Message, Code: event.Code}
|
||||
}
|
||||
rejected := types.WithOpenAIError(*rejection.Error, rejection.Status, types.ErrOptionWithSkipRetry())
|
||||
if rejected.StatusCode < 400 || rejected.StatusCode > 599 {
|
||||
rejected.StatusCode = http.StatusBadRequest
|
||||
}
|
||||
return rejected
|
||||
}
|
||||
if strings.HasPrefix(event.Type, "response.") {
|
||||
if !accepted {
|
||||
// Like HTTP, bind the session only once upstream accepted the request.
|
||||
service.RecordChannelAffinity(c, s.lockedChannelID)
|
||||
}
|
||||
accepted = true
|
||||
if event.Response != nil && event.Response.ID != "" {
|
||||
responseID = event.Response.ID
|
||||
}
|
||||
}
|
||||
accumulator.Observe(&event.ResponsesStreamResponse)
|
||||
}
|
||||
switch event.Type {
|
||||
case "response.completed", "response.done", "response.incomplete", "response.failed", "response.cancelled", "response.canceled":
|
||||
if event.Response != nil {
|
||||
s.lastResponseID = event.Response.ID
|
||||
}
|
||||
info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonDone, nil)
|
||||
state.terminal = &incoming
|
||||
ConsumeResponsesQuota(c, info, accumulator.Finish())
|
||||
return nil
|
||||
}
|
||||
if err := s.writeClient(incoming.kind, incoming.body); err != nil {
|
||||
s.shutdown()
|
||||
}
|
||||
if accepted && pendingControl != nil {
|
||||
if err := s.writeTarget(websocket.TextMessage, pendingControl); err != nil {
|
||||
s.shutdown()
|
||||
}
|
||||
sentControl = pendingControl
|
||||
pendingControl = nil
|
||||
}
|
||||
case control := <-state.controls:
|
||||
if pendingControl != nil || sentControl != nil {
|
||||
s.sendError(control.eventID, control.streamID, newResponsesWSInvalidRequestError(errors.New("a response control event is already pending")))
|
||||
continue
|
||||
}
|
||||
if !accepted {
|
||||
pendingControl = control.body
|
||||
continue
|
||||
}
|
||||
if err := s.writeTarget(websocket.TextMessage, control.body); err != nil {
|
||||
s.shutdown()
|
||||
}
|
||||
sentControl = control.body
|
||||
case <-idle.C:
|
||||
info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonTimeout, context.DeadlineExceeded)
|
||||
state.closeAfter = true
|
||||
ConsumeResponsesQuota(c, info, accumulator.Finish())
|
||||
return nil
|
||||
case <-s.ctx.Done():
|
||||
info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonClientGone, s.ctx.Err())
|
||||
ConsumeResponsesQuota(c, info, accumulator.Finish())
|
||||
return nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// A single generation is active, but a cancel can fail independently. Correlate
|
||||
// available IDs first; uncorrelated control/server errors require a bounded
|
||||
// connection close because continuing could leave the generation stuck forever.
|
||||
func responsesWSErrorEndsRequest(event responsesWSErrorEvent, streamID, responseID string, control []byte) (terminal, ambiguous, controlError bool) {
|
||||
if len(control) > 0 {
|
||||
var pending struct {
|
||||
EventID string `json:"event_id"`
|
||||
StreamID string `json:"stream_id"`
|
||||
ResponseID string `json:"response_id"`
|
||||
}
|
||||
_ = common.Unmarshal(control, &pending)
|
||||
if event.EventID != "" && event.EventID == pending.EventID {
|
||||
return false, false, true
|
||||
}
|
||||
if event.ResponseID != "" && event.ResponseID == pending.ResponseID && event.ResponseID != responseID {
|
||||
return false, false, true
|
||||
}
|
||||
if (event.StreamID == "" || event.StreamID == pending.StreamID) && event.Error != nil && event.Error.Type == "invalid_request_error" {
|
||||
switch event.Error.Code {
|
||||
case "response_not_found", "response_not_active", "response_already_completed":
|
||||
return false, false, true
|
||||
}
|
||||
}
|
||||
}
|
||||
if event.StreamID != "" && event.StreamID != streamID || responseID != "" && event.ResponseID != "" && event.ResponseID != responseID {
|
||||
return false, false, false
|
||||
}
|
||||
return true, len(control) > 0 && event.ResponseID == "", false
|
||||
}
|
||||
|
||||
func (s *responsesWSSession) restoreConnectionContext(c *gin.Context, model string) *types.NewAPIError {
|
||||
// The channel type is not rechecked: a locked channel keeps its type, and
|
||||
// the FilterResponsesWebSocket rules were applied when it was selected.
|
||||
channel, err := appmodel.CacheGetChannel(s.lockedChannelID)
|
||||
if err != nil || channel == nil || channel.Status != common.ChannelStatusEnabled || !channel.GetSetting().ResponsesWebSocketEnabled {
|
||||
return types.NewErrorWithStatusCode(errors.New("Responses WebSocket is disabled for this channel"), types.ErrorCode(appdto.FilterResponsesWebSocket), http.StatusForbidden, types.ErrOptionWithSkipRetry())
|
||||
}
|
||||
keyEnabled := channel.Key == s.lockedKey
|
||||
if channel.ChannelInfo.IsMultiKey {
|
||||
keys := channel.GetKeys()
|
||||
status := channel.ChannelInfo.MultiKeyStatusList[s.lockedKeyIndex]
|
||||
keyEnabled = s.lockedKeyIndex >= 0 && s.lockedKeyIndex < len(keys) && keys[s.lockedKeyIndex] == s.lockedKey && (status == 0 || status == common.ChannelStatusEnabled)
|
||||
}
|
||||
if !keyEnabled {
|
||||
return types.NewErrorWithStatusCode(errors.New("the upstream connection credential is no longer enabled"), types.ErrorCodeAccessDenied, http.StatusForbidden, types.ErrOptionWithSkipRetry())
|
||||
}
|
||||
// Changes that alter the physical upstream connection require a new
|
||||
// handshake. Request-level settings can be refreshed without rotating keys.
|
||||
if channel.GetBaseURL() != s.lockedContext[appconstant.ContextKeyChannelBaseUrl] ||
|
||||
!reflect.DeepEqual(channel.GetHeaderOverride(), s.lockedContext[appconstant.ContextKeyChannelHeaderOverride]) {
|
||||
return types.NewErrorWithStatusCode(errors.New("upstream connection settings changed; reconnect required"), types.ErrorCodeAccessDenied, http.StatusForbidden, types.ErrOptionWithSkipRetry())
|
||||
}
|
||||
if previous, ok := s.lockedContext[appconstant.ContextKeyChannelSetting].(dto.ChannelSettings); ok && previous.Proxy != channel.GetSetting().Proxy {
|
||||
return types.NewErrorWithStatusCode(errors.New("upstream proxy changed; reconnect required"), types.ErrorCodeAccessDenied, http.StatusForbidden, types.ErrOptionWithSkipRetry())
|
||||
}
|
||||
// An advanced custom connection was dialed through the route matched for
|
||||
// this model. Only the fields that shape the upstream handshake force a
|
||||
// reconnect: target path, protocol conversion and credential placement.
|
||||
// Other channel types match the zero route on both sides.
|
||||
route, _ := channel.GetOtherSettings().AdvancedCustom.MatchPathForModel(c.Request.URL.Path, model)
|
||||
if strings.TrimSpace(route.UpstreamPath) != strings.TrimSpace(s.lockedRoute.UpstreamPath) ||
|
||||
route.IsNative() != s.lockedRoute.IsNative() ||
|
||||
!reflect.DeepEqual(route.Auth, s.lockedRoute.Auth) {
|
||||
return types.NewErrorWithStatusCode(errors.New("upstream route changed; reconnect required"), types.ErrorCodeAccessDenied, http.StatusForbidden, types.ErrOptionWithSkipRetry())
|
||||
}
|
||||
if pin, found, _ := service.GetChannelConstraints(c).ResolvedPin(); found {
|
||||
if pin.ChannelId != s.lockedChannelID {
|
||||
return types.NewErrorWithStatusCode(errors.New("channel pin changed; reconnect required"), types.ErrorCodeAccessDenied, http.StatusForbidden, types.ErrOptionWithSkipRetry())
|
||||
}
|
||||
} else {
|
||||
group := common.GetContextKeyString(c, appconstant.ContextKeyUsingGroup)
|
||||
if group == "auto" {
|
||||
if !slices.Contains(service.GetRequestAutoGroups(c, common.GetContextKeyString(c, appconstant.ContextKeyUserGroup)), s.lockedGroup) {
|
||||
return types.NewErrorWithStatusCode(errors.New("the connection group is no longer allowed"), types.ErrorCodeAccessDenied, http.StatusForbidden, types.ErrOptionWithSkipRetry())
|
||||
}
|
||||
group = s.lockedGroup
|
||||
common.SetContextKey(c, appconstant.ContextKeyAutoGroup, group)
|
||||
}
|
||||
if !appmodel.IsChannelEnabledForGroupModel(group, model, s.lockedChannelID) {
|
||||
return types.NewErrorWithStatusCode(errors.New("the connection channel is no longer allowed for this group and model"), types.ErrorCodeAccessDenied, http.StatusForbidden, types.ErrOptionWithSkipRetry())
|
||||
}
|
||||
}
|
||||
for key, value := range s.lockedContext {
|
||||
c.Set(string(key), value)
|
||||
}
|
||||
common.SetContextKey(c, appconstant.ContextKeyChannelSetting, channel.GetSetting())
|
||||
common.SetContextKey(c, appconstant.ContextKeyChannelOtherSetting, channel.GetOtherSettings())
|
||||
common.SetContextKey(c, appconstant.ContextKeyChannelParamOverride, channel.GetParamOverride())
|
||||
common.SetContextKey(c, appconstant.ContextKeyChannelModelMapping, channel.GetModelMapping())
|
||||
common.SetContextKey(c, appconstant.ContextKeyChannelStatusCodeMapping, channel.GetStatusCodeMapping())
|
||||
policy := service.RequestPolicy(c)
|
||||
policy.BeginAttempt(channel, s.lockedGroup)
|
||||
policy.AddEvent(service.PolicyEvent{ChannelID: channel.Id, Decision: service.PolicyDecision{Action: "select", Reason: "pinned_channel", Source: "channel_constraint"}})
|
||||
return nil
|
||||
}
|
||||
|
||||
func buildResponsesWSCreatePayload(c *gin.Context, info *relaycommon.RelayInfo, req dto.OpenAIResponsesRequest, generate common.RawMessage, streamID string) ([]byte, *types.NewAPIError) {
|
||||
_, body, closer, apiErr := PrepareResponsesRequest(c, info, &req)
|
||||
if apiErr != nil {
|
||||
return nil, apiErr
|
||||
}
|
||||
defer closer.Close()
|
||||
jsonData, err := io.ReadAll(body)
|
||||
if err != nil {
|
||||
return nil, types.NewError(err, types.ErrorCodeReadRequestBodyFailed, types.ErrOptionWithSkipRetry())
|
||||
}
|
||||
event, err := buildResponsesWSCreateEvent(jsonData, generate, streamID)
|
||||
if err != nil {
|
||||
return nil, types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
|
||||
}
|
||||
return event, nil
|
||||
}
|
||||
|
||||
func (s *responsesWSSession) startTargetReader(target *websocket.Conn) {
|
||||
limit := int64(helper.DefaultMaxScannerBufferSize)
|
||||
if appconstant.StreamScannerMaxBufferMB > 0 {
|
||||
limit = int64(appconstant.StreamScannerMaxBufferMB) << 20
|
||||
}
|
||||
target.SetReadLimit(limit)
|
||||
s.workers.Go(func() {
|
||||
for {
|
||||
kind, body, err := target.ReadMessage()
|
||||
incoming := responsesWSMessage{kind: kind, body: body, err: err}
|
||||
if state := s.getCurrent(); state != nil {
|
||||
select {
|
||||
case state.inbox <- incoming:
|
||||
case <-state.done:
|
||||
if err == nil {
|
||||
if writeErr := s.writeClient(kind, body); writeErr != nil {
|
||||
s.shutdown()
|
||||
return
|
||||
}
|
||||
} else {
|
||||
s.shutdown()
|
||||
}
|
||||
case <-s.ctx.Done():
|
||||
return
|
||||
}
|
||||
} else if err == nil {
|
||||
if writeErr := s.writeClient(kind, body); writeErr != nil {
|
||||
s.shutdown()
|
||||
return
|
||||
}
|
||||
} else {
|
||||
s.shutdown()
|
||||
}
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func (s *responsesWSSession) getCurrent() *responsesWSCallState {
|
||||
s.stateMu.Lock()
|
||||
defer s.stateMu.Unlock()
|
||||
return s.current
|
||||
}
|
||||
|
||||
func (s *responsesWSSession) tryReserveCurrent(state *responsesWSCallState) bool {
|
||||
s.stateMu.Lock()
|
||||
defer s.stateMu.Unlock()
|
||||
if s.current != nil {
|
||||
return false
|
||||
}
|
||||
s.current = state
|
||||
return true
|
||||
}
|
||||
|
||||
func (s *responsesWSSession) getTarget() *websocket.Conn {
|
||||
s.connectionMu.Lock()
|
||||
defer s.connectionMu.Unlock()
|
||||
return s.target
|
||||
}
|
||||
|
||||
func (s *responsesWSSession) setTarget(target *websocket.Conn) bool {
|
||||
s.connectionMu.Lock()
|
||||
defer s.connectionMu.Unlock()
|
||||
if s.ctx.Err() != nil {
|
||||
_ = target.Close()
|
||||
return false
|
||||
}
|
||||
s.target = target
|
||||
return true
|
||||
}
|
||||
|
||||
func (s *responsesWSSession) writeTarget(kind int, message []byte) error {
|
||||
s.targetWriteMu.Lock()
|
||||
defer s.targetWriteMu.Unlock()
|
||||
target := s.getTarget()
|
||||
if target == nil {
|
||||
return errors.New("responses websocket upstream is not connected")
|
||||
}
|
||||
if err := target.SetWriteDeadline(time.Now().Add(responsesWSWriteTimeout)); err != nil {
|
||||
return err
|
||||
}
|
||||
return target.WriteMessage(kind, message)
|
||||
}
|
||||
|
||||
func (s *responsesWSSession) writeClient(kind int, message []byte) error {
|
||||
s.clientWriteMu.Lock()
|
||||
defer s.clientWriteMu.Unlock()
|
||||
if err := s.client.SetWriteDeadline(time.Now().Add(responsesWSWriteTimeout)); err != nil {
|
||||
return err
|
||||
}
|
||||
return s.client.WriteMessage(kind, message)
|
||||
}
|
||||
|
||||
func (s *responsesWSSession) sendError(eventID, streamID string, apiErr *types.NewAPIError) {
|
||||
if apiErr == nil {
|
||||
return
|
||||
}
|
||||
payload, err := buildResponsesWSErrorPayload(eventID, streamID, apiErr)
|
||||
if err == nil {
|
||||
_ = s.writeClient(websocket.TextMessage, payload)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *responsesWSSession) closeTarget() {
|
||||
s.connectionMu.Lock()
|
||||
target, unregister := s.target, s.unregister
|
||||
s.target, s.unregister = nil, nil
|
||||
s.connectionMu.Unlock()
|
||||
if unregister != nil {
|
||||
unregister()
|
||||
}
|
||||
if target != nil {
|
||||
_ = target.Close()
|
||||
}
|
||||
}
|
||||
|
||||
func (s *responsesWSSession) shutdown() {
|
||||
s.cancel()
|
||||
s.closeTarget()
|
||||
_ = s.client.Close()
|
||||
}
|
||||
|
||||
func (s *responsesWSSession) registerChannelClose(channelID int) {
|
||||
unregister := wsmanager.Register(channelID, wsmanager.KindResponses, s.closeForPolicy)
|
||||
s.connectionMu.Lock()
|
||||
if s.ctx.Err() != nil {
|
||||
s.connectionMu.Unlock()
|
||||
unregister()
|
||||
return
|
||||
}
|
||||
s.unregister = unregister
|
||||
s.connectionMu.Unlock()
|
||||
}
|
||||
|
||||
// closeForPolicy tells both peers why the connection ends, like the realtime
|
||||
// relay does; WriteControl is safe alongside a blocked data writer.
|
||||
func (s *responsesWSSession) closeForPolicy(reason string) {
|
||||
closeMessage := websocket.FormatCloseMessage(websocket.ClosePolicyViolation, reason)
|
||||
deadline := time.Now().Add(time.Second)
|
||||
_ = s.client.WriteControl(websocket.CloseMessage, closeMessage, deadline)
|
||||
if target := s.getTarget(); target != nil {
|
||||
_ = target.WriteControl(websocket.CloseMessage, closeMessage, deadline)
|
||||
}
|
||||
s.shutdown()
|
||||
}
|
||||
|
||||
// Stream identity belongs to the WebSocket envelope. For the legacy wrapped
|
||||
// input, the top-level field takes precedence over response.stream_id.
|
||||
func parseResponsesWSEnvelope(message []byte) (responsesWSCreateEvent, string, error) {
|
||||
var event responsesWSCreateEvent
|
||||
if err := common.Unmarshal(message, &event); err != nil {
|
||||
return event, "", fmt.Errorf("invalid websocket event json: %w", err)
|
||||
}
|
||||
streamRaw := event.StreamID
|
||||
if len(streamRaw) == 0 && len(event.Request) > 0 {
|
||||
var wrapped struct {
|
||||
StreamID common.RawMessage `json:"stream_id"`
|
||||
}
|
||||
if err := common.Unmarshal(event.Request, &wrapped); err == nil {
|
||||
streamRaw = wrapped.StreamID
|
||||
}
|
||||
}
|
||||
var streamID string
|
||||
if len(streamRaw) > 0 {
|
||||
if err := common.Unmarshal(streamRaw, &streamID); err != nil || len(streamID) < 1 || len(streamID) > 256 {
|
||||
return event, "", errors.New("stream_id must contain 1-256 ASCII letters, digits, underscores, hyphens, or periods")
|
||||
}
|
||||
for _, char := range streamID {
|
||||
if !(char >= 'a' && char <= 'z' || char >= 'A' && char <= 'Z' || char >= '0' && char <= '9' || char == '_' || char == '-' || char == '.') {
|
||||
return event, "", errors.New("stream_id must contain only ASCII letters, digits, underscores, hyphens, or periods")
|
||||
}
|
||||
}
|
||||
}
|
||||
if strings.TrimSpace(event.Type) == "" {
|
||||
return event, streamID, errors.New("websocket event type is required")
|
||||
}
|
||||
return event, streamID, nil
|
||||
}
|
||||
|
||||
func newResponsesWSInvalidRequestError(err error) *types.NewAPIError {
|
||||
return types.NewErrorWithStatusCode(err, types.ErrorCodeInvalidRequest, http.StatusBadRequest, types.ErrOptionWithSkipRetry())
|
||||
}
|
||||
|
||||
// normalizeResponsesWSCreateEvent builds the HTTP-shaped request body from a
|
||||
// response.create whose envelope the read loop already parsed and validated.
|
||||
// The message is decoded once more as a raw map (the body source) and once as
|
||||
// the typed request; the envelope is not parsed again.
|
||||
func normalizeResponsesWSCreateEvent(message []byte, event responsesWSCreateEvent, streamID string) (responsesWSCreateRequest, error) {
|
||||
create := responsesWSCreateRequest{StreamID: streamID, Generate: event.Generate}
|
||||
// Decode the wrapped request on its own so outer fields cannot enter the HTTP body.
|
||||
source := message
|
||||
if len(event.Request) > 0 {
|
||||
source = event.Request
|
||||
}
|
||||
var raw map[string]common.RawMessage
|
||||
if err := common.Unmarshal(source, &raw); err != nil {
|
||||
return create, err
|
||||
}
|
||||
if len(event.Request) > 0 {
|
||||
if len(create.Generate) == 0 {
|
||||
create.Generate = raw["generate"]
|
||||
}
|
||||
} else {
|
||||
for _, key := range []string{"type", "event_id", "background", "stream", "stream_options"} {
|
||||
delete(raw, key)
|
||||
}
|
||||
}
|
||||
delete(raw, "generate")
|
||||
delete(raw, "stream_id")
|
||||
payload, err := common.Marshal(raw)
|
||||
if err != nil {
|
||||
return create, err
|
||||
}
|
||||
if err := common.Unmarshal(payload, &create.Request); err != nil {
|
||||
return create, err
|
||||
}
|
||||
if helper.ExceedsMaxTokensLimit(create.Request.MaxOutputTokens) {
|
||||
return create, errors.New("max_output_tokens is invalid")
|
||||
}
|
||||
create.Request.Stream = nil
|
||||
create.Request.StreamOptions = nil
|
||||
create.Body = payload
|
||||
return create, nil
|
||||
}
|
||||
|
||||
func buildResponsesWSCreateEvent(jsonData []byte, generate common.RawMessage, streamID string) ([]byte, error) {
|
||||
var event map[string]common.RawMessage
|
||||
if err := common.Unmarshal(jsonData, &event); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
typeData, err := common.Marshal(responsesWSEventTypeResponseCreate)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
event["type"] = typeData
|
||||
// The normalized body never carries stream_id; a channel parameter
|
||||
// override may inject one, and it must not shadow the envelope's value.
|
||||
delete(event, "stream_id")
|
||||
if streamID != "" {
|
||||
streamData, err := common.Marshal(streamID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
event["stream_id"] = streamData
|
||||
}
|
||||
delete(event, "event_id")
|
||||
delete(event, "background")
|
||||
delete(event, "stream")
|
||||
delete(event, "stream_options")
|
||||
if len(generate) > 0 {
|
||||
event["generate"] = generate
|
||||
}
|
||||
return common.Marshal(event)
|
||||
}
|
||||
|
||||
func buildResponsesWSErrorPayload(eventID, streamID string, apiErr *types.NewAPIError) ([]byte, error) {
|
||||
if apiErr == nil {
|
||||
return nil, errors.New("api error is nil")
|
||||
}
|
||||
status := apiErr.StatusCode
|
||||
if status == 0 {
|
||||
status = http.StatusInternalServerError
|
||||
}
|
||||
openaiErr := apiErr.ToOpenAIError()
|
||||
return common.Marshal(&responsesWSErrorEvent{
|
||||
Type: "error",
|
||||
Status: status,
|
||||
EventID: eventID,
|
||||
StreamID: streamID,
|
||||
Error: &openaiErr,
|
||||
})
|
||||
}
|
||||
|
||||
// checkResponsesWSModelAccess applies the token model limit with the same name
|
||||
// matching as the HTTP distributor. Unlike HTTP, it also runs for requests
|
||||
// pinned to a channel: the check is stricter there on purpose, because a
|
||||
// persistent connection keeps serving the model after the pin was resolved.
|
||||
func checkResponsesWSModelAccess(c *gin.Context, modelName string) *types.NewAPIError {
|
||||
if !common.GetContextKeyBool(c, appconstant.ContextKeyTokenModelLimitEnabled) {
|
||||
return nil
|
||||
}
|
||||
raw, ok := common.GetContextKey(c, appconstant.ContextKeyTokenModelLimit)
|
||||
if !ok {
|
||||
return types.NewErrorWithStatusCode(errors.New("token has no model access"), types.ErrorCodeAccessDenied, http.StatusForbidden, types.ErrOptionWithSkipRetry())
|
||||
}
|
||||
tokenModelLimit, ok := raw.(map[string]bool)
|
||||
if !ok {
|
||||
tokenModelLimit = map[string]bool{}
|
||||
}
|
||||
if !middleware.TokenModelLimitAllows(tokenModelLimit, modelName) {
|
||||
return types.NewErrorWithStatusCode(fmt.Errorf("token is not allowed to use model %s", modelName), types.ErrorCodeAccessDenied, http.StatusForbidden, types.ErrOptionWithSkipRetry())
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// selectResponsesWSChannel narrows the shared HTTP selection to channels that
|
||||
// speak the Responses WebSocket protocol and renders its outcome as a
|
||||
// non-retryable NewAPIError.
|
||||
func selectResponsesWSChannel(c *gin.Context, modelName string, retryParam *service.RetryParam) (*appmodel.Channel, *types.NewAPIError) {
|
||||
constraints := service.GetChannelConstraints(c)
|
||||
if !slices.ContainsFunc(constraints.Filters, func(filter appdto.ChannelFilter) bool {
|
||||
return filter.Kind == appdto.FilterResponsesWebSocket
|
||||
}) {
|
||||
constraints.AddFilter(appdto.ChannelFilter{Kind: appdto.FilterResponsesWebSocket})
|
||||
}
|
||||
channel, _, selectErr := service.SelectChannelForRequest(c, modelName, retryParam)
|
||||
if selectErr != nil {
|
||||
message := selectErr.Message
|
||||
if selectErr.MessageID != "" {
|
||||
message = i18n.T(c, selectErr.MessageID, selectErr.Params)
|
||||
}
|
||||
code := selectErr.Code
|
||||
if code == "" {
|
||||
code = types.ErrorCodeGetChannelFailed
|
||||
}
|
||||
return nil, types.NewErrorWithStatusCode(errors.New(message), code, selectErr.StatusCode, types.ErrOptionWithSkipRetry())
|
||||
}
|
||||
if err := middleware.SetupContextForSelectedChannel(c, channel, modelName); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return channel, nil
|
||||
}
|
||||
@@ -0,0 +1,682 @@
|
||||
package relay
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
appdto "github.com/QuantumNous/new-api/dto"
|
||||
"github.com/QuantumNous/new-api/i18n"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
"github.com/QuantumNous/new-api/relaykit/types"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/QuantumNous/new-api/setting/operation_setting"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/gorilla/websocket"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// normalizeResponsesWSTestMessage runs the read-loop envelope parse followed by
|
||||
// request normalization, exactly as the session does for one response.create.
|
||||
func normalizeResponsesWSTestMessage(message []byte) (responsesWSCreateRequest, string, error) {
|
||||
envelope, streamID, err := parseResponsesWSEnvelope(message)
|
||||
if err != nil {
|
||||
return responsesWSCreateRequest{StreamID: streamID}, envelope.EventID, err
|
||||
}
|
||||
create, err := normalizeResponsesWSCreateEvent(message, envelope, streamID)
|
||||
return create, envelope.EventID, err
|
||||
}
|
||||
|
||||
func TestNormalizeResponsesWSMaxOutputTokens(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
value string
|
||||
valid bool
|
||||
}{
|
||||
{value: "0", valid: true},
|
||||
{value: "1073741823", valid: true},
|
||||
{value: "1073741824"},
|
||||
{value: "18446744073686646784"},
|
||||
{value: "-1"},
|
||||
} {
|
||||
for _, wrapped := range []bool{false, true} {
|
||||
t.Run(fmt.Sprintf("%s/wrapped=%t", tc.value, wrapped), func(t *testing.T) {
|
||||
fields := `"model":"gpt-5.1","input":"hi","max_output_tokens":` + tc.value
|
||||
payload := `{"type":"response.create",` + fields + `}`
|
||||
if wrapped {
|
||||
payload = `{"type":"response.create","response":{` + fields + `}}`
|
||||
}
|
||||
create, _, err := normalizeResponsesWSTestMessage([]byte(payload))
|
||||
if !tc.valid {
|
||||
require.Error(t, err)
|
||||
assert.Equal(t, http.StatusBadRequest, newResponsesWSInvalidRequestError(err).StatusCode)
|
||||
return
|
||||
}
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, create.Request.MaxOutputTokens)
|
||||
assert.Equal(t, tc.value, fmt.Sprint(*create.Request.MaxOutputTokens))
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSelectResponsesWSChannelHonorsPinsAndFilters(t *testing.T) {
|
||||
require.NoError(t, i18n.Init())
|
||||
database := setupRelayChannelDB(t)
|
||||
enabled := &model.Channel{Name: "enabled", Key: "sk-test", Status: common.ChannelStatusEnabled, Type: constant.ChannelTypeOpenAI}
|
||||
enabled.SetSetting(dto.ChannelSettings{ResponsesWebSocketEnabled: true})
|
||||
wsDisabled := &model.Channel{Name: "ws-disabled", Key: "sk-test", Status: common.ChannelStatusEnabled, Type: constant.ChannelTypeOpenAI}
|
||||
disabled := &model.Channel{Name: "disabled", Key: "sk-test", Status: common.ChannelStatusManuallyDisabled, Type: constant.ChannelTypeOpenAI}
|
||||
filtered := &model.Channel{Name: "filtered", Key: "sk-test", Status: common.ChannelStatusEnabled, Type: constant.ChannelTypeAdvancedCustom}
|
||||
for _, channel := range []*model.Channel{enabled, disabled, filtered, wsDisabled} {
|
||||
require.NoError(t, database.Create(channel).Error)
|
||||
}
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
channelID int
|
||||
status int
|
||||
}{
|
||||
{name: "token pin overrides origin pin", channelID: enabled.Id},
|
||||
{name: "disabled pin rejects", channelID: disabled.Id, status: http.StatusForbidden},
|
||||
{name: "pin cannot bypass websocket switch", channelID: wsDisabled.Id, status: http.StatusBadRequest},
|
||||
{name: "pin cannot bypass path filter", channelID: filtered.Id, status: http.StatusBadRequest},
|
||||
{name: "missing pin rejects", channelID: 99999, status: http.StatusBadRequest},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil)
|
||||
constraints := service.GetChannelConstraints(c)
|
||||
constraints.AddPin(appdto.ChannelPin{ChannelId: disabled.Id, Source: appdto.PinSourceOriginTask, Rank: appdto.PinRankOriginTask, RetryMode: appdto.PinRetrySameChannel})
|
||||
constraints.AddPin(appdto.ChannelPin{ChannelId: tc.channelID, Source: appdto.PinSourceToken, Rank: appdto.PinRankToken, RetryMode: appdto.PinRetrySingleAttempt})
|
||||
constraints.AddFilter(appdto.ChannelFilter{Kind: appdto.FilterRequestPath, RequestPath: c.Request.URL.Path})
|
||||
channel, apiErr := selectResponsesWSChannel(c, "gpt-5.1", &service.RetryParam{Ctx: c, ModelName: "gpt-5.1", TokenGroup: "default"})
|
||||
if tc.status != 0 {
|
||||
require.NotNil(t, apiErr)
|
||||
assert.Equal(t, tc.status, apiErr.StatusCode)
|
||||
assert.Nil(t, channel)
|
||||
assert.False(t, service.ShouldRetryRelayError(c, apiErr, 2))
|
||||
return
|
||||
}
|
||||
require.Nil(t, apiErr)
|
||||
require.NotNil(t, channel)
|
||||
assert.Equal(t, enabled.Id, channel.Id)
|
||||
assert.Equal(t, enabled.Id, common.GetContextKeyInt(c, constant.ContextKeyChannelId))
|
||||
assert.False(t, service.ShouldRetryRelayError(c, types.NewErrorWithStatusCode(errors.New("upstream failed"), types.ErrorCodeDoRequestFailed, 503), 2))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// The WebSocket relay must admit the same model names as the HTTP distributor
|
||||
// when a token restricts models (reasoning suffixes and @modifiers included).
|
||||
func TestCheckResponsesWSModelAccessMatchesHTTPTokenLimits(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
model string
|
||||
status int
|
||||
}{
|
||||
{model: "gpt-5.1"},
|
||||
{model: "gpt-5.1-high"},
|
||||
{model: "gpt-5.1@thinking:on"},
|
||||
{model: "gpt-4o", status: http.StatusForbidden},
|
||||
} {
|
||||
t.Run(tc.model, func(t *testing.T) {
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil)
|
||||
common.SetContextKey(c, constant.ContextKeyTokenModelLimitEnabled, true)
|
||||
common.SetContextKey(c, constant.ContextKeyTokenModelLimit, map[string]bool{"gpt-5.1": true})
|
||||
apiErr := checkResponsesWSModelAccess(c, tc.model)
|
||||
if tc.status == 0 {
|
||||
assert.Nil(t, apiErr)
|
||||
return
|
||||
}
|
||||
require.NotNil(t, apiErr)
|
||||
assert.Equal(t, tc.status, apiErr.StatusCode)
|
||||
assert.False(t, service.ShouldRetryRelayError(c, apiErr, 2))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSelectResponsesWSChannelAcceptsNativeResponsesChannelTypes(t *testing.T) {
|
||||
require.NoError(t, i18n.Init())
|
||||
database := setupRelayChannelDB(t)
|
||||
nativeRoute := &dto.AdvancedCustomConfig{Routes: []dto.AdvancedCustomRoute{{IncomingPath: "/v1/responses", UpstreamPath: "/v1/responses"}}}
|
||||
noneRoute := &dto.AdvancedCustomConfig{Routes: []dto.AdvancedCustomRoute{{IncomingPath: "/v1/responses", UpstreamPath: "/v1/responses", Converter: "none"}}}
|
||||
convertedRoute := &dto.AdvancedCustomConfig{Routes: []dto.AdvancedCustomRoute{{IncomingPath: "/v1/responses", UpstreamPath: "/v1/chat/completions", Converter: "openai_responses_to_openai_chat_completions"}}}
|
||||
chatRoute := &dto.AdvancedCustomConfig{Routes: []dto.AdvancedCustomRoute{{IncomingPath: "/v1/chat/completions", UpstreamPath: "/v1/chat/completions"}}}
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
channelType int
|
||||
advanced *dto.AdvancedCustomConfig
|
||||
rejectedBy appdto.ChannelFilterKind
|
||||
}{
|
||||
{name: "new api", channelType: constant.ChannelTypeNewAPI},
|
||||
{name: "sub2api", channelType: constant.ChannelTypeSub2API},
|
||||
{name: "advanced custom native responses route", channelType: constant.ChannelTypeAdvancedCustom, advanced: nativeRoute},
|
||||
{name: "advanced custom explicit none converter", channelType: constant.ChannelTypeAdvancedCustom, advanced: noneRoute},
|
||||
{name: "advanced custom converter route", channelType: constant.ChannelTypeAdvancedCustom, advanced: convertedRoute, rejectedBy: appdto.FilterResponsesWebSocket},
|
||||
{name: "advanced custom without responses route", channelType: constant.ChannelTypeAdvancedCustom, advanced: chatRoute, rejectedBy: appdto.FilterRequestPath},
|
||||
{name: "anthropic", channelType: constant.ChannelTypeAnthropic, rejectedBy: appdto.FilterResponsesWebSocket},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
channel := &model.Channel{Name: tc.name, Key: "sk-test", Status: common.ChannelStatusEnabled, Type: tc.channelType}
|
||||
channel.SetSetting(dto.ChannelSettings{ResponsesWebSocketEnabled: true})
|
||||
channel.SetOtherSettings(dto.ChannelOtherSettings{AdvancedCustom: tc.advanced})
|
||||
require.NoError(t, database.Create(channel).Error)
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil)
|
||||
constraints := service.GetChannelConstraints(c)
|
||||
constraints.AddPin(appdto.ChannelPin{ChannelId: channel.Id, Source: appdto.PinSourceToken, Rank: appdto.PinRankToken, RetryMode: appdto.PinRetrySingleAttempt})
|
||||
constraints.AddFilter(appdto.ChannelFilter{Kind: appdto.FilterRequestPath, RequestPath: c.Request.URL.Path})
|
||||
selected, apiErr := selectResponsesWSChannel(c, "ws-model", &service.RetryParam{Ctx: c, ModelName: "ws-model", TokenGroup: "default"})
|
||||
if tc.rejectedBy != "" {
|
||||
require.NotNil(t, apiErr)
|
||||
assert.Nil(t, selected)
|
||||
assert.Equal(t, http.StatusBadRequest, apiErr.StatusCode)
|
||||
assert.Equal(t, types.ErrorCode(tc.rejectedBy), apiErr.GetErrorCode())
|
||||
return
|
||||
}
|
||||
require.Nil(t, apiErr)
|
||||
require.NotNil(t, selected)
|
||||
assert.Equal(t, channel.Id, selected.Id)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestRestoreConnectionContextRejectsChangedAdvancedCustomRoute(t *testing.T) {
|
||||
database := setupRelayChannelDB(t)
|
||||
baseURL := "http://upstream.example"
|
||||
channel := &model.Channel{Name: "advanced", Key: "sk-test", Status: common.ChannelStatusEnabled, Type: constant.ChannelTypeAdvancedCustom, BaseURL: &baseURL}
|
||||
channel.SetSetting(dto.ChannelSettings{ResponsesWebSocketEnabled: true})
|
||||
channel.SetOtherSettings(dto.ChannelOtherSettings{AdvancedCustom: &dto.AdvancedCustomConfig{Routes: []dto.AdvancedCustomRoute{{IncomingPath: "/v1/responses", UpstreamPath: "/v1/responses"}}}})
|
||||
require.NoError(t, database.Create(channel).Error)
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil)
|
||||
service.GetChannelConstraints(c).AddPin(appdto.ChannelPin{ChannelId: channel.Id, Source: appdto.PinSourceToken, Rank: appdto.PinRankToken, RetryMode: appdto.PinRetrySingleAttempt})
|
||||
route, ok := channel.GetOtherSettings().AdvancedCustom.MatchPathForModel("/v1/responses", "ws-model")
|
||||
require.True(t, ok)
|
||||
session := &responsesWSSession{lockedChannelID: channel.Id, lockedModel: "ws-model", lockedKey: "sk-test", lockedRoute: route,
|
||||
lockedContext: map[constant.ContextKey]any{constant.ContextKeyChannelBaseUrl: channel.GetBaseURL(), constant.ContextKeyChannelHeaderOverride: channel.GetHeaderOverride()}}
|
||||
require.Nil(t, session.restoreConnectionContext(c, "ws-model"))
|
||||
|
||||
// Request-level edits (model list, explicit none converter) keep the connection.
|
||||
channel.SetOtherSettings(dto.ChannelOtherSettings{AdvancedCustom: &dto.AdvancedCustomConfig{Routes: []dto.AdvancedCustomRoute{{IncomingPath: "/v1/responses", UpstreamPath: "/v1/responses", Converter: "none", Models: []string{"ws-model", "other-model"}}}}})
|
||||
require.NoError(t, database.Model(channel).Update("settings", channel.OtherSettings).Error)
|
||||
require.Nil(t, session.restoreConnectionContext(c, "ws-model"))
|
||||
|
||||
channel.SetOtherSettings(dto.ChannelOtherSettings{AdvancedCustom: &dto.AdvancedCustomConfig{Routes: []dto.AdvancedCustomRoute{{IncomingPath: "/v1/responses", UpstreamPath: "/v2/responses"}}}})
|
||||
require.NoError(t, database.Model(channel).Update("settings", channel.OtherSettings).Error)
|
||||
apiErr := session.restoreConnectionContext(c, "ws-model")
|
||||
require.NotNil(t, apiErr)
|
||||
assert.Equal(t, http.StatusForbidden, apiErr.StatusCode)
|
||||
assert.ErrorContains(t, apiErr, "upstream route changed")
|
||||
}
|
||||
|
||||
func TestResponsesWSChannelRoutingRequiresExplicitOptIn(t *testing.T) {
|
||||
require.NoError(t, i18n.Init())
|
||||
database := setupRelayChannelDB(t)
|
||||
require.NoError(t, database.AutoMigrate(&model.Ability{}))
|
||||
legacy := &model.Channel{Name: "legacy-http", Key: "sk-test", Type: constant.ChannelTypeOpenAI, Status: common.ChannelStatusEnabled, Group: "default", Models: "ws-model", Priority: common.GetPointer(int64(10))}
|
||||
enabled := &model.Channel{Name: "websocket", Key: "sk-test", Type: constant.ChannelTypeCodex, Status: common.ChannelStatusEnabled, Group: "default", Models: "ws-model", Priority: common.GetPointer(int64(0))}
|
||||
unsupported := &model.Channel{Name: "unsupported", Key: "sk-test", Type: constant.ChannelTypeAnthropic, Status: common.ChannelStatusEnabled, Group: "default", Models: "ws-model", Priority: common.GetPointer(int64(5))}
|
||||
enabled.SetSetting(dto.ChannelSettings{ResponsesWebSocketEnabled: true})
|
||||
unsupported.SetSetting(dto.ChannelSettings{ResponsesWebSocketEnabled: true})
|
||||
for _, channel := range []*model.Channel{legacy, enabled, unsupported} {
|
||||
require.NoError(t, database.Create(channel).Error)
|
||||
require.NoError(t, database.Create(&model.Ability{ChannelId: channel.Id, Model: "ws-model", Group: "default", Enabled: true, Priority: channel.Priority}).Error)
|
||||
}
|
||||
previousCache := common.MemoryCacheEnabled
|
||||
t.Cleanup(func() {
|
||||
defer func() { common.MemoryCacheEnabled = previousCache }()
|
||||
ids := []int{legacy.Id, enabled.Id, unsupported.Id}
|
||||
require.NoError(t, database.Where("channel_id IN ?", ids).Delete(&model.Ability{}).Error)
|
||||
require.NoError(t, database.Where("id IN ?", ids).Delete(&model.Channel{}).Error)
|
||||
model.InitChannelCache()
|
||||
})
|
||||
common.MemoryCacheEnabled = true
|
||||
model.InitChannelCache()
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil)
|
||||
params := &service.RetryParam{Ctx: c, ModelName: "ws-model", TokenGroup: "default"}
|
||||
channel, apiErr := selectResponsesWSChannel(c, "ws-model", params)
|
||||
require.Nil(t, apiErr)
|
||||
require.NotNil(t, channel)
|
||||
assert.Equal(t, enabled.Id, channel.Id)
|
||||
httpChannel, err := model.GetRandomSatisfiedChannel("default", "ws-model", 0, nil)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, httpChannel)
|
||||
assert.Equal(t, legacy.Id, httpChannel.Id)
|
||||
|
||||
// Disabling the saved setting takes effect for the next create on an existing session.
|
||||
enabled.SetSetting(dto.ChannelSettings{ResponsesWebSocketEnabled: false})
|
||||
require.NoError(t, database.Model(enabled).Update("setting", enabled.Setting).Error)
|
||||
model.InitChannelCache()
|
||||
session := &responsesWSSession{lockedChannelID: enabled.Id, lockedModel: "ws-model"}
|
||||
apiErr = session.restoreConnectionContext(c, "ws-model")
|
||||
require.NotNil(t, apiErr)
|
||||
assert.Equal(t, http.StatusForbidden, apiErr.StatusCode)
|
||||
channel, apiErr = selectResponsesWSChannel(c, "ws-model", params)
|
||||
require.NotNil(t, apiErr)
|
||||
assert.Nil(t, channel)
|
||||
}
|
||||
|
||||
// Session affinity follows the HTTP distributor: a strict binding to an
|
||||
// unusable channel fails the request, a prefer binding falls back, and an
|
||||
// unusable binding is dropped from the cache either way.
|
||||
func TestSelectResponsesWSChannelHonorsStrictSessionBinding(t *testing.T) {
|
||||
require.NoError(t, i18n.Init())
|
||||
database := setupRelayChannelDB(t)
|
||||
require.NoError(t, database.AutoMigrate(&model.Ability{}))
|
||||
bound := &model.Channel{Name: "bound", Key: "sk-test", Type: constant.ChannelTypeOpenAI, Status: common.ChannelStatusAutoDisabled, Group: "default", Models: "ws-model", Priority: common.GetPointer(int64(10))}
|
||||
fallback := &model.Channel{Name: "fallback", Key: "sk-test", Type: constant.ChannelTypeOpenAI, Status: common.ChannelStatusEnabled, Group: "default", Models: "ws-model", Priority: common.GetPointer(int64(0))}
|
||||
for _, channel := range []*model.Channel{bound, fallback} {
|
||||
channel.SetSetting(dto.ChannelSettings{ResponsesWebSocketEnabled: true})
|
||||
require.NoError(t, database.Create(channel).Error)
|
||||
require.NoError(t, database.Create(&model.Ability{ChannelId: channel.Id, Model: "ws-model", Group: "default", Enabled: channel.Status == common.ChannelStatusEnabled, Priority: channel.Priority}).Error)
|
||||
}
|
||||
previousCache := common.MemoryCacheEnabled
|
||||
t.Cleanup(func() {
|
||||
defer func() { common.MemoryCacheEnabled = previousCache }()
|
||||
ids := []int{bound.Id, fallback.Id}
|
||||
require.NoError(t, database.Where("channel_id IN ?", ids).Delete(&model.Ability{}).Error)
|
||||
require.NoError(t, database.Where("id IN ?", ids).Delete(&model.Channel{}).Error)
|
||||
model.InitChannelCache()
|
||||
})
|
||||
common.MemoryCacheEnabled = true
|
||||
model.InitChannelCache()
|
||||
affinity := operation_setting.GetChannelAffinitySetting()
|
||||
previousAffinity := *affinity
|
||||
t.Cleanup(func() { *affinity = previousAffinity })
|
||||
snapshot, err := model.BuildRequestPolicy(map[string]string{
|
||||
"channel_affinity_setting.enabled": "true",
|
||||
"channel_affinity_setting.session_mode": "strict",
|
||||
"channel_affinity_setting.rules": `[{"name":"session","model_regex":[".*"],"key_sources":[{"type":"request_header","key":"X-Session"}],"session_mode":"inherit"}]`,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
*affinity = snapshot.Affinity
|
||||
newSessionContext := func() *gin.Context {
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil)
|
||||
c.Request.Header.Set("X-Session", t.Name())
|
||||
return c
|
||||
}
|
||||
seed := newSessionContext()
|
||||
_, found := service.GetPreferredChannelByAffinity(seed, "ws-model", "default")
|
||||
require.False(t, found)
|
||||
seed.Set("channel_id", bound.Id)
|
||||
service.RecordChannelAffinity(seed, bound.Id)
|
||||
t.Cleanup(func() { service.ClearCurrentChannelAffinityCache(seed) })
|
||||
|
||||
strict := newSessionContext()
|
||||
channel, apiErr := selectResponsesWSChannel(strict, "ws-model", &service.RetryParam{Ctx: strict, ModelName: "ws-model", TokenGroup: "default", Retry: common.GetPointer(0)})
|
||||
require.NotNil(t, apiErr)
|
||||
assert.Nil(t, channel)
|
||||
assert.Equal(t, http.StatusServiceUnavailable, apiErr.StatusCode)
|
||||
assert.False(t, service.ShouldRetryRelayError(strict, apiErr, 2))
|
||||
_, found = service.GetPreferredChannelByAffinity(seed, "ws-model", "default")
|
||||
assert.False(t, found, "an unusable binding is cleared unless keep_on_channel_disabled is set")
|
||||
|
||||
affinity.SessionMode = "prefer"
|
||||
service.RecordChannelAffinity(seed, bound.Id)
|
||||
prefer := newSessionContext()
|
||||
channel, apiErr = selectResponsesWSChannel(prefer, "ws-model", &service.RetryParam{Ctx: prefer, ModelName: "ws-model", TokenGroup: "default", Retry: common.GetPointer(0)})
|
||||
require.Nil(t, apiErr)
|
||||
require.NotNil(t, channel)
|
||||
assert.Equal(t, fallback.Id, channel.Id)
|
||||
}
|
||||
|
||||
func TestNormalizeResponsesWSCreateEventWrapper(t *testing.T) {
|
||||
message := []byte(`{
|
||||
"type": "response.create",
|
||||
"event_id": "evt_1",
|
||||
"generate": false,
|
||||
"response": {
|
||||
"model": "gpt-5.3-codex-spark",
|
||||
"input": "hi",
|
||||
"store": false,
|
||||
"stream": true,
|
||||
"stream_options": {"include_usage": true}
|
||||
}
|
||||
}`)
|
||||
|
||||
create, eventID, err := normalizeResponsesWSTestMessage(message)
|
||||
if err != nil {
|
||||
t.Fatalf("normalizeResponsesWSTestMessage() error = %v", err)
|
||||
}
|
||||
req := create.Request
|
||||
if eventID != "evt_1" {
|
||||
t.Fatalf("eventID = %q, want evt_1", eventID)
|
||||
}
|
||||
if req.Model != "gpt-5.3-codex-spark" {
|
||||
t.Fatalf("model = %q", req.Model)
|
||||
}
|
||||
if strings.TrimSpace(string(create.Generate)) != "false" {
|
||||
t.Fatalf("generate = %s, want false", create.Generate)
|
||||
}
|
||||
if req.Stream != nil {
|
||||
t.Fatalf("stream = %v, want nil", req.Stream)
|
||||
}
|
||||
if req.StreamOptions != nil {
|
||||
t.Fatalf("stream_options = %#v, want nil", req.StreamOptions)
|
||||
}
|
||||
if strings.TrimSpace(string(req.Store)) != "false" {
|
||||
t.Fatalf("store = %s, want false", req.Store)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNormalizeResponsesWSCreateEventFlat(t *testing.T) {
|
||||
message := []byte(`{
|
||||
"type": "response.create",
|
||||
"event_id": "evt_2",
|
||||
"model": "gpt-5.3-codex-spark",
|
||||
"input": "hi",
|
||||
"generate": false,
|
||||
"stream": true,
|
||||
"background": true,
|
||||
"stream_options": {"include_usage": true}
|
||||
}`)
|
||||
|
||||
create, eventID, err := normalizeResponsesWSTestMessage(message)
|
||||
if err != nil {
|
||||
t.Fatalf("normalizeResponsesWSTestMessage() error = %v", err)
|
||||
}
|
||||
req := create.Request
|
||||
if eventID != "evt_2" {
|
||||
t.Fatalf("eventID = %q, want evt_2", eventID)
|
||||
}
|
||||
if req.Model != "gpt-5.3-codex-spark" {
|
||||
t.Fatalf("model = %q", req.Model)
|
||||
}
|
||||
if strings.TrimSpace(string(create.Generate)) != "false" {
|
||||
t.Fatalf("generate = %s, want false", create.Generate)
|
||||
}
|
||||
if req.Stream != nil {
|
||||
t.Fatalf("stream = %v, want nil", req.Stream)
|
||||
}
|
||||
if req.StreamOptions != nil {
|
||||
t.Fatalf("stream_options = %#v, want nil", req.StreamOptions)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildResponsesWSCreateEventIsFlat(t *testing.T) {
|
||||
payload := []byte(`{
|
||||
"model": "gpt-5.3-codex-spark",
|
||||
"input": "hi",
|
||||
"store": false,
|
||||
"event_id": "evt_upstream",
|
||||
"stream": true,
|
||||
"background": true,
|
||||
"stream_options": {"include_usage": true}
|
||||
}`)
|
||||
|
||||
got, err := buildResponsesWSCreateEvent(payload, common.RawMessage(`false`), "")
|
||||
if err != nil {
|
||||
t.Fatalf("buildResponsesWSCreateEvent() error = %v", err)
|
||||
}
|
||||
var data map[string]any
|
||||
if err := common.Unmarshal(got, &data); err != nil {
|
||||
t.Fatalf("unmarshal result: %v", err)
|
||||
}
|
||||
if data["type"] != responsesWSEventTypeResponseCreate {
|
||||
t.Fatalf("type = %#v", data["type"])
|
||||
}
|
||||
if data["model"] != "gpt-5.3-codex-spark" || data["input"] != "hi" || data["store"] != false {
|
||||
t.Fatalf("unexpected flat event fields: %s", got)
|
||||
}
|
||||
if data["generate"] != false {
|
||||
t.Fatalf("generate = %#v, want false", data["generate"])
|
||||
}
|
||||
for _, key := range []string{"response", "event_id", "stream", "background", "stream_options"} {
|
||||
if _, ok := data[key]; ok {
|
||||
t.Fatalf("field %q should not be present in upstream event: %s", key, got)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestHTTPResponsesRequestOmitsWebSocketMetadata(t *testing.T) {
|
||||
var req dto.OpenAIResponsesRequest
|
||||
require.NoError(t, common.Unmarshal([]byte(`{"model":"gpt-5.3-codex-spark","input":"hi","generate":false,"stream_id":"planner"}`), &req))
|
||||
got, err := common.Marshal(req)
|
||||
require.NoError(t, err)
|
||||
var data map[string]any
|
||||
require.NoError(t, common.Unmarshal(got, &data))
|
||||
assert.NotContains(t, data, "stream_id")
|
||||
assert.NotContains(t, data, "generate")
|
||||
}
|
||||
|
||||
func TestBuildResponsesWSErrorPayloadIncludesStatus(t *testing.T) {
|
||||
payload, err := buildResponsesWSErrorPayload("evt_err", "", types.NewErrorWithStatusCode(
|
||||
errors.New("model is required"),
|
||||
types.ErrorCodeInvalidRequest,
|
||||
http.StatusBadRequest,
|
||||
types.ErrOptionWithSkipRetry(),
|
||||
))
|
||||
if err != nil {
|
||||
t.Fatalf("buildResponsesWSErrorPayload() error = %v", err)
|
||||
}
|
||||
var data struct {
|
||||
Type string `json:"type"`
|
||||
Status int `json:"status"`
|
||||
EventID string `json:"event_id"`
|
||||
Error *types.OpenAIError `json:"error"`
|
||||
}
|
||||
if err := common.Unmarshal(payload, &data); err != nil {
|
||||
t.Fatalf("unmarshal result: %v", err)
|
||||
}
|
||||
if data.Type != "error" || data.Status != http.StatusBadRequest || data.EventID != "evt_err" {
|
||||
t.Fatalf("unexpected error event: %s", payload)
|
||||
}
|
||||
if data.Error == nil || data.Error.Code != string(types.ErrorCodeInvalidRequest) {
|
||||
t.Fatalf("unexpected error body: %#v", data.Error)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResponsesWSInvalidRequestErrorUsesBadRequestStatus(t *testing.T) {
|
||||
payload, err := buildResponsesWSErrorPayload("", "", newResponsesWSInvalidRequestError(errors.New("bad event")))
|
||||
if err != nil {
|
||||
t.Fatalf("buildResponsesWSErrorPayload() error = %v", err)
|
||||
}
|
||||
var data struct {
|
||||
Status int `json:"status"`
|
||||
}
|
||||
if err := common.Unmarshal(payload, &data); err != nil {
|
||||
t.Fatalf("unmarshal result: %v", err)
|
||||
}
|
||||
if data.Status != http.StatusBadRequest {
|
||||
t.Fatalf("status = %d, want %d", data.Status, http.StatusBadRequest)
|
||||
}
|
||||
}
|
||||
|
||||
func newTestResponsesWSTarget(t *testing.T) (*websocket.Conn, func()) {
|
||||
t.Helper()
|
||||
target, _, cleanup := newTestWebSocketPair(t)
|
||||
return target, cleanup
|
||||
}
|
||||
|
||||
func newTestWebSocketPair(t *testing.T) (*websocket.Conn, *websocket.Conn, func()) {
|
||||
t.Helper()
|
||||
upgrader := websocket.Upgrader{}
|
||||
serverConnCh := make(chan *websocket.Conn, 1)
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
conn, err := upgrader.Upgrade(w, r, nil)
|
||||
if err != nil {
|
||||
t.Errorf("upgrade websocket: %v", err)
|
||||
return
|
||||
}
|
||||
serverConnCh <- conn
|
||||
}))
|
||||
|
||||
targetURL := "ws" + strings.TrimPrefix(server.URL, "http")
|
||||
target, _, err := websocket.DefaultDialer.Dial(targetURL, nil)
|
||||
if err != nil {
|
||||
server.Close()
|
||||
t.Fatalf("dial websocket: %v", err)
|
||||
}
|
||||
serverConn := <-serverConnCh
|
||||
cleanup := func() {
|
||||
_ = target.Close()
|
||||
_ = serverConn.Close()
|
||||
server.Close()
|
||||
}
|
||||
return target, serverConn, cleanup
|
||||
}
|
||||
|
||||
func TestResponsesWSMessageSizeLimit(t *testing.T) {
|
||||
previous := constant.MaxRequestBodyMB
|
||||
constant.MaxRequestBodyMB = 1
|
||||
t.Cleanup(func() { constant.MaxRequestBodyMB = previous })
|
||||
client, server, cleanup := newTestWebSocketPair(t)
|
||||
defer cleanup()
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil)
|
||||
var admitted atomic.Bool
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
defer close(done)
|
||||
ResponsesWebSocketHelper(c, server, func(*http.Request, string, func(*gin.Context) *types.NewAPIError) *types.NewAPIError {
|
||||
admitted.Store(true)
|
||||
return nil
|
||||
})
|
||||
}()
|
||||
_ = client.WriteMessage(websocket.TextMessage, []byte(strings.Repeat("x", (1<<20)+1)))
|
||||
require.NoError(t, client.SetReadDeadline(time.Now().Add(5*time.Second)))
|
||||
_, _, err := client.ReadMessage()
|
||||
assert.True(t, websocket.IsCloseError(err, websocket.CloseMessageTooBig), "oversized message must be rejected before admission: %v", err)
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("websocket request did not exit after oversized message")
|
||||
}
|
||||
assert.False(t, admitted.Load())
|
||||
}
|
||||
|
||||
func TestResponsesWSShutdownInterruptsBusyWriter(t *testing.T) {
|
||||
client, server, cleanupClient := newTestWebSocketPair(t)
|
||||
defer cleanupClient()
|
||||
target, peer, cleanupTarget := newTestWebSocketPair(t)
|
||||
defer cleanupTarget()
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
s := &responsesWSSession{ctx: ctx, cancel: cancel, client: server, target: target}
|
||||
// A blocked network writer owns this lock. Closing the connection must
|
||||
// remain possible so that writer can be interrupted.
|
||||
s.targetWriteMu.Lock()
|
||||
defer s.targetWriteMu.Unlock()
|
||||
done := make(chan struct{})
|
||||
go func() { s.shutdown(); close(done) }()
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("shutdown waited for the network writer")
|
||||
}
|
||||
require.NoError(t, peer.SetReadDeadline(time.Now().Add(time.Second)))
|
||||
_, _, err := peer.ReadMessage()
|
||||
assert.Error(t, err)
|
||||
require.NoError(t, client.SetReadDeadline(time.Now().Add(time.Second)))
|
||||
_, _, err = client.ReadMessage()
|
||||
assert.Error(t, err)
|
||||
assert.Nil(t, s.getTarget())
|
||||
}
|
||||
|
||||
func TestResponsesWSPassthroughPreservesRawPricingParameters(t *testing.T) {
|
||||
create, _, err := normalizeResponsesWSTestMessage([]byte(`{"type":"response.create","generate":false,"response":{"model":"gpt-5.1","input":"hi","vendor":{"tier":"premium"},"stream":true}}`))
|
||||
require.NoError(t, err)
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(string(create.Body)))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
common.SetContextKey(c, constant.ContextKeyOriginalModel, create.Request.Model)
|
||||
common.SetContextKey(c, constant.ContextKeyChannelType, constant.ChannelTypeOpenAI)
|
||||
common.SetContextKey(c, constant.ContextKeyChannelSetting, dto.ChannelSettings{PassThroughBodyEnabled: true})
|
||||
info := relaycommon.GenRelayInfoResponses(c, &create.Request)
|
||||
payload, apiErr := buildResponsesWSCreatePayload(c, info, create.Request, create.Generate, create.StreamID)
|
||||
require.Nil(t, apiErr)
|
||||
assert.JSONEq(t, `{"type":"response.create","generate":false,"model":"gpt-5.1","input":"hi","vendor":{"tier":"premium"}}`, string(payload))
|
||||
storage, err := common.GetBodyStorage(c)
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, storage.Close())
|
||||
}
|
||||
|
||||
func TestResponsesWSStreamIdentity(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name, fields, want string
|
||||
invalid bool
|
||||
}{
|
||||
{name: "default", fields: `"model":"gpt-4o"`},
|
||||
{name: "named", fields: `"model":"gpt-4o","stream_id":"planner.A-1_2"`, want: "planner.A-1_2"},
|
||||
{name: "wrapped", fields: `"response":{"model":"gpt-4o","stream_id":"wrapped"}`, want: "wrapped"},
|
||||
{name: "top-level-wins", fields: `"stream_id":"outer","response":{"model":"gpt-4o","stream_id":"inner"}`, want: "outer"},
|
||||
{name: "maximum", fields: `"model":"gpt-4o","stream_id":"` + strings.Repeat("a", 256) + `"`, want: strings.Repeat("a", 256)},
|
||||
{name: "too-long", fields: `"stream_id":"` + strings.Repeat("a", 257) + `"`, invalid: true},
|
||||
{name: "empty", fields: `"stream_id":""`, invalid: true},
|
||||
{name: "null", fields: `"stream_id":null`, invalid: true},
|
||||
{name: "number", fields: `"stream_id":123`, invalid: true},
|
||||
{name: "unicode", fields: `"stream_id":"计划"`, invalid: true},
|
||||
{name: "spaces", fields: `"stream_id":"a b"`, invalid: true},
|
||||
{name: "invalid-top-level", fields: `"stream_id":"","response":{"stream_id":"inner"}`, invalid: true},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
create, eventID, err := normalizeResponsesWSTestMessage([]byte(`{"type":"response.create","event_id":"test",` + tc.fields + `}`))
|
||||
assert.Equal(t, "test", eventID)
|
||||
if tc.invalid {
|
||||
require.ErrorContains(t, err, "stream_id")
|
||||
return
|
||||
}
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, tc.want, create.StreamID)
|
||||
assert.NotContains(t, string(create.Body), "stream_id")
|
||||
for _, passthrough := range []bool{false, true} {
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(string(create.Body)))
|
||||
common.SetContextKey(c, constant.ContextKeyOriginalModel, create.Request.Model)
|
||||
common.SetContextKey(c, constant.ContextKeyChannelType, constant.ChannelTypeOpenAI)
|
||||
common.SetContextKey(c, constant.ContextKeyChannelSetting, dto.ChannelSettings{PassThroughBodyEnabled: passthrough})
|
||||
info := relaycommon.GenRelayInfoResponses(c, &create.Request)
|
||||
payload, apiErr := buildResponsesWSCreatePayload(c, info, create.Request, create.Generate, create.StreamID)
|
||||
require.Nil(t, apiErr)
|
||||
var event map[string]any
|
||||
require.NoError(t, common.Unmarshal(payload, &event))
|
||||
if tc.want == "" {
|
||||
assert.NotContains(t, event, "stream_id")
|
||||
} else {
|
||||
assert.Equal(t, tc.want, event["stream_id"])
|
||||
}
|
||||
if storage, err := common.GetBodyStorage(c); err == nil {
|
||||
require.NoError(t, storage.Close())
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestResponsesWSErrorAttribution(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name, payload, control string
|
||||
terminal, ambiguous, controlError bool
|
||||
}{
|
||||
{name: "active-server-failure", payload: `{"status":500,"error":{"type":"server_error"}}`, terminal: true},
|
||||
{name: "other-stream", payload: `{"stream_id":"other","status":500}`},
|
||||
{name: "previous-response", payload: `{"response_id":"previous","status":500}`},
|
||||
{name: "cancel-event-id", payload: `{"event_id":"cancel","status":500}`, control: `{"event_id":"cancel"}`, controlError: true},
|
||||
{name: "cancel-target", payload: `{"response_id":"missing","status":400}`, control: `{"response_id":"missing"}`, controlError: true},
|
||||
{name: "cancel-code", payload: `{"error":{"type":"invalid_request_error","code":"response_not_found"}}`, control: `{"response_id":"missing"}`, controlError: true},
|
||||
{name: "ambiguous-after-cancel", payload: `{"status":500}`, control: `{"response_id":"active"}`, terminal: true, ambiguous: true},
|
||||
{name: "explicit-generation-failure", payload: `{"response_id":"active","status":500}`, control: `{"response_id":"active"}`, terminal: true},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
var event responsesWSErrorEvent
|
||||
require.NoError(t, common.Unmarshal([]byte(tc.payload), &event))
|
||||
terminal, ambiguous, controlError := responsesWSErrorEndsRequest(event, "planner", "active", []byte(tc.control))
|
||||
assert.Equal(t, tc.terminal, terminal)
|
||||
assert.Equal(t, tc.ambiguous, ambiguous)
|
||||
assert.Equal(t, tc.controlError, controlError)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -2,7 +2,10 @@ package relay
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/pkg/wsmanager"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
"github.com/QuantumNous/new-api/relaykit/types"
|
||||
@@ -33,6 +36,18 @@ func WssHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *types.
|
||||
if resp != nil {
|
||||
info.TargetWs = resp.(*websocket.Conn)
|
||||
defer info.TargetWs.Close()
|
||||
var closeOnce sync.Once
|
||||
unregister := wsmanager.Register(info.ChannelId, wsmanager.KindRealtime, func(reason string) {
|
||||
closeOnce.Do(func() {
|
||||
deadline := time.Now().Add(time.Second)
|
||||
closeMessage := websocket.FormatCloseMessage(websocket.ClosePolicyViolation, reason)
|
||||
_ = info.ClientWs.WriteControl(websocket.CloseMessage, closeMessage, deadline)
|
||||
_ = info.TargetWs.WriteControl(websocket.CloseMessage, closeMessage, deadline)
|
||||
_ = info.ClientWs.Close()
|
||||
_ = info.TargetWs.Close()
|
||||
})
|
||||
})
|
||||
defer unregister()
|
||||
}
|
||||
|
||||
usage, newAPIError := adaptor.DoResponse(c, nil, info)
|
||||
|
||||
Reference in new issue
Block a user