| @@ -52,7 +52,9 @@ func OaiResponsesToChatHandler(c *gin.Context, info *relaycommon.RelayInfo, resp | |||
| } | |||
| 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 != "" { | |||
| @@ -64,7 +66,9 @@ func OaiResponsesToChatHandler(c *gin.Context, info *relaycommon.RelayInfo, resp | |||
| chatId := helper.GetResponseID(c) | |||
| chatResp, usage, err := service.ResponsesResponseToChatCompletionsResponse(&responsesResp, chatId) | |||
| 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 { | |||
| @@ -28,7 +28,9 @@ func OaiResponsesHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http | |||
| } | |||
| err = common.Unmarshal(responseBody, &responsesResponse) | |||
| 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 != "" { | |||
| apiErr := types.WithOpenAIError(*oaiError, resp.StatusCode) | |||
| @@ -1,10 +1,16 @@ | |||
| package openai | |||
| import ( | |||
| "io" | |||
| "net/http" | |||
| "net/http/httptest" | |||
| "strings" | |||
| "testing" | |||
| "github.com/QuantumNous/new-api/common" | |||
| "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/require" | |||
| ) | |||
| @@ -62,3 +68,37 @@ func TestResponsesStreamEventErrorParsing(t *testing.T) { | |||
| 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 | |||
| import ( | |||
| "io" | |||
| "net/http" | |||
| "strings" | |||
| @@ -53,14 +54,24 @@ func ExtractTopLevelChatID(c *gin.Context) (string, error) { | |||
| if c == 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) | |||
| if err != nil { | |||
| return "", err | |||
| } | |||
| bodyBytes, err := bodyStorage.Bytes() | |||
| if err != nil || len(bodyBytes) == 0 { | |||
| if _, seekErr := bodyStorage.Seek(0, io.SeekStart); seekErr == nil { | |||
| c.Request.Body = io.NopCloser(bodyStorage) | |||
| } | |||
| return "", err | |||
| } | |||
| if _, seekErr := bodyStorage.Seek(0, io.SeekStart); seekErr != nil { | |||
| return "", seekErr | |||
| } | |||
| c.Request.Body = io.NopCloser(bodyStorage) | |||
| var payload map[string]interface{} | |||
| if err := commonpkg.Unmarshal(bodyBytes, &payload); err != nil { | |||
| return "", nil | |||
| @@ -2,6 +2,7 @@ package common | |||
| import ( | |||
| "bytes" | |||
| "io" | |||
| "net/http" | |||
| "net/http/httptest" | |||
| "testing" | |||
| @@ -14,12 +15,17 @@ func TestExtractChatIDFromReusableBody(t *testing.T) { | |||
| gin.SetMode(gin.TestMode) | |||
| w := httptest.NewRecorder() | |||
| 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") | |||
| chatID, err := ExtractTopLevelChatID(c) | |||
| require.NoError(t, err) | |||
| 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) { | |||
| @@ -87,3 +93,20 @@ func TestCaptureUpstreamIDFromDialResponseHeaders(t *testing.T) { | |||
| CaptureUpstreamIDFromHTTPResponse(info, resp) | |||
| 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)) | |||
| } | |||