Skip to content

Commit cdff8d2

Browse files
committed
fix: gate Titan embedding dimensions to v2
Addresses PR review (Codex P2). Titan v1 (G1) rejects any InvokeModel body field other than inputText, so serializing `dimensions` for it — e.g. from an embedding config shared with v2 — would fail the request. Only set dimensions for amazon.titan-embed-text-v2 and ignore it for v1 (which returns its fixed-size embedding). Adds a regression test. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01WnKgCRfMvFcZAkoVbr8k79
1 parent fbe454e commit cdff8d2

2 files changed

Lines changed: 31 additions & 1 deletion

File tree

packages/engine/internal/connector/provider_bedrock.go

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -261,9 +261,14 @@ func (p *BedrockEmbeddingProvider) Embeddings(ctx context.Context, req *Embeddin
261261
out := make([][]float64, len(req.Inputs))
262262
totalTokens := 0
263263

264+
// dimensions is only accepted by Titan v2; Titan v1 (G1) rejects any field
265+
// other than inputText, so ignore it there rather than failing a request
266+
// that reuses a shared embedding config.
267+
supportsDimensions := strings.Contains(req.Model, "titan-embed-text-v2")
268+
264269
for i, text := range req.Inputs {
265270
body := titanEmbedRequest{InputText: text}
266-
if req.Dimensions > 0 {
271+
if req.Dimensions > 0 && supportsDimensions {
267272
body.Dimensions = req.Dimensions
268273
}
269274
bodyJSON, err := json.Marshal(body)

packages/engine/internal/connector/provider_bedrock_embed_test.go

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,31 @@ func TestBedrockEmbeddingProvider_Titan(t *testing.T) {
6969
}
7070
}
7171

72+
func TestBedrockEmbeddingProvider_V1IgnoresDimensions(t *testing.T) {
73+
mock := &mockBedrockInvokeClient{
74+
InvokeFunc: func(ctx context.Context, input *bedrockruntime.InvokeModelInput, opts ...func(*bedrockruntime.Options)) (*bedrockruntime.InvokeModelOutput, error) {
75+
var body titanEmbedRequest
76+
if err := json.Unmarshal(input.Body, &body); err != nil {
77+
t.Fatalf("unmarshaling request body: %v", err)
78+
}
79+
// Titan v1 must not receive dimensions even when requested.
80+
if body.Dimensions != 0 {
81+
t.Errorf("dimensions = %d, want 0 (omitted) for titan v1", body.Dimensions)
82+
}
83+
resp, _ := json.Marshal(titanEmbedResponse{Embedding: []float64{1}, InputTextTokenCount: 1})
84+
return &bedrockruntime.InvokeModelOutput{Body: resp}, nil
85+
},
86+
}
87+
p := &BedrockEmbeddingProvider{Client: mock}
88+
if _, err := p.Embeddings(context.Background(), &EmbeddingRequest{
89+
Model: "amazon.titan-embed-text-v1",
90+
Inputs: []string{"x"},
91+
Dimensions: 512,
92+
}); err != nil {
93+
t.Fatalf("Embeddings() error: %v", err)
94+
}
95+
}
96+
7297
func TestBedrockEmbeddingProvider_UnsupportedModel(t *testing.T) {
7398
p := &BedrockEmbeddingProvider{Client: &mockBedrockInvokeClient{}}
7499
_, err := p.Embeddings(context.Background(), &EmbeddingRequest{

0 commit comments

Comments
 (0)