Skip to content

phsl / new-api

Go to a project
Commit 59a93cf5 by CaIon
Options

fix(openai): align image streaming relay governance 

Route OpenAI image streaming through shared stream handling, split image/realtime/usage helpers for maintainability, and include the related image request and rate limit updates.
parent 867d8acf
Showing with 784 additions and 163 deletions
common/init.go
...@@ -112,11 +112,11 @@ func InitEnv() { ...@@ -112,11 +112,11 @@ func InitEnv() {
// Initialize rate limit variables // Initialize rate limit variables
GlobalApiRateLimitEnable = GetEnvOrDefaultBool("GLOBAL_API_RATE_LIMIT_ENABLE", true) GlobalApiRateLimitEnable = GetEnvOrDefaultBool("GLOBAL_API_RATE_LIMIT_ENABLE", true)
GlobalApiRateLimitNum = GetEnvOrDefault("GLOBAL_API_RATE_LIMIT", 180) GlobalApiRateLimitNum = GetEnvOrDefault("GLOBAL_API_RATE_LIMIT", 360)
GlobalApiRateLimitDuration = int64(GetEnvOrDefault("GLOBAL_API_RATE_LIMIT_DURATION", 180)) GlobalApiRateLimitDuration = int64(GetEnvOrDefault("GLOBAL_API_RATE_LIMIT_DURATION", 180))
GlobalWebRateLimitEnable = GetEnvOrDefaultBool("GLOBAL_WEB_RATE_LIMIT_ENABLE", true) GlobalWebRateLimitEnable = GetEnvOrDefaultBool("GLOBAL_WEB_RATE_LIMIT_ENABLE", true)
GlobalWebRateLimitNum = GetEnvOrDefault("GLOBAL_WEB_RATE_LIMIT", 60) GlobalWebRateLimitNum = GetEnvOrDefault("GLOBAL_WEB_RATE_LIMIT", 120)
GlobalWebRateLimitDuration = int64(GetEnvOrDefault("GLOBAL_WEB_RATE_LIMIT_DURATION", 180)) GlobalWebRateLimitDuration = int64(GetEnvOrDefault("GLOBAL_WEB_RATE_LIMIT_DURATION", 180))
CriticalRateLimitEnable = GetEnvOrDefaultBool("CRITICAL_RATE_LIMIT_ENABLE", true) CriticalRateLimitEnable = GetEnvOrDefaultBool("CRITICAL_RATE_LIMIT_ENABLE", true)
......
dto/openai_image.go
...@@ -26,7 +26,7 @@ type ImageRequest struct { ...@@ -26,7 +26,7 @@ type ImageRequest struct {
OutputFormat json.RawMessage `json:"output_format,omitempty"` OutputFormat json.RawMessage `json:"output_format,omitempty"`
OutputCompression json.RawMessage `json:"output_compression,omitempty"` OutputCompression json.RawMessage `json:"output_compression,omitempty"`
PartialImages json.RawMessage `json:"partial_images,omitempty"` PartialImages json.RawMessage `json:"partial_images,omitempty"`
Stream bool `json:"stream,omitempty"` Stream *bool `json:"stream,omitempty"`
Images json.RawMessage `json:"images,omitempty"` Images json.RawMessage `json:"images,omitempty"`
Mask json.RawMessage `json:"mask,omitempty"` Mask json.RawMessage `json:"mask,omitempty"`
InputFidelity json.RawMessage `json:"input_fidelity,omitempty"` InputFidelity json.RawMessage `json:"input_fidelity,omitempty"`
...@@ -163,7 +163,7 @@ func (i *ImageRequest) GetTokenCountMeta() *types.TokenCountMeta { ...@@ -163,7 +163,7 @@ func (i *ImageRequest) GetTokenCountMeta() *types.TokenCountMeta {
} }
func (i *ImageRequest) IsStream(c *gin.Context) bool { func (i *ImageRequest) IsStream(c *gin.Context) bool {
return i.Stream return i.Stream != nil && *i.Stream
} }
func (i *ImageRequest) SetModelName(modelName string) { func (i *ImageRequest) SetModelName(modelName string) {
......
dto/openai_image_test.go deleted 100644 → 0
package dto
import (
"testing"
"github.com/stretchr/testify/require"
)
// TestImageRequestStreamJSON verifies that image requests preserve stream=true.
func TestImageRequestStreamJSON(t *testing.T) {
var req ImageRequest
require.NoError(t, req.UnmarshalJSON([]byte(`{"model":"gpt-image-1","prompt":"draw a cat","stream":true}`)))
require.True(t, req.Stream)
require.True(t, req.IsStream(nil))
}
relay/channel/openai/adaptor.go
...@@ -632,7 +632,7 @@ func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycom ...@@ -632,7 +632,7 @@ func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycom
if info.IsStream { if info.IsStream {
usage, err = OpenaiImageStreamHandler(c, info, resp) usage, err = OpenaiImageStreamHandler(c, info, resp)
} else { } else {
usage, err = OpenaiHandlerWithUsage(c, info, resp) usage, err = OpenaiImageHandler(c, info, resp)
} }
case relayconstant.RelayModeRerank: case relayconstant.RelayModeRerank:
usage, err = common_handler.RerankHandler(c, info, resp) usage, err = common_handler.RerankHandler(c, info, resp)
......
relay/channel/openai/image_edit_test.go
...@@ -16,106 +16,83 @@ import ( ...@@ -16,106 +16,83 @@ import (
"github.com/stretchr/testify/require" "github.com/stretchr/testify/require"
) )
// TestConvertImageEditRequestKeepsValidMultipartStreamFields verifies multipart replay. // TestConvertImageEditRequestMultipart verifies that ConvertImageRequest
func TestConvertImageEditRequestKeepsValidMultipartStreamFields(t *testing.T) { // re-serializes multipart image edit requests with all fields (including
// stream) and the file intact, both when the form was already parsed and when
// it must be re-parsed from the reusable body.
func TestConvertImageEditRequestMultipart(t *testing.T) {
gin.SetMode(gin.TestMode) gin.SetMode(gin.TestMode)
var body bytes.Buffer newMultipartContext := func(t *testing.T, prompt string) *gin.Context {
writer := multipart.NewWriter(&body) var body bytes.Buffer
require.NoError(t, writer.WriteField("model", "gpt-image-1")) writer := multipart.NewWriter(&body)
require.NoError(t, writer.WriteField("prompt", "edit this image")) require.NoError(t, writer.WriteField("model", "gpt-image-1"))
require.NoError(t, writer.WriteField("stream", "true")) require.NoError(t, writer.WriteField("prompt", prompt))
require.NoError(t, writer.WriteField("partial_images", "3")) require.NoError(t, writer.WriteField("stream", "true"))
part, err := writer.CreateFormFile("image", "input.png") require.NoError(t, writer.WriteField("partial_images", "3"))
require.NoError(t, err) part, err := writer.CreateFormFile("image", "input.png")
_, err = part.Write([]byte("fake image")) require.NoError(t, err)
require.NoError(t, err) _, err = part.Write([]byte("fake image"))
require.NoError(t, writer.Close()) require.NoError(t, err)
require.NoError(t, writer.Close())
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder) c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/edits", &body) c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/edits", &body)
c.Request.Header.Set("Content-Type", writer.FormDataContentType()) c.Request.Header.Set("Content-Type", writer.FormDataContentType())
require.NoError(t, c.Request.ParseMultipartForm(32<<20)) return c
info := &relaycommon.RelayInfo{
RelayMode: relayconstant.RelayModeImagesEdits,
} }
request := dto.ImageRequest{
Model: "gpt-image-1",
Prompt: "edit this image",
Stream: true,
}
converted, err := (&Adaptor{}).ConvertImageRequest(c, info, request)
require.NoError(t, err)
convertedBody, ok := converted.(*bytes.Buffer)
require.True(t, ok)
contentType := c.Request.Header.Get("Content-Type")
replayedRequest := httptest.NewRequest(http.MethodPost, "/v1/images/edits", bytes.NewReader(convertedBody.Bytes()))
replayedRequest.Header.Set("Content-Type", contentType)
require.NoError(t, replayedRequest.ParseMultipartForm(32<<20))
require.Equal(t, "gpt-image-1", replayedRequest.PostForm.Get("model"))
require.Equal(t, "edit this image", replayedRequest.PostForm.Get("prompt"))
require.Equal(t, "true", replayedRequest.PostForm.Get("stream"))
require.Equal(t, "3", replayedRequest.PostForm.Get("partial_images"))
require.Len(t, replayedRequest.MultipartForm.File["image"], 1)
file, err := replayedRequest.MultipartForm.File["image"][0].Open() convertAndReplay := func(t *testing.T, c *gin.Context, prompt string) {
require.NoError(t, err) info := &relaycommon.RelayInfo{
defer file.Close() RelayMode: relayconstant.RelayModeImagesEdits,
fileBytes, err := io.ReadAll(file) }
require.NoError(t, err) request := dto.ImageRequest{
require.Equal(t, []byte("fake image"), fileBytes) Model: "gpt-image-1",
} Prompt: prompt,
Stream: common.GetPointer(true),
// TestConvertImageEditRequestParsesReusableMultipartWhenFormIsMissing verifies fallback parsing. }
func TestConvertImageEditRequestParsesReusableMultipartWhenFormIsMissing(t *testing.T) {
gin.SetMode(gin.TestMode) converted, err := (&Adaptor{}).ConvertImageRequest(c, info, request)
require.NoError(t, err)
var body bytes.Buffer convertedBody, ok := converted.(*bytes.Buffer)
writer := multipart.NewWriter(&body) require.True(t, ok)
require.NoError(t, writer.WriteField("model", "gpt-image-1"))
require.NoError(t, writer.WriteField("prompt", "edit without pre-parsed form")) replayedRequest := httptest.NewRequest(http.MethodPost, "/v1/images/edits", bytes.NewReader(convertedBody.Bytes()))
require.NoError(t, writer.WriteField("stream", "true")) replayedRequest.Header.Set("Content-Type", c.Request.Header.Get("Content-Type"))
part, err := writer.CreateFormFile("image", "input.png") require.NoError(t, replayedRequest.ParseMultipartForm(32<<20))
require.NoError(t, err)
_, err = part.Write([]byte("fake image")) require.Equal(t, "gpt-image-1", replayedRequest.PostForm.Get("model"))
require.NoError(t, err) require.Equal(t, prompt, replayedRequest.PostForm.Get("prompt"))
require.NoError(t, writer.Close()) require.Equal(t, "true", replayedRequest.PostForm.Get("stream"))
require.Equal(t, "3", replayedRequest.PostForm.Get("partial_images"))
require.Len(t, replayedRequest.MultipartForm.File["image"], 1)
file, err := replayedRequest.MultipartForm.File["image"][0].Open()
require.NoError(t, err)
defer file.Close()
fileBytes, err := io.ReadAll(file)
require.NoError(t, err)
require.Equal(t, []byte("fake image"), fileBytes)
}
recorder := httptest.NewRecorder() t.Run("with pre-parsed form", func(t *testing.T) {
c, _ := gin.CreateTestContext(recorder) prompt := "edit this image"
c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/edits", &body) c := newMultipartContext(t, prompt)
c.Request.Header.Set("Content-Type", writer.FormDataContentType()) require.NoError(t, c.Request.ParseMultipartForm(32<<20))
storage, err := common.GetBodyStorage(c) convertAndReplay(t, c, prompt)
require.NoError(t, err) })
c.Request.Body = io.NopCloser(storage)
c.Request.MultipartForm = nil
c.Request.PostForm = nil
info := &relaycommon.RelayInfo{ t.Run("re-parses reusable body when form is missing", func(t *testing.T) {
RelayMode: relayconstant.RelayModeImagesEdits, prompt := "edit without pre-parsed form"
} c := newMultipartContext(t, prompt)
request := dto.ImageRequest{
Model: "gpt-image-1",
Prompt: "edit without pre-parsed form",
Stream: true,
}
converted, err := (&Adaptor{}).ConvertImageRequest(c, info, request) storage, err := common.GetBodyStorage(c)
require.NoError(t, err) require.NoError(t, err)
c.Request.Body = io.NopCloser(storage)
c.Request.MultipartForm = nil
c.Request.PostForm = nil
convertedBody, ok := converted.(*bytes.Buffer) convertAndReplay(t, c, prompt)
require.True(t, ok) })
replayedRequest := httptest.NewRequest(http.MethodPost, "/v1/images/edits", bytes.NewReader(convertedBody.Bytes()))
replayedRequest.Header.Set("Content-Type", c.Request.Header.Get("Content-Type"))
require.NoError(t, replayedRequest.ParseMultipartForm(32<<20))
require.Equal(t, "edit without pre-parsed form", replayedRequest.PostForm.Get("prompt"))
require.Equal(t, "true", replayedRequest.PostForm.Get("stream"))
require.Len(t, replayedRequest.MultipartForm.File["image"], 1)
} }
relay/channel/openai/relay_realtime.go 0 → 100644
package openai
import (
"fmt"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/logger"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relay/helper"
"github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/types"
"github.com/bytedance/gopkg/util/gopool"
"github.com/gin-gonic/gin"
"github.com/gorilla/websocket"
)
func OpenaiRealtimeHandler(c *gin.Context, info *relaycommon.RelayInfo) (*types.NewAPIError, *dto.RealtimeUsage) {
if info == nil || info.ClientWs == nil || info.TargetWs == nil {
return types.NewError(fmt.Errorf("invalid websocket connection"), types.ErrorCodeBadResponse), nil
}
info.IsStream = true
clientConn := info.ClientWs
targetConn := info.TargetWs
clientClosed := make(chan struct{})
targetClosed := make(chan struct{})
sendChan := make(chan []byte, 100)
receiveChan := make(chan []byte, 100)
errChan := make(chan error, 2)
usage := &dto.RealtimeUsage{}
localUsage := &dto.RealtimeUsage{}
sumUsage := &dto.RealtimeUsage{}
gopool.Go(func() {
defer func() {
if r := recover(); r != nil {
errChan <- fmt.Errorf("panic in client reader: %v", r)
}
}()
for {
select {
case <-c.Done():
return
default:
_, message, err := clientConn.ReadMessage()
if err != nil {
if !websocket.IsCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway) {
errChan <- fmt.Errorf("error reading from client: %v", err)
}
close(clientClosed)
return
}
realtimeEvent := &dto.RealtimeEvent{}
err = common.Unmarshal(message, realtimeEvent)
if err != nil {
errChan <- fmt.Errorf("error unmarshalling message: %v", err)
return
}
if realtimeEvent.Type == dto.RealtimeEventTypeSessionUpdate {
if realtimeEvent.Session != nil {
if realtimeEvent.Session.Tools != nil {
info.RealtimeTools = realtimeEvent.Session.Tools
}
}
}
textToken, audioToken, err := service.CountTokenRealtime(info, *realtimeEvent, info.UpstreamModelName)
if err != nil {
errChan <- fmt.Errorf("error counting text token: %v", err)
return
}
logger.LogInfo(c, fmt.Sprintf("type: %s, textToken: %d, audioToken: %d", realtimeEvent.Type, textToken, audioToken))
localUsage.TotalTokens += textToken + audioToken
localUsage.InputTokens += textToken + audioToken
localUsage.InputTokenDetails.TextTokens += textToken
localUsage.InputTokenDetails.AudioTokens += audioToken
err = helper.WssString(c, targetConn, string(message))
if err != nil {
errChan <- fmt.Errorf("error writing to target: %v", err)
return
}
select {
case sendChan <- message:
default:
}
}
}
})
gopool.Go(func() {
defer func() {
if r := recover(); r != nil {
errChan <- fmt.Errorf("panic in target reader: %v", r)
}
}()
for {
select {
case <-c.Done():
return
default:
_, message, err := targetConn.ReadMessage()
if err != nil {
if !websocket.IsCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway) {
errChan <- fmt.Errorf("error reading from target: %v", err)
}
close(targetClosed)
return
}
info.SetFirstResponseTime()
realtimeEvent := &dto.RealtimeEvent{}
err = common.Unmarshal(message, realtimeEvent)
if err != nil {
errChan <- fmt.Errorf("error unmarshalling message: %v", err)
return
}
if realtimeEvent.Type == dto.RealtimeEventTypeResponseDone {
realtimeUsage := realtimeEvent.Response.Usage
if realtimeUsage != nil {
usage.TotalTokens += realtimeUsage.TotalTokens
usage.InputTokens += realtimeUsage.InputTokens
usage.OutputTokens += realtimeUsage.OutputTokens
usage.InputTokenDetails.AudioTokens += realtimeUsage.InputTokenDetails.AudioTokens
usage.InputTokenDetails.CachedTokens += realtimeUsage.InputTokenDetails.CachedTokens
usage.InputTokenDetails.TextTokens += realtimeUsage.InputTokenDetails.TextTokens
usage.OutputTokenDetails.AudioTokens += realtimeUsage.OutputTokenDetails.AudioTokens
usage.OutputTokenDetails.TextTokens += realtimeUsage.OutputTokenDetails.TextTokens
err := preConsumeUsage(c, info, usage, sumUsage)
if err != nil {
errChan <- fmt.Errorf("error consume usage: %v", err)
return
}
// 本次计费完成，清除
usage = &dto.RealtimeUsage{}
localUsage = &dto.RealtimeUsage{}
} else {
textToken, audioToken, err := service.CountTokenRealtime(info, *realtimeEvent, info.UpstreamModelName)
if err != nil {
errChan <- fmt.Errorf("error counting text token: %v", err)
return
}
logger.LogInfo(c, fmt.Sprintf("type: %s, textToken: %d, audioToken: %d", realtimeEvent.Type, textToken, audioToken))
localUsage.TotalTokens += textToken + audioToken
info.IsFirstRequest = false
localUsage.InputTokens += textToken + audioToken
localUsage.InputTokenDetails.TextTokens += textToken
localUsage.InputTokenDetails.AudioTokens += audioToken
err = preConsumeUsage(c, info, localUsage, sumUsage)
if err != nil {
errChan <- fmt.Errorf("error consume usage: %v", err)
return
}
// 本次计费完成，清除
localUsage = &dto.RealtimeUsage{}
// print now usage
}
logger.LogInfo(c, fmt.Sprintf("realtime streaming sumUsage: %v", sumUsage))
logger.LogInfo(c, fmt.Sprintf("realtime streaming localUsage: %v", localUsage))
logger.LogInfo(c, fmt.Sprintf("realtime streaming localUsage: %v", localUsage))
} else if realtimeEvent.Type == dto.RealtimeEventTypeSessionUpdated || realtimeEvent.Type == dto.RealtimeEventTypeSessionCreated {
realtimeSession := realtimeEvent.Session
if realtimeSession != nil {
// update audio format
info.InputAudioFormat = common.GetStringIfEmpty(realtimeSession.InputAudioFormat, info.InputAudioFormat)
info.OutputAudioFormat = common.GetStringIfEmpty(realtimeSession.OutputAudioFormat, info.OutputAudioFormat)
}
} else {
textToken, audioToken, err := service.CountTokenRealtime(info, *realtimeEvent, info.UpstreamModelName)
if err != nil {
errChan <- fmt.Errorf("error counting text token: %v", err)
return
}
logger.LogInfo(c, fmt.Sprintf("type: %s, textToken: %d, audioToken: %d", realtimeEvent.Type, textToken, audioToken))
localUsage.TotalTokens += textToken + audioToken
localUsage.OutputTokens += textToken + audioToken
localUsage.OutputTokenDetails.TextTokens += textToken
localUsage.OutputTokenDetails.AudioTokens += audioToken
}
err = helper.WssString(c, clientConn, string(message))
if err != nil {
errChan <- fmt.Errorf("error writing to client: %v", err)
return
}
select {
case receiveChan <- message:
default:
}
}
}
})
select {
case <-clientClosed:
case <-targetClosed:
case err := <-errChan:
//return service.OpenAIErrorWrapper(err, "realtime_error", http.StatusInternalServerError), nil
logger.LogError(c, "realtime error: "+err.Error())
case <-c.Done():
}
if usage.TotalTokens != 0 {
_ = preConsumeUsage(c, info, usage, sumUsage)
}
if localUsage.TotalTokens != 0 {
_ = preConsumeUsage(c, info, localUsage, sumUsage)
}
// check usage total tokens, if 0, use local usage
return nil, sumUsage
}
func preConsumeUsage(ctx *gin.Context, info *relaycommon.RelayInfo, usage *dto.RealtimeUsage, totalUsage *dto.RealtimeUsage) error {
if usage == nil || totalUsage == nil {
return fmt.Errorf("invalid usage pointer")
}
totalUsage.TotalTokens += usage.TotalTokens
totalUsage.InputTokens += usage.InputTokens
totalUsage.OutputTokens += usage.OutputTokens
totalUsage.InputTokenDetails.CachedTokens += usage.InputTokenDetails.CachedTokens
totalUsage.InputTokenDetails.TextTokens += usage.InputTokenDetails.TextTokens
totalUsage.InputTokenDetails.AudioTokens += usage.InputTokenDetails.AudioTokens
totalUsage.OutputTokenDetails.TextTokens += usage.OutputTokenDetails.TextTokens
totalUsage.OutputTokenDetails.AudioTokens += usage.OutputTokenDetails.AudioTokens
// clear usage
err := service.PreWssConsumeQuota(ctx, info, usage)
return err
}
relay/channel/openai/usage.go 0 → 100644
package openai
import (
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/dto"
relaycommon "github.com/QuantumNous/new-api/relay/common"
)
func applyUsagePostProcessing(info *relaycommon.RelayInfo, usage *dto.Usage, responseBody []byte) {
if info == nil || usage == nil {
return
}
switch info.ChannelType {
case constant.ChannelTypeDeepSeek:
if usage.PromptTokensDetails.CachedTokens == 0 && usage.PromptCacheHitTokens != 0 {
usage.PromptTokensDetails.CachedTokens = usage.PromptCacheHitTokens
}
case constant.ChannelTypeZhipu_v4:
// 智普的cached_tokens在标准位置: usage.prompt_tokens_details.cached_tokens
if usage.PromptTokensDetails.CachedTokens == 0 {
if usage.InputTokensDetails != nil && usage.InputTokensDetails.CachedTokens > 0 {
usage.PromptTokensDetails.CachedTokens = usage.InputTokensDetails.CachedTokens
} else if cachedTokens, ok := extractCachedTokensFromBody(responseBody); ok {
usage.PromptTokensDetails.CachedTokens = cachedTokens
} else if usage.PromptCacheHitTokens > 0 {
usage.PromptTokensDetails.CachedTokens = usage.PromptCacheHitTokens
}
}
case constant.ChannelTypeMoonshot:
// Moonshot的cached_tokens在非标准位置: choices[].usage.cached_tokens
if usage.PromptTokensDetails.CachedTokens == 0 {
if usage.InputTokensDetails != nil && usage.InputTokensDetails.CachedTokens > 0 {
usage.PromptTokensDetails.CachedTokens = usage.InputTokensDetails.CachedTokens
} else if cachedTokens, ok := extractMoonshotCachedTokensFromBody(responseBody); ok {
usage.PromptTokensDetails.CachedTokens = cachedTokens
} else if cachedTokens, ok := extractCachedTokensFromBody(responseBody); ok {
usage.PromptTokensDetails.CachedTokens = cachedTokens
} else if usage.PromptCacheHitTokens > 0 {
usage.PromptTokensDetails.CachedTokens = usage.PromptCacheHitTokens
}
}
case constant.ChannelTypeOpenAI:
if usage.PromptTokensDetails.CachedTokens == 0 {
if cachedTokens, ok := extractLlamaCachedTokensFromBody(responseBody); ok {
usage.PromptTokensDetails.CachedTokens = cachedTokens
}
}
}
}
func extractCachedTokensFromBody(body []byte) (int, bool) {
if len(body) == 0 {
return 0, false
}
var payload struct {
Usage struct {
PromptTokensDetails struct {
CachedTokens *int `json:"cached_tokens"`
} `json:"prompt_tokens_details"`
CachedTokens *int `json:"cached_tokens"`
PromptCacheHitTokens *int `json:"prompt_cache_hit_tokens"`
} `json:"usage"`
}
if err := common.Unmarshal(body, &payload); err != nil {
return 0, false
}
if payload.Usage.PromptTokensDetails.CachedTokens != nil {
return *payload.Usage.PromptTokensDetails.CachedTokens, true
}
if payload.Usage.CachedTokens != nil {
return *payload.Usage.CachedTokens, true
}
if payload.Usage.PromptCacheHitTokens != nil {
return *payload.Usage.PromptCacheHitTokens, true
}
return 0, false
}
// extractMoonshotCachedTokensFromBody 从Moonshot的非标准位置提取cached_tokens
// Moonshot的流式响应格式: {"choices":[{"usage":{"cached_tokens":111}}]}
func extractMoonshotCachedTokensFromBody(body []byte) (int, bool) {
if len(body) == 0 {
return 0, false
}
var payload struct {
Choices []struct {
Usage struct {
CachedTokens *int `json:"cached_tokens"`
} `json:"usage"`
} `json:"choices"`
}
if err := common.Unmarshal(body, &payload); err != nil {
return 0, false
}
// 遍历choices查找cached_tokens
for _, choice := range payload.Choices {
if choice.Usage.CachedTokens != nil && *choice.Usage.CachedTokens > 0 {
return *choice.Usage.CachedTokens, true
}
}
return 0, false
}
// extractLlamaCachedTokensFromBody 从llama.cpp的非标准位置提取cache_n
func extractLlamaCachedTokensFromBody(body []byte) (int, bool) {
if len(body) == 0 {
return 0, false
}
var payload struct {
Timings struct {
CachedTokens *int `json:"cache_n"`
} `json:"timings"`
}
if err := common.Unmarshal(body, &payload); err != nil {
return 0, false
}
if payload.Timings.CachedTokens == nil {
return 0, false
}
return *payload.Timings.CachedTokens, true
}
relay/channel/xai/adaptor.go
...@@ -114,7 +114,7 @@ func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, request ...@@ -114,7 +114,7 @@ func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, request
func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *types.NewAPIError) { func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *types.NewAPIError) {
switch info.RelayMode { switch info.RelayMode {
case constant.RelayModeImagesGenerations, constant.RelayModeImagesEdits: case constant.RelayModeImagesGenerations, constant.RelayModeImagesEdits:
usage, err = openai.OpenaiHandlerWithUsage(c, info, resp) usage, err = openai.OpenaiImageHandler(c, info, resp)
case constant.RelayModeResponses: case constant.RelayModeResponses:
if info.IsStream { if info.IsStream {
usage, err = openai.OaiResponsesStreamHandler(c, info, resp) usage, err = openai.OaiResponsesStreamHandler(c, info, resp)
......
relay/helper/openai_image_request_test.go
...@@ -15,59 +15,57 @@ import ( ...@@ -15,59 +15,57 @@ import (
"github.com/stretchr/testify/require" "github.com/stretchr/testify/require"
) )
// TestGetAndValidOpenAIImageRequestMultipartStream verifies reusable image edit parsing. // 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) { func TestGetAndValidOpenAIImageRequestMultipartStream(t *testing.T) {
gin.SetMode(gin.TestMode) gin.SetMode(gin.TestMode)
var body bytes.Buffer newContext := func(t *testing.T, streamValue string, withImage bool) (*gin.Context, string) {
writer := multipart.NewWriter(&body) var body bytes.Buffer
require.NoError(t, writer.WriteField("model", "gpt-image-1")) writer := multipart.NewWriter(&body)
require.NoError(t, writer.WriteField("prompt", "edit this image")) require.NoError(t, writer.WriteField("model", "gpt-image-1"))
require.NoError(t, writer.WriteField("stream", "true")) require.NoError(t, writer.WriteField("prompt", "edit this image"))
require.NoError(t, writer.WriteField("n", "1")) require.NoError(t, writer.WriteField("stream", streamValue))
part, err := writer.CreateFormFile("image", "input.png") if withImage {
require.NoError(t, err) part, err := writer.CreateFormFile("image", "input.png")
_, err = part.Write([]byte("fake image")) require.NoError(t, err)
require.NoError(t, err) _, err = part.Write([]byte("fake image"))
require.NoError(t, writer.Close()) require.NoError(t, err)
originalBody := body.String() }
require.NoError(t, writer.Close())
originalBody := body.String()
recorder := httptest.NewRecorder() c, _ := gin.CreateTestContext(httptest.NewRecorder())
c, _ := gin.CreateTestContext(recorder) c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/edits", &body)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/edits", &body) c.Request.Header.Set("Content-Type", writer.FormDataContentType())
c.Request.Header.Set("Content-Type", writer.FormDataContentType()) return c, originalBody
}
req, err := GetAndValidOpenAIImageRequest(c, relayconstant.RelayModeImagesEdits) t.Run("valid stream value keeps body replayable", func(t *testing.T) {
require.NoError(t, err) c, originalBody := newContext(t, "true", true)
require.True(t, req.Stream)
require.True(t, req.IsStream(c))
bodyAfterValidation, err := io.ReadAll(c.Request.Body) req, err := GetAndValidOpenAIImageRequest(c, relayconstant.RelayModeImagesEdits)
require.NoError(t, err) require.NoError(t, err)
require.Equal(t, originalBody, string(bodyAfterValidation)) require.NotNil(t, req.Stream)
require.True(t, *req.Stream)
require.True(t, req.IsStream(c))
form, err := common.ParseMultipartFormReusable(c) bodyAfterValidation, err := io.ReadAll(c.Request.Body)
require.NoError(t, err) require.NoError(t, err)
require.Equal(t, "true", url.Values(form.Value).Get("stream")) require.Equal(t, originalBody, string(bodyAfterValidation))
require.Len(t, form.File["image"], 1)
}
// TestGetAndValidOpenAIImageRequestMultipartStreamInvalidValue verifies stream validation.
func TestGetAndValidOpenAIImageRequestMultipartStreamInvalidValue(t *testing.T) {
gin.SetMode(gin.TestMode)
var body bytes.Buffer form, err := common.ParseMultipartFormReusable(c)
writer := multipart.NewWriter(&body) require.NoError(t, err)
require.NoError(t, writer.WriteField("model", "gpt-image-1")) require.Equal(t, "true", url.Values(form.Value).Get("stream"))
require.NoError(t, writer.WriteField("stream", "notabool")) require.Len(t, form.File["image"], 1)
require.NoError(t, writer.Close()) })
recorder := httptest.NewRecorder() t.Run("invalid stream value is rejected", func(t *testing.T) {
c, _ := gin.CreateTestContext(recorder) c, _ := newContext(t, "notabool", false)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/edits", &body)
c.Request.Header.Set("Content-Type", writer.FormDataContentType())
_, err := GetAndValidOpenAIImageRequest(c, relayconstant.RelayModeImagesEdits) _, err := GetAndValidOpenAIImageRequest(c, relayconstant.RelayModeImagesEdits)
require.Error(t, err) require.Error(t, err)
require.Contains(t, err.Error(), "invalid stream value") require.Contains(t, err.Error(), "invalid stream value")
})
} }
relay/helper/stream_scanner.go
...@@ -22,8 +22,8 @@ import ( ...@@ -22,8 +22,8 @@ import (
) )
const ( const (
InitialScannerBufferSize = 64 << 10 // 64KB (64*1024) InitialScannerBufferSize = 64 << 10 // 64KB (64*1024)
DefaultMaxScannerBufferSize = 64 << 20 // 64MB (64*1024*1024) default SSE buffer size DefaultMaxScannerBufferSize = 128 << 20 // 64MB (64*1024*1024) default SSE buffer size
DefaultPingInterval = 10 * time.Second DefaultPingInterval = 10 * time.Second
) )
......
relay/helper/valid_request.go
...@@ -163,7 +163,7 @@ func GetAndValidOpenAIImageRequest(c *gin.Context, relayMode int) (*dto.ImageReq ...@@ -163,7 +163,7 @@ func GetAndValidOpenAIImageRequest(c *gin.Context, relayMode int) (*dto.ImageReq
if err != nil { if err != nil {
return nil, fmt.Errorf("invalid stream value: %w", err) return nil, fmt.Errorf("invalid stream value: %w", err)
} }
imageRequest.Stream = stream imageRequest.Stream = common.GetPointer(stream)
} }
if imageValue := formData.Get("image"); imageValue != "" { if imageValue := formData.Get("image"); imageValue != "" {
imageRequest.Image, _ = common.Marshal(imageValue) imageRequest.Image, _ = common.Marshal(imageValue)
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or sign in to comment