Skip to content

Commit f4c6221

Browse files
committed
more
Signed-off-by: Takeshi Yoneda <t.y.mathetake@gmail.com>
1 parent 713456b commit f4c6221

File tree

6 files changed

+517
-813
lines changed

6 files changed

+517
-813
lines changed

internal/extproc/translator/openai_awsbedrock.go

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,6 @@ import (
1616
"strings"
1717

1818
"github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream"
19-
extprocv3 "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
2019
"k8s.io/utils/ptr"
2120

2221
"github.com/envoyproxy/ai-gateway/internal/apischema/awsbedrock"
@@ -102,8 +101,8 @@ func (o *openAIToAWSBedrockTranslatorV1ChatCompletion) RequestBody(_ []byte, ope
102101
}
103102
}
104103

105-
mut := &extprocv3.BodyMutation_Body{}
106-
if mut.Body, err = json.Marshal(bedrockReq); err != nil {
104+
newBody, err = json.Marshal(bedrockReq)
105+
if err != nil {
107106
return nil, nil, fmt.Errorf("failed to marshal body: %w", err)
108107
}
109108
newHeaders = []internalapi.Header{
@@ -480,10 +479,10 @@ func (o *openAIToAWSBedrockTranslatorV1ChatCompletion) ResponseHeaders(headers m
480479
contentType := headers["content-type"]
481480
if contentType == "application/vnd.amazon.eventstream" {
482481
// We need to change the content-type to text/event-stream for streaming responses.
483-
newHeaders = []internalapi.Header{{"content-type", "text/event-stream"}}
482+
newHeaders = []internalapi.Header{{contentType, "text/event-stream"}}
484483
}
485484
}
486-
return nil, nil
485+
return
487486
}
488487

489488
func (o *openAIToAWSBedrockTranslatorV1ChatCompletion) bedrockStopReasonToOpenAIStopReason(

internal/extproc/translator/openai_awsbedrock_test.go

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1293,7 +1293,6 @@ func TestOpenAIToAWSBedrockTranslatorV1ChatCompletion_Streaming_ResponseBody(t *
12931293
hm, newBody, tokenUsage, _, err := o.ResponseBody(nil, bytes.NewBuffer([]byte{buf[i]}), i == len(buf)-1, nil)
12941294
require.NoError(t, err)
12951295
require.Nil(t, hm)
1296-
require.NotNil(t, newBody)
12971296
if len(newBody) > 0 {
12981297
results = append(results, string(newBody))
12991298
}
@@ -1407,10 +1406,11 @@ func TestOpenAIToAWSBedrockTranslator_ResponseError(t *testing.T) {
14071406
require.NoError(t, err)
14081407
require.NotNil(t, newBody)
14091408
require.NotNil(t, hm)
1410-
require.NotNil(t, hm)
1411-
require.Len(t, hm, 1)
1412-
require.Equal(t, contentLengthHeaderName, hm[0].Key())
1413-
require.Equal(t, strconv.Itoa(len(newBody)), string(hm[0].Value()))
1409+
require.Len(t, hm, 2)
1410+
require.Equal(t, contentTypeHeaderName, hm[0].Key())
1411+
require.Equal(t, jsonContentType, hm[0].Value())
1412+
require.Equal(t, contentLengthHeaderName, hm[1].Key())
1413+
require.Equal(t, strconv.Itoa(len(newBody)), hm[1].Value())
14141414

14151415
var openAIError openai.Error
14161416
err = json.Unmarshal(newBody, &openAIError)

internal/extproc/translator/openai_completions_test.go

Lines changed: 6 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -190,11 +190,8 @@ func TestOpenAIToOpenAITranslatorV1CompletionResponseBody(t *testing.T) {
190190
require.NoError(t, err)
191191
require.Equal(t, tc.expTokenUsage, tokenUsage)
192192
require.Equal(t, tc.expModel, responseModel)
193-
194-
// For non-streaming responses, body should be passed through.
195-
require.NotNil(t, bodyMutation)
196-
require.NotNil(t, headerMutation)
197-
require.Equal(t, contentLengthHeaderName, headerMutation[0].Key())
193+
require.Nil(t, bodyMutation)
194+
require.Nil(t, headerMutation)
198195
})
199196
}
200197
}
@@ -277,11 +274,8 @@ func TestOpenAIToOpenAITranslatorV1CompletionResponseError(t *testing.T) {
277274

278275
headerMutation, bodyMutation, err := translator.ResponseError(respHeaders, strings.NewReader(errorBody))
279276
require.NoError(t, err)
280-
281-
// For passthrough, errors should be returned as-is.
282-
require.NotNil(t, bodyMutation)
283-
require.NotNil(t, headerMutation)
284-
require.Equal(t, errorBody, string(bodyMutation))
277+
require.Nil(t, bodyMutation)
278+
require.Nil(t, headerMutation)
285279
})
286280

287281
t.Run("non_json_error", func(t *testing.T) {
@@ -293,10 +287,7 @@ func TestOpenAIToOpenAITranslatorV1CompletionResponseError(t *testing.T) {
293287

294288
headerMutation, bodyMutation, err := translator.ResponseError(respHeaders, strings.NewReader(errorBody))
295289
require.NoError(t, err)
296-
require.NotNil(t, headerMutation)
297-
require.NotNil(t, bodyMutation)
298-
299-
// Should still pass through the error.
300-
require.Equal(t, errorBody, string(bodyMutation))
290+
require.Nil(t, headerMutation)
291+
require.Nil(t, bodyMutation)
301292
})
302293
}

internal/extproc/translator/openai_gcpanthropic_test.go

Lines changed: 6 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -894,15 +894,14 @@ func TestOpenAIToGCPAnthropicTranslator_ResponseError(t *testing.T) {
894894
require.NoError(t, err)
895895
require.NotNil(t, bm)
896896
require.NotNil(t, hm)
897-
898-
newBody := bm
899-
require.NotNil(t, newBody)
900-
require.Len(t, hm, 1)
901-
require.Equal(t, contentLengthHeaderName, hm[0].Key())
902-
require.Equal(t, strconv.Itoa(len(newBody)), string(hm[0].Value()))
897+
require.Len(t, hm, 2)
898+
require.Equal(t, contentTypeHeaderName, hm[0].Key())
899+
require.Equal(t, jsonContentType, hm[0].Value())
900+
require.Equal(t, contentLengthHeaderName, hm[1].Key())
901+
require.Equal(t, strconv.Itoa(len(bm)), hm[1].Value())
903902

904903
var gotError openai.Error
905-
err = json.Unmarshal(newBody, &gotError)
904+
err = json.Unmarshal(bm, &gotError)
906905
require.NoError(t, err)
907906

908907
if diff := cmp.Diff(tt.expectedOutput, gotError); diff != "" {

0 commit comments

Comments
 (0)