Skip to content

Commit 87eadc7

Browse files
Python: Batch Google AI embedding requests to the 100-input limit (#14700)
### Motivation and Context `GoogleAITextEmbedding.generate_raw_embeddings` sent every input text in a single `embed_content` request. The Gemini `batchEmbedContents` endpoint accepts at most 100 inputs per request (the same limit `langchain-google-genai` batches around), so calling `generate_embeddings` with more than 100 texts, e.g. when indexing a document set into a vector store, fails with a 400 instead of returning the embeddings. ### Description Split `texts` into chunks of at most 100 (`MAX_EMBEDDING_BATCH_SIZE`), call `embed_content` once per chunk and concatenate the embeddings in input order. The three client paths (user-supplied client, Vertex AI, API key) share the same helper. Requests with 100 or fewer texts are unchanged (a single call). The Vertex AI embedding service (`vertex_ai_text_embedding.py`) is left untouched: its limit was not verified. ### Tests New `test_raw_embedding_batches_large_inputs`: 250 texts produce calls of 100/100/50 and the results come back in order. It fails on main and passes with the fix; the rest of the file still passes. ### Contribution Checklist - [x] The code builds clean without any errors or warnings - [x] The PR follows the SK Contribution Guidelines - [x] All unit tests pass, and I have added new tests where possible This change was prepared with the assistance of an AI coding tool (Claude Code) and reviewed before submission.
1 parent 356acb6 commit 87eadc7

2 files changed

Lines changed: 41 additions & 13 deletions

File tree

‎python/semantic_kernel/connectors/ai/google/google_ai/services/google_ai_text_embedding.py‎

Lines changed: 21 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,9 @@
2222
else:
2323
from typing_extensions import override # pragma: no cover
2424

25+
# The Gemini batchEmbedContents endpoint accepts at most 100 inputs per request.
26+
MAX_EMBEDDING_BATCH_SIZE = 100
27+
2528

2629
class GoogleAITextEmbedding(GoogleAIBase, EmbeddingGeneratorBase):
2730
"""Google AI Text Embedding Service."""
@@ -119,31 +122,36 @@ async def generate_raw_embeddings(
119122
if not self.service_settings.embedding_model_id:
120123
raise ServiceInitializationError("The Google AI embedding model ID is required.")
121124

122-
async def _embed_content(client: Client) -> EmbedContentResponse:
125+
async def _embed_content(client: Client, contents: list[str]) -> EmbedContentResponse:
123126
return await client.aio.models.embed_content(
124127
model=self.service_settings.embedding_model_id, # type: ignore[arg-type]
125-
contents=texts, # type: ignore[arg-type]
128+
contents=contents, # type: ignore[arg-type]
126129
config=EmbedContentConfigDict(output_dimensionality=settings.output_dimensionality),
127130
)
128131

132+
async def _embed_all(client: Client) -> list[list[float]]:
133+
# The batch embedding endpoint accepts at most MAX_EMBEDDING_BATCH_SIZE inputs per request.
134+
embeddings: list[list[float]] = []
135+
for start in range(0, len(texts), MAX_EMBEDDING_BATCH_SIZE):
136+
response = await _embed_content(client, texts[start : start + MAX_EMBEDDING_BATCH_SIZE])
137+
embeddings.extend(embedding.values for embedding in response.embeddings) # type: ignore
138+
return embeddings
139+
129140
if self.client:
130-
response: EmbedContentResponse = await _embed_content(self.client)
131-
elif self.service_settings.use_vertexai:
141+
return await _embed_all(self.client)
142+
if self.service_settings.use_vertexai:
132143
with Client(
133144
vertexai=True,
134145
project=self.service_settings.cloud_project_id,
135146
location=self.service_settings.cloud_region,
136147
http_options=self._get_http_options(),
137148
) as client:
138-
response: EmbedContentResponse = await _embed_content(client) # type: ignore[no-redef]
139-
else:
140-
with Client(
141-
api_key=self.service_settings.api_key.get_secret_value(), # type: ignore[union-attr]
142-
http_options=self._get_http_options(),
143-
) as client:
144-
response: EmbedContentResponse = await _embed_content(client) # type: ignore[no-redef]
145-
146-
return [embedding.values for embedding in response.embeddings] # type: ignore
149+
return await _embed_all(client)
150+
with Client(
151+
api_key=self.service_settings.api_key.get_secret_value(), # type: ignore[union-attr]
152+
http_options=self._get_http_options(),
153+
) as client:
154+
return await _embed_all(client)
147155

148156
@override
149157
def get_prompt_execution_settings_class(

‎python/tests/unit/connectors/ai/google/google_ai/services/test_google_ai_text_embedding.py‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -235,3 +235,23 @@ async def test_embedding_with_custom_client(google_ai_unit_test_env, prompt) ->
235235

236236
# Verify that the custom client's method was called
237237
mock_embed_content.assert_called_once()
238+
239+
240+
@patch.object(AsyncModels, "embed_content", new_callable=AsyncMock)
241+
async def test_raw_embedding_batches_large_inputs(mock_google_model_embed_content, google_ai_unit_test_env):
242+
"""Inputs above the 100-per-request limit are split into several requests and returned in order."""
243+
model_id = google_ai_unit_test_env["GOOGLE_AI_EMBEDDING_MODEL_ID"]
244+
texts = [f"text {i}" for i in range(250)]
245+
246+
async def _fake_embed(*, model, contents, config):
247+
assert len(contents) <= 100
248+
return EmbedContentResponse(embeddings=[ContentEmbedding(values=[float(t.split()[1])]) for t in contents])
249+
250+
mock_google_model_embed_content.side_effect = _fake_embed
251+
settings = GoogleAIEmbeddingPromptExecutionSettings()
252+
253+
response = await GoogleAITextEmbedding().generate_raw_embeddings(texts, settings=settings)
254+
255+
assert [len(call.kwargs["contents"]) for call in mock_google_model_embed_content.call_args_list] == [100, 100, 50]
256+
assert all(call.kwargs["model"] == model_id for call in mock_google_model_embed_content.call_args_list)
257+
assert response == [[float(i)] for i in range(250)]

0 commit comments

Comments
 (0)