diff --git a/relay/channel/openai/relay-openai.go b/relay/channel/openai/relay-openai.go index 9a0619eb27f5..1e5d4db97e4f 100644 --- a/relay/channel/openai/relay-openai.go +++ b/relay/channel/openai/relay-openai.go @@ -4,6 +4,7 @@ import ( "fmt" "io" "net/http" + "sort" "strings" "github.com/QuantumNous/new-api/common" @@ -121,6 +122,7 @@ func OaiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Re var secondLastStreamData string // 存储倒数第二个stream data，用于音频模型 seenStreamToolCalls := make(map[string]struct{}) var streamFunctionCallNames []string + streamChoiceStates := make(map[int]openAIStreamChoiceState) // 检查是否为音频模型 isAudioModel := strings.Contains(strings.ToLower(model), "audio") @@ -133,6 +135,19 @@ func OaiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Re } } if len(data) > 0 { + var streamResponse dto.ChatCompletionsStreamResponse + if err := common.UnmarshalJsonStr(data, &streamResponse); err == nil { + for _, choice := range streamResponse.Choices { + state := streamChoiceStates[choice.Index] + if choice.FinishReason != nil && *choice.FinishReason != "" { + state.finished = true + } + if len(choice.Delta.ToolCalls) > 0 { + state.toolCalls = true + } + streamChoiceStates[choice.Index] = state + } + } // 对音频模型，保存倒数第二个stream data if isAudioModel && lastStreamData != "" { secondLastStreamData = lastStreamData @@ -165,17 +180,69 @@ func OaiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Re } } + // A finish_reason is an explicit terminal signal even when an upstream + // omits [DONE]. EOF/timeout without either signal is only a partial stream + // and must not be normalized into a successful downstream completion. + streamCompleted := openAIStreamCompleted(info, streamChoiceStates) + var synthesizedFinishResponse *dto.ChatCompletionsStreamResponse + if streamCompleted && info.RelayFormat == types.RelayFormatOpenAI && + info.StreamStatus != nil && info.StreamStatus.EndReason == relaycommon.StreamEndReasonDone { + missingChoiceIndexes := make([]int, 0) + for choiceIndex, state := range streamChoiceStates { + if !state.finished { + missingChoiceIndexes = append(missingChoiceIndexes, choiceIndex) + } + } + if len(missingChoiceIndexes) > 0 { + sort.Ints(missingChoiceIndexes) + synthesizedFinishResponse = &dto.ChatCompletionsStreamResponse{ + Id: responseId, + Object: "chat.completion.chunk", + Created: createAt, + Model: model, + } + synthesizedFinishResponse.SetSystemFingerprint(systemFingerprint) + for _, choiceIndex := range missingChoiceIndexes { + finishReason := types.FinishReasonStop + if streamChoiceStates[choiceIndex].toolCalls { + finishReason = types.FinishReasonToolCalls + } + synthesizedFinishResponse.Choices = append(synthesizedFinishResponse.Choices, + dto.ChatCompletionsStreamResponseChoice{ + FinishReason: common.GetPointer(finishReason), + Index: choiceIndex, + }) + } + } + } + // 处理最后的响应 shouldSendLastResp := true - if err := handleLastResponse(lastStreamData, &responseId, &createAt, &systemFingerprint, &model, &usage, - &containStreamUsage, info, &shouldSendLastResp); err != nil { - logger.LogError(c, fmt.Sprintf("error handling last response: %s, lastStreamData: [%s]", err.Error(), lastStreamData)) + if streamCompleted { + if err := handleLastResponse(lastStreamData, &responseId, &createAt, &systemFingerprint, &model, &usage, + &containStreamUsage, info, &shouldSendLastResp); err != nil { + logger.LogError(c, fmt.Sprintf("error handling last response: %s, lastStreamData: [%s]", err.Error(), lastStreamData)) + } } - if info.RelayFormat == types.RelayFormatOpenAI { + if info.RelayFormat == types.RelayFormatOpenAI && streamCompleted { + // [DONE] is an explicit upstream completion signal. Some compatible + // providers omit the terminal JSON chunk, so emit one before a trailing + // usage-only chunk for clients that require finish_reason. + lastStreamHasChoices := true + var lastStreamResponse dto.ChatCompletionsStreamResponse + if err := common.UnmarshalJsonStr(lastStreamData, &lastStreamResponse); err == nil { + lastStreamHasChoices = len(lastStreamResponse.Choices) > 0 + } + if synthesizedFinishResponse != nil && !lastStreamHasChoices { + _ = helper.ObjectData(c, synthesizedFinishResponse) + } if shouldSendLastResp { _ = sendStreamData(c, info, lastStreamData, info.ChannelSetting.ForceFormat, info.ChannelSetting.ThinkingToContent) } + if synthesizedFinishResponse != nil && lastStreamHasChoices { + _ = helper.ObjectData(c, synthesizedFinishResponse) + } } if !containStreamUsage { @@ -189,11 +256,46 @@ func OaiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Re info.CountBillableToolCall(dto.BuildInCallFunctionCall, name) } + if !streamCompleted { + if info.StreamStatus != nil { + info.StreamStatus.RecordError("openai stream ended without finish_reason or [DONE]") + logger.LogError(c, fmt.Sprintf("openai stream ended before terminal chunk: %s", info.StreamStatus.Summary())) + } + return usage, nil + } + HandleFinalResponse(c, info, lastStreamData, responseId, createAt, model, systemFingerprint, usage, containStreamUsage) return usage, nil } +type openAIStreamChoiceState struct { + finished bool + toolCalls bool +} + +func openAIStreamCompleted(info *relaycommon.RelayInfo, choiceStates map[int]openAIStreamChoiceState) bool { + if info == nil || info.StreamStatus == nil { + return true + } + switch info.StreamStatus.EndReason { + case relaycommon.StreamEndReasonDone: + return true + case relaycommon.StreamEndReasonEOF: + if len(choiceStates) == 0 { + return false + } + for _, state := range choiceStates { + if !state.finished { + return false + } + } + return true + default: + return false + } +} + func collectStreamFunctionCallNames(data string, seen map[string]struct{}, names *[]string) { var streamResponse dto.ChatCompletionsStreamResponse if err := common.UnmarshalJsonStr(data, &streamResponse); err != nil { diff --git a/relay/channel/openai/relay_openai_stream_test.go b/relay/channel/openai/relay_openai_stream_test.go new file mode 100644 index 000000000000..dd7bca433d32 --- /dev/null +++ b/relay/channel/openai/relay_openai_stream_test.go @@ -0,0 +1,295 @@ +package openai + +import ( + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/QuantumNous/new-api/constant" + relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/QuantumNous/new-api/relaykit/types" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +func newOpenAIStreamTestContext(t *testing.T, body string) (*gin.Context, *httptest.ResponseRecorder, *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{ + StatusCode: http.StatusOK, + Body: io.NopCloser(strings.NewReader(body)), + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + } + info := &relaycommon.RelayInfo{ + ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "gpt-test"}, + IsStream: true, + RelayFormat: types.RelayFormatOpenAI, + ShouldIncludeUsage: true, + DisablePing: true, + } + return c, recorder, resp, info +} + +func TestOaiStreamHandlerDoesNotFinalizeOnEOFWithoutFinishReason(t *testing.T) { + oldMode := gin.Mode() + gin.SetMode(gin.TestMode) + t.Cleanup(func() { gin.SetMode(oldMode) }) + + oldTimeout := constant.StreamingTimeout + constant.StreamingTimeout = 30 + t.Cleanup(func() { constant.StreamingTimeout = oldTimeout }) + + body := strings.Join([]string{ + `data: {"id":"chatcmpl-test","created":1710000000,"model":"gpt-test","choices":[{"index":0,"delta":{"role":"assistant","content":"partial"}}]}`, + ``, + `data: {"id":"chatcmpl-test","created":1710000000,"model":"gpt-test","choices":[],"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2}}`, + ``, + }, "\n") + c, recorder, resp, info := newOpenAIStreamTestContext(t, body) + + usage, err := OaiStreamHandler(c, info, resp) + + require.Nil(t, err) + require.NotNil(t, usage) + require.NotNil(t, info.StreamStatus) + require.Equal(t, relaycommon.StreamEndReasonEOF, info.StreamStatus.EndReason) + require.True(t, info.StreamStatus.HasErrors()) + require.Contains(t, info.StreamStatus.Errors[0].Message, "finish_reason") + + got := recorder.Body.String() + require.Contains(t, got, `"content":"partial"`) + require.NotContains(t, got, `"usage"`) + require.NotContains(t, got, "data: [DONE]") +} + +func TestOaiStreamHandlerDoesNotFinalizeAfterPartialTimeout(t *testing.T) { + oldMode := gin.Mode() + gin.SetMode(gin.TestMode) + t.Cleanup(func() { gin.SetMode(oldMode) }) + + oldTimeout := constant.StreamingTimeout + constant.StreamingTimeout = 1 + t.Cleanup(func() { constant.StreamingTimeout = oldTimeout }) + + body := strings.Join([]string{ + `data: {"id":"chatcmpl-test","created":1710000000,"model":"gpt-test","choices":[{"index":0,"delta":{"role":"assistant","content":"partial"}}]}`, + ``, + `data: {"id":"chatcmpl-test","created":1710000000,"model":"gpt-test","choices":[],"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2}}`, + ``, + }, "\n") + c, recorder, resp, info := newOpenAIStreamTestContext(t, "") + reader, writer := io.Pipe() + resp.Body = reader + writerDone := make(chan struct{}) + go func() { + defer close(writerDone) + _, _ = io.WriteString(writer, body) + }() + + usage, err := OaiStreamHandler(c, info, resp) + _ = writer.Close() + <-writerDone + + require.Nil(t, err) + require.NotNil(t, usage) + require.NotNil(t, info.StreamStatus) + require.Equal(t, relaycommon.StreamEndReasonTimeout, info.StreamStatus.EndReason) + require.True(t, info.StreamStatus.HasErrors()) + require.Contains(t, info.StreamStatus.Errors[0].Message, "finish_reason") + + got := recorder.Body.String() + require.Contains(t, got, `"content":"partial"`) + require.NotContains(t, got, `"usage"`) + require.NotContains(t, got, "data: [DONE]") +} + +func TestOaiStreamHandlerFinalizesAfterFinishReason(t *testing.T) { + oldMode := gin.Mode() + gin.SetMode(gin.TestMode) + t.Cleanup(func() { gin.SetMode(oldMode) }) + + oldTimeout := constant.StreamingTimeout + constant.StreamingTimeout = 30 + t.Cleanup(func() { constant.StreamingTimeout = oldTimeout }) + + body := strings.Join([]string{ + `data: {"id":"chatcmpl-test","created":1710000000,"model":"gpt-test","choices":[{"index":0,"delta":{"role":"assistant","content":"complete"}}]}`, + ``, + `data: {"id":"chatcmpl-test","created":1710000000,"model":"gpt-test","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}`, + ``, + `data: [DONE]`, + ``, + }, "\n") + c, recorder, resp, info := newOpenAIStreamTestContext(t, body) + + _, err := OaiStreamHandler(c, info, resp) + + require.Nil(t, err) + require.NotNil(t, info.StreamStatus) + require.Equal(t, relaycommon.StreamEndReasonDone, info.StreamStatus.EndReason) + require.False(t, info.StreamStatus.HasErrors()) + + got := recorder.Body.String() + require.Contains(t, got, `"content":"complete"`) + require.Contains(t, got, `"finish_reason":"stop"`) + require.Contains(t, got, "data: [DONE]") +} + +func TestOaiStreamHandlerSynthesizesStopAfterDoneWithoutFinishReason(t *testing.T) { + oldMode := gin.Mode() + gin.SetMode(gin.TestMode) + t.Cleanup(func() { gin.SetMode(oldMode) }) + + oldTimeout := constant.StreamingTimeout + constant.StreamingTimeout = 30 + t.Cleanup(func() { constant.StreamingTimeout = oldTimeout }) + + body := strings.Join([]string{ + `data: {"id":"chatcmpl-test","created":1710000000,"model":"gpt-test","choices":[{"index":0,"delta":{"role":"assistant","content":"complete"},"finish_reason":null}]}`, + ``, + `data: [DONE]`, + ``, + }, "\n") + c, recorder, resp, info := newOpenAIStreamTestContext(t, body) + + _, err := OaiStreamHandler(c, info, resp) + + require.Nil(t, err) + require.NotNil(t, info.StreamStatus) + require.Equal(t, relaycommon.StreamEndReasonDone, info.StreamStatus.EndReason) + require.False(t, info.StreamStatus.HasErrors()) + + got := recorder.Body.String() + require.Contains(t, got, `"content":"complete"`) + require.Contains(t, got, `"finish_reason":"stop"`) + require.Less(t, strings.Index(got, `"finish_reason":"stop"`), strings.Index(got, "data: [DONE]")) +} + +func TestOaiStreamHandlerSynthesizesFinishReasonBeforeTrailingUsage(t *testing.T) { + oldMode := gin.Mode() + gin.SetMode(gin.TestMode) + t.Cleanup(func() { gin.SetMode(oldMode) }) + + oldTimeout := constant.StreamingTimeout + constant.StreamingTimeout = 30 + t.Cleanup(func() { constant.StreamingTimeout = oldTimeout }) + + body := strings.Join([]string{ + `data: {"id":"chatcmpl-test","created":1710000000,"model":"gpt-test","choices":[{"index":0,"delta":{"role":"assistant","content":"complete"},"finish_reason":null}]}`, + ``, + `data: {"id":"chatcmpl-test","created":1710000000,"model":"gpt-test","choices":[],"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2}}`, + ``, + `data: [DONE]`, + ``, + }, "\n") + c, recorder, resp, info := newOpenAIStreamTestContext(t, body) + + _, err := OaiStreamHandler(c, info, resp) + + require.Nil(t, err) + require.NotNil(t, info.StreamStatus) + require.Equal(t, relaycommon.StreamEndReasonDone, info.StreamStatus.EndReason) + require.False(t, info.StreamStatus.HasErrors()) + + got := recorder.Body.String() + finishIndex := strings.Index(got, `"finish_reason":"stop"`) + usageIndex := strings.Index(got, `"total_tokens":2`) + doneIndex := strings.Index(got, "data: [DONE]") + require.NotEqual(t, -1, finishIndex) + require.NotEqual(t, -1, usageIndex) + require.NotEqual(t, -1, doneIndex) + require.Less(t, finishIndex, usageIndex) + require.Less(t, usageIndex, doneIndex) +} + +func TestOaiStreamHandlerSynthesizesToolCallsAfterDone(t *testing.T) { + oldMode := gin.Mode() + gin.SetMode(gin.TestMode) + t.Cleanup(func() { gin.SetMode(oldMode) }) + + oldTimeout := constant.StreamingTimeout + constant.StreamingTimeout = 30 + t.Cleanup(func() { constant.StreamingTimeout = oldTimeout }) + + body := strings.Join([]string{ + `data: {"id":"chatcmpl-test","created":1710000000,"model":"gpt-test","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"id":"call-1","type":"function","function":{"name":"lookup","arguments":"{}"}}]},"finish_reason":null}]}`, + ``, + `data: [DONE]`, + ``, + }, "\n") + c, recorder, resp, info := newOpenAIStreamTestContext(t, body) + + _, err := OaiStreamHandler(c, info, resp) + + require.Nil(t, err) + require.NotNil(t, info.StreamStatus) + require.Equal(t, relaycommon.StreamEndReasonDone, info.StreamStatus.EndReason) + require.False(t, info.StreamStatus.HasErrors()) + + got := recorder.Body.String() + require.Contains(t, got, `"finish_reason":"tool_calls"`) + require.Less(t, strings.Index(got, `"finish_reason":"tool_calls"`), strings.Index(got, "data: [DONE]")) +} + +func TestOaiStreamHandlerDoesNotFinalizeEOFWhenAnyChoiceIsIncomplete(t *testing.T) { + oldMode := gin.Mode() + gin.SetMode(gin.TestMode) + t.Cleanup(func() { gin.SetMode(oldMode) }) + + oldTimeout := constant.StreamingTimeout + constant.StreamingTimeout = 30 + t.Cleanup(func() { constant.StreamingTimeout = oldTimeout }) + + body := strings.Join([]string{ + `data: {"id":"chatcmpl-test","created":1710000000,"model":"gpt-test","choices":[{"index":0,"delta":{"content":"complete"},"finish_reason":"stop"},{"index":1,"delta":{"content":"partial"},"finish_reason":null}]}`, + ``, + }, "\n") + c, recorder, resp, info := newOpenAIStreamTestContext(t, body) + + _, err := OaiStreamHandler(c, info, resp) + + require.Nil(t, err) + require.NotNil(t, info.StreamStatus) + require.Equal(t, relaycommon.StreamEndReasonEOF, info.StreamStatus.EndReason) + require.True(t, info.StreamStatus.HasErrors()) + + got := recorder.Body.String() + require.NotContains(t, got, `"finish_reason":"stop"`) + require.NotContains(t, got, "data: [DONE]") +} + +func TestOaiStreamHandlerFinalizesOnEOFAfterFinishReason(t *testing.T) { + oldMode := gin.Mode() + gin.SetMode(gin.TestMode) + t.Cleanup(func() { gin.SetMode(oldMode) }) + + oldTimeout := constant.StreamingTimeout + constant.StreamingTimeout = 30 + t.Cleanup(func() { constant.StreamingTimeout = oldTimeout }) + + body := strings.Join([]string{ + `data: {"id":"chatcmpl-test","created":1710000000,"model":"gpt-test","choices":[{"index":0,"delta":{"role":"assistant","content":"complete"}}]}`, + ``, + `data: {"id":"chatcmpl-test","created":1710000000,"model":"gpt-test","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}`, + ``, + }, "\n") + c, recorder, resp, info := newOpenAIStreamTestContext(t, body) + + _, err := OaiStreamHandler(c, info, resp) + + require.Nil(t, err) + require.NotNil(t, info.StreamStatus) + require.Equal(t, relaycommon.StreamEndReasonEOF, info.StreamStatus.EndReason) + require.False(t, info.StreamStatus.HasErrors()) + + got := recorder.Body.String() + require.Contains(t, got, `"content":"complete"`) + require.Contains(t, got, `"finish_reason":"stop"`) + require.Contains(t, got, "data: [DONE]") +}