重构前
This commit is contained in:
@@ -0,0 +1,91 @@
|
||||
package helper
|
||||
|
||||
import (
|
||||
"strings"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/pkg/billingexpr"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func ResolveIncomingBillingExprRequestInput(c *gin.Context, info *relaycommon.RelayInfo) (billingexpr.RequestInput, error) {
|
||||
if info != nil && info.BillingRequestInput != nil {
|
||||
input := cloneRequestInput(*info.BillingRequestInput)
|
||||
merged := cloneStringMap(info.RequestHeaders)
|
||||
for k, v := range input.Headers {
|
||||
merged[k] = v
|
||||
}
|
||||
input.Headers = merged
|
||||
return input, nil
|
||||
}
|
||||
|
||||
input := billingexpr.RequestInput{}
|
||||
if info != nil {
|
||||
input.Headers = cloneStringMap(info.RequestHeaders)
|
||||
}
|
||||
|
||||
bodyBytes, err := readIncomingBillingExprBody(c)
|
||||
if err != nil {
|
||||
return billingexpr.RequestInput{}, err
|
||||
}
|
||||
input.Body = bodyBytes
|
||||
return input, nil
|
||||
}
|
||||
|
||||
func BuildBillingExprRequestInputFromRequest(request dto.Request, headers map[string]string) (billingexpr.RequestInput, error) {
|
||||
input := billingexpr.RequestInput{
|
||||
Headers: cloneStringMap(headers),
|
||||
}
|
||||
if request == nil {
|
||||
return input, nil
|
||||
}
|
||||
|
||||
bodyBytes, err := common.Marshal(request)
|
||||
if err != nil {
|
||||
return billingexpr.RequestInput{}, err
|
||||
}
|
||||
input.Body = bodyBytes
|
||||
return input, nil
|
||||
}
|
||||
|
||||
func readIncomingBillingExprBody(c *gin.Context) ([]byte, error) {
|
||||
if c == nil || c.Request == nil || !isJSONContentType(c.Request.Header.Get("Content-Type")) {
|
||||
return nil, nil
|
||||
}
|
||||
storage, err := common.GetBodyStorage(c)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return storage.Bytes()
|
||||
}
|
||||
|
||||
func cloneRequestInput(src billingexpr.RequestInput) billingexpr.RequestInput {
|
||||
input := billingexpr.RequestInput{
|
||||
Headers: cloneStringMap(src.Headers),
|
||||
}
|
||||
if len(src.Body) > 0 {
|
||||
input.Body = append([]byte(nil), src.Body...)
|
||||
}
|
||||
return input
|
||||
}
|
||||
|
||||
func isJSONContentType(contentType string) bool {
|
||||
contentType = strings.ToLower(strings.TrimSpace(contentType))
|
||||
return strings.HasPrefix(contentType, "application/json")
|
||||
}
|
||||
|
||||
func cloneStringMap(src map[string]string) map[string]string {
|
||||
if len(src) == 0 {
|
||||
return map[string]string{}
|
||||
}
|
||||
dst := make(map[string]string, len(src))
|
||||
for key, value := range src {
|
||||
if strings.TrimSpace(key) == "" {
|
||||
continue
|
||||
}
|
||||
dst[key] = value
|
||||
}
|
||||
return dst
|
||||
}
|
||||
@@ -0,0 +1,63 @@
|
||||
package helper
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/samber/lo"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
func TestResolveIncomingBillingExprRequestInput(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
recorder := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(recorder)
|
||||
ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
|
||||
ctx.Request.Header.Set("Content-Type", "application/json")
|
||||
|
||||
body := []byte(`{"service_tier":"fast"}`)
|
||||
ctx.Request.Body = io.NopCloser(bytes.NewReader(body))
|
||||
ctx.Set(common.KeyRequestBody, body)
|
||||
|
||||
info := &relaycommon.RelayInfo{
|
||||
RequestHeaders: map[string]string{"Content-Type": "application/json"},
|
||||
}
|
||||
|
||||
input, err := ResolveIncomingBillingExprRequestInput(ctx, info)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, body, input.Body)
|
||||
require.Equal(t, "application/json", input.Headers["Content-Type"])
|
||||
}
|
||||
|
||||
func TestBuildBillingExprRequestInputFromRequest(t *testing.T) {
|
||||
request := &dto.GeneralOpenAIRequest{
|
||||
Model: "gemini-3.1-pro-preview",
|
||||
Stream: lo.ToPtr(true),
|
||||
Messages: []dto.Message{
|
||||
{
|
||||
Role: "user",
|
||||
Content: "hi",
|
||||
},
|
||||
},
|
||||
MaxTokens: lo.ToPtr(uint(3000)),
|
||||
}
|
||||
|
||||
input, err := BuildBillingExprRequestInputFromRequest(request, map[string]string{
|
||||
"Content-Type": "application/json",
|
||||
"X-Test": "1",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "application/json", input.Headers["Content-Type"])
|
||||
require.Equal(t, "1", input.Headers["X-Test"])
|
||||
require.True(t, gjson.GetBytes(input.Body, "stream").Bool())
|
||||
require.Equal(t, "user", gjson.GetBytes(input.Body, "messages.0.role").String())
|
||||
require.Equal(t, float64(3000), gjson.GetBytes(input.Body, "max_tokens").Float())
|
||||
}
|
||||
@@ -0,0 +1,227 @@
|
||||
package helper
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/logger"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
"github.com/QuantumNous/new-api/relaykit/types"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/gorilla/websocket"
|
||||
)
|
||||
|
||||
func FlushWriter(c *gin.Context) (err error) {
|
||||
defer func() {
|
||||
if r := recover(); r != nil {
|
||||
err = fmt.Errorf("flush panic recovered: %v", r)
|
||||
}
|
||||
}()
|
||||
|
||||
if c == nil || c.Writer == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
if requestContextDone(c) {
|
||||
return fmt.Errorf("request context done: %w", c.Request.Context().Err())
|
||||
}
|
||||
|
||||
flusher, ok := c.Writer.(http.Flusher)
|
||||
if !ok {
|
||||
return errors.New("streaming error: flusher not found")
|
||||
}
|
||||
|
||||
flusher.Flush()
|
||||
return nil
|
||||
}
|
||||
|
||||
func requestContextDone(c *gin.Context) bool {
|
||||
return c != nil && c.Request != nil && c.Request.Context().Err() != nil
|
||||
}
|
||||
|
||||
func SetEventStreamHeaders(c *gin.Context) {
|
||||
// 检查是否已经设置过头部
|
||||
if _, exists := c.Get("event_stream_headers_set"); exists {
|
||||
return
|
||||
}
|
||||
|
||||
// 设置标志，表示头部已经设置过
|
||||
c.Set("event_stream_headers_set", true)
|
||||
|
||||
c.Writer.Header().Set("Content-Type", "text/event-stream")
|
||||
c.Writer.Header().Set("Cache-Control", "no-cache")
|
||||
c.Writer.Header().Set("Connection", "keep-alive")
|
||||
c.Writer.Header().Set("Transfer-Encoding", "chunked")
|
||||
c.Writer.Header().Set("X-Accel-Buffering", "no")
|
||||
}
|
||||
|
||||
func ClaudeData(c *gin.Context, resp dto.ClaudeResponse) error {
|
||||
if requestContextDone(c) {
|
||||
return nil
|
||||
}
|
||||
|
||||
jsonData, err := common.Marshal(resp)
|
||||
if err != nil {
|
||||
common.SysError("error marshalling stream response: " + err.Error())
|
||||
} else {
|
||||
c.Render(-1, common.CustomEvent{Data: fmt.Sprintf("event: %s\n", resp.Type)})
|
||||
c.Render(-1, common.CustomEvent{Data: "data: " + string(jsonData)})
|
||||
}
|
||||
_ = FlushWriter(c)
|
||||
return nil
|
||||
}
|
||||
|
||||
func ClaudeChunkData(c *gin.Context, resp dto.ClaudeResponse, data string) {
|
||||
if requestContextDone(c) {
|
||||
return
|
||||
}
|
||||
|
||||
c.Render(-1, common.CustomEvent{Data: fmt.Sprintf("event: %s\n", resp.Type)})
|
||||
c.Render(-1, common.CustomEvent{Data: fmt.Sprintf("data: %s\n", data)})
|
||||
_ = FlushWriter(c)
|
||||
}
|
||||
|
||||
func ResponseChunkData(c *gin.Context, resp dto.ResponsesStreamResponse, data string) error {
|
||||
if requestContextDone(c) {
|
||||
return fmt.Errorf("request context done: %w", c.Request.Context().Err())
|
||||
}
|
||||
|
||||
c.Render(-1, common.CustomEvent{Data: fmt.Sprintf("event: %s\n", resp.Type)})
|
||||
c.Render(-1, common.CustomEvent{Data: fmt.Sprintf("data: %s", data)})
|
||||
return FlushWriter(c)
|
||||
}
|
||||
|
||||
func StringData(c *gin.Context, str string) error {
|
||||
if c == nil || c.Writer == nil {
|
||||
return errors.New("context or writer is nil")
|
||||
}
|
||||
|
||||
if requestContextDone(c) {
|
||||
return fmt.Errorf("request context done: %w", c.Request.Context().Err())
|
||||
}
|
||||
|
||||
c.Render(-1, common.CustomEvent{Data: "data: " + str})
|
||||
return FlushWriter(c)
|
||||
}
|
||||
|
||||
func PingData(c *gin.Context) error {
|
||||
if c == nil || c.Writer == nil {
|
||||
return errors.New("context or writer is nil")
|
||||
}
|
||||
|
||||
if requestContextDone(c) {
|
||||
return fmt.Errorf("request context done: %w", c.Request.Context().Err())
|
||||
}
|
||||
|
||||
if _, err := c.Writer.Write([]byte(": PING\n\n")); err != nil {
|
||||
return fmt.Errorf("write ping data failed: %w", err)
|
||||
}
|
||||
return FlushWriter(c)
|
||||
}
|
||||
|
||||
func ObjectData(c *gin.Context, object interface{}) error {
|
||||
if object == nil {
|
||||
return errors.New("object is nil")
|
||||
}
|
||||
jsonData, err := common.Marshal(object)
|
||||
if err != nil {
|
||||
return fmt.Errorf("error marshalling object: %w", err)
|
||||
}
|
||||
return StringData(c, string(jsonData))
|
||||
}
|
||||
|
||||
func Done(c *gin.Context) {
|
||||
_ = StringData(c, "[DONE]")
|
||||
}
|
||||
|
||||
func WssString(c *gin.Context, ws *websocket.Conn, str string) error {
|
||||
if ws == nil {
|
||||
logger.LogError(c, "websocket connection is nil")
|
||||
return errors.New("websocket connection is nil")
|
||||
}
|
||||
//common.LogInfo(c, fmt.Sprintf("sending message: %s", str))
|
||||
return ws.WriteMessage(1, []byte(str))
|
||||
}
|
||||
|
||||
func WssObject(c *gin.Context, ws *websocket.Conn, object interface{}) error {
|
||||
jsonData, err := common.Marshal(object)
|
||||
if err != nil {
|
||||
return fmt.Errorf("error marshalling object: %w", err)
|
||||
}
|
||||
if ws == nil {
|
||||
logger.LogError(c, "websocket connection is nil")
|
||||
return errors.New("websocket connection is nil")
|
||||
}
|
||||
//common.LogInfo(c, fmt.Sprintf("sending message: %s", jsonData))
|
||||
return ws.WriteMessage(1, jsonData)
|
||||
}
|
||||
|
||||
func WssError(c *gin.Context, ws *websocket.Conn, openaiError types.OpenAIError) {
|
||||
if ws == nil {
|
||||
return
|
||||
}
|
||||
errorObj := &dto.RealtimeEvent{
|
||||
Type: "error",
|
||||
EventId: GetLocalRealtimeID(c),
|
||||
Error: &openaiError,
|
||||
}
|
||||
_ = WssObject(c, ws, errorObj)
|
||||
}
|
||||
|
||||
func GetResponseID(c *gin.Context) string {
|
||||
logID := c.GetString(common.RequestIdKey)
|
||||
return fmt.Sprintf("chatcmpl-%s", logID)
|
||||
}
|
||||
|
||||
func GetLocalRealtimeID(c *gin.Context) string {
|
||||
logID := c.GetString(common.RequestIdKey)
|
||||
return fmt.Sprintf("evt_%s", logID)
|
||||
}
|
||||
|
||||
func GenerateStartEmptyResponse(id string, createAt int64, model string, systemFingerprint *string) *dto.ChatCompletionsStreamResponse {
|
||||
return &dto.ChatCompletionsStreamResponse{
|
||||
Id: id,
|
||||
Object: "chat.completion.chunk",
|
||||
Created: createAt,
|
||||
Model: model,
|
||||
SystemFingerprint: systemFingerprint,
|
||||
Choices: []dto.ChatCompletionsStreamResponseChoice{
|
||||
{
|
||||
Delta: dto.ChatCompletionsStreamResponseChoiceDelta{
|
||||
Role: "assistant",
|
||||
Content: common.GetPointer(""),
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func GenerateStopResponse(id string, createAt int64, model string, finishReason string) *dto.ChatCompletionsStreamResponse {
|
||||
return &dto.ChatCompletionsStreamResponse{
|
||||
Id: id,
|
||||
Object: "chat.completion.chunk",
|
||||
Created: createAt,
|
||||
Model: model,
|
||||
SystemFingerprint: nil,
|
||||
Choices: []dto.ChatCompletionsStreamResponseChoice{
|
||||
{
|
||||
FinishReason: &finishReason,
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func GenerateFinalUsageResponse(id string, createAt int64, model string, usage dto.Usage) *dto.ChatCompletionsStreamResponse {
|
||||
return &dto.ChatCompletionsStreamResponse{
|
||||
Id: id,
|
||||
Object: "chat.completion.chunk",
|
||||
Created: createAt,
|
||||
Model: model,
|
||||
SystemFingerprint: nil,
|
||||
Choices: make([]dto.ChatCompletionsStreamResponseChoice, 0),
|
||||
Usage: &usage,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,72 @@
|
||||
package helper
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
relayconstant "github.com/QuantumNous/new-api/relay/constant"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// TestMaxTokensBounds guards the billing invariant that user-supplied max
|
||||
// token fields are bounded on every relay format. These values feed
|
||||
// pre-consume quota math (preConsumedTokens * ratio); a huge or
|
||||
// wrapped-negative value (e.g. 18446744073686646784 parsed into *uint) must
|
||||
// be rejected at validation instead of corrupting the pre-charge.
|
||||
func TestMaxTokensBounds(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
newJSONContext := func(t *testing.T, body string) *gin.Context {
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/relay", bytes.NewBufferString(body))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
return c
|
||||
}
|
||||
|
||||
const hugeN = "18446744073686646784"
|
||||
|
||||
t.Run("openai max_tokens overflow rejected", func(t *testing.T) {
|
||||
c := newJSONContext(t, `{"model":"gpt-4o","messages":[{"role":"user","content":"hi"}],"max_tokens":`+hugeN+`}`)
|
||||
_, err := GetAndValidateTextRequest(c, relayconstant.RelayModeChatCompletions)
|
||||
require.Error(t, err)
|
||||
require.Contains(t, err.Error(), "max_tokens is invalid")
|
||||
})
|
||||
|
||||
t.Run("openai max_completion_tokens overflow rejected", func(t *testing.T) {
|
||||
c := newJSONContext(t, `{"model":"gpt-4o","messages":[{"role":"user","content":"hi"}],"max_completion_tokens":`+hugeN+`}`)
|
||||
_, err := GetAndValidateTextRequest(c, relayconstant.RelayModeChatCompletions)
|
||||
require.Error(t, err)
|
||||
require.Contains(t, err.Error(), "max_tokens is invalid")
|
||||
})
|
||||
|
||||
t.Run("claude max_tokens overflow rejected", func(t *testing.T) {
|
||||
c := newJSONContext(t, `{"model":"claude-sonnet-4","messages":[{"role":"user","content":"hi"}],"max_tokens":`+hugeN+`}`)
|
||||
_, err := GetAndValidateClaudeRequest(c)
|
||||
require.Error(t, err)
|
||||
require.Contains(t, err.Error(), "max_tokens is invalid")
|
||||
})
|
||||
|
||||
t.Run("claude normal max_tokens accepted", func(t *testing.T) {
|
||||
c := newJSONContext(t, `{"model":"claude-sonnet-4","messages":[{"role":"user","content":"hi"}],"max_tokens":8192}`)
|
||||
req, err := GetAndValidateClaudeRequest(c)
|
||||
require.NoError(t, err)
|
||||
require.EqualValues(t, 8192, *req.MaxTokens)
|
||||
})
|
||||
|
||||
t.Run("gemini maxOutputTokens overflow rejected", func(t *testing.T) {
|
||||
c := newJSONContext(t, `{"contents":[{"parts":[{"text":"hi"}]}],"generationConfig":{"maxOutputTokens":`+hugeN+`}}`)
|
||||
_, err := GetAndValidateGeminiRequest(c)
|
||||
require.Error(t, err)
|
||||
require.Contains(t, err.Error(), "maxOutputTokens is invalid")
|
||||
})
|
||||
|
||||
t.Run("responses max_output_tokens overflow rejected", func(t *testing.T) {
|
||||
c := newJSONContext(t, `{"model":"gpt-4o","input":"hi","max_output_tokens":`+hugeN+`}`)
|
||||
_, err := GetAndValidateResponsesRequest(c)
|
||||
require.Error(t, err)
|
||||
require.Contains(t, err.Error(), "max_output_tokens is invalid")
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,63 @@
|
||||
package helper
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
|
||||
"github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func ModelMappedHelper(c *gin.Context, info *common.RelayInfo, request dto.Request) error {
|
||||
if info.ChannelMeta == nil {
|
||||
info.ChannelMeta = &common.ChannelMeta{}
|
||||
}
|
||||
|
||||
// map model name
|
||||
modelMapping := c.GetString("model_mapping")
|
||||
if modelMapping != "" && modelMapping != "{}" {
|
||||
modelMap := make(map[string]string)
|
||||
err := json.Unmarshal([]byte(modelMapping), &modelMap)
|
||||
if err != nil {
|
||||
return fmt.Errorf("unmarshal_model_mapping_failed")
|
||||
}
|
||||
|
||||
// 支持链式模型重定向，最终使用链尾的模型
|
||||
currentModel := info.OriginModelName
|
||||
visitedModels := map[string]bool{
|
||||
currentModel: true,
|
||||
}
|
||||
for {
|
||||
if mappedModel, exists := modelMap[currentModel]; exists && mappedModel != "" {
|
||||
// 模型重定向循环检测，避免无限循环
|
||||
if visitedModels[mappedModel] {
|
||||
if mappedModel == currentModel {
|
||||
if currentModel == info.OriginModelName {
|
||||
info.IsModelMapped = false
|
||||
return nil
|
||||
} else {
|
||||
info.IsModelMapped = true
|
||||
break
|
||||
}
|
||||
}
|
||||
return errors.New("model_mapping_contains_cycle")
|
||||
}
|
||||
visitedModels[mappedModel] = true
|
||||
currentModel = mappedModel
|
||||
info.IsModelMapped = true
|
||||
} else {
|
||||
break
|
||||
}
|
||||
}
|
||||
if info.IsModelMapped {
|
||||
info.UpstreamModelName = currentModel
|
||||
}
|
||||
}
|
||||
|
||||
if request != nil {
|
||||
request.SetModelName(info.UpstreamModelName)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,160 @@
|
||||
package helper
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"fmt"
|
||||
"io"
|
||||
"mime/multipart"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/url"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/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/require"
|
||||
)
|
||||
|
||||
// TestGetAndValidOpenAIImageRequestMultipartStream verifies multipart image
|
||||
// edit parsing: the stream field is parsed and validated, and the request body
|
||||
// stays replayable for the upstream request.
|
||||
func TestGetAndValidOpenAIImageRequestMultipartStream(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
newContext := func(t *testing.T, streamValue string, withImage bool) (*gin.Context, string) {
|
||||
var body bytes.Buffer
|
||||
writer := multipart.NewWriter(&body)
|
||||
require.NoError(t, writer.WriteField("model", "gpt-image-1"))
|
||||
require.NoError(t, writer.WriteField("prompt", "edit this image"))
|
||||
require.NoError(t, writer.WriteField("stream", streamValue))
|
||||
if withImage {
|
||||
part, err := writer.CreateFormFile("image", "input.png")
|
||||
require.NoError(t, err)
|
||||
_, err = part.Write([]byte("fake image"))
|
||||
require.NoError(t, err)
|
||||
}
|
||||
require.NoError(t, writer.Close())
|
||||
originalBody := body.String()
|
||||
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/edits", &body)
|
||||
c.Request.Header.Set("Content-Type", writer.FormDataContentType())
|
||||
return c, originalBody
|
||||
}
|
||||
|
||||
t.Run("valid stream value keeps body replayable", func(t *testing.T) {
|
||||
c, originalBody := newContext(t, "true", true)
|
||||
|
||||
req, err := GetAndValidOpenAIImageRequest(c, relayconstant.RelayModeImagesEdits)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, req.Stream)
|
||||
require.True(t, *req.Stream)
|
||||
require.True(t, req.IsStream(c.Request))
|
||||
|
||||
bodyAfterValidation, err := io.ReadAll(c.Request.Body)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, originalBody, string(bodyAfterValidation))
|
||||
|
||||
form, err := common.ParseMultipartFormReusable(c)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "true", url.Values(form.Value).Get("stream"))
|
||||
require.Len(t, form.File["image"], 1)
|
||||
})
|
||||
|
||||
t.Run("invalid stream value is rejected", func(t *testing.T) {
|
||||
c, _ := newContext(t, "notabool", false)
|
||||
|
||||
_, err := GetAndValidOpenAIImageRequest(c, relayconstant.RelayModeImagesEdits)
|
||||
require.Error(t, err)
|
||||
require.Contains(t, err.Error(), "invalid stream value")
|
||||
})
|
||||
}
|
||||
|
||||
// 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.
|
||||
func TestGetAndValidOpenAIImageRequestNBounds(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
newJSONContext := func(t *testing.T, body string) *gin.Context {
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewBufferString(body))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
return c
|
||||
}
|
||||
|
||||
boundErr := fmt.Sprintf("n must be an integer between 1 and %d", dto.MaxImageN)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
body string
|
||||
wantErr string
|
||||
wantN uint
|
||||
}{
|
||||
{
|
||||
name: "overflowed uint64 n is rejected",
|
||||
body: `{"model":"gpt-image-1","prompt":"a cat","n":18446744073686646784}`,
|
||||
wantErr: boundErr,
|
||||
},
|
||||
{
|
||||
name: "n above max is rejected",
|
||||
body: fmt.Sprintf(`{"model":"gpt-image-1","prompt":"a cat","n":%d}`, dto.MaxImageN+1),
|
||||
wantErr: boundErr,
|
||||
},
|
||||
{
|
||||
name: "n at max is accepted",
|
||||
body: fmt.Sprintf(`{"model":"gpt-image-1","prompt":"a cat","n":%d}`, dto.MaxImageN),
|
||||
wantN: dto.MaxImageN,
|
||||
},
|
||||
{
|
||||
name: "explicit n is accepted",
|
||||
body: `{"model":"gpt-image-1","prompt":"a cat","n":3}`,
|
||||
wantN: 3,
|
||||
},
|
||||
{
|
||||
name: "zero n defaults to 1",
|
||||
body: `{"model":"gpt-image-1","prompt":"a cat","n":0}`,
|
||||
wantN: 1,
|
||||
},
|
||||
{
|
||||
name: "absent n defaults to 1",
|
||||
body: `{"model":"gpt-image-1","prompt":"a cat"}`,
|
||||
wantN: 1,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
c := newJSONContext(t, tt.body)
|
||||
req, err := GetAndValidOpenAIImageRequest(c, relayconstant.RelayModeImagesGenerations)
|
||||
if tt.wantErr != "" {
|
||||
require.Error(t, err)
|
||||
require.Contains(t, err.Error(), tt.wantErr)
|
||||
return
|
||||
}
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, req.N)
|
||||
require.Equal(t, tt.wantN, *req.N)
|
||||
require.Equal(t, float64(tt.wantN), req.GetTokenCountMeta().BillingRatios["n"])
|
||||
})
|
||||
}
|
||||
|
||||
t.Run("negative multipart n is rejected", func(t *testing.T) {
|
||||
var body bytes.Buffer
|
||||
writer := multipart.NewWriter(&body)
|
||||
require.NoError(t, writer.WriteField("model", "gpt-image-1"))
|
||||
require.NoError(t, writer.WriteField("prompt", "edit this image"))
|
||||
require.NoError(t, writer.WriteField("n", "-22904832"))
|
||||
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())
|
||||
|
||||
_, err := GetAndValidOpenAIImageRequest(c, relayconstant.RelayModeImagesEdits)
|
||||
require.Error(t, err)
|
||||
require.Contains(t, err.Error(), boundErr)
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,337 @@
|
||||
package helper
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"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/types"
|
||||
"github.com/QuantumNous/new-api/setting/billing_setting"
|
||||
"github.com/QuantumNous/new-api/setting/operation_setting"
|
||||
"github.com/QuantumNous/new-api/setting/ratio_setting"
|
||||
hosttypes "github.com/QuantumNous/new-api/types"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func modelPriceNotConfiguredError(modelName string, userId int) error {
|
||||
if model.IsAdmin(userId) {
|
||||
return fmt.Errorf(
|
||||
"模型 %s 的价格未配置。请前往「系统设置 → 运营设置」开启自用模式，或在「系统设置 → 分组与模型定价设置」中为该模型配置价格；"+
|
||||
"Model %s price not configured. Go to System Settings → Operation Settings to enable self-use mode, or configure the model price in System Settings → Group & Model Pricing.",
|
||||
modelName, modelName,
|
||||
)
|
||||
}
|
||||
return fmt.Errorf(
|
||||
"模型 %s 的价格尚未由管理员配置，暂时无法使用，请联系站点管理员开启该模型；"+
|
||||
"Model %s has not been priced by the administrator yet. Please contact the site administrator to enable this model.",
|
||||
modelName, modelName,
|
||||
)
|
||||
}
|
||||
|
||||
// https://docs.claude.com/en/docs/build-with-claude/prompt-caching#1-hour-cache-duration
|
||||
const claudeCacheCreation1hMultiplier = 6 / 3.75
|
||||
|
||||
// defaultTieredPreConsumeMaxTokens is the fallback completion-token estimate
|
||||
// used for tiered expression pre-consume when the client omits max_tokens, so
|
||||
// the pre-consumed quota still reflects a plausible output cost in paid groups.
|
||||
const defaultTieredPreConsumeMaxTokens = 8192
|
||||
|
||||
// HandleGroupRatio checks for "auto_group" in the context and updates the group ratio and relayInfo.UsingGroup if present
|
||||
func HandleGroupRatio(ctx *gin.Context, relayInfo *relaycommon.RelayInfo) hosttypes.GroupRatioInfo {
|
||||
groupRatioInfo := hosttypes.GroupRatioInfo{
|
||||
GroupRatio: 1.0, // default ratio
|
||||
GroupSpecialRatio: -1,
|
||||
}
|
||||
|
||||
// check auto group
|
||||
autoGroup, exists := ctx.Get("auto_group")
|
||||
if exists {
|
||||
logger.LogDebug(ctx, "final group: %s", autoGroup)
|
||||
relayInfo.UsingGroup = autoGroup.(string)
|
||||
}
|
||||
|
||||
// check user group special ratio
|
||||
userGroupRatio, ok := ratio_setting.GetGroupGroupRatio(relayInfo.UserGroup, relayInfo.UsingGroup)
|
||||
if ok {
|
||||
// user group special ratio
|
||||
groupRatioInfo.GroupSpecialRatio = userGroupRatio
|
||||
groupRatioInfo.GroupRatio = userGroupRatio
|
||||
groupRatioInfo.HasSpecialRatio = true
|
||||
} else {
|
||||
// normal group ratio
|
||||
groupRatioInfo.GroupRatio = ratio_setting.GetGroupRatio(relayInfo.UsingGroup)
|
||||
}
|
||||
|
||||
return groupRatioInfo
|
||||
}
|
||||
|
||||
func ModelPriceHelper(c *gin.Context, info *relaycommon.RelayInfo, promptTokens int, meta *types.TokenCountMeta) (hosttypes.PriceData, error) {
|
||||
modelPrice, usePrice := ratio_setting.GetModelPrice(info.OriginModelName, false)
|
||||
|
||||
groupRatioInfo := HandleGroupRatio(c, info)
|
||||
|
||||
// Check if this model uses tiered_expr billing
|
||||
if billing_setting.GetBillingMode(info.OriginModelName) == billing_setting.BillingModeTieredExpr {
|
||||
return modelPriceHelperTiered(c, info, promptTokens, meta, groupRatioInfo)
|
||||
}
|
||||
|
||||
var preConsumedQuota int
|
||||
var modelRatio float64
|
||||
var completionRatio float64
|
||||
var cacheRatio float64
|
||||
var imageRatio float64
|
||||
var cacheCreationRatio float64
|
||||
var cacheCreationRatio5m float64
|
||||
var cacheCreationRatio1h float64
|
||||
var audioRatio float64
|
||||
var audioCompletionRatio float64
|
||||
var freeModel bool
|
||||
if !usePrice {
|
||||
preConsumedTokens := common.Max(promptTokens, common.PreConsumedQuota)
|
||||
if meta.MaxTokens != 0 {
|
||||
preConsumedTokens += meta.MaxTokens
|
||||
}
|
||||
var success bool
|
||||
var matchName string
|
||||
modelRatio, success, matchName = ratio_setting.GetModelRatio(info.OriginModelName)
|
||||
if !success {
|
||||
acceptUnsetRatio := false
|
||||
if info.UserSetting.AcceptUnsetRatioModel {
|
||||
acceptUnsetRatio = true
|
||||
}
|
||||
if !acceptUnsetRatio {
|
||||
return hosttypes.PriceData{}, modelPriceNotConfiguredError(matchName, info.UserId)
|
||||
}
|
||||
}
|
||||
completionRatio = ratio_setting.GetCompletionRatio(info.OriginModelName)
|
||||
cacheRatio, _ = ratio_setting.GetCacheRatio(info.OriginModelName)
|
||||
cacheCreationRatio, _ = ratio_setting.GetCreateCacheRatio(info.OriginModelName)
|
||||
cacheCreationRatio5m = cacheCreationRatio
|
||||
// 固定1h和5min缓存写入价格的比例
|
||||
cacheCreationRatio1h = cacheCreationRatio * claudeCacheCreation1hMultiplier
|
||||
imageRatio, _ = ratio_setting.GetImageRatio(info.OriginModelName)
|
||||
audioRatio = ratio_setting.GetAudioRatio(info.OriginModelName)
|
||||
audioCompletionRatio = ratio_setting.GetAudioCompletionRatio(info.OriginModelName)
|
||||
ratio := modelRatio * groupRatioInfo.GroupRatio
|
||||
quota, err := common.QuotaFromFloatStrict(float64(preConsumedTokens) * ratio)
|
||||
if err != nil {
|
||||
return hosttypes.PriceData{}, err
|
||||
}
|
||||
preConsumedQuota = quota
|
||||
} else {
|
||||
if meta.ImagePriceRatio != 0 {
|
||||
modelPrice = modelPrice * meta.ImagePriceRatio
|
||||
}
|
||||
}
|
||||
|
||||
// check if free model pre-consume is disabled
|
||||
if !operation_setting.GetQuotaSetting().EnableFreeModelPreConsume {
|
||||
// if model price or ratio is 0, do not pre-consume quota
|
||||
if groupRatioInfo.GroupRatio == 0 {
|
||||
preConsumedQuota = 0
|
||||
freeModel = true
|
||||
} else if usePrice {
|
||||
if modelPrice == 0 {
|
||||
preConsumedQuota = 0
|
||||
freeModel = true
|
||||
}
|
||||
} else {
|
||||
if modelRatio == 0 {
|
||||
preConsumedQuota = 0
|
||||
freeModel = true
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
priceData := hosttypes.PriceData{
|
||||
FreeModel: freeModel,
|
||||
ModelPrice: modelPrice,
|
||||
ModelRatio: modelRatio,
|
||||
CompletionRatio: completionRatio,
|
||||
GroupRatioInfo: groupRatioInfo,
|
||||
UsePrice: usePrice,
|
||||
CacheRatio: cacheRatio,
|
||||
ImageRatio: imageRatio,
|
||||
AudioRatio: audioRatio,
|
||||
AudioCompletionRatio: audioCompletionRatio,
|
||||
CacheCreationRatio: cacheCreationRatio,
|
||||
CacheCreation5mRatio: cacheCreationRatio5m,
|
||||
CacheCreation1hRatio: cacheCreationRatio1h,
|
||||
QuotaToPreConsume: preConsumedQuota,
|
||||
}
|
||||
if usePrice {
|
||||
for name, ratio := range meta.BillingRatios {
|
||||
priceData.AddOtherRatio(name, ratio)
|
||||
}
|
||||
quotaToPreConsume := priceData.ApplyOtherRatiosToFloat(modelPrice * common.QuotaPerUnit * groupRatioInfo.GroupRatio)
|
||||
quota, err := common.QuotaFromFloatStrict(quotaToPreConsume)
|
||||
if err != nil {
|
||||
return hosttypes.PriceData{}, err
|
||||
}
|
||||
priceData.QuotaToPreConsume = quota
|
||||
}
|
||||
|
||||
if common.DebugEnabled {
|
||||
logger.LogDebug(c, "model_price_helper result: %s", priceData.ToSetting())
|
||||
}
|
||||
info.PriceData = priceData
|
||||
return priceData, nil
|
||||
}
|
||||
|
||||
// ModelPriceHelperPerCall 按次/按量计费的 PriceHelper (MJ、Task)
|
||||
func ModelPriceHelperPerCall(c *gin.Context, info *relaycommon.RelayInfo) (hosttypes.PriceData, error) {
|
||||
groupRatioInfo := HandleGroupRatio(c, info)
|
||||
|
||||
modelPrice, success := ratio_setting.GetModelPrice(info.OriginModelName, true)
|
||||
usePrice := success
|
||||
var modelRatio float64
|
||||
|
||||
if !success {
|
||||
defaultPrice, ok := ratio_setting.GetDefaultModelPriceMap()[info.OriginModelName]
|
||||
if ok {
|
||||
modelPrice = defaultPrice
|
||||
usePrice = true
|
||||
} else {
|
||||
var ratioSuccess bool
|
||||
var matchName string
|
||||
modelRatio, ratioSuccess, matchName = ratio_setting.GetModelRatio(info.OriginModelName)
|
||||
acceptUnsetRatio := false
|
||||
if info.UserSetting.AcceptUnsetRatioModel {
|
||||
acceptUnsetRatio = true
|
||||
}
|
||||
if !ratioSuccess && !acceptUnsetRatio {
|
||||
return hosttypes.PriceData{}, modelPriceNotConfiguredError(matchName, info.UserId)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
var quota int
|
||||
freeModel := false
|
||||
|
||||
if usePrice {
|
||||
var err error
|
||||
quota, err = common.QuotaFromFloatStrict(modelPrice * common.QuotaPerUnit * groupRatioInfo.GroupRatio)
|
||||
if err != nil {
|
||||
return hosttypes.PriceData{}, err
|
||||
}
|
||||
if !operation_setting.GetQuotaSetting().EnableFreeModelPreConsume {
|
||||
if groupRatioInfo.GroupRatio == 0 || modelPrice == 0 {
|
||||
quota = 0
|
||||
freeModel = true
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// 按量计费：以模型倍率的一半作为预扣额度
|
||||
var err error
|
||||
quota, err = common.QuotaFromFloatStrict(modelRatio / 2 * common.QuotaPerUnit * groupRatioInfo.GroupRatio)
|
||||
if err != nil {
|
||||
return hosttypes.PriceData{}, err
|
||||
}
|
||||
modelPrice = -1
|
||||
if !operation_setting.GetQuotaSetting().EnableFreeModelPreConsume {
|
||||
if groupRatioInfo.GroupRatio == 0 || modelRatio == 0 {
|
||||
quota = 0
|
||||
freeModel = true
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
priceData := hosttypes.PriceData{
|
||||
FreeModel: freeModel,
|
||||
ModelPrice: modelPrice,
|
||||
ModelRatio: modelRatio,
|
||||
UsePrice: usePrice,
|
||||
Quota: quota,
|
||||
GroupRatioInfo: groupRatioInfo,
|
||||
}
|
||||
return priceData, nil
|
||||
}
|
||||
|
||||
func HasModelBillingConfig(modelName string) bool {
|
||||
if _, ok := ratio_setting.GetModelPrice(modelName, false); ok {
|
||||
return true
|
||||
}
|
||||
if _, ok, _ := ratio_setting.GetModelRatio(modelName); ok {
|
||||
return true
|
||||
}
|
||||
if billing_setting.GetBillingMode(modelName) != billing_setting.BillingModeTieredExpr {
|
||||
return false
|
||||
}
|
||||
expr, ok := billing_setting.GetBillingExpr(modelName)
|
||||
return ok && strings.TrimSpace(expr) != ""
|
||||
}
|
||||
|
||||
func modelPriceHelperTiered(c *gin.Context, info *relaycommon.RelayInfo, promptTokens int, meta *types.TokenCountMeta, groupRatioInfo hosttypes.GroupRatioInfo) (hosttypes.PriceData, error) {
|
||||
exprStr, ok := billing_setting.GetBillingExpr(info.OriginModelName)
|
||||
if !ok {
|
||||
return hosttypes.PriceData{}, fmt.Errorf("model %s is configured as tiered_expr but has no billing expression", info.OriginModelName)
|
||||
}
|
||||
|
||||
estimatedCompletionTokens := meta.MaxTokens
|
||||
if estimatedCompletionTokens == 0 && groupRatioInfo.GroupRatio != 0 {
|
||||
estimatedCompletionTokens = defaultTieredPreConsumeMaxTokens
|
||||
}
|
||||
|
||||
requestInput, err := ResolveIncomingBillingExprRequestInput(c, info)
|
||||
if err != nil {
|
||||
return hosttypes.PriceData{}, err
|
||||
}
|
||||
|
||||
rawCost, trace, err := billingexpr.RunExprWithRequest(exprStr, billingexpr.TokenParams{
|
||||
P: float64(promptTokens),
|
||||
C: float64(estimatedCompletionTokens),
|
||||
Len: float64(promptTokens),
|
||||
}, requestInput)
|
||||
if err != nil {
|
||||
return hosttypes.PriceData{}, fmt.Errorf("model %s tiered expr run failed: %w", info.OriginModelName, err)
|
||||
}
|
||||
|
||||
// Expression coefficients are $/1M tokens prices; convert to quota the same way per-call billing does.
|
||||
quotaBeforeGroup := rawCost / 1_000_000 * common.QuotaPerUnit
|
||||
preConsumedQuota, err := billingexpr.QuotaRoundStrict(quotaBeforeGroup * groupRatioInfo.GroupRatio)
|
||||
if err != nil {
|
||||
return hosttypes.PriceData{}, err
|
||||
}
|
||||
|
||||
freeModel := false
|
||||
if !operation_setting.GetQuotaSetting().EnableFreeModelPreConsume {
|
||||
if groupRatioInfo.GroupRatio == 0 {
|
||||
preConsumedQuota = 0
|
||||
freeModel = true
|
||||
}
|
||||
}
|
||||
|
||||
exprHash := billingexpr.ExprHashString(exprStr)
|
||||
snapshot := &billingexpr.BillingSnapshot{
|
||||
BillingMode: billing_setting.BillingModeTieredExpr,
|
||||
ModelName: info.OriginModelName,
|
||||
ExprString: exprStr,
|
||||
ExprHash: exprHash,
|
||||
GroupRatio: groupRatioInfo.GroupRatio,
|
||||
EstimatedPromptTokens: promptTokens,
|
||||
EstimatedCompletionTokens: estimatedCompletionTokens,
|
||||
EstimatedQuotaBeforeGroup: quotaBeforeGroup,
|
||||
EstimatedQuotaAfterGroup: preConsumedQuota,
|
||||
EstimatedTier: trace.MatchedTier,
|
||||
QuotaPerUnit: common.QuotaPerUnit,
|
||||
ExprVersion: billingexpr.ExprVersion(exprStr),
|
||||
}
|
||||
info.TieredBillingSnapshot = snapshot
|
||||
info.BillingRequestInput = &requestInput
|
||||
|
||||
priceData := hosttypes.PriceData{
|
||||
FreeModel: freeModel,
|
||||
GroupRatioInfo: groupRatioInfo,
|
||||
QuotaToPreConsume: preConsumedQuota,
|
||||
}
|
||||
|
||||
logger.LogDebug(c, "model_price_helper_tiered result: model=%s preConsume=%d quotaBeforeGroup=%.2f groupRatio=%.2f tier=%s", info.OriginModelName, preConsumedQuota, quotaBeforeGroup, groupRatioInfo.GroupRatio, trace.MatchedTier)
|
||||
|
||||
info.PriceData = priceData
|
||||
return priceData, nil
|
||||
}
|
||||
@@ -0,0 +1,274 @@
|
||||
package helper
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/pkg/billingexpr"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/relaykit/types"
|
||||
"github.com/QuantumNous/new-api/setting/billing_setting"
|
||||
"github.com/QuantumNous/new-api/setting/config"
|
||||
"github.com/QuantumNous/new-api/setting/ratio_setting"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestModelPriceHelperTieredUsesPreloadedRequestInput(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
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": `{"tiered-test-model":"tiered_expr"}`,
|
||||
"billing_setting.billing_expr": `{"tiered-test-model":"param(\"stream\") == true ? tier(\"stream\", p * 3) : tier(\"base\", p * 2)"}`,
|
||||
}))
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(recorder)
|
||||
req := httptest.NewRequest(http.MethodPost, "/api/channel/test/1", nil)
|
||||
req.Body = nil
|
||||
req.ContentLength = 0
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
ctx.Request = req
|
||||
ctx.Set("group", "default")
|
||||
|
||||
info := &relaycommon.RelayInfo{
|
||||
OriginModelName: "tiered-test-model",
|
||||
UserGroup: "default",
|
||||
UsingGroup: "default",
|
||||
RequestHeaders: map[string]string{"Content-Type": "application/json"},
|
||||
BillingRequestInput: &billingexpr.RequestInput{
|
||||
Headers: map[string]string{"Content-Type": "application/json"},
|
||||
Body: []byte(`{"stream":true}`),
|
||||
},
|
||||
}
|
||||
|
||||
priceData, err := ModelPriceHelper(ctx, info, 1000, &types.TokenCountMeta{
|
||||
BillingRatios: map[string]float64{"n": 3},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 1500, priceData.QuotaToPreConsume)
|
||||
require.NotNil(t, info.TieredBillingSnapshot)
|
||||
require.Equal(t, "stream", info.TieredBillingSnapshot.EstimatedTier)
|
||||
require.Equal(t, billing_setting.BillingModeTieredExpr, info.TieredBillingSnapshot.BillingMode)
|
||||
require.Equal(t, common.QuotaPerUnit, info.TieredBillingSnapshot.QuotaPerUnit)
|
||||
}
|
||||
|
||||
func TestModelPriceHelperTieredPreConsumeMaxTokensFallback(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
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": `{"tiered-fallback-model":"tiered_expr"}`,
|
||||
"billing_setting.billing_expr": `{"tiered-fallback-model":"tier(\"base\", p * 3 + c * 15)"}`,
|
||||
"group_ratio_setting.group_ratio": `{"default":1,"free":0}`,
|
||||
}))
|
||||
|
||||
const promptTokens = 1000
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
group string
|
||||
maxTokens int
|
||||
expected int
|
||||
}{
|
||||
{
|
||||
// max_tokens omitted in a paid group -> fall back to 8192 completion tokens.
|
||||
// p*3 + c*15 = 1000*3 + 8192*15 = 125880 -> /1e6 * 500000 = 62940
|
||||
name: "non-free group falls back to 8192 completion tokens",
|
||||
group: "default",
|
||||
maxTokens: 0,
|
||||
expected: 62940,
|
||||
},
|
||||
{
|
||||
// explicit max_tokens is used verbatim, no fallback.
|
||||
// 1000*3 + 100*15 = 4500 -> /1e6 * 500000 = 2250
|
||||
name: "explicit max_tokens is used verbatim",
|
||||
group: "default",
|
||||
maxTokens: 100,
|
||||
expected: 2250,
|
||||
},
|
||||
{
|
||||
// free group (ratio 0) stays zero; fallback is gated on non-zero group ratio.
|
||||
name: "free group stays zero without fallback",
|
||||
group: "free",
|
||||
maxTokens: 0,
|
||||
expected: 0,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
recorder := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(recorder)
|
||||
req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
ctx.Request = req
|
||||
ctx.Set("group", tc.group)
|
||||
|
||||
info := &relaycommon.RelayInfo{
|
||||
OriginModelName: "tiered-fallback-model",
|
||||
UserGroup: tc.group,
|
||||
UsingGroup: tc.group,
|
||||
RequestHeaders: map[string]string{"Content-Type": "application/json"},
|
||||
BillingRequestInput: &billingexpr.RequestInput{
|
||||
Headers: map[string]string{"Content-Type": "application/json"},
|
||||
Body: []byte(`{}`),
|
||||
},
|
||||
}
|
||||
|
||||
priceData, err := ModelPriceHelper(ctx, info, promptTokens, &types.TokenCountMeta{MaxTokens: tc.maxTokens})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, tc.expected, priceData.QuotaToPreConsume)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestModelPriceHelperTieredRejectsPreConsumeOverflow(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
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": `{"tiered-overflow-model":"tiered_expr"}`,
|
||||
"billing_setting.billing_expr": `{"tiered-overflow-model":"tier(\"overflow\", p * 1000000000000000)"}`,
|
||||
"group_ratio_setting.group_ratio": `{"default":1}`,
|
||||
}))
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(recorder)
|
||||
ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
|
||||
ctx.Set("group", "default")
|
||||
info := &relaycommon.RelayInfo{
|
||||
OriginModelName: "tiered-overflow-model",
|
||||
UserGroup: "default",
|
||||
UsingGroup: "default",
|
||||
BillingRequestInput: &billingexpr.RequestInput{
|
||||
Body: []byte(`{}`),
|
||||
},
|
||||
}
|
||||
|
||||
_, err := ModelPriceHelper(ctx, info, 1000, &types.TokenCountMeta{})
|
||||
|
||||
var clamp *common.QuotaClamp
|
||||
require.ErrorAs(t, err, &clamp)
|
||||
require.Equal(t, "QuotaRound", clamp.Op)
|
||||
require.Equal(t, common.QuotaClampOverflow, clamp.Kind)
|
||||
}
|
||||
|
||||
func TestModelPriceHelperRequestBillingRatiosOnlyApplyToFixedPrice(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
savedModelPrices := ratio_setting.ModelPrice2JSONString()
|
||||
savedModelRatios := ratio_setting.ModelRatio2JSONString()
|
||||
t.Cleanup(func() {
|
||||
require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(savedModelPrices))
|
||||
require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(savedModelRatios))
|
||||
})
|
||||
|
||||
modelPrices, err := common.Marshal(map[string]float64{
|
||||
"fixed-image-price": 0.04,
|
||||
"fractional-image-price": 0.0000012,
|
||||
"overflow-image-price": float64(common.MaxQuota) / common.QuotaPerUnit / 2,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(string(modelPrices)))
|
||||
modelRatios, err := common.Marshal(map[string]float64{"ratio-image-price": 15})
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(string(modelRatios)))
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
model string
|
||||
wantQuota int
|
||||
wantUsePrice bool
|
||||
wantImageCount bool
|
||||
}{
|
||||
{
|
||||
name: "fixed price applies image count",
|
||||
model: "fixed-image-price",
|
||||
wantQuota: 180000,
|
||||
wantUsePrice: true,
|
||||
wantImageCount: true,
|
||||
},
|
||||
{
|
||||
name: "ratio price ignores request billing ratios",
|
||||
model: "ratio-image-price",
|
||||
wantQuota: 15000,
|
||||
wantUsePrice: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
ctx.Set("group", "default")
|
||||
info := &relaycommon.RelayInfo{
|
||||
OriginModelName: tt.model,
|
||||
UserGroup: "default",
|
||||
UsingGroup: "default",
|
||||
}
|
||||
meta := &types.TokenCountMeta{
|
||||
ImagePriceRatio: 3,
|
||||
BillingRatios: map[string]float64{"n": 3},
|
||||
}
|
||||
|
||||
priceData, err := ModelPriceHelper(ctx, info, 1000, meta)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, tt.wantQuota, priceData.QuotaToPreConsume)
|
||||
require.Equal(t, tt.wantUsePrice, priceData.UsePrice)
|
||||
require.Equal(t, tt.wantImageCount, priceData.HasOtherRatio("n"))
|
||||
require.Equal(t, priceData.OtherRatios(), info.PriceData.OtherRatios())
|
||||
})
|
||||
}
|
||||
|
||||
newInfo := func(model string) (*gin.Context, *relaycommon.RelayInfo) {
|
||||
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
ctx.Set("group", "default")
|
||||
return ctx, &relaycommon.RelayInfo{
|
||||
OriginModelName: model,
|
||||
UserGroup: "default",
|
||||
UsingGroup: "default",
|
||||
}
|
||||
}
|
||||
meta := &types.TokenCountMeta{BillingRatios: map[string]float64{"n": 3}}
|
||||
|
||||
ctx, info := newInfo("fractional-image-price")
|
||||
priceData, err := ModelPriceHelper(ctx, info, 0, meta)
|
||||
require.NoError(t, err)
|
||||
// 0.0000012 * 500000 * 3 = 1.8, then truncate once to 1.
|
||||
require.Equal(t, 1, priceData.QuotaToPreConsume)
|
||||
|
||||
ctx, info = newInfo("overflow-image-price")
|
||||
_, err = ModelPriceHelper(ctx, info, 0, meta)
|
||||
var clamp *common.QuotaClamp
|
||||
require.ErrorAs(t, err, &clamp)
|
||||
require.Equal(t, "QuotaFromFloat", clamp.Op)
|
||||
require.Equal(t, common.QuotaClampOverflow, clamp.Kind)
|
||||
require.Nil(t, info.Billing)
|
||||
}
|
||||
@@ -0,0 +1,52 @@
|
||||
package helper
|
||||
|
||||
import (
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
)
|
||||
|
||||
// StreamResult is passed to each dataHandler invocation, providing methods
|
||||
// to record soft errors, signal fatal stops, or mark normal completion.
|
||||
// StreamScannerHandler checks IsStopped() after each callback invocation.
|
||||
type StreamResult struct {
|
||||
status *relaycommon.StreamStatus
|
||||
stopped bool
|
||||
}
|
||||
|
||||
func newStreamResult(status *relaycommon.StreamStatus) *StreamResult {
|
||||
return &StreamResult{status: status}
|
||||
}
|
||||
|
||||
// Error records a soft error. The stream continues processing.
|
||||
// Can be called multiple times per chunk.
|
||||
func (r *StreamResult) Error(err error) {
|
||||
if err == nil {
|
||||
return
|
||||
}
|
||||
r.status.RecordError(err.Error())
|
||||
}
|
||||
|
||||
// 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.status.SetEndReason(relaycommon.StreamEndReasonHandlerStop, err)
|
||||
r.stopped = true
|
||||
}
|
||||
|
||||
// Done signals that the handler has finished processing normally
|
||||
// (e.g., Dify "message_end"). The stream stops after this chunk.
|
||||
func (r *StreamResult) Done() {
|
||||
r.status.SetEndReason(relaycommon.StreamEndReasonDone, nil)
|
||||
r.stopped = true
|
||||
}
|
||||
|
||||
// IsStopped returns whether Stop() or Done() was called during this chunk.
|
||||
func (r *StreamResult) IsStopped() bool {
|
||||
return r.stopped
|
||||
}
|
||||
|
||||
// reset clears the per-chunk stopped flag so the object can be reused.
|
||||
func (r *StreamResult) reset() {
|
||||
r.stopped = false
|
||||
}
|
||||
@@ -0,0 +1,310 @@
|
||||
package helper
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"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/service"
|
||||
"github.com/QuantumNous/new-api/setting/operation_setting"
|
||||
|
||||
"github.com/bytedance/gopkg/util/gopool"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
const (
|
||||
InitialScannerBufferSize = 64 << 10 // 64KB (64*1024)
|
||||
DefaultMaxScannerBufferSize = 128 << 20 // 64MB (64*1024*1024) default SSE buffer size
|
||||
DefaultPingInterval = 10 * time.Second
|
||||
// streamWriteTimeout bounds a single blocked write to a slow client so the
|
||||
// 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
|
||||
)
|
||||
|
||||
func getScannerBufferSize() int {
|
||||
if constant.StreamScannerMaxBufferMB > 0 {
|
||||
return constant.StreamScannerMaxBufferMB << 20
|
||||
}
|
||||
return DefaultMaxScannerBufferSize
|
||||
}
|
||||
|
||||
func NewStreamScanner(reader io.Reader) *bufio.Scanner {
|
||||
scanner := bufio.NewScanner(reader)
|
||||
scanner.Buffer(make([]byte, InitialScannerBufferSize), getScannerBufferSize())
|
||||
return scanner
|
||||
}
|
||||
|
||||
func copyCodexSSEHeaders(c *gin.Context, resp *http.Response) {
|
||||
if c == nil || c.Writer == nil || resp == nil {
|
||||
return
|
||||
}
|
||||
// codex
|
||||
for _, name := range []string{"X-Reasoning-Included", "X-Codex-Turn-State"} {
|
||||
values := resp.Header.Values(name)
|
||||
if !service.ShouldCopyUpstreamHeader(c, name, values) {
|
||||
continue
|
||||
}
|
||||
for _, value := range values {
|
||||
if value != "" {
|
||||
c.Writer.Header().Add(name, value)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ExtendWriteDeadline pushes the connection write deadline forward before each
|
||||
// stream write. Best-effort: writers that don't support deadlines (e.g.
|
||||
// httptest recorders) are silently ignored.
|
||||
func ExtendWriteDeadline(c *gin.Context) {
|
||||
if c == nil || c.Writer == nil {
|
||||
return
|
||||
}
|
||||
_ = http.NewResponseController(c.Writer).SetWriteDeadline(time.Now().Add(streamWriteTimeout))
|
||||
}
|
||||
|
||||
func StreamScannerHandler(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo, dataHandler func(data string, sr *StreamResult)) {
|
||||
|
||||
if resp == nil || dataHandler == nil {
|
||||
return
|
||||
}
|
||||
|
||||
// 无条件新建 StreamStatus
|
||||
info.StreamStatus = relaycommon.NewStreamStatus()
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
|
||||
streamingTimeout := time.Duration(constant.StreamingTimeout) * time.Second
|
||||
|
||||
var (
|
||||
stopChan = make(chan bool, 3) // 增加缓冲区避免阻塞
|
||||
scanner = NewStreamScanner(resp.Body)
|
||||
ticker = time.NewTicker(streamingTimeout)
|
||||
pingTicker *time.Ticker
|
||||
writeMutex sync.Mutex // Mutex to protect concurrent writes
|
||||
wg sync.WaitGroup // 用于等待所有 goroutine 退出
|
||||
cleanupOnce sync.Once
|
||||
stopOnce sync.Once
|
||||
)
|
||||
|
||||
stop := func() {
|
||||
stopOnce.Do(func() {
|
||||
close(stopChan)
|
||||
})
|
||||
}
|
||||
|
||||
generalSettings := operation_setting.GetGeneralSetting()
|
||||
pingEnabled := generalSettings.PingIntervalEnabled && !info.DisablePing
|
||||
pingInterval := time.Duration(generalSettings.PingIntervalSeconds) * time.Second
|
||||
if pingInterval <= 0 {
|
||||
pingInterval = DefaultPingInterval
|
||||
}
|
||||
|
||||
if pingEnabled {
|
||||
pingTicker = time.NewTicker(pingInterval)
|
||||
}
|
||||
|
||||
logger.LogDebug(c, "relay timeout seconds: %d", common.RelayTimeout)
|
||||
logger.LogDebug(c, "relay max idle conns: %d", common.RelayMaxIdleConns)
|
||||
logger.LogDebug(c, "relay max idle conns per host: %d", common.RelayMaxIdleConnsPerHost)
|
||||
logger.LogDebug(c, "streaming timeout seconds: %d", int64(streamingTimeout.Seconds()))
|
||||
logger.LogDebug(c, "ping interval seconds: %d", int64(pingInterval.Seconds()))
|
||||
|
||||
cleanup := func() {
|
||||
cleanupOnce.Do(func() {
|
||||
cancel()
|
||||
stop()
|
||||
if resp.Body != nil {
|
||||
_ = resp.Body.Close()
|
||||
}
|
||||
|
||||
ticker.Stop()
|
||||
if pingTicker != nil {
|
||||
pingTicker.Stop()
|
||||
}
|
||||
|
||||
wg.Wait()
|
||||
})
|
||||
}
|
||||
// Ensure gin.Context is not returned to Gin's pool while any stream goroutine can still use it.
|
||||
defer cleanup()
|
||||
|
||||
scanner.Split(bufio.ScanLines)
|
||||
copyCodexSSEHeaders(c, resp)
|
||||
SetEventStreamHeaders(c)
|
||||
|
||||
ctx = context.WithValue(ctx, "stop_chan", stopChan)
|
||||
|
||||
// Handle ping data sending with improved error handling
|
||||
if pingEnabled && pingTicker != nil {
|
||||
wg.Add(1)
|
||||
gopool.Go(func() {
|
||||
defer func() {
|
||||
if r := recover(); r != nil {
|
||||
logger.LogError(c, fmt.Sprintf("ping goroutine panic: %v", r))
|
||||
info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonPanic, fmt.Errorf("ping panic: %v", r))
|
||||
stop()
|
||||
}
|
||||
logger.LogDebug(c, "ping goroutine exited")
|
||||
wg.Done()
|
||||
}()
|
||||
|
||||
// 添加超时保护，防止 goroutine 无限运行
|
||||
maxPingDuration := 30 * time.Minute // 最大 ping 持续时间
|
||||
pingTimeout := time.NewTimer(maxPingDuration)
|
||||
defer pingTimeout.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-pingTicker.C:
|
||||
var err error
|
||||
func() {
|
||||
writeMutex.Lock()
|
||||
defer writeMutex.Unlock()
|
||||
ExtendWriteDeadline(c)
|
||||
err = PingData(c)
|
||||
}()
|
||||
if err != nil {
|
||||
logger.LogError(c, "ping data error: "+err.Error())
|
||||
info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonPingFail, err)
|
||||
return
|
||||
}
|
||||
logger.LogDebug(c, "ping data sent")
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-stopChan:
|
||||
return
|
||||
case <-c.Request.Context().Done():
|
||||
// 监听客户端断开连接
|
||||
return
|
||||
case <-pingTimeout.C:
|
||||
logger.LogError(c, "ping goroutine max duration reached")
|
||||
return
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
dataChan := make(chan string, 10)
|
||||
|
||||
wg.Add(1)
|
||||
gopool.Go(func() {
|
||||
defer func() {
|
||||
if r := recover(); r != nil {
|
||||
logger.LogError(c, fmt.Sprintf("data handler goroutine panic: %v", r))
|
||||
info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonPanic, fmt.Errorf("handler panic: %v", r))
|
||||
}
|
||||
stop()
|
||||
wg.Done()
|
||||
}()
|
||||
sr := newStreamResult(info.StreamStatus)
|
||||
for data := range dataChan {
|
||||
sr.reset()
|
||||
func() {
|
||||
writeMutex.Lock()
|
||||
defer writeMutex.Unlock()
|
||||
ExtendWriteDeadline(c)
|
||||
dataHandler(data, sr)
|
||||
}()
|
||||
if sr.IsStopped() {
|
||||
return
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
// Scanner goroutine with improved error handling
|
||||
wg.Add(1)
|
||||
common.RelayCtxGo(ctx, func() {
|
||||
defer func() {
|
||||
close(dataChan)
|
||||
if r := recover(); r != nil {
|
||||
logger.LogError(c, fmt.Sprintf("scanner goroutine panic: %v", r))
|
||||
info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonPanic, fmt.Errorf("scanner panic: %v", r))
|
||||
}
|
||||
stop()
|
||||
logger.LogDebug(c, "scanner goroutine exited")
|
||||
wg.Done()
|
||||
}()
|
||||
|
||||
for scanner.Scan() {
|
||||
// 检查是否需要停止
|
||||
select {
|
||||
case <-stopChan:
|
||||
return
|
||||
case <-ctx.Done():
|
||||
return
|
||||
default:
|
||||
}
|
||||
|
||||
ticker.Reset(streamingTimeout)
|
||||
data := scanner.Text()
|
||||
logger.LogDebug(c, "stream scanner data: %s", data)
|
||||
|
||||
if len(data) < 6 {
|
||||
continue
|
||||
}
|
||||
if data[:5] != "data:" && data[:6] != "[DONE]" {
|
||||
continue
|
||||
}
|
||||
data = data[5:]
|
||||
data = strings.TrimSpace(data)
|
||||
if data == "" {
|
||||
continue
|
||||
}
|
||||
if !strings.HasPrefix(data, "[DONE]") {
|
||||
info.SetFirstResponseTime()
|
||||
info.ReceivedResponseCount++
|
||||
|
||||
select {
|
||||
case dataChan <- data:
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-stopChan:
|
||||
return
|
||||
}
|
||||
} else {
|
||||
info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonDone, nil)
|
||||
logger.LogDebug(c, "received [DONE], stopping scanner")
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
if err := scanner.Err(); err != nil {
|
||||
if err != io.EOF {
|
||||
logger.LogError(c, "scanner error: "+err.Error())
|
||||
info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonScannerErr, err)
|
||||
}
|
||||
}
|
||||
info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonEOF, nil)
|
||||
})
|
||||
|
||||
// 主循环等待完成或超时
|
||||
select {
|
||||
case <-ticker.C:
|
||||
info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonTimeout, nil)
|
||||
case <-stopChan:
|
||||
// EndReason already set by the goroutine that triggered stopChan
|
||||
case <-c.Request.Context().Done():
|
||||
// 客户端断开：立即 cleanup 关闭上游 resp.Body，解除 scanner 阻塞并让上游停止生成，
|
||||
// 避免为已放弃的请求继续消费上游 token。
|
||||
info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonClientGone, c.Request.Context().Err())
|
||||
}
|
||||
|
||||
cleanup()
|
||||
if info.StreamStatus.IsNormalEnd() && !info.StreamStatus.HasErrors() {
|
||||
logger.LogInfo(c, fmt.Sprintf("stream ended: %s", info.StreamStatus.Summary()))
|
||||
} else {
|
||||
logger.LogError(c, fmt.Sprintf("stream ended: %s, received=%d", info.StreamStatus.Summary(), info.ReceivedResponseCount))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,573 @@
|
||||
package helper
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/setting/operation_setting"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func init() {
|
||||
gin.SetMode(gin.TestMode)
|
||||
if constant.StreamingTimeout == 0 {
|
||||
constant.StreamingTimeout = 30
|
||||
}
|
||||
}
|
||||
|
||||
func setupStreamTest(t *testing.T, body io.Reader) (*gin.Context, *http.Response, *relaycommon.RelayInfo) {
|
||||
t.Helper()
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
|
||||
|
||||
resp := &http.Response{
|
||||
Body: io.NopCloser(body),
|
||||
}
|
||||
|
||||
info := &relaycommon.RelayInfo{
|
||||
ChannelMeta: &relaycommon.ChannelMeta{},
|
||||
}
|
||||
|
||||
return c, resp, info
|
||||
}
|
||||
|
||||
func buildSSEBody(n int) string {
|
||||
var b strings.Builder
|
||||
for i := 0; i < n; i++ {
|
||||
fmt.Fprintf(&b, "data: {\"id\":%d,\"choices\":[{\"delta\":{\"content\":\"token_%d\"}}]}\n", i, i)
|
||||
}
|
||||
b.WriteString("data: [DONE]\n")
|
||||
return b.String()
|
||||
}
|
||||
|
||||
// ---------- Basic correctness ----------
|
||||
|
||||
func TestStreamScannerHandler_NilInputs(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||||
|
||||
info := &relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{}}
|
||||
|
||||
StreamScannerHandler(c, nil, info, func(data string, sr *StreamResult) {})
|
||||
StreamScannerHandler(c, &http.Response{Body: io.NopCloser(strings.NewReader(""))}, info, nil)
|
||||
}
|
||||
|
||||
func TestNewStreamScanner_AllowsLargeStreamLine(t *testing.T) {
|
||||
oldBufferMB := constant.StreamScannerMaxBufferMB
|
||||
constant.StreamScannerMaxBufferMB = 1
|
||||
t.Cleanup(func() {
|
||||
constant.StreamScannerMaxBufferMB = oldBufferMB
|
||||
})
|
||||
|
||||
payload := strings.Repeat("x", 128<<10)
|
||||
scanner := NewStreamScanner(strings.NewReader("data: " + payload + "\n"))
|
||||
scanner.Split(bufio.ScanLines)
|
||||
|
||||
require.True(t, scanner.Scan())
|
||||
assert.Equal(t, "data: "+payload, scanner.Text())
|
||||
require.NoError(t, scanner.Err())
|
||||
}
|
||||
|
||||
func TestStreamScannerHandler_EmptyBody(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
c, resp, info := setupStreamTest(t, strings.NewReader(""))
|
||||
|
||||
var called atomic.Bool
|
||||
StreamScannerHandler(c, resp, info, func(data string, sr *StreamResult) {
|
||||
called.Store(true)
|
||||
})
|
||||
|
||||
assert.False(t, called.Load(), "handler should not be called for empty body")
|
||||
}
|
||||
|
||||
func TestStreamScannerHandler_1000Chunks(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
const numChunks = 1000
|
||||
body := buildSSEBody(numChunks)
|
||||
c, resp, info := setupStreamTest(t, strings.NewReader(body))
|
||||
|
||||
var count atomic.Int64
|
||||
StreamScannerHandler(c, resp, info, func(data string, sr *StreamResult) {
|
||||
count.Add(1)
|
||||
})
|
||||
|
||||
assert.Equal(t, int64(numChunks), count.Load())
|
||||
assert.Equal(t, numChunks, info.ReceivedResponseCount)
|
||||
}
|
||||
|
||||
func TestStreamScannerHandler_OrderPreserved(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
const numChunks = 500
|
||||
body := buildSSEBody(numChunks)
|
||||
c, resp, info := setupStreamTest(t, strings.NewReader(body))
|
||||
|
||||
var mu sync.Mutex
|
||||
received := make([]string, 0, numChunks)
|
||||
|
||||
StreamScannerHandler(c, resp, info, func(data string, sr *StreamResult) {
|
||||
mu.Lock()
|
||||
received = append(received, data)
|
||||
mu.Unlock()
|
||||
})
|
||||
|
||||
require.Equal(t, numChunks, len(received))
|
||||
for i := 0; i < numChunks; i++ {
|
||||
expected := fmt.Sprintf("{\"id\":%d,\"choices\":[{\"delta\":{\"content\":\"token_%d\"}}]}", i, i)
|
||||
assert.Equal(t, expected, received[i], "chunk %d out of order", i)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStreamScannerHandler_DoneStopsScanner(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
body := buildSSEBody(50) + "data: should_not_appear\n"
|
||||
c, resp, info := setupStreamTest(t, strings.NewReader(body))
|
||||
|
||||
var count atomic.Int64
|
||||
StreamScannerHandler(c, resp, info, func(data string, sr *StreamResult) {
|
||||
count.Add(1)
|
||||
})
|
||||
|
||||
assert.Equal(t, int64(50), count.Load(), "data after [DONE] must not be processed")
|
||||
}
|
||||
|
||||
func TestStreamScannerHandler_StopStopsStream(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
const numChunks = 200
|
||||
body := buildSSEBody(numChunks)
|
||||
c, resp, info := setupStreamTest(t, strings.NewReader(body))
|
||||
|
||||
const stopAt int64 = 50
|
||||
var count atomic.Int64
|
||||
StreamScannerHandler(c, resp, info, func(data string, sr *StreamResult) {
|
||||
n := count.Add(1)
|
||||
if n >= stopAt {
|
||||
sr.Stop(fmt.Errorf("fatal at %d", n))
|
||||
}
|
||||
})
|
||||
|
||||
assert.Equal(t, stopAt, count.Load())
|
||||
require.NotNil(t, info.StreamStatus)
|
||||
assert.Equal(t, relaycommon.StreamEndReasonHandlerStop, info.StreamStatus.EndReason)
|
||||
}
|
||||
|
||||
func TestStreamScannerHandler_SkipsNonDataLines(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
var b strings.Builder
|
||||
b.WriteString(": comment line\n")
|
||||
b.WriteString("event: message\n")
|
||||
b.WriteString("id: 12345\n")
|
||||
b.WriteString("retry: 5000\n")
|
||||
for i := 0; i < 100; i++ {
|
||||
fmt.Fprintf(&b, "data: payload_%d\n", i)
|
||||
b.WriteString(": interleaved comment\n")
|
||||
}
|
||||
b.WriteString("data: [DONE]\n")
|
||||
|
||||
c, resp, info := setupStreamTest(t, strings.NewReader(b.String()))
|
||||
|
||||
var count atomic.Int64
|
||||
StreamScannerHandler(c, resp, info, func(data string, sr *StreamResult) {
|
||||
count.Add(1)
|
||||
})
|
||||
|
||||
assert.Equal(t, int64(100), count.Load())
|
||||
}
|
||||
|
||||
func TestStreamScannerHandler_DataWithExtraSpaces(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
body := "data: {\"trimmed\":true} \ndata: [DONE]\n"
|
||||
c, resp, info := setupStreamTest(t, strings.NewReader(body))
|
||||
|
||||
var got string
|
||||
StreamScannerHandler(c, resp, info, func(data string, sr *StreamResult) {
|
||||
got = data
|
||||
})
|
||||
|
||||
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) {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
|
||||
pr, pw := io.Pipe()
|
||||
t.Cleanup(func() {
|
||||
_ = pr.Close()
|
||||
_ = pw.Close()
|
||||
})
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil).WithContext(ctx)
|
||||
|
||||
resp := &http.Response{Body: pr}
|
||||
info := &relaycommon.RelayInfo{
|
||||
DisablePing: true,
|
||||
ChannelMeta: &relaycommon.ChannelMeta{},
|
||||
}
|
||||
|
||||
var count atomic.Int64
|
||||
firstHandled := make(chan struct{})
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
StreamScannerHandler(c, resp, info, func(data string, sr *StreamResult) {
|
||||
count.Add(1)
|
||||
_ = StringData(c, data)
|
||||
if data == "first" {
|
||||
close(firstHandled)
|
||||
}
|
||||
})
|
||||
close(done)
|
||||
}()
|
||||
|
||||
_, err := fmt.Fprint(pw, "data: first\n")
|
||||
require.NoError(t, err)
|
||||
|
||||
select {
|
||||
case <-firstHandled:
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("timed out waiting for first chunk")
|
||||
}
|
||||
|
||||
cancel()
|
||||
|
||||
// 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")
|
||||
}
|
||||
|
||||
// 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")
|
||||
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
|
||||
oldSeconds := setting.PingIntervalSeconds
|
||||
setting.PingIntervalEnabled = true
|
||||
setting.PingIntervalSeconds = 1
|
||||
t.Cleanup(func() {
|
||||
setting.PingIntervalEnabled = oldEnabled
|
||||
setting.PingIntervalSeconds = oldSeconds
|
||||
})
|
||||
|
||||
pr, pw := io.Pipe()
|
||||
go func() {
|
||||
defer pw.Close()
|
||||
for i := 0; i < 4; i++ {
|
||||
fmt.Fprintf(pw, "data: chunk_%d\n", i)
|
||||
time.Sleep(400 * time.Millisecond)
|
||||
}
|
||||
fmt.Fprint(pw, "data: [DONE]\n")
|
||||
}()
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
|
||||
|
||||
resp := &http.Response{Body: pr}
|
||||
info := &relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{}}
|
||||
|
||||
var count atomic.Int64
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
StreamScannerHandler(c, resp, info, func(data string, sr *StreamResult) {
|
||||
count.Add(1)
|
||||
})
|
||||
close(done)
|
||||
}()
|
||||
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("timed out waiting for stream to finish")
|
||||
}
|
||||
|
||||
assert.Equal(t, int64(4), count.Load())
|
||||
|
||||
body := recorder.Body.String()
|
||||
pingCount := strings.Count(body, ": PING")
|
||||
assert.GreaterOrEqual(t, pingCount, 1,
|
||||
"expected at least 1 ping during slow stream with 1s interval; got %d", pingCount)
|
||||
}
|
||||
|
||||
func TestStreamScannerHandler_PingDisabledByRelayInfo(t *testing.T) {
|
||||
setting := operation_setting.GetGeneralSetting()
|
||||
oldEnabled := setting.PingIntervalEnabled
|
||||
oldSeconds := setting.PingIntervalSeconds
|
||||
setting.PingIntervalEnabled = true
|
||||
setting.PingIntervalSeconds = 1
|
||||
t.Cleanup(func() {
|
||||
setting.PingIntervalEnabled = oldEnabled
|
||||
setting.PingIntervalSeconds = oldSeconds
|
||||
})
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
|
||||
|
||||
resp := &http.Response{Body: io.NopCloser(strings.NewReader(buildSSEBody(5)))}
|
||||
info := &relaycommon.RelayInfo{
|
||||
DisablePing: true,
|
||||
ChannelMeta: &relaycommon.ChannelMeta{},
|
||||
}
|
||||
|
||||
var count atomic.Int64
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
StreamScannerHandler(c, resp, info, func(data string, sr *StreamResult) {
|
||||
count.Add(1)
|
||||
})
|
||||
close(done)
|
||||
}()
|
||||
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("timed out")
|
||||
}
|
||||
|
||||
assert.Equal(t, int64(5), count.Load())
|
||||
|
||||
body := recorder.Body.String()
|
||||
pingCount := strings.Count(body, ": PING")
|
||||
assert.Equal(t, 0, pingCount, "pings should be disabled when DisablePing=true")
|
||||
}
|
||||
|
||||
// ---------- StreamStatus integration ----------
|
||||
|
||||
func TestStreamScannerHandler_StreamStatus_DoneReason(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
body := buildSSEBody(10)
|
||||
c, resp, info := setupStreamTest(t, strings.NewReader(body))
|
||||
|
||||
StreamScannerHandler(c, resp, info, func(data string, sr *StreamResult) {})
|
||||
|
||||
require.NotNil(t, info.StreamStatus)
|
||||
assert.Equal(t, relaycommon.StreamEndReasonDone, info.StreamStatus.EndReason)
|
||||
assert.Nil(t, info.StreamStatus.EndError)
|
||||
assert.True(t, info.StreamStatus.IsNormalEnd())
|
||||
assert.False(t, info.StreamStatus.HasErrors())
|
||||
}
|
||||
|
||||
func TestStreamScannerHandler_StreamStatus_EOFWithoutDone(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
var b strings.Builder
|
||||
for i := 0; i < 5; i++ {
|
||||
fmt.Fprintf(&b, "data: {\"id\":%d}\n", i)
|
||||
}
|
||||
c, resp, info := setupStreamTest(t, strings.NewReader(b.String()))
|
||||
|
||||
StreamScannerHandler(c, resp, info, func(data string, sr *StreamResult) {})
|
||||
|
||||
require.NotNil(t, info.StreamStatus)
|
||||
assert.Equal(t, relaycommon.StreamEndReasonEOF, info.StreamStatus.EndReason)
|
||||
assert.True(t, info.StreamStatus.IsNormalEnd())
|
||||
}
|
||||
|
||||
func TestStreamScannerHandler_StreamStatus_HandlerStop(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
body := buildSSEBody(100)
|
||||
c, resp, info := setupStreamTest(t, strings.NewReader(body))
|
||||
|
||||
var count atomic.Int64
|
||||
StreamScannerHandler(c, resp, info, func(data string, sr *StreamResult) {
|
||||
n := count.Add(1)
|
||||
if n >= 10 {
|
||||
sr.Stop(fmt.Errorf("stop at 10"))
|
||||
}
|
||||
})
|
||||
|
||||
require.NotNil(t, info.StreamStatus)
|
||||
assert.Equal(t, relaycommon.StreamEndReasonHandlerStop, info.StreamStatus.EndReason)
|
||||
assert.True(t, info.StreamStatus.HasErrors())
|
||||
}
|
||||
|
||||
func TestStreamScannerHandler_StreamStatus_HandlerDone(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
body := buildSSEBody(20)
|
||||
c, resp, info := setupStreamTest(t, strings.NewReader(body))
|
||||
|
||||
var count atomic.Int64
|
||||
StreamScannerHandler(c, resp, info, func(data string, sr *StreamResult) {
|
||||
n := count.Add(1)
|
||||
if n >= 5 {
|
||||
sr.Done()
|
||||
}
|
||||
})
|
||||
|
||||
assert.Equal(t, int64(5), count.Load())
|
||||
require.NotNil(t, info.StreamStatus)
|
||||
assert.Equal(t, relaycommon.StreamEndReasonDone, info.StreamStatus.EndReason)
|
||||
assert.False(t, info.StreamStatus.HasErrors())
|
||||
}
|
||||
|
||||
func TestStreamScannerHandler_StreamStatus_Timeout(t *testing.T) {
|
||||
// Not parallel: modifies global constant.StreamingTimeout
|
||||
oldTimeout := constant.StreamingTimeout
|
||||
constant.StreamingTimeout = 1
|
||||
t.Cleanup(func() { constant.StreamingTimeout = oldTimeout })
|
||||
|
||||
pr, pw := io.Pipe()
|
||||
go func() {
|
||||
fmt.Fprint(pw, "data: {\"id\":1}\n")
|
||||
time.Sleep(2 * time.Second)
|
||||
pw.Close()
|
||||
}()
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
|
||||
|
||||
resp := &http.Response{Body: pr}
|
||||
info := &relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{}}
|
||||
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
StreamScannerHandler(c, resp, info, func(data string, sr *StreamResult) {})
|
||||
close(done)
|
||||
}()
|
||||
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("timed out waiting for stream timeout")
|
||||
}
|
||||
|
||||
require.NotNil(t, info.StreamStatus)
|
||||
assert.Equal(t, relaycommon.StreamEndReasonTimeout, info.StreamStatus.EndReason)
|
||||
assert.False(t, info.StreamStatus.IsNormalEnd())
|
||||
}
|
||||
|
||||
func TestStreamScannerHandler_StreamStatus_SoftErrors(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
body := buildSSEBody(10)
|
||||
c, resp, info := setupStreamTest(t, strings.NewReader(body))
|
||||
|
||||
StreamScannerHandler(c, resp, info, func(data string, sr *StreamResult) {
|
||||
sr.Error(fmt.Errorf("soft error for chunk"))
|
||||
})
|
||||
|
||||
require.NotNil(t, info.StreamStatus)
|
||||
assert.Equal(t, relaycommon.StreamEndReasonDone, info.StreamStatus.EndReason)
|
||||
assert.True(t, info.StreamStatus.HasErrors())
|
||||
assert.Equal(t, 10, info.StreamStatus.TotalErrorCount())
|
||||
}
|
||||
|
||||
func TestStreamScannerHandler_StreamStatus_MultipleErrorsPerChunk(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
body := buildSSEBody(5)
|
||||
c, resp, info := setupStreamTest(t, strings.NewReader(body))
|
||||
|
||||
StreamScannerHandler(c, resp, info, func(data string, sr *StreamResult) {
|
||||
sr.Error(fmt.Errorf("error A"))
|
||||
sr.Error(fmt.Errorf("error B"))
|
||||
})
|
||||
|
||||
require.NotNil(t, info.StreamStatus)
|
||||
assert.Equal(t, relaycommon.StreamEndReasonDone, info.StreamStatus.EndReason)
|
||||
assert.Equal(t, 10, info.StreamStatus.TotalErrorCount())
|
||||
}
|
||||
|
||||
func TestStreamScannerHandler_StreamStatus_ErrorThenStop(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
// Use a large body without [DONE] to avoid race between scanner's [DONE]
|
||||
// and handler's Stop on the sync.Once EndReason.
|
||||
var b strings.Builder
|
||||
for i := 0; i < 100; i++ {
|
||||
fmt.Fprintf(&b, "data: {\"id\":%d}\n", i)
|
||||
}
|
||||
c, resp, info := setupStreamTest(t, strings.NewReader(b.String()))
|
||||
|
||||
var count atomic.Int64
|
||||
StreamScannerHandler(c, resp, info, func(data string, sr *StreamResult) {
|
||||
count.Add(1)
|
||||
sr.Error(fmt.Errorf("soft error"))
|
||||
sr.Stop(fmt.Errorf("fatal"))
|
||||
})
|
||||
|
||||
assert.Equal(t, int64(1), count.Load())
|
||||
require.NotNil(t, info.StreamStatus)
|
||||
assert.Equal(t, relaycommon.StreamEndReasonHandlerStop, info.StreamStatus.EndReason)
|
||||
assert.Equal(t, 2, info.StreamStatus.TotalErrorCount())
|
||||
}
|
||||
|
||||
func TestStreamScannerHandler_StreamStatus_InitializedIfNil(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
body := buildSSEBody(1)
|
||||
c, resp, info := setupStreamTest(t, strings.NewReader(body))
|
||||
|
||||
assert.Nil(t, info.StreamStatus)
|
||||
|
||||
StreamScannerHandler(c, resp, info, func(data string, sr *StreamResult) {})
|
||||
|
||||
assert.NotNil(t, info.StreamStatus)
|
||||
}
|
||||
|
||||
func TestStreamScannerHandler_StreamStatus_ReplacesPreInitialized(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
body := buildSSEBody(5)
|
||||
c, resp, info := setupStreamTest(t, strings.NewReader(body))
|
||||
|
||||
info.StreamStatus = relaycommon.NewStreamStatus()
|
||||
info.StreamStatus.RecordError("pre-existing error")
|
||||
|
||||
StreamScannerHandler(c, resp, info, func(data string, sr *StreamResult) {})
|
||||
|
||||
assert.Equal(t, relaycommon.StreamEndReasonDone, info.StreamStatus.EndReason)
|
||||
assert.Equal(t, 0, info.StreamStatus.TotalErrorCount())
|
||||
}
|
||||
@@ -0,0 +1,406 @@
|
||||
package helper
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"math"
|
||||
"net/url"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/logger"
|
||||
relayconstant "github.com/QuantumNous/new-api/relay/constant"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
"github.com/QuantumNous/new-api/relaykit/types"
|
||||
"github.com/samber/lo"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func GetAndValidateRequest(c *gin.Context, format types.RelayFormat) (request dto.Request, err error) {
|
||||
relayMode := relayconstant.Path2RelayMode(c.Request.URL.Path)
|
||||
|
||||
switch format {
|
||||
case types.RelayFormatOpenAI:
|
||||
request, err = GetAndValidateTextRequest(c, relayMode)
|
||||
case types.RelayFormatGemini:
|
||||
if strings.Contains(c.Request.URL.Path, ":embedContent") {
|
||||
request, err = GetAndValidateGeminiEmbeddingRequest(c)
|
||||
} else if strings.Contains(c.Request.URL.Path, ":batchEmbedContents") {
|
||||
request, err = GetAndValidateGeminiBatchEmbeddingRequest(c)
|
||||
} else {
|
||||
request, err = GetAndValidateGeminiRequest(c)
|
||||
}
|
||||
case types.RelayFormatClaude:
|
||||
request, err = GetAndValidateClaudeRequest(c)
|
||||
case types.RelayFormatOpenAIResponses:
|
||||
request, err = GetAndValidateResponsesRequest(c)
|
||||
case types.RelayFormatOpenAIResponsesCompaction:
|
||||
request, err = GetAndValidateResponsesCompactionRequest(c)
|
||||
case types.RelayFormatOpenAIAlphaSearch:
|
||||
request, err = GetAndValidateAlphaSearchRequest(c)
|
||||
|
||||
case types.RelayFormatOpenAIImage:
|
||||
request, err = GetAndValidOpenAIImageRequest(c, relayMode)
|
||||
case types.RelayFormatEmbedding:
|
||||
request, err = GetAndValidateEmbeddingRequest(c, relayMode)
|
||||
case types.RelayFormatRerank:
|
||||
request, err = GetAndValidateRerankRequest(c)
|
||||
case types.RelayFormatOpenAIAudio:
|
||||
request, err = GetAndValidAudioRequest(c, relayMode)
|
||||
case types.RelayFormatOpenAIRealtime:
|
||||
request = &dto.BaseRequest{}
|
||||
default:
|
||||
return nil, fmt.Errorf("unsupported relay format: %s", format)
|
||||
}
|
||||
return request, err
|
||||
}
|
||||
|
||||
func GetAndValidAudioRequest(c *gin.Context, relayMode int) (*dto.AudioRequest, error) {
|
||||
audioRequest := &dto.AudioRequest{}
|
||||
err := common.UnmarshalBodyReusable(c, audioRequest)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
switch relayMode {
|
||||
case relayconstant.RelayModeAudioSpeech:
|
||||
if audioRequest.Model == "" {
|
||||
return nil, errors.New("model is required")
|
||||
}
|
||||
default:
|
||||
if audioRequest.Model == "" {
|
||||
return nil, errors.New("model is required")
|
||||
}
|
||||
if audioRequest.ResponseFormat == "" {
|
||||
audioRequest.ResponseFormat = "json"
|
||||
}
|
||||
}
|
||||
return audioRequest, nil
|
||||
}
|
||||
|
||||
func GetAndValidateRerankRequest(c *gin.Context) (*dto.RerankRequest, error) {
|
||||
var rerankRequest *dto.RerankRequest
|
||||
err := common.UnmarshalBodyReusable(c, &rerankRequest)
|
||||
if err != nil {
|
||||
logger.LogError(c, fmt.Sprintf("getAndValidateTextRequest failed: %s", err.Error()))
|
||||
return nil, types.NewError(err, types.ErrorCodeInvalidRequest, types.ErrOptionWithSkipRetry())
|
||||
}
|
||||
|
||||
if rerankRequest.Query == "" {
|
||||
return nil, types.NewError(fmt.Errorf("query is empty"), types.ErrorCodeInvalidRequest, types.ErrOptionWithSkipRetry())
|
||||
}
|
||||
if len(rerankRequest.Documents) == 0 {
|
||||
return nil, types.NewError(fmt.Errorf("documents is empty"), types.ErrorCodeInvalidRequest, types.ErrOptionWithSkipRetry())
|
||||
}
|
||||
return rerankRequest, nil
|
||||
}
|
||||
|
||||
func GetAndValidateEmbeddingRequest(c *gin.Context, relayMode int) (*dto.EmbeddingRequest, error) {
|
||||
var embeddingRequest *dto.EmbeddingRequest
|
||||
err := common.UnmarshalBodyReusable(c, &embeddingRequest)
|
||||
if err != nil {
|
||||
logger.LogError(c, fmt.Sprintf("getAndValidateTextRequest failed: %s", err.Error()))
|
||||
return nil, types.NewError(err, types.ErrorCodeInvalidRequest, types.ErrOptionWithSkipRetry())
|
||||
}
|
||||
|
||||
if embeddingRequest.Input == nil {
|
||||
return nil, fmt.Errorf("input is empty")
|
||||
}
|
||||
if relayMode == relayconstant.RelayModeModerations && embeddingRequest.Model == "" {
|
||||
embeddingRequest.Model = "omni-moderation-latest"
|
||||
}
|
||||
if relayMode == relayconstant.RelayModeEmbeddings && embeddingRequest.Model == "" {
|
||||
embeddingRequest.Model = c.Param("model")
|
||||
}
|
||||
return embeddingRequest, nil
|
||||
}
|
||||
|
||||
// maxTokensLimit bounds user-supplied max token fields. These values feed
|
||||
// pre-consume quota math (preConsumedTokens * ratio); an unbounded value can
|
||||
// overflow the conversion and corrupt billing.
|
||||
const maxTokensLimit = math.MaxInt32 / 2
|
||||
|
||||
func exceedsMaxTokensLimit(values ...*uint) bool {
|
||||
for _, v := range values {
|
||||
if lo.FromPtrOr(v, uint(0)) > maxTokensLimit {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func GetAndValidateResponsesRequest(c *gin.Context) (*dto.OpenAIResponsesRequest, error) {
|
||||
request := &dto.OpenAIResponsesRequest{}
|
||||
err := common.UnmarshalBodyReusable(c, request)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if request.Model == "" {
|
||||
return nil, errors.New("model is required")
|
||||
}
|
||||
if request.Input == nil {
|
||||
return nil, errors.New("input is required")
|
||||
}
|
||||
if exceedsMaxTokensLimit(request.MaxOutputTokens) {
|
||||
return nil, errors.New("max_output_tokens is invalid")
|
||||
}
|
||||
return request, nil
|
||||
}
|
||||
|
||||
func GetAndValidateAlphaSearchRequest(c *gin.Context) (*dto.AlphaSearchRequest, error) {
|
||||
request := &dto.AlphaSearchRequest{}
|
||||
if err := common.UnmarshalBodyReusable(c, request); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if request.Model == "" {
|
||||
return nil, errors.New("model is required")
|
||||
}
|
||||
storage, err := common.GetBodyStorage(c)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
rawBody, err := storage.Bytes()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
request.RawBody = rawBody
|
||||
return request, nil
|
||||
}
|
||||
|
||||
func GetAndValidateResponsesCompactionRequest(c *gin.Context) (*dto.OpenAIResponsesCompactionRequest, error) {
|
||||
request := &dto.OpenAIResponsesCompactionRequest{}
|
||||
if err := common.UnmarshalBodyReusable(c, request); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if request.Model == "" {
|
||||
return nil, errors.New("model is required")
|
||||
}
|
||||
return request, nil
|
||||
}
|
||||
|
||||
func GetAndValidOpenAIImageRequest(c *gin.Context, relayMode int) (*dto.ImageRequest, error) {
|
||||
imageRequest := &dto.ImageRequest{}
|
||||
|
||||
switch relayMode {
|
||||
case relayconstant.RelayModeImagesEdits:
|
||||
if strings.Contains(c.Request.Header.Get("Content-Type"), "multipart/form-data") {
|
||||
form, err := common.ParseMultipartFormReusable(c)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to parse image edit form request: %w", err)
|
||||
}
|
||||
formData := url.Values(form.Value)
|
||||
c.Request.MultipartForm = form
|
||||
c.Request.PostForm = formData
|
||||
imageRequest.Prompt = formData.Get("prompt")
|
||||
imageRequest.Model = formData.Get("model")
|
||||
if nValue := strings.TrimSpace(formData.Get("n")); nValue != "" {
|
||||
n, err := strconv.Atoi(nValue)
|
||||
if err != nil || n < 0 || n > dto.MaxImageN {
|
||||
return nil, fmt.Errorf("n must be an integer between 1 and %d", dto.MaxImageN)
|
||||
}
|
||||
imageRequest.N = common.GetPointer(uint(n))
|
||||
}
|
||||
imageRequest.Quality = formData.Get("quality")
|
||||
imageRequest.Size = formData.Get("size")
|
||||
if streamValue := strings.TrimSpace(formData.Get("stream")); streamValue != "" {
|
||||
stream, err := strconv.ParseBool(streamValue)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("invalid stream value: %w", err)
|
||||
}
|
||||
imageRequest.Stream = common.GetPointer(stream)
|
||||
}
|
||||
if imageValue := formData.Get("image"); imageValue != "" {
|
||||
imageRequest.Image, _ = common.Marshal(imageValue)
|
||||
}
|
||||
|
||||
if imageRequest.Model == "gpt-image-1" {
|
||||
if imageRequest.Quality == "" {
|
||||
imageRequest.Quality = "standard"
|
||||
}
|
||||
}
|
||||
if imageRequest.N == nil || *imageRequest.N == 0 {
|
||||
imageRequest.N = common.GetPointer(uint(1))
|
||||
}
|
||||
|
||||
hasWatermark := formData.Has("watermark")
|
||||
if hasWatermark {
|
||||
watermark := formData.Get("watermark") == "true"
|
||||
imageRequest.Watermark = &watermark
|
||||
}
|
||||
break
|
||||
}
|
||||
fallthrough
|
||||
default:
|
||||
err := common.UnmarshalBodyReusable(c, imageRequest)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if imageRequest.Model == "" {
|
||||
//imageRequest.Model = "dall-e-3"
|
||||
return nil, errors.New("model is required")
|
||||
}
|
||||
|
||||
if strings.Contains(imageRequest.Size, "×") {
|
||||
return nil, errors.New("size an unexpected error occurred in the parameter, please use 'x' instead of the multiplication sign '×'")
|
||||
}
|
||||
|
||||
if imageRequest.N != nil && *imageRequest.N > dto.MaxImageN {
|
||||
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 imageRequest.Quality == "" {
|
||||
imageRequest.Quality = "auto"
|
||||
}
|
||||
}
|
||||
|
||||
//if imageRequest.Prompt == "" {
|
||||
// return nil, errors.New("prompt is required")
|
||||
//}
|
||||
|
||||
if imageRequest.N == nil || *imageRequest.N == 0 {
|
||||
imageRequest.N = common.GetPointer(uint(1))
|
||||
}
|
||||
}
|
||||
|
||||
return imageRequest, nil
|
||||
}
|
||||
|
||||
func GetAndValidateClaudeRequest(c *gin.Context) (textRequest *dto.ClaudeRequest, err error) {
|
||||
textRequest = &dto.ClaudeRequest{}
|
||||
err = common.UnmarshalBodyReusable(c, textRequest)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if textRequest.Messages == nil || len(textRequest.Messages) == 0 {
|
||||
return nil, errors.New("field messages is required")
|
||||
}
|
||||
if textRequest.Model == "" {
|
||||
return nil, errors.New("field model is required")
|
||||
}
|
||||
if exceedsMaxTokensLimit(textRequest.MaxTokens, textRequest.MaxTokensToSample) {
|
||||
return nil, errors.New("max_tokens is invalid")
|
||||
}
|
||||
|
||||
//if textRequest.Stream {
|
||||
// relayInfo.IsStream = true
|
||||
//}
|
||||
|
||||
return textRequest, nil
|
||||
}
|
||||
|
||||
func GetAndValidateTextRequest(c *gin.Context, relayMode int) (*dto.GeneralOpenAIRequest, error) {
|
||||
textRequest := &dto.GeneralOpenAIRequest{}
|
||||
err := common.UnmarshalBodyReusable(c, textRequest)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if relayMode == relayconstant.RelayModeModerations && textRequest.Model == "" {
|
||||
textRequest.Model = "text-moderation-latest"
|
||||
}
|
||||
if relayMode == relayconstant.RelayModeEmbeddings && textRequest.Model == "" {
|
||||
textRequest.Model = c.Param("model")
|
||||
}
|
||||
|
||||
if exceedsMaxTokensLimit(textRequest.MaxTokens, textRequest.MaxCompletionTokens) {
|
||||
return nil, errors.New("max_tokens is invalid")
|
||||
}
|
||||
if textRequest.Model == "" {
|
||||
return nil, errors.New("model is required")
|
||||
}
|
||||
if textRequest.WebSearchOptions != nil {
|
||||
if textRequest.WebSearchOptions.SearchContextSize != "" {
|
||||
validSizes := map[string]bool{
|
||||
"high": true,
|
||||
"medium": true,
|
||||
"low": true,
|
||||
}
|
||||
if !validSizes[textRequest.WebSearchOptions.SearchContextSize] {
|
||||
return nil, errors.New("invalid search_context_size, must be one of: high, medium, low")
|
||||
}
|
||||
} else {
|
||||
textRequest.WebSearchOptions.SearchContextSize = "medium"
|
||||
}
|
||||
}
|
||||
switch relayMode {
|
||||
case relayconstant.RelayModeCompletions:
|
||||
if textRequest.Prompt == "" {
|
||||
return nil, errors.New("field prompt is required")
|
||||
}
|
||||
case relayconstant.RelayModeChatCompletions:
|
||||
// For FIM (Fill-in-the-middle) requests with prefix/suffix, messages is optional
|
||||
// It will be filled by provider-specific adaptors if needed (e.g., SiliconFlow)。Or it is allowed by model vendor(s) (e.g., DeepSeek)
|
||||
if len(textRequest.Messages) == 0 && textRequest.Prefix == nil && textRequest.Suffix == nil {
|
||||
return nil, errors.New("field messages is required")
|
||||
}
|
||||
case relayconstant.RelayModeEmbeddings:
|
||||
case relayconstant.RelayModeModerations:
|
||||
if textRequest.Input == nil || textRequest.Input == "" {
|
||||
return nil, errors.New("field input is required")
|
||||
}
|
||||
case relayconstant.RelayModeEdits:
|
||||
if textRequest.Instruction == "" {
|
||||
return nil, errors.New("field instruction is required")
|
||||
}
|
||||
}
|
||||
return textRequest, nil
|
||||
}
|
||||
|
||||
func GetAndValidateGeminiRequest(c *gin.Context) (*dto.GeminiChatRequest, error) {
|
||||
request := &dto.GeminiChatRequest{}
|
||||
err := common.UnmarshalBodyReusable(c, request)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(request.Contents) == 0 && len(request.Requests) == 0 {
|
||||
return nil, errors.New("contents is required")
|
||||
}
|
||||
if exceedsMaxTokensLimit(request.GenerationConfig.MaxOutputTokens) {
|
||||
return nil, errors.New("maxOutputTokens is invalid")
|
||||
}
|
||||
|
||||
//if c.Query("alt") == "sse" {
|
||||
// relayInfo.IsStream = true
|
||||
//}
|
||||
|
||||
return request, nil
|
||||
}
|
||||
|
||||
func GetAndValidateGeminiEmbeddingRequest(c *gin.Context) (*dto.GeminiEmbeddingRequest, error) {
|
||||
request := &dto.GeminiEmbeddingRequest{}
|
||||
err := common.UnmarshalBodyReusable(c, request)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return request, nil
|
||||
}
|
||||
|
||||
func GetAndValidateGeminiBatchEmbeddingRequest(c *gin.Context) (*dto.GeminiBatchEmbeddingRequest, error) {
|
||||
request := &dto.GeminiBatchEmbeddingRequest{}
|
||||
err := common.UnmarshalBodyReusable(c, request)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return request, nil
|
||||
}
|
||||
Reference in New Issue
Block a user