Merge fda1156952 into 9677f3d80c

src/embedding_client.py

@@ -218,8 +218,16 @@ class _EmbeddingClient:
                 api_key=config.api_key,
                 http_options=http_options,
             )
-            # Gemini has a 2048 token limit
-            self.max_embedding_tokens: int = min(max_input_tokens, 2048)
+            # Gemini's embedding models have model-specific input token caps
+            # (shared across all modalities):
+            #   - gemini-embedding-001: 2048 tokens
+            #   - gemini-embedding-2: 8192 tokens
+            # Unknown models default conservatively to 2048.
+            model_id = self.model.removeprefix("models/")
+            gemini_model_token_cap = 8192 if model_id == "gemini-embedding-2" else 2048
+            self.max_embedding_tokens: int = min(
+                max_input_tokens, gemini_model_token_cap
+            )
             # Gemini batch size is not documented, using conservative estimate
             self.max_batch_size: int = config.max_batch_size or 100
         else:  # openai

tests/llm/test_embedding_client.py

@@ -139,6 +139,164 @@ async def test_openai_embedding_client_rejects_dimension_mismatch(
         await client.embed("hello world")
 
 
+def test_gemini_embedding_client_gemini_001_caps_at_2048(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """gemini-embedding-001 should cap max_embedding_tokens at 2048."""
+
+    class FakeGeminiClient:
+        def __init__(self, *, api_key: str, http_options: Any) -> None:
+            self.api_key: str = api_key
+
+    monkeypatch.setattr("src.embedding_client.genai.Client", FakeGeminiClient)
+
+    # When max_input_tokens is above the model cap, it should be clamped to 2048.
+    client_above_cap = _EmbeddingClient(
+        EmbeddingModelConfig(
+            transport="gemini",
+            model="gemini-embedding-001",
+            api_key="test-key",
+        ),
+        vector_dimensions=8,
+        max_input_tokens=20_000,
+        max_tokens_per_request=300_000,
+        send_dimensions=True,
+    )
+    assert client_above_cap.max_embedding_tokens == 2048
+
+    # When max_input_tokens is below the model cap, it should pass through unchanged.
+    client_below_cap = _EmbeddingClient(
+        EmbeddingModelConfig(
+            transport="gemini",
+            model="gemini-embedding-001",
+            api_key="test-key",
+        ),
+        vector_dimensions=8,
+        max_input_tokens=1024,
+        max_tokens_per_request=300_000,
+        send_dimensions=True,
+    )
+    assert client_below_cap.max_embedding_tokens == 1024
+
+
+def test_gemini_embedding_client_gemini_2_caps_at_8192(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """gemini-embedding-2 should cap max_embedding_tokens at 8192."""
+
+    class FakeGeminiClient:
+        def __init__(self, *, api_key: str, http_options: Any) -> None:
+            self.api_key: str = api_key
+
+    monkeypatch.setattr("src.embedding_client.genai.Client", FakeGeminiClient)
+
+    # When max_input_tokens is above the model cap, it should be clamped to 8192.
+    client_above_cap = _EmbeddingClient(
+        EmbeddingModelConfig(
+            transport="gemini",
+            model="gemini-embedding-2",
+            api_key="test-key",
+        ),
+        vector_dimensions=8,
+        max_input_tokens=20_000,
+        max_tokens_per_request=300_000,
+        send_dimensions=True,
+    )
+    assert client_above_cap.max_embedding_tokens == 8192
+
+    # When max_input_tokens is below the model cap, it should pass through unchanged.
+    client_below_cap = _EmbeddingClient(
+        EmbeddingModelConfig(
+            transport="gemini",
+            model="gemini-embedding-2",
+            api_key="test-key",
+        ),
+        vector_dimensions=8,
+        max_input_tokens=4096,
+        max_tokens_per_request=300_000,
+        send_dimensions=True,
+    )
+    assert client_below_cap.max_embedding_tokens == 4096
+
+
+def test_gemini_embedding_client_models_prefixed_gemini_2_caps_at_8192(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """The canonical 'models/gemini-embedding-2' form (as returned by Gemini's
+    API) must still be recognized and granted the 8192 cap."""
+
+    class FakeGeminiClient:
+        def __init__(self, *, api_key: str, http_options: Any) -> None:
+            self.api_key: str = api_key
+
+    monkeypatch.setattr("src.embedding_client.genai.Client", FakeGeminiClient)
+
+    client = _EmbeddingClient(
+        EmbeddingModelConfig(
+            transport="gemini",
+            model="models/gemini-embedding-2",
+            api_key="test-key",
+        ),
+        vector_dimensions=8,
+        max_input_tokens=20_000,
+        max_tokens_per_request=300_000,
+        send_dimensions=True,
+    )
+    assert client.max_embedding_tokens == 8192
+
+
+def test_gemini_embedding_client_unknown_model_defaults_to_2048(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Unknown Gemini model names should conservatively default to a 2048 cap."""
+
+    class FakeGeminiClient:
+        def __init__(self, *, api_key: str, http_options: Any) -> None:
+            self.api_key: str = api_key
+
+    monkeypatch.setattr("src.embedding_client.genai.Client", FakeGeminiClient)
+
+    client = _EmbeddingClient(
+        EmbeddingModelConfig(
+            transport="gemini",
+            model="gemini-embedding-unknown-future-model",
+            api_key="test-key",
+        ),
+        vector_dimensions=8,
+        max_input_tokens=20_000,
+        max_tokens_per_request=300_000,
+        send_dimensions=True,
+    )
+    assert client.max_embedding_tokens == 2048
+
+
+def test_gemini_embedding_client_near_miss_model_id_defaults_to_2048(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """A model id that merely contains 'gemini-embedding-2' as a substring
+    (e.g. 'gemini-embedding-20') must not be granted the 8192 cap reserved
+    for the exact 'gemini-embedding-2' model id."""
+
+    class FakeGeminiClient:
+        def __init__(self, *, api_key: str, http_options: Any) -> None:
+            self.api_key: str = api_key
+
+    monkeypatch.setattr("src.embedding_client.genai.Client", FakeGeminiClient)
+
+    client = _EmbeddingClient(
+        EmbeddingModelConfig(
+            transport="gemini",
+            model="gemini-embedding-20",
+            api_key="test-key",
+        ),
+        vector_dimensions=8,
+        max_input_tokens=20_000,
+        max_tokens_per_request=300_000,
+        send_dimensions=True,
+    )
+    assert client.max_embedding_tokens == 2048
+
+
 @pytest.mark.asyncio
 async def test_gemini_embedding_client_uses_output_dimensionality(
     monkeypatch: pytest.MonkeyPatch,