fix: normalize models/ prefix before exact-match model cap comparison

Signed-off-by: Corey Hemminger <hemminger@hotmail.com>
This commit is contained in:
Corey Hemminger 2026-07-22 12:28:48 -05:00
parent ef3a50ba24
commit fda1156952
2 changed files with 28 additions and 4 deletions

View File

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

View File

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