diff --git a/python/semantic_kernel/connectors/ai/google/google_ai/services/google_ai_text_embedding.py b/python/semantic_kernel/connectors/ai/google/google_ai/services/google_ai_text_embedding.py index 4e74830e960f..8b8613b0386e 100644 --- a/python/semantic_kernel/connectors/ai/google/google_ai/services/google_ai_text_embedding.py +++ b/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( diff --git a/python/tests/unit/connectors/ai/google/google_ai/services/test_google_ai_text_embedding.py b/python/tests/unit/connectors/ai/google/google_ai/services/test_google_ai_text_embedding.py index 9b734d145f4a..5b1ab28fdafd 100644 --- a/python/tests/unit/connectors/ai/google/google_ai/services/test_google_ai_text_embedding.py +++ b/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)]