diff --git a/go/adk/pkg/embedding/embedding.go b/go/adk/pkg/embedding/embedding.go index 371a1de9e..44b5ab000 100644 --- a/go/adk/pkg/embedding/embedding.go +++ b/go/adk/pkg/embedding/embedding.go @@ -115,9 +115,8 @@ func generateEmbeddings(ctx context.Context, client openai.Client, cfg *adk.Embe log := logr.FromContextOrDiscard(ctx) resp, err := client.Embeddings.New(ctx, openai.EmbeddingNewParams{ - Model: openai.EmbeddingModel(cfg.Model), - Input: openai.EmbeddingNewParamsInputUnion{OfArrayOfStrings: texts}, - Dimensions: openai.Int(int64(TargetDimension)), + Model: openai.EmbeddingModel(cfg.Model), + Input: openai.EmbeddingNewParamsInputUnion{OfArrayOfStrings: texts}, }, embeddingPassthroughOpts(ctx, cfg, isAzureFamily)...) if err != nil { return nil, fmt.Errorf("%s embeddings request failed: %w", provider, err) diff --git a/go/adk/pkg/embedding/embedding_test.go b/go/adk/pkg/embedding/embedding_test.go index 63ce20c38..220d36ad8 100644 --- a/go/adk/pkg/embedding/embedding_test.go +++ b/go/adk/pkg/embedding/embedding_test.go @@ -29,9 +29,8 @@ func TestOpenAIProvider_UsesAPIKeyNotKagentToken(t *testing.T) { t.Fatalf("read body: %v", err) } var req struct { - Model string `json:"model"` - Input []string `json:"input"` - Dimensions int `json:"dimensions"` + Model string `json:"model"` + Input []string `json:"input"` } if err := json.Unmarshal(body, &req); err != nil { t.Fatalf("unmarshal request: %v", err) @@ -39,8 +38,18 @@ func TestOpenAIProvider_UsesAPIKeyNotKagentToken(t *testing.T) { if req.Model != "text-embedding-3-small" { t.Errorf("model = %q, want text-embedding-3-small", req.Model) } - if req.Dimensions != TargetDimension { - t.Errorf("dimensions = %d, want %d", req.Dimensions, TargetDimension) + // dimensions must not be sent at all: providers that don't offer + // exactly TargetDimension (e.g. Bedrock Titan behind an + // OpenAI-compatible gateway) would reject the request outright. + // processEmbeddings truncates whatever the provider returns instead. + // Checked against the raw body, not a decoded field, since decoding + // into an int can't tell an absent field from an explicit 0. + var rawFields map[string]json.RawMessage + if err := json.Unmarshal(body, &rawFields); err != nil { + t.Fatalf("unmarshal request as map: %v", err) + } + if _, present := rawFields["dimensions"]; present { + t.Errorf("request body has a dimensions field, want it omitted: %s", body) } if len(req.Input) != 1 || req.Input[0] != "hello" { t.Errorf("input = %v, want [hello]", req.Input) diff --git a/python/packages/kagent-adk/src/kagent/adk/models/_embedding.py b/python/packages/kagent-adk/src/kagent/adk/models/_embedding.py index deccc2978..e7481d0bc 100644 --- a/python/packages/kagent-adk/src/kagent/adk/models/_embedding.py +++ b/python/packages/kagent-adk/src/kagent/adk/models/_embedding.py @@ -174,7 +174,6 @@ async def _embed_openai(self, texts: List[str]) -> List[List[float]]: response = await client.embeddings.create( model=self.config.model, input=texts, - dimensions=self.TARGET_DIMENSION, ) return [item.embedding for item in response.data] diff --git a/python/packages/kagent-adk/tests/unittests/test_embedding.py b/python/packages/kagent-adk/tests/unittests/test_embedding.py index 0144f6fdf..cc4363a01 100644 --- a/python/packages/kagent-adk/tests/unittests/test_embedding.py +++ b/python/packages/kagent-adk/tests/unittests/test_embedding.py @@ -53,6 +53,24 @@ async def test_generate_single_text(self): result = await client.generate("hello world") assert result == vec + @pytest.mark.asyncio + async def test_generate_single_text_omits_dimensions(self): + # Providers that don't offer exactly 768 dimensions (e.g. Bedrock Titan + # behind an OpenAI-compatible gateway) reject the request outright when + # dimensions is set. _process_embeddings truncates whatever comes back. + client = make_client(provider="openai", model="text-embedding-3-small") + vec = [0.1] * 768 + mock_response = make_openai_embedding_response([vec]) + with mock.patch("openai.AsyncOpenAI") as mock_cls: + instance = mock.AsyncMock() + instance.embeddings.create = mock.AsyncMock(return_value=mock_response) + mock_cls.return_value = instance + await client.generate("hello world") + instance.embeddings.create.assert_called_once_with( + model="text-embedding-3-small", + input=["hello world"], + ) + @pytest.mark.asyncio async def test_generate_batch_texts(self): client = make_client(provider="openai", model="text-embedding-3-small")