From 6620a3c68fb9d1f894385c02a09e7172cee3f049 Mon Sep 17 00:00:00 2001 From: Corey Hemminger Date: Wed, 22 Jul 2026 09:06:10 -0500 Subject: [PATCH 1/4] fix(embedding): raise Gemini max_embedding_tokens cap from 2048 to 8192 Gemini's embedding models actually support up to 8192 input tokens, not 2048. The artificial 2048 cap was truncating/chunking text more aggressively than necessary for Gemini-backed embedding configs. Added a unit test asserting the client honors the 8192 cap (and still respects a lower max_input_tokens when configured below the cap). Signed-off-by: Corey Hemminger --- src/embedding_client.py | 4 +-- tests/llm/test_embedding_client.py | 40 ++++++++++++++++++++++++++++++ 2 files changed, 42 insertions(+), 2 deletions(-) diff --git a/src/embedding_client.py b/src/embedding_client.py index 07197f43..c319fd53 100644 --- a/src/embedding_client.py +++ b/src/embedding_client.py @@ -191,8 +191,8 @@ 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 support up to 8192 input tokens. + self.max_embedding_tokens: int = min(max_input_tokens, 8192) # Gemini batch size is not documented, using conservative estimate self.max_batch_size: int = 100 else: # openai diff --git a/tests/llm/test_embedding_client.py b/tests/llm/test_embedding_client.py index a4642159..ed2751d3 100644 --- a/tests/llm/test_embedding_client.py +++ b/tests/llm/test_embedding_client.py @@ -92,6 +92,46 @@ async def test_openai_embedding_client_rejects_dimension_mismatch( await client.embed("hello world") +def test_gemini_embedding_client_honors_8192_token_cap( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Gemini's max_embedding_tokens should be capped at 8192, not 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 provider cap, it should be clamped to 8192. + 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 == 8192 + + # When max_input_tokens is below the provider 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 + + @pytest.mark.asyncio async def test_gemini_embedding_client_uses_output_dimensionality( monkeypatch: pytest.MonkeyPatch, From 0eec91f26c0dab209c071f137b2ff5a31b013e3b Mon Sep 17 00:00:00 2001 From: Corey Hemminger Date: Wed, 22 Jul 2026 11:38:36 -0500 Subject: [PATCH 2/4] fix: make Gemini embedding token limit model-aware (gemini-embedding-001: 2048, gemini-embedding-2: 8192) Signed-off-by: Corey Hemminger --- src/embedding_client.py | 16 ++++++- tests/llm/test_embedding_client.py | 75 ++++++++++++++++++++++++++++-- 2 files changed, 84 insertions(+), 7 deletions(-) diff --git a/src/embedding_client.py b/src/embedding_client.py index c319fd53..c75ed345 100644 --- a/src/embedding_client.py +++ b/src/embedding_client.py @@ -191,8 +191,20 @@ class _EmbeddingClient: api_key=config.api_key, http_options=http_options, ) - # Gemini's embedding models support up to 8192 input tokens. - self.max_embedding_tokens: int = min(max_input_tokens, 8192) + # 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. + if "gemini-embedding-2" in self.model: + gemini_model_token_cap = 8192 + elif "gemini-embedding-001" in self.model: + gemini_model_token_cap = 2048 + else: + gemini_model_token_cap = 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 = 100 else: # openai diff --git a/tests/llm/test_embedding_client.py b/tests/llm/test_embedding_client.py index ed2751d3..ce156e4c 100644 --- a/tests/llm/test_embedding_client.py +++ b/tests/llm/test_embedding_client.py @@ -92,10 +92,10 @@ async def test_openai_embedding_client_rejects_dimension_mismatch( await client.embed("hello world") -def test_gemini_embedding_client_honors_8192_token_cap( +def test_gemini_embedding_client_gemini_001_caps_at_2048( monkeypatch: pytest.MonkeyPatch, ) -> None: - """Gemini's max_embedding_tokens should be capped at 8192, not 2048.""" + """gemini-embedding-001 should cap max_embedding_tokens at 2048.""" class FakeGeminiClient: def __init__(self, *, api_key: str, http_options: Any) -> None: @@ -103,7 +103,7 @@ def test_gemini_embedding_client_honors_8192_token_cap( monkeypatch.setattr("src.embedding_client.genai.Client", FakeGeminiClient) - # When max_input_tokens is above the provider cap, it should be clamped to 8192. + # When max_input_tokens is above the model cap, it should be clamped to 2048. client_above_cap = _EmbeddingClient( EmbeddingModelConfig( transport="gemini", @@ -115,9 +115,9 @@ def test_gemini_embedding_client_honors_8192_token_cap( max_tokens_per_request=300_000, send_dimensions=True, ) - assert client_above_cap.max_embedding_tokens == 8192 + assert client_above_cap.max_embedding_tokens == 2048 - # When max_input_tokens is below the provider cap, it should pass through unchanged. + # When max_input_tokens is below the model cap, it should pass through unchanged. client_below_cap = _EmbeddingClient( EmbeddingModelConfig( transport="gemini", @@ -132,6 +132,71 @@ def test_gemini_embedding_client_honors_8192_token_cap( 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_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 + + @pytest.mark.asyncio async def test_gemini_embedding_client_uses_output_dimensionality( monkeypatch: pytest.MonkeyPatch, From ef3a50ba24aba887244d6f3d70bb5c7477d3f06c Mon Sep 17 00:00:00 2001 From: Corey Hemminger Date: Wed, 22 Jul 2026 12:25:45 -0500 Subject: [PATCH 3/4] fix: use exact match instead of substring for gemini-embedding-2 model cap Signed-off-by: Corey Hemminger --- src/embedding_client.py | 4 +--- tests/llm/test_embedding_client.py | 27 +++++++++++++++++++++++++++ 2 files changed, 28 insertions(+), 3 deletions(-) diff --git a/src/embedding_client.py b/src/embedding_client.py index c75ed345..8b6493d2 100644 --- a/src/embedding_client.py +++ b/src/embedding_client.py @@ -196,10 +196,8 @@ class _EmbeddingClient: # - gemini-embedding-001: 2048 tokens # - gemini-embedding-2: 8192 tokens # Unknown models default conservatively to 2048. - if "gemini-embedding-2" in self.model: + if self.model == "gemini-embedding-2": gemini_model_token_cap = 8192 - elif "gemini-embedding-001" in self.model: - gemini_model_token_cap = 2048 else: gemini_model_token_cap = 2048 self.max_embedding_tokens: int = min( diff --git a/tests/llm/test_embedding_client.py b/tests/llm/test_embedding_client.py index ce156e4c..fe516102 100644 --- a/tests/llm/test_embedding_client.py +++ b/tests/llm/test_embedding_client.py @@ -197,6 +197,33 @@ def test_gemini_embedding_client_unknown_model_defaults_to_2048( 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, From fda1156952f8eb87040078d719cbca965a871ff3 Mon Sep 17 00:00:00 2001 From: Corey Hemminger Date: Wed, 22 Jul 2026 12:28:48 -0500 Subject: [PATCH 4/4] fix: normalize models/ prefix before exact-match model cap comparison Signed-off-by: Corey Hemminger --- src/embedding_client.py | 6 ++---- tests/llm/test_embedding_client.py | 26 ++++++++++++++++++++++++++ 2 files changed, 28 insertions(+), 4 deletions(-) diff --git a/src/embedding_client.py b/src/embedding_client.py index 8b6493d2..d7812602 100644 --- a/src/embedding_client.py +++ b/src/embedding_client.py @@ -196,10 +196,8 @@ class _EmbeddingClient: # - gemini-embedding-001: 2048 tokens # - gemini-embedding-2: 8192 tokens # Unknown models default conservatively to 2048. - if self.model == "gemini-embedding-2": - gemini_model_token_cap = 8192 - else: - gemini_model_token_cap = 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 ) diff --git a/tests/llm/test_embedding_client.py b/tests/llm/test_embedding_client.py index fe516102..a53188be 100644 --- a/tests/llm/test_embedding_client.py +++ b/tests/llm/test_embedding_client.py @@ -172,6 +172,32 @@ def test_gemini_embedding_client_gemini_2_caps_at_8192( 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: