Change error response to json and improve e2e stability (#669)

* Change error response to json and improve e2e stability

1. For all failure request, change the Content-Type to application/json for readbility and integration
2. Use Fatal to replace Error and return unexpected error immediately. In this case, we won’t see invalid memory access issues
3. Use status code for comparision instead of text which is not reliable.

Signed-off-by: Jiaxin Shan <seedjeffwan@gmail.com>

* Fix the lint issue

* Force the openai client not to retry for errors

Signed-off-by: Jiaxin Shan <seedjeffwan@gmail.com>

---------

Signed-off-by: Jiaxin Shan <seedjeffwan@gmail.com>
Signed-off-by: Varun Gupta <varungup90@gmail.com>

pkg/plugins/gateway/gateway.go

@@ -215,7 +215,7 @@ func (s *Server) Process(srv extProcPb.ExternalProcessor_ProcessServer) error {
 }
 
 func (s *Server) HandleRequestHeaders(ctx context.Context, requestID string, req *extProcPb.ProcessingRequest) (*extProcPb.ProcessingResponse, utils.User, int64, string) {
-	klog.Info("\n\n")
+	klog.Info("\n")
 	klog.InfoS("-- In RequestHeaders processing ...", "requestID", requestID)
 	var username string
 	var user utils.User
@@ -482,15 +482,19 @@ func (s *Server) HandleResponseBody(ctx context.Context, requestID string, req *
 				}}},
 				err.Error()), complete
 		} else if len(res.Model) == 0 {
-			err = ErrorUnknownResponse
-			klog.ErrorS(err, "unexpected response", "requestID", requestID, "responseBody", string(b.ResponseBody.GetBody()))
+			msg := ErrorUnknownResponse.Error()
+			responseBodyContent := string(b.ResponseBody.GetBody())
+			if len(responseBodyContent) != 0 {
+				msg = responseBodyContent
+			}
+			klog.ErrorS(err, "unexpected response", "requestID", requestID, "responseBody", responseBodyContent)
 			complete = true
 			return generateErrorResponse(
 				envoyTypePb.StatusCode_InternalServerError,
 				[]*configPb.HeaderValueOption{{Header: &configPb.HeaderValue{
 					Key: HeaderErrorResponseUnknown, RawValue: []byte("true"),
 				}}},
-				err.Error()), complete
+				msg), complete
 		}
 		// Do not overwrite model, res can be empty.
 		usage = res.Usage
