diff --git a/hindsight-api-slim/hindsight_api/engine/embeddings.py b/hindsight-api-slim/hindsight_api/engine/embeddings.py index d4e01cd5af..293817ac36 100644 --- a/hindsight-api-slim/hindsight_api/engine/embeddings.py +++ b/hindsight-api-slim/hindsight_api/engine/embeddings.py @@ -2019,7 +2019,9 @@ def encode(self, texts: list[str]) -> list[list[float]]: # batches too, which is why RetryBudget takes a lock. budget = self.retry_policy.new_budget() - all_embeddings = self._encode_batched(texts, lambda batch: self._embed_batch(batch, budget)) + all_embeddings = self._encode_batched( + texts, lambda batch: self._embed_batch(batch, budget), batch_size=self._effective_batch_size() + ) # L2-normalize when output_dimensionality is set — Gemini only returns # normalized vectors at full 3072 dims; truncated dims need re-normalization @@ -2034,6 +2036,30 @@ def encode(self, texts: list[str]) -> list[list[float]]: return all_embeddings + def _effective_batch_size(self) -> int: + """How many texts may share one ``embed_content`` call. + + ``batch_size`` everywhere except the Vertex models that take exactly one + Content per request: on Vertex the SDK routes every embedding model whose + name contains ``gemini`` — bar ``gemini-embedding-001`` — and every ``maas`` + model to the single-content ``embedContent`` endpoint, and raises + ``ValueError("The embedContent API for this model only supports one content + at a time.")`` client-side for anything longer. We cannot batch around that: + each text has to be its own Content to come back as its own vector (#4001), + so for these models one request per text is the only shape that returns the + 1:1 alignment ``_embed_batch`` asserts. The Gemini API (non-Vertex) path has + no such limit and keeps the configured batch size. + + Mirrored from ``google.genai._transformers.t_is_vertex_embed_content_model`` + rather than imported: it is private, and a copy that drifts fails loudly + here (the SDK raises) instead of silently sending batches that never worked. + """ + if not self._is_vertexai: + return self.batch_size + model = self.model.removeprefix("google/") + single_content_only = ("gemini" in model and model != "gemini-embedding-001") or "maas" in model + return 1 if single_content_only else self.batch_size + def _embed_batch(self, batch: list[str], budget: "RetryBudget") -> list[list[float]]: """Embed one batch-sized slice through the google.genai sync client.""" from google.genai import types as genai_types diff --git a/hindsight-api-slim/tests/test_gemini_embeddings.py b/hindsight-api-slim/tests/test_gemini_embeddings.py index 17bff85331..c15eba207e 100644 --- a/hindsight-api-slim/tests/test_gemini_embeddings.py +++ b/hindsight-api-slim/tests/test_gemini_embeddings.py @@ -110,6 +110,57 @@ async def test_initialization_vertexai_success(self): location="us-central1", ) + @pytest.mark.parametrize( + "model,vertexai,expected_requests", + [ + # Vertex routes these to the single-content embedContent endpoint, which + # rejects a second Content client-side — one request per text is the only + # shape that comes back 1:1. + ("gemini-embedding-2-preview", True, 3), + ("google/gemini-embedding-2-preview", True, 3), + ("text-multilingual-maas-002", True, 3), + # The one Vertex gemini model the SDK still batches, and everything on the + # Gemini API, keep the configured batch size: one request for all three. + ("gemini-embedding-001", True, 1), + ("text-embedding-005", True, 1), + ("gemini-embedding-2-preview", False, 1), + ], + ) + async def test_vertex_single_content_models_get_one_request_per_text(self, model, vertexai, expected_requests): + """Batching is capped at one text where the API takes one Content (#4001 follow-up). + + Each text is already sent as its own Content, which is what keeps a + multimodal model from fusing a batch into a single vector. On Vertex the + SDK then refuses more than one Content for these models outright — + ``ValueError: The embedContent API for this model only supports one content + at a time.`` — so every encode() against gemini-embedding-2 raised before + reaching the network. Counting requests rather than asserting on the + exception is deliberate: the fix is the request shape, and a regression + would show up here as three texts back in one call. + """ + mock_genai = _make_mock_genai() + embed_content = mock_genai.Client.return_value.models.embed_content + emb = GeminiEmbeddings( + model=model, + api_key=None if vertexai else "test-key", + vertexai_project_id="test-project" if vertexai else None, + ) + with _patch_google_import(mock_genai): + await emb.initialize() + + texts = ["a", "b", "c"] + # One vector per text in the batch the call actually carries, so the 1:1 + # check inside _embed_batch passes for either shape. + embed_content.side_effect = lambda **kw: _make_mock_embed_result([[0.1] * 768] * len(kw["contents"])) + embed_content.reset_mock() + vectors = emb.encode(texts) + + assert len(vectors) == len(texts) + assert embed_content.call_count == expected_requests + assert [len(c.kwargs["contents"]) for c in embed_content.call_args_list] == ( + [1, 1, 1] if expected_requests == 3 else [3] + ) + async def test_initialization_missing_api_key(self): """Test that missing API key raises ValueError when no vertexai_project_id.""" mock_genai = _make_mock_genai()