more

Signed-off-by: Takeshi Yoneda <t.y.mathetake@gmail.com>

Author: mathetake

internal/extproc/translator/openai_awsbedrock.go

@@ -16,7 +16,6 @@ import (
 	"strings"
 
 	"github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream"
-	extprocv3 "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
 	"k8s.io/utils/ptr"
 
 	"github.com/envoyproxy/ai-gateway/internal/apischema/awsbedrock"
@@ -102,8 +101,8 @@ func (o *openAIToAWSBedrockTranslatorV1ChatCompletion) RequestBody(_ []byte, ope
 		}
 	}
 
-	mut := &extprocv3.BodyMutation_Body{}
-	if mut.Body, err = json.Marshal(bedrockReq); err != nil {
+	newBody, err = json.Marshal(bedrockReq)
+	if err != nil {
 		return nil, nil, fmt.Errorf("failed to marshal body: %w", err)
 	}
 	newHeaders = []internalapi.Header{
@@ -480,10 +479,10 @@ func (o *openAIToAWSBedrockTranslatorV1ChatCompletion) ResponseHeaders(headers m
 		contentType := headers["content-type"]
 		if contentType == "application/vnd.amazon.eventstream" {
 			// We need to change the content-type to text/event-stream for streaming responses.
-			newHeaders = []internalapi.Header{{"content-type", "text/event-stream"}}
+			newHeaders = []internalapi.Header{{contentType, "text/event-stream"}}
 		}
 	}
-	return nil, nil
+	return
 }
 
 func (o *openAIToAWSBedrockTranslatorV1ChatCompletion) bedrockStopReasonToOpenAIStopReason(

internal/extproc/translator/openai_awsbedrock_test.go

@@ -1293,7 +1293,6 @@ func TestOpenAIToAWSBedrockTranslatorV1ChatCompletion_Streaming_ResponseBody(t *
 			hm, newBody, tokenUsage, _, err := o.ResponseBody(nil, bytes.NewBuffer([]byte{buf[i]}), i == len(buf)-1, nil)
 			require.NoError(t, err)
 			require.Nil(t, hm)
-			require.NotNil(t, newBody)
 			if len(newBody) > 0 {
 				results = append(results, string(newBody))
 			}
@@ -1407,10 +1406,11 @@ func TestOpenAIToAWSBedrockTranslator_ResponseError(t *testing.T) {
 			require.NoError(t, err)
 			require.NotNil(t, newBody)
 			require.NotNil(t, hm)
-			require.NotNil(t, hm)
-			require.Len(t, hm, 1)
-			require.Equal(t, contentLengthHeaderName, hm[0].Key())
-			require.Equal(t, strconv.Itoa(len(newBody)), string(hm[0].Value()))
+			require.Len(t, hm, 2)
+			require.Equal(t, contentTypeHeaderName, hm[0].Key())
+			require.Equal(t, jsonContentType, hm[0].Value())
+			require.Equal(t, contentLengthHeaderName, hm[1].Key())
+			require.Equal(t, strconv.Itoa(len(newBody)), hm[1].Value())
 
 			var openAIError openai.Error
 			err = json.Unmarshal(newBody, &openAIError)

internal/extproc/translator/openai_completions_test.go

@@ -190,11 +190,8 @@ func TestOpenAIToOpenAITranslatorV1CompletionResponseBody(t *testing.T) {
 			require.NoError(t, err)
 			require.Equal(t, tc.expTokenUsage, tokenUsage)
 			require.Equal(t, tc.expModel, responseModel)
-
-			// For non-streaming responses, body should be passed through.
-			require.NotNil(t, bodyMutation)
-			require.NotNil(t, headerMutation)
-			require.Equal(t, contentLengthHeaderName, headerMutation[0].Key())
+			require.Nil(t, bodyMutation)
+			require.Nil(t, headerMutation)
 		})
 	}
 }
@@ -277,11 +274,8 @@ func TestOpenAIToOpenAITranslatorV1CompletionResponseError(t *testing.T) {
 
 		headerMutation, bodyMutation, err := translator.ResponseError(respHeaders, strings.NewReader(errorBody))
 		require.NoError(t, err)
-
-		// For passthrough, errors should be returned as-is.
-		require.NotNil(t, bodyMutation)
-		require.NotNil(t, headerMutation)
-		require.Equal(t, errorBody, string(bodyMutation))
+		require.Nil(t, bodyMutation)
+		require.Nil(t, headerMutation)
 	})
 
 	t.Run("non_json_error", func(t *testing.T) {
@@ -293,10 +287,7 @@ func TestOpenAIToOpenAITranslatorV1CompletionResponseError(t *testing.T) {
 
 		headerMutation, bodyMutation, err := translator.ResponseError(respHeaders, strings.NewReader(errorBody))
 		require.NoError(t, err)
-		require.NotNil(t, headerMutation)
-		require.NotNil(t, bodyMutation)
-
-		// Should still pass through the error.
-		require.Equal(t, errorBody, string(bodyMutation))
+		require.Nil(t, headerMutation)
+		require.Nil(t, bodyMutation)
 	})
 }

internal/extproc/translator/openai_gcpanthropic_test.go

@@ -894,15 +894,14 @@ func TestOpenAIToGCPAnthropicTranslator_ResponseError(t *testing.T) {
 			require.NoError(t, err)
 			require.NotNil(t, bm)
 			require.NotNil(t, hm)
-
-			newBody := bm
-			require.NotNil(t, newBody)
-			require.Len(t, hm, 1)
-			require.Equal(t, contentLengthHeaderName, hm[0].Key())
-			require.Equal(t, strconv.Itoa(len(newBody)), string(hm[0].Value()))
+			require.Len(t, hm, 2)
+			require.Equal(t, contentTypeHeaderName, hm[0].Key())
+			require.Equal(t, jsonContentType, hm[0].Value())
+			require.Equal(t, contentLengthHeaderName, hm[1].Key())
+			require.Equal(t, strconv.Itoa(len(bm)), hm[1].Value())
 
 			var gotError openai.Error
-			err = json.Unmarshal(newBody, &gotError)
+			err = json.Unmarshal(bm, &gotError)
 			require.NoError(t, err)
 
 			if diff := cmp.Diff(tt.expectedOutput, gotError); diff != "" {