@@ -663,6 +667,14 @@ func validateRoutingStrategy(routingStrategy string) bool {
 }
 
 func generateErrorResponse(statusCode envoyTypePb.StatusCode, headers []*configPb.HeaderValueOption, body string) *extProcPb.ProcessingResponse {
+	// Set the Content-Type header to application/json
+	headers = append(headers, &configPb.HeaderValueOption{
+		Header: &configPb.HeaderValue{
+			Key:   "Content-Type",
+			Value: "application/json",
+		},
+	})
+
 	return &extProcPb.ProcessingResponse{
 		Response: &extProcPb.ProcessingResponse_ImmediateResponse{
 			ImmediateResponse: &extProcPb.ImmediateResponse{
@@ -672,7 +684,7 @@ func generateErrorResponse(statusCode envoyTypePb.StatusCode, headers []*configP
 				Headers: &extProcPb.HeaderMutation{
 					SetHeaders: headers,
 				},
-				Body: body,
+				Body: generateErrorMessage(body, int(statusCode)),
 			},
 		},
 	}
@@ -719,3 +731,15 @@ func GetRoutingStrategy(headers []*configPb.HeaderValue) (string, bool) {
 
 	return routingStrategy, routingStrategyEnabled
 }
+
+// generateErrorMessage constructs a JSON error message using fmt.Sprintf
+func generateErrorMessage(message string, code int) string {
+	errorStruct := map[string]interface{}{
+		"error": map[string]interface{}{
+			"message": message,
+			"code":    code,
+		},
+	}
+	jsonData, _ := json.Marshal(errorStruct)
+	return string(jsonData)
+}

test/e2e/e2e_test.go

@@ -18,6 +18,8 @@ package e2e
 
 import (
 	"context"
+	"errors"
+	"net/http"
 	"testing"
 
 	"github.com/openai/openai-go"
@@ -43,52 +45,70 @@ func TestBaseModelInference(t *testing.T) {
 		Model: openai.F(modelName),
 	})
 	if err != nil {
-		t.Error("chat completions failed", err)
+		t.Fatalf("chat completions failed: %v", err)
 	}
+
 	assert.Equal(t, modelName, chatCompletion.Model)
+	assert.NotEmpty(t, chatCompletion.Choices, "chat completion has no choices returned")
+	assert.NotNil(t, chatCompletion.Choices[0].Message.Content, "chat completion has no message returned")
 }
 
 func TestBaseModelInferenceFailures(t *testing.T) {
-	// error on invalid api key
-	client := createOpenAIClient(baseURL, "fake-api-key")
-	_, err := client.Chat.Completions.New(context.TODO(), openai.ChatCompletionNewParams{
-		Messages: openai.F([]openai.ChatCompletionMessageParamUnion{
-			openai.UserMessage("Say this is a test"),
-		}),
-		Model: openai.F(modelName),
-	})
-	if err == nil {
-		t.Error("500 Internal Server Error expected for invalid api-key")
+	testCases := []struct {
+		name            string
+		apiKey          string
+		modelName       string
+		routingStrategy string
+		expectErrCode   int
+	}{
+		{
+			name:      "Invalid API Key",
+			apiKey:    "fake-api-key",
+			modelName: modelName,
+			// TODO: it is supposed to be 401. Let's handle such case and fix this.
+			expectErrCode: 500,
+		},
+		{
+			name:          "Invalid Model Name",
+			apiKey:        apiKey,
+			modelName:     "fake-model-name",
+			expectErrCode: 400,
+		},
+		{
+			name:            "Invalid Routing Strategy",
+			apiKey:          apiKey,
+			modelName:       modelName,
+			routingStrategy: "invalid-routing-strategy",
+			expectErrCode:   400,
+		},
 	}
 
-	// error on invalid model name
-	client = createOpenAIClient(baseURL, apiKey)
-	_, err = client.Chat.Completions.New(context.TODO(), openai.ChatCompletionNewParams{
-		Messages: openai.F([]openai.ChatCompletionMessageParamUnion{
-			openai.UserMessage("Say this is a test"),
-		}),
-		Model: openai.F("fake-model-name"),
-	})
-	assert.Contains(t, err.Error(), "400 Bad Request")
-	if err == nil {
-		t.Error("400 Bad Request expected for invalid api-key")
-	}
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			var client *openai.Client
+			if tc.routingStrategy != "" {
+				var dst *http.Response
+				client = createOpenAIClientWithRoutingStrategy(baseURL, tc.apiKey,
+					tc.routingStrategy, option.WithResponseInto(&dst))
+			} else {
+				client = createOpenAIClient(baseURL, tc.apiKey)
+			}
 
-	// invalid routing strategy
-	client = openai.NewClient(
-		option.WithBaseURL(baseURL),
-		option.WithAPIKey(apiKey),
-		option.WithHeader("routing-strategy", "invalid-routing-strategy"),
-	)
-	client.Options = append(client.Options, option.WithHeader("routing-strategy", "invalid-routing-strategy"))
-	_, err = client.Chat.Completions.New(context.TODO(), openai.ChatCompletionNewParams{
-		Messages: openai.F([]openai.ChatCompletionMessageParamUnion{
-			openai.UserMessage("Say this is a test"),
-		}),
-		Model: openai.F(modelName),
-	})
-	if err == nil {
-		t.Error("400 Bad Request expected for invalid routing-strategy")
+			_, err := client.Chat.Completions.New(context.TODO(), openai.ChatCompletionNewParams{
+				Messages: openai.F([]openai.ChatCompletionMessageParamUnion{
+					openai.UserMessage("Say this is a test"),
+				}),
+				Model: openai.F(tc.modelName),
+			})
+
+			assert.Error(t, err)
+			var apiErr *openai.Error
+			if !errors.As(err, &apiErr) {
+				t.Fatalf("Error is not an APIError: %+v", err)
+			}
+			if assert.ErrorAs(t, err, &apiErr) {
+				assert.Equal(t, apiErr.StatusCode, tc.expectErrCode)
+			}
+		})
 	}
-	assert.Contains(t, err.Error(), "400 Bad Request")
 }

test/e2e/util.go

@@ -83,6 +83,7 @@ func createOpenAIClient(baseURL, apiKey string) *openai.Client {
 			r.URL.Path = "/v1" + r.URL.Path
 			return mn(r)
 		}),
+		option.WithMaxRetries(0),
 	)
 }
 
@@ -96,6 +97,7 @@ func createOpenAIClientWithRoutingStrategy(baseURL, apiKey, routingStrategy stri
 			return mn(r)
 		}),
 		option.WithHeader("routing-strategy", routingStrategy),
+		option.WithMaxRetries(0),
 		respOpt,
 	)
 }

test/run-e2e-tests.sh

@@ -66,7 +66,7 @@ if [ -n "$INSTALL_AIBRIX" ]; then
   cd ../..
 
   kubectl port-forward svc/llama2-7b 8000:8000 &
-  kubectl -n envoy-gateway-system port-forward service/envoy-aibrix-system-aibrix-eg-903790dc  8888:80 &
+  kubectl -n envoy-gateway-system port-forward service/envoy-aibrix-system-aibrix-eg-903790dc 8888:80 &
 
   function cleanup {
     echo "Cleaning up..."