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:
parent
4f9a41360a
commit
6620a3c68f
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
Loading…
Reference in New Issue