Python: Batch Google AI embedding requests to the 100-input limit

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

@@ -22,6 +22,9 @@
 else:
     from typing_extensions import override  # pragma: no cover
 
+# The Gemini batchEmbedContents endpoint accepts at most 100 inputs per request.
+MAX_EMBEDDING_BATCH_SIZE = 100
+
 
 class GoogleAITextEmbedding(GoogleAIBase, EmbeddingGeneratorBase):
     """Google AI Text Embedding Service."""
@@ -119,31 +122,36 @@ async def generate_raw_embeddings(
         if not self.service_settings.embedding_model_id:
             raise ServiceInitializationError("The Google AI embedding model ID is required.")
 
-        async def _embed_content(client: Client) -> EmbedContentResponse:
+        async def _embed_content(client: Client, contents: list[str]) -> EmbedContentResponse:
             return await client.aio.models.embed_content(
                 model=self.service_settings.embedding_model_id,  # type: ignore[arg-type]
-                contents=texts,  # type: ignore[arg-type]
+                contents=contents,  # type: ignore[arg-type]
                 config=EmbedContentConfigDict(output_dimensionality=settings.output_dimensionality),
             )
 
+        async def _embed_all(client: Client) -> list[list[float]]:
+            # The batch embedding endpoint accepts at most MAX_EMBEDDING_BATCH_SIZE inputs per request.
+            embeddings: list[list[float]] = []
+            for start in range(0, len(texts), MAX_EMBEDDING_BATCH_SIZE):
+                response = await _embed_content(client, texts[start : start + MAX_EMBEDDING_BATCH_SIZE])
+                embeddings.extend(embedding.values for embedding in response.embeddings)  # type: ignore
+            return embeddings
+
         if self.client:
-            response: EmbedContentResponse = await _embed_content(self.client)
-        elif self.service_settings.use_vertexai:
+            return await _embed_all(self.client)
+        if self.service_settings.use_vertexai:
             with Client(
                 vertexai=True,
                 project=self.service_settings.cloud_project_id,
                 location=self.service_settings.cloud_region,
                 http_options=self._get_http_options(),
             ) as client:
-                response: EmbedContentResponse = await _embed_content(client)  # type: ignore[no-redef]
-        else:
-            with Client(
-                api_key=self.service_settings.api_key.get_secret_value(),  # type: ignore[union-attr]
-                http_options=self._get_http_options(),
-            ) as client:
-                response: EmbedContentResponse = await _embed_content(client)  # type: ignore[no-redef]
-
-        return [embedding.values for embedding in response.embeddings]  # type: ignore
+                return await _embed_all(client)
+        with Client(
+            api_key=self.service_settings.api_key.get_secret_value(),  # type: ignore[union-attr]
+            http_options=self._get_http_options(),
+        ) as client:
+            return await _embed_all(client)
 
     @override
     def get_prompt_execution_settings_class(

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

@@ -235,3 +235,23 @@ async def test_embedding_with_custom_client(google_ai_unit_test_env, prompt) ->
 
     # Verify that the custom client's method was called
     mock_embed_content.assert_called_once()
+
+
+@patch.object(AsyncModels, "embed_content", new_callable=AsyncMock)
+async def test_raw_embedding_batches_large_inputs(mock_google_model_embed_content, google_ai_unit_test_env):
+    """Inputs above the 100-per-request limit are split into several requests and returned in order."""
+    model_id = google_ai_unit_test_env["GOOGLE_AI_EMBEDDING_MODEL_ID"]
+    texts = [f"text {i}" for i in range(250)]
+
+    async def _fake_embed(*, model, contents, config):
+        assert len(contents) <= 100
+        return EmbedContentResponse(embeddings=[ContentEmbedding(values=[float(t.split()[1])]) for t in contents])
+
+    mock_google_model_embed_content.side_effect = _fake_embed
+    settings = GoogleAIEmbeddingPromptExecutionSettings()
+
+    response = await GoogleAITextEmbedding().generate_raw_embeddings(texts, settings=settings)
+
+    assert [len(call.kwargs["contents"]) for call in mock_google_model_embed_content.call_args_list] == [100, 100, 50]
+    assert all(call.kwargs["model"] == model_id for call in mock_google_model_embed_content.call_args_list)
+    assert response == [[float(i)] for i in range(250)]