Skip to content
Open
Show file tree
Hide file tree
Changes from 5 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 2 additions & 3 deletions go/adk/pkg/embedding/embedding.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
19 changes: 14 additions & 5 deletions go/adk/pkg/embedding/embedding_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -29,18 +29,27 @@ 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)
}
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)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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]

Expand Down
18 changes: 18 additions & 0 deletions python/packages/kagent-adk/tests/unittests/test_embedding.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down
Loading