new-api/relay/channel/openai/upstream_body_test.go

package openai

import (
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/constant"
	"github.com/QuantumNous/new-api/dto"
	relaycommon "github.com/QuantumNous/new-api/relay/common"
	"github.com/QuantumNous/new-api/relay/helper"
	"github.com/QuantumNous/new-api/types"
	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestResponsesStreamEventErrorParsing(t *testing.T) {
	t.Run("standalone error event", func(t *testing.T) {
		data := `{"type":"error","error":{"type":"too_many_requests","code":"too_many_requests","message":"Too Many Requests","param":null}}`

		var streamResp dto.ResponsesStreamResponse
		err := common.UnmarshalJsonStr(data, &streamResp)
		require.NoError(t, err)
		assert.Equal(t, "error", streamResp.Type)
		assert.NotNil(t, streamResp.Error)

		oaiErr := dto.GetOpenAIError(streamResp.Error)
		require.NotNil(t, oaiErr)
		assert.Equal(t, "too_many_requests", oaiErr.Type)
		assert.Equal(t, "Too Many Requests", oaiErr.Message)
	})

	t.Run("response.failed event", func(t *testing.T) {
		data := `{"type":"response.failed","response":{"id":"test-id","status":"failed","error":{"type":"server_error","message":"Internal server error"}}}`

		var streamResp dto.ResponsesStreamResponse
		err := common.UnmarshalJsonStr(data, &streamResp)
		require.NoError(t, err)
		assert.Equal(t, "response.failed", streamResp.Type)
		require.NotNil(t, streamResp.Response)

		oaiErr := streamResp.Response.GetOpenAIError()
		require.NotNil(t, oaiErr)
		assert.Equal(t, "server_error", oaiErr.Type)
	})

	t.Run("response.error event with error in response object", func(t *testing.T) {
		data := `{"type":"response.error","response":{"id":"test-id","status":"failed","error":{"type":"invalid_request_error","message":"Invalid model"}}}`

		var streamResp dto.ResponsesStreamResponse
		err := common.UnmarshalJsonStr(data, &streamResp)
		require.NoError(t, err)
		assert.Equal(t, "response.error", streamResp.Type)

		oaiErr := streamResp.Response.GetOpenAIError()
		require.NotNil(t, oaiErr)
		assert.Equal(t, "invalid_request_error", oaiErr.Type)
	})

	t.Run("normal event has no error", func(t *testing.T) {
		data := `{"type":"response.output_text.delta","delta":"hello"}`

		var streamResp dto.ResponsesStreamResponse
		err := common.UnmarshalJsonStr(data, &streamResp)
		require.NoError(t, err)
		assert.Equal(t, "response.output_text.delta", streamResp.Type)
		assert.Nil(t, streamResp.Error)
	})
}

func TestOaiResponsesHandlerAttachesUpstreamBodyOnInvalidJSON(t *testing.T) {
	gin.SetMode(gin.TestMode)
	w := httptest.NewRecorder()
	c, _ := gin.CreateTestContext(w)

	resp := &http.Response{
		StatusCode: http.StatusOK,
		Body:       io.NopCloser(strings.NewReader(`{"incomplete":`)),
	}

	usage, err := OaiResponsesHandler(c, &relaycommon.RelayInfo{}, resp)
	require.Nil(t, usage)
	require.Error(t, err)
	assert.NotEmpty(t, err.UpstreamBody)
	assert.Contains(t, err.UpstreamBody, `{"incomplete":`)
}

func TestOaiResponsesStreamHandlerDoesNotForwardErrorEvent(t *testing.T) {
	gin.SetMode(gin.TestMode)
	oldTimeout := constant.StreamingTimeout
	constant.StreamingTimeout = 30
	t.Cleanup(func() {
		constant.StreamingTimeout = oldTimeout
	})

	w := httptest.NewRecorder()
	c, _ := gin.CreateTestContext(w)
	c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)

	errorEvent := `{"type":"error","error":{"type":"tokens","code":"rate_limit_exceeded","message":"Request too large","param":null},"sequence_number":2}`
	resp := &http.Response{
		StatusCode: http.StatusOK,
		Body:       io.NopCloser(strings.NewReader("data: " + errorEvent + "\n\n")),
	}

	usage, err := OaiResponsesStreamHandler(c, &relaycommon.RelayInfo{}, resp)
	require.Nil(t, usage)
	require.Error(t, err)
	assert.Equal(t, http.StatusInternalServerError, err.StatusCode)
	assert.Contains(t, err.UpstreamBody, "rate_limit_exceeded")
	assert.NotContains(t, w.Body.String(), "rate_limit_exceeded")
	assert.Empty(t, w.Body.String())
	assert.Empty(t, w.Header().Get("Content-Type"))
}

func TestOaiStreamHandlerDoesNotForwardInitialErrorEvent(t *testing.T) {
	gin.SetMode(gin.TestMode)
	oldTimeout := constant.StreamingTimeout
	constant.StreamingTimeout = 30
	t.Cleanup(func() {
		constant.StreamingTimeout = oldTimeout
	})

	w := httptest.NewRecorder()
	c, _ := gin.CreateTestContext(w)
	c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)

	errorEvent := `{"error":{"type":"tokens","code":"rate_limit_exceeded","message":"Request too large","param":null}}`
	resp := &http.Response{
		StatusCode: http.StatusOK,
		Body:       io.NopCloser(strings.NewReader("data: " + errorEvent + "\n\n")),
	}
	info := &relaycommon.RelayInfo{
		RelayFormat: types.RelayFormatOpenAI,
		ChannelMeta: &relaycommon.ChannelMeta{
			UpstreamModelName: "gpt-test",
		},
	}

	usage, err := OaiStreamHandler(c, info, resp)
	require.Nil(t, usage)
	require.Error(t, err)
	assert.Contains(t, err.UpstreamBody, "rate_limit_exceeded")
	assert.NotContains(t, w.Body.String(), "rate_limit_exceeded")
	assert.Empty(t, w.Body.String())
	assert.Empty(t, w.Header().Get("Content-Type"))
}

func TestOaiStreamHandlerReturnsInitialErrorAfterPing(t *testing.T) {
	gin.SetMode(gin.TestMode)
	oldTimeout := constant.StreamingTimeout
	constant.StreamingTimeout = 30
	t.Cleanup(func() {
		constant.StreamingTimeout = oldTimeout
	})

	w := httptest.NewRecorder()
	c, _ := gin.CreateTestContext(w)
	c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)

	require.NoError(t, helper.PingData(c))

	errorEvent := `{"error":{"type":"tokens","code":"rate_limit_exceeded","message":"Request too large","param":null}}`
	resp := &http.Response{
		StatusCode: http.StatusOK,
		Body:       io.NopCloser(strings.NewReader("data: " + errorEvent + "\n\n")),
	}
	info := &relaycommon.RelayInfo{
		RelayFormat: types.RelayFormatOpenAI,
		ChannelMeta: &relaycommon.ChannelMeta{
			UpstreamModelName: "gpt-test",
		},
	}

	usage, err := OaiStreamHandler(c, info, resp)
	require.Nil(t, usage)
	require.Error(t, err)
	assert.Equal(t, http.StatusInternalServerError, err.StatusCode)
	assert.Contains(t, err.UpstreamBody, "rate_limit_exceeded")
	assert.Contains(t, w.Body.String(), ": PING")
	assert.NotContains(t, w.Body.String(), "rate_limit_exceeded")
}

func TestOaiStreamHandlerDoesNotRetryAfterNormalStreamData(t *testing.T) {
	gin.SetMode(gin.TestMode)
	oldTimeout := constant.StreamingTimeout
	constant.StreamingTimeout = 30
	t.Cleanup(func() {
		constant.StreamingTimeout = oldTimeout
	})

	w := httptest.NewRecorder()
	c, _ := gin.CreateTestContext(w)
	c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)

	normalEvent := `{"id":"chatcmpl-test","object":"chat.completion.chunk","created":1,"model":"gpt-test","choices":[{"index":0,"delta":{"content":"hi"}}]}`
	errorEvent := `{"error":{"type":"tokens","code":"rate_limit_exceeded","message":"Request too large","param":null}}`
	resp := &http.Response{
		StatusCode: http.StatusOK,
		Body:       io.NopCloser(strings.NewReader("data: " + normalEvent + "\n\ndata: " + errorEvent + "\n\n")),
	}
	info := &relaycommon.RelayInfo{
		RelayFormat: types.RelayFormatOpenAI,
		ChannelMeta: &relaycommon.ChannelMeta{
			UpstreamModelName: "gpt-test",
		},
	}

	usage, err := OaiStreamHandler(c, info, resp)
	require.Nil(t, err)
	require.NotNil(t, usage)
	assert.Contains(t, w.Body.String(), "hi")
	assert.NotContains(t, w.Body.String(), "rate_limit_exceeded")
}

func TestOaiResponsesStreamHandlerDoesNotRetryAfterNormalStreamData(t *testing.T) {
	gin.SetMode(gin.TestMode)
	oldTimeout := constant.StreamingTimeout
	constant.StreamingTimeout = 30
	t.Cleanup(func() {
		constant.StreamingTimeout = oldTimeout
	})

	w := httptest.NewRecorder()
	c, _ := gin.CreateTestContext(w)
	c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)

	normalEvent := `{"type":"response.completed","response":{"id":"resp_test","usage":{"input_tokens":1,"output_tokens":1,"total_tokens":2}}}`
	errorEvent := `{"type":"error","error":{"type":"tokens","code":"rate_limit_exceeded","message":"Request too large","param":null},"sequence_number":2}`
	resp := &http.Response{
		StatusCode: http.StatusOK,
		Body:       io.NopCloser(strings.NewReader("data: " + normalEvent + "\n\ndata: " + errorEvent + "\n\n")),
	}
	info := &relaycommon.RelayInfo{
		ChannelMeta: &relaycommon.ChannelMeta{
			UpstreamModelName: "gpt-test",
		},
	}

	usage, err := OaiResponsesStreamHandler(c, info, resp)
	require.Nil(t, err)
	require.NotNil(t, usage)
	assert.Contains(t, w.Body.String(), "response.completed")
	assert.NotContains(t, w.Body.String(), "rate_limit_exceeded")
}

func TestOaiResponsesToChatHandlerAttachesUpstreamBodyOnInvalidJSON(t *testing.T) {
	gin.SetMode(gin.TestMode)
	w := httptest.NewRecorder()
	c, _ := gin.CreateTestContext(w)

	resp := &http.Response{
		StatusCode: http.StatusOK,
		Body:       io.NopCloser(strings.NewReader(`{"incomplete":`)),
	}

	usage, err := OaiResponsesToChatHandler(c, &relaycommon.RelayInfo{}, resp)
	require.Nil(t, usage)
	require.Error(t, err)
	assert.NotEmpty(t, err.UpstreamBody)
	assert.Contains(t, err.UpstreamBody, `{"incomplete":`)
}