Skip to content

Commit 713456b

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

18 files changed

+198
-368
lines changed

‎internal/extproc/translator/anthropic_anthropic_test.go‎

Lines changed: 9 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -10,8 +10,6 @@ import (
1010
"strings"
1111
"testing"
1212

13-
corev3 "github.com/envoyproxy/go-control-plane/envoy/config/core/v3"
14-
extprocv3 "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
1513
"github.com/stretchr/testify/require"
1614

1715
anthropicschema "github.com/envoyproxy/ai-gateway/internal/apischema/anthropic"
@@ -27,7 +25,7 @@ func TestAnthropicToAnthropic_RequestBody(t *testing.T) {
2725
modelNameOverride string
2826

2927
expRequestModel internalapi.RequestModel
30-
expBodyMutation *extprocv3.BodyMutation
28+
expNewBody []byte
3129
}{
3230
{
3331
name: "no mutation",
@@ -36,7 +34,7 @@ func TestAnthropicToAnthropic_RequestBody(t *testing.T) {
3634
forceBodyMutation: false,
3735
modelNameOverride: "",
3836
expRequestModel: "claude-2",
39-
expBodyMutation: nil,
37+
expNewBody: nil,
4038
},
4139
{
4240
name: "model override",
@@ -45,11 +43,7 @@ func TestAnthropicToAnthropic_RequestBody(t *testing.T) {
4543
forceBodyMutation: false,
4644
modelNameOverride: "claude-100.1",
4745
expRequestModel: "claude-100.1",
48-
expBodyMutation: &extprocv3.BodyMutation{
49-
Mutation: &extprocv3.BodyMutation_Body{
50-
Body: []byte(`{"model":"claude-100.1","messages":[{"role":"user","content":"Hello!"}], "stream": true}`),
51-
},
52-
},
46+
expNewBody: []byte(`{"model":"claude-100.1","messages":[{"role":"user","content":"Hello!"}], "stream": true}`),
5347
},
5448
{
5549
name: "force mutation",
@@ -58,11 +52,7 @@ func TestAnthropicToAnthropic_RequestBody(t *testing.T) {
5852
forceBodyMutation: true,
5953
modelNameOverride: "",
6054
expRequestModel: "claude-2",
61-
expBodyMutation: &extprocv3.BodyMutation{
62-
Mutation: &extprocv3.BodyMutation_Body{
63-
Body: []byte(`{"model":"claude-2","messages":[{"role":"user","content":"Hello!"}]}`),
64-
},
65-
},
55+
expNewBody: []byte(`{"model":"claude-2","messages":[{"role":"user","content":"Hello!"}]}`),
6656
},
6757
} {
6858
t.Run(tc.name, func(t *testing.T) {
@@ -71,24 +61,14 @@ func TestAnthropicToAnthropic_RequestBody(t *testing.T) {
7161

7262
headerMutation, bodyMutation, err := translator.RequestBody(tc.original, &tc.body, tc.forceBodyMutation)
7363
require.NoError(t, err)
74-
expHeaderMutation := &extprocv3.HeaderMutation{
75-
SetHeaders: []*corev3.HeaderValueOption{
76-
{
77-
Header: &corev3.HeaderValue{
78-
Key: pathHeaderName,
79-
RawValue: []byte("/v1/messages"),
80-
},
81-
},
82-
},
64+
expHeaders := []internalapi.Header{
65+
{pathHeaderName, "/v1/messages"},
8366
}
8467
if bodyMutation != nil {
85-
expHeaderMutation.SetHeaders = append(expHeaderMutation.SetHeaders, &corev3.HeaderValueOption{Header: &corev3.HeaderValue{
86-
Key: contentLengthHeaderName,
87-
RawValue: []byte(strconv.Itoa(len(bodyMutation.GetBody()))),
88-
}})
68+
expHeaders = append(expHeaders, internalapi.Header{contentLengthHeaderName, strconv.Itoa(len(bodyMutation))})
8969
}
90-
require.Equal(t, expHeaderMutation, headerMutation)
91-
require.Equal(t, tc.expBodyMutation, bodyMutation)
70+
require.Equal(t, expHeaders, headerMutation)
71+
require.Equal(t, tc.expNewBody, bodyMutation)
9272

9373
require.Equal(t, tc.expRequestModel, translator.(*anthropicToAnthropicTranslator).requestModel)
9474
require.Equal(t, tc.body.Stream, translator.(*anthropicToAnthropicTranslator).stream)

‎internal/extproc/translator/anthropic_awsanthropic_test.go‎

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -76,14 +76,14 @@ func TestAnthropicToAWSAnthropicTranslator_RequestBody_ModelNameOverride(t *test
7676

7777
// Check path header contains expected model (URL encoded).
7878
// Use the last element as it takes precedence when multiple headers are set.
79-
pathHeader := headerMutation.SetHeaders[len(headerMutation.SetHeaders)-2]
80-
require.Equal(t, pathHeaderName, pathHeader.Header.Key)
79+
pathHeader := headerMutation[len(headerMutation)-2]
80+
require.Equal(t, pathHeaderName, pathHeader.Key())
8181
expectedPath := "/model/" + tt.expectedInPath + "/invoke"
82-
assert.Equal(t, expectedPath, string(pathHeader.Header.RawValue))
82+
assert.Equal(t, expectedPath, pathHeader.Value())
8383

8484
// Check that model field is removed from body (since it's in the path).
8585
var modifiedReq map[string]any
86-
err = json.Unmarshal(bodyMutation.GetBody(), &modifiedReq)
86+
err = json.Unmarshal(bodyMutation, &modifiedReq)
8787
require.NoError(t, err)
8888
_, hasModel := modifiedReq["model"]
8989
assert.False(t, hasModel, "model field should be removed from request body")
@@ -144,9 +144,9 @@ func TestAnthropicToAWSAnthropicTranslator_RequestBody_StreamingPaths(t *testing
144144

145145
// Check path contains expected suffix.
146146
// Use the last element as it takes precedence when multiple headers are set.
147-
pathHeader := headerMutation.SetHeaders[len(headerMutation.SetHeaders)-2]
147+
pathHeader := headerMutation[len(headerMutation)-2]
148148
expectedPath := "/model/anthropic.claude-3-sonnet-20240229-v1:0" + tt.expectedPathSuffix
149-
assert.Equal(t, expectedPath, string(pathHeader.Header.RawValue))
149+
assert.Equal(t, expectedPath, pathHeader.Value())
150150
})
151151
}
152152
}
@@ -188,8 +188,9 @@ func TestAnthropicToAWSAnthropicTranslator_URLEncoding(t *testing.T) {
188188
require.NotNil(t, headerMutation)
189189

190190
// Use the last element as it takes precedence when multiple headers are set.
191-
pathHeader := headerMutation.SetHeaders[len(headerMutation.SetHeaders)-2]
192-
assert.Equal(t, tt.expectedPath, string(pathHeader.Header.RawValue))
191+
pathHeader := headerMutation[len(headerMutation)-2]
192+
assert.Equal(t, pathHeaderName, pathHeader.Key())
193+
assert.Equal(t, tt.expectedPath, pathHeader.Value())
193194
})
194195
}
195196
}

‎internal/extproc/translator/anthropic_gcpanthropic_test.go‎

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -60,14 +60,14 @@ func TestAnthropicToGCPAnthropicTranslator_RequestBody_ModelNameOverride(t *test
6060
require.NotNil(t, bodyMutation)
6161

6262
// Check path header contains expected model.
63-
pathHeader := headerMutation.SetHeaders[0]
64-
require.Equal(t, pathHeaderName, pathHeader.Header.Key)
63+
pathHeader := headerMutation[0]
64+
require.Equal(t, pathHeaderName, pathHeader.Key())
6565
expectedPath := "publishers/anthropic/models/" + tt.expectedInPath + ":rawPredict"
66-
assert.Equal(t, expectedPath, string(pathHeader.Header.RawValue))
66+
assert.Equal(t, expectedPath, string(pathHeader.Value()))
6767

6868
// Check that model field is removed from body (since it's in the path).
6969
var modifiedReq map[string]any
70-
err = json.Unmarshal(bodyMutation.GetBody(), &modifiedReq)
70+
err = json.Unmarshal(bodyMutation, &modifiedReq)
7171
require.NoError(t, err)
7272
_, hasModel := modifiedReq["model"]
7373
assert.False(t, hasModel, "model field should be removed from request body")
@@ -120,7 +120,7 @@ func TestAnthropicToGCPAnthropicTranslator_BackendVersionHandling(t *testing.T)
120120
require.NotNil(t, bodyMutation)
121121

122122
var outputReq map[string]any
123-
err = json.Unmarshal(bodyMutation.GetBody(), &outputReq)
123+
err = json.Unmarshal(bodyMutation, &outputReq)
124124
require.NoError(t, err)
125125

126126
require.Contains(t, outputReq, "anthropic_version")
@@ -181,9 +181,10 @@ func TestAnthropicToGCPAnthropicTranslator_RequestBody_StreamingPaths(t *testing
181181
require.NotNil(t, headerMutation)
182182

183183
// Check path contains expected specifier.
184-
pathHeader := headerMutation.SetHeaders[0]
184+
pathHeader := headerMutation[0]
185185
expectedPath := "publishers/anthropic/models/claude-3-sonnet-20240229:" + tt.expectedSpecifier
186-
assert.Equal(t, expectedPath, string(pathHeader.Header.RawValue))
186+
assert.Equal(t, pathHeaderName, pathHeader.Key())
187+
assert.Equal(t, expectedPath, pathHeader.Value())
187188
})
188189
}
189190
}

‎internal/extproc/translator/cohere_rerank_v2_test.go‎

Lines changed: 15 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -62,24 +62,24 @@ func TestCohereToCohereTranslatorV2Rerank_RequestBody(t *testing.T) {
6262
headerMutation, bodyMutation, err := translator.RequestBody([]byte(originalBody), &req, tc.onRetry)
6363
require.NoError(t, err)
6464
require.NotNil(t, headerMutation)
65-
require.GreaterOrEqual(t, len(headerMutation.SetHeaders), 1)
66-
require.Equal(t, pathHeaderName, headerMutation.SetHeaders[0].Header.Key)
67-
require.Equal(t, tc.expPath, string(headerMutation.SetHeaders[0].Header.RawValue))
65+
require.GreaterOrEqual(t, len(headerMutation), 1)
66+
require.Equal(t, pathHeaderName, headerMutation[0].Key())
67+
require.Equal(t, tc.expPath, headerMutation[0].Value())
6868

6969
switch {
7070
case tc.expBodyContains != "":
7171
require.NotNil(t, bodyMutation)
72-
require.Contains(t, string(bodyMutation.GetBody()), tc.expBodyContains)
72+
require.Contains(t, string(bodyMutation), tc.expBodyContains)
7373
// Verify content-length header is set.
74-
require.Len(t, headerMutation.SetHeaders, 2)
75-
require.Equal(t, contentLengthHeaderName, headerMutation.SetHeaders[1].Header.Key)
74+
require.Len(t, headerMutation, 2)
75+
require.Equal(t, contentLengthHeaderName, headerMutation[1].Key())
7676
case bodyMutation != nil:
7777
// If there's a body mutation (like on retry), content-length header should be set.
78-
require.Len(t, headerMutation.SetHeaders, 2)
79-
require.Equal(t, contentLengthHeaderName, headerMutation.SetHeaders[1].Header.Key)
78+
require.Len(t, headerMutation, 2)
79+
require.Equal(t, contentLengthHeaderName, headerMutation[1].Key())
8080
default:
8181
// No body mutation, only path header.
82-
require.Len(t, headerMutation.SetHeaders, 1)
82+
require.Len(t, headerMutation, 1)
8383
}
8484
})
8585
}
@@ -95,12 +95,12 @@ func TestCohereToCohereTranslatorV2Rerank_RequestBody_InvalidJSONCreatesBodyWith
9595
require.NotNil(t, headerMutation)
9696
require.NotNil(t, bodyMutation)
9797
// Body should contain the override model
98-
require.Contains(t, string(bodyMutation.GetBody()), `"model":"override-model"`)
98+
require.Contains(t, string(bodyMutation), `"model":"override-model"`)
9999
// Verify content-length header is set alongside :path
100-
require.GreaterOrEqual(t, len(headerMutation.SetHeaders), 2)
101-
require.Equal(t, pathHeaderName, headerMutation.SetHeaders[0].Header.Key)
102-
require.Equal(t, "/v2/rerank", string(headerMutation.SetHeaders[0].Header.RawValue))
103-
require.Equal(t, contentLengthHeaderName, headerMutation.SetHeaders[1].Header.Key)
100+
require.GreaterOrEqual(t, len(headerMutation), 2)
101+
require.Equal(t, pathHeaderName, headerMutation[0].Key())
102+
require.Equal(t, "/v2/rerank", string(headerMutation[0].Value()))
103+
require.Equal(t, contentLengthHeaderName, headerMutation[1].Key())
104104
}
105105

106106
func TestCohereToCohereTranslatorV2Rerank_RequestBody_SetModelNameError(t *testing.T) {
@@ -197,7 +197,7 @@ func TestCohereToCohereTranslatorV2Rerank_ResponseError(t *testing.T) {
197197
require.NotNil(t, bodyMutation)
198198

199199
var cohereErr cohereschema.RerankV2Error
200-
require.NoError(t, json.Unmarshal(bodyMutation.GetBody(), &cohereErr))
200+
require.NoError(t, json.Unmarshal(bodyMutation, &cohereErr))
201201
require.NotNil(t, cohereErr.Message)
202202
require.Equal(t, errorBody, *cohereErr.Message)
203203
})

‎internal/extproc/translator/imagegeneration_openai_openai_test.go‎

Lines changed: 11 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -25,15 +25,14 @@ func TestOpenAIToOpenAIImageTranslator_RequestBody_ModelOverrideAndPath(t *testi
2525
hm, bm, err := tr.RequestBody(original, req, false)
2626
require.NoError(t, err)
2727
require.NotNil(t, hm)
28-
require.Len(t, hm.SetHeaders, 2) // path and content-length headers
29-
require.Equal(t, pathHeaderName, hm.SetHeaders[0].Header.Key)
30-
require.Equal(t, "/v1/images/generations", string(hm.SetHeaders[0].Header.RawValue))
31-
require.Equal(t, contentLengthHeaderName, hm.SetHeaders[1].Header.Key)
28+
require.Len(t, hm, 2) // path and content-length headers
29+
require.Equal(t, pathHeaderName, hm[0].Key())
30+
require.Equal(t, "/v1/images/generations", string(hm[0].Value()))
31+
require.Equal(t, contentLengthHeaderName, hm[1].Key())
3232

3333
require.NotNil(t, bm)
34-
mutated := bm.GetBody()
3534
var got openaisdk.ImageGenerateParams
36-
require.NoError(t, json.Unmarshal(mutated, &got))
35+
require.NoError(t, json.Unmarshal(bm, &got))
3736
require.Equal(t, "gpt-image-1", got.Model)
3837
}
3938

@@ -47,15 +46,15 @@ func TestOpenAIToOpenAIImageTranslator_RequestBody_ForceMutation(t *testing.T) {
4746
require.NotNil(t, hm)
4847
// Content-Length is set only when body mutated; with force it should be mutated to original.
4948
foundCL := false
50-
for _, h := range hm.SetHeaders {
51-
if h.Header.Key == contentLengthHeaderName {
49+
for _, h := range hm {
50+
if h.Key() == contentLengthHeaderName {
5251
foundCL = true
5352
break
5453
}
5554
}
5655
require.True(t, foundCL)
5756
require.NotNil(t, bm)
58-
require.Equal(t, original, bm.GetBody())
57+
require.Equal(t, original, bm)
5958
}
6059

6160
func TestOpenAIToOpenAIImageTranslator_ResponseError_NonJSON(t *testing.T) {
@@ -70,7 +69,7 @@ func TestOpenAIToOpenAIImageTranslator_ResponseError_NonJSON(t *testing.T) {
7069
var got struct {
7170
Error openai.ErrorType `json:"error"`
7271
}
73-
require.NoError(t, json.Unmarshal(bm.GetBody(), &got))
72+
require.NoError(t, json.Unmarshal(bm, &got))
7473
require.Equal(t, openAIBackendError, got.Error.Type)
7574
}
7675

@@ -96,8 +95,8 @@ func TestOpenAIToOpenAIImageTranslator_RequestBody_NoOverrideNoForce(t *testing.
9695
require.NoError(t, err)
9796
require.NotNil(t, hm)
9897
// Only path header present; content-length should not be set when no mutation
99-
require.Len(t, hm.SetHeaders, 1)
100-
require.Equal(t, pathHeaderName, hm.SetHeaders[0].Header.Key)
98+
require.Len(t, hm, 1)
99+
require.Equal(t, pathHeaderName, hm[0].Key())
101100
require.Nil(t, bm)
102101
}
103102

0 commit comments

Comments
 (0)