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 <hemminger@hotmail.com>
This commit is contained in:
Corey Hemminger 2026-07-22 09:06:10 -05:00
parent 4f9a41360a
commit 6620a3c68f
2 changed files with 42 additions and 2 deletions

View File

@ -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

View File

@ -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,