From fda1156952f8eb87040078d719cbca965a871ff3 Mon Sep 17 00:00:00 2001
From: Corey Hemminger <hemminger@hotmail.com>
Date: Wed, 22 Jul 2026 12:28:48 -0500
Subject: [PATCH] fix: normalize models/ prefix before exact-match model cap
 comparison

Signed-off-by: Corey Hemminger <hemminger@hotmail.com>
---
 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: