| @@ -52,7 +52,9 @@ func OaiResponsesToChatHandler(c *gin.Context, info *relaycommon.RelayInfo, resp | |||||
| } | } | ||||
| if err := common.Unmarshal(body, &responsesResp); err != nil { | if err := common.Unmarshal(body, &responsesResp); err != nil { | ||||
| return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError) | |||||
| apiErr := types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError) | |||||
| apiErr.UpstreamBody = service.TruncateBody(string(body)) | |||||
| return nil, apiErr | |||||
| } | } | ||||
| if oaiError := responsesResp.GetOpenAIError(); oaiError != nil && oaiError.Type != "" { | if oaiError := responsesResp.GetOpenAIError(); oaiError != nil && oaiError.Type != "" { | ||||
| @@ -64,7 +66,9 @@ func OaiResponsesToChatHandler(c *gin.Context, info *relaycommon.RelayInfo, resp | |||||
| chatId := helper.GetResponseID(c) | chatId := helper.GetResponseID(c) | ||||
| chatResp, usage, err := service.ResponsesResponseToChatCompletionsResponse(&responsesResp, chatId) | chatResp, usage, err := service.ResponsesResponseToChatCompletionsResponse(&responsesResp, chatId) | ||||
| if err != nil { | if err != nil { | ||||
| return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError) | |||||
| apiErr := types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError) | |||||
| apiErr.UpstreamBody = service.TruncateBody(string(body)) | |||||
| return nil, apiErr | |||||
| } | } | ||||
| if usage == nil || usage.TotalTokens == 0 { | if usage == nil || usage.TotalTokens == 0 { | ||||
| @@ -28,7 +28,9 @@ func OaiResponsesHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http | |||||
| } | } | ||||
| err = common.Unmarshal(responseBody, &responsesResponse) | err = common.Unmarshal(responseBody, &responsesResponse) | ||||
| if err != nil { | if err != nil { | ||||
| return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError) | |||||
| apiErr := types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError) | |||||
| apiErr.UpstreamBody = service.TruncateBody(string(responseBody)) | |||||
| return nil, apiErr | |||||
| } | } | ||||
| if oaiError := responsesResponse.GetOpenAIError(); oaiError != nil && oaiError.Type != "" { | if oaiError := responsesResponse.GetOpenAIError(); oaiError != nil && oaiError.Type != "" { | ||||
| apiErr := types.WithOpenAIError(*oaiError, resp.StatusCode) | apiErr := types.WithOpenAIError(*oaiError, resp.StatusCode) | ||||
| @@ -1,10 +1,16 @@ | |||||
| package openai | package openai | ||||
| import ( | import ( | ||||
| "io" | |||||
| "net/http" | |||||
| "net/http/httptest" | |||||
| "strings" | |||||
| "testing" | "testing" | ||||
| "github.com/QuantumNous/new-api/common" | "github.com/QuantumNous/new-api/common" | ||||
| "github.com/QuantumNous/new-api/dto" | "github.com/QuantumNous/new-api/dto" | ||||
| relaycommon "github.com/QuantumNous/new-api/relay/common" | |||||
| "github.com/gin-gonic/gin" | |||||
| "github.com/stretchr/testify/assert" | "github.com/stretchr/testify/assert" | ||||
| "github.com/stretchr/testify/require" | "github.com/stretchr/testify/require" | ||||
| ) | ) | ||||
| @@ -62,3 +68,37 @@ func TestResponsesStreamEventErrorParsing(t *testing.T) { | |||||
| assert.Nil(t, streamResp.Error) | 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 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":`) | |||||
| } | |||||
| @@ -1,6 +1,7 @@ | |||||
| package common | package common | ||||
| import ( | import ( | ||||
| "io" | |||||
| "net/http" | "net/http" | ||||
| "strings" | "strings" | ||||
| @@ -53,14 +54,24 @@ func ExtractTopLevelChatID(c *gin.Context) (string, error) { | |||||
| if c == nil { | if c == nil { | ||||
| return "", nil | return "", nil | ||||
| } | } | ||||
| if !strings.HasPrefix(strings.ToLower(strings.TrimSpace(c.Request.Header.Get("Content-Type"))), "application/json") { | |||||
| return "", nil | |||||
| } | |||||
| bodyStorage, err := commonpkg.GetBodyStorage(c) | bodyStorage, err := commonpkg.GetBodyStorage(c) | ||||
| if err != nil { | if err != nil { | ||||
| return "", err | return "", err | ||||
| } | } | ||||
| bodyBytes, err := bodyStorage.Bytes() | bodyBytes, err := bodyStorage.Bytes() | ||||
| if err != nil || len(bodyBytes) == 0 { | if err != nil || len(bodyBytes) == 0 { | ||||
| if _, seekErr := bodyStorage.Seek(0, io.SeekStart); seekErr == nil { | |||||
| c.Request.Body = io.NopCloser(bodyStorage) | |||||
| } | |||||
| return "", err | return "", err | ||||
| } | } | ||||
| if _, seekErr := bodyStorage.Seek(0, io.SeekStart); seekErr != nil { | |||||
| return "", seekErr | |||||
| } | |||||
| c.Request.Body = io.NopCloser(bodyStorage) | |||||
| var payload map[string]interface{} | var payload map[string]interface{} | ||||
| if err := commonpkg.Unmarshal(bodyBytes, &payload); err != nil { | if err := commonpkg.Unmarshal(bodyBytes, &payload); err != nil { | ||||
| return "", nil | return "", nil | ||||
| @@ -2,6 +2,7 @@ package common | |||||
| import ( | import ( | ||||
| "bytes" | "bytes" | ||||
| "io" | |||||
| "net/http" | "net/http" | ||||
| "net/http/httptest" | "net/http/httptest" | ||||
| "testing" | "testing" | ||||
| @@ -14,12 +15,17 @@ func TestExtractChatIDFromReusableBody(t *testing.T) { | |||||
| gin.SetMode(gin.TestMode) | gin.SetMode(gin.TestMode) | ||||
| w := httptest.NewRecorder() | w := httptest.NewRecorder() | ||||
| c, _ := gin.CreateTestContext(w) | c, _ := gin.CreateTestContext(w) | ||||
| c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewBufferString(`{"model":"gpt-4o-mini","chat_id":"chat_abc","messages":[{"role":"user","content":"hi"}]}`)) | |||||
| body := `{"model":"gpt-4o-mini","chat_id":"chat_abc","messages":[{"role":"user","content":"hi"}]}` | |||||
| c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewBufferString(body)) | |||||
| c.Request.Header.Set("Content-Type", "application/json") | c.Request.Header.Set("Content-Type", "application/json") | ||||
| chatID, err := ExtractTopLevelChatID(c) | chatID, err := ExtractTopLevelChatID(c) | ||||
| require.NoError(t, err) | require.NoError(t, err) | ||||
| require.Equal(t, "chat_abc", chatID) | require.Equal(t, "chat_abc", chatID) | ||||
| bodyAfter, err := io.ReadAll(c.Request.Body) | |||||
| require.NoError(t, err) | |||||
| require.Equal(t, body, string(bodyAfter)) | |||||
| } | } | ||||
| func TestExtractChatIDReturnsEmptyWhenMissing(t *testing.T) { | func TestExtractChatIDReturnsEmptyWhenMissing(t *testing.T) { | ||||
| @@ -87,3 +93,20 @@ func TestCaptureUpstreamIDFromDialResponseHeaders(t *testing.T) { | |||||
| CaptureUpstreamIDFromHTTPResponse(info, resp) | CaptureUpstreamIDFromHTTPResponse(info, resp) | ||||
| require.Equal(t, "ws_upstream_1", info.UpstreamID) | require.Equal(t, "ws_upstream_1", info.UpstreamID) | ||||
| } | } | ||||
| func TestExtractTopLevelChatIDSkipsNonJSONBody(t *testing.T) { | |||||
| gin.SetMode(gin.TestMode) | |||||
| w := httptest.NewRecorder() | |||||
| c, _ := gin.CreateTestContext(w) | |||||
| body := "--boundary\r\nContent-Disposition: form-data; name=\"chat_id\"\r\n\r\nchat_abc\r\n--boundary--\r\n" | |||||
| c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/edits", bytes.NewBufferString(body)) | |||||
| c.Request.Header.Set("Content-Type", "multipart/form-data; boundary=boundary") | |||||
| chatID, err := ExtractTopLevelChatID(c) | |||||
| require.NoError(t, err) | |||||
| require.Empty(t, chatID) | |||||
| bodyAfter, err := io.ReadAll(c.Request.Body) | |||||
| require.NoError(t, err) | |||||
| require.Equal(t, body, string(bodyAfter)) | |||||
| } | |||||