fix(router): keep embedding images out of the pre-call count when any backend in the group embeds images

Mark Cohere embed-english-light-v3.0 as embedding images, as its v3 siblings are.

litellm/litellm_core_utils/token_counter.py

@@ -740,15 +740,16 @@ def _embeds_media(model: str) -> bool:
 
 
 def embedding_text_input(value: _EmbeddingInputT, media_models: Callable[[], Sequence[str]]) -> _EmbeddingInputT:
-    """``value`` without its media when every backend in ``media_models`` declares image embedding
-    input. A text-only backend reads a data URI as text, so for it the media stays in."""
+    """``value`` without its media when a backend in ``media_models`` declares image embedding input.
+    The pre-call check holds one count for the whole group, so an image backend is never refused over
+    media it embeds, and only a group of text-only backends, which read a data URI as text, counts it."""
     from litellm.litellm_core_utils.prompt_templates.tpm_media_utils import split_embedding_media_for_tpm
 
     text, media = split_embedding_media_for_tpm(value)
     if not media:
         return value
     models: Final = tuple(media_models())
-    if not models or not all(map(_embeds_media, models)):
+    if not any(map(_embeds_media, models)):
         return value
     return cast(_EmbeddingInputT, text)  # cast-ok: the projection keeps the input's container shape
 

litellm/model_prices_and_context_window_backup.json

@@ -25620,7 +25620,8 @@
         "litellm_provider": "cohere",
         "max_input_tokens": 512,
         "mode": "embedding",
-        "output_cost_per_token": 0.0
+        "output_cost_per_token": 0.0,
+        "supports_embedding_image_input": true
     },
     "embed-english-v3.0": {
         "input_cost_per_image": 0.0001,

litellm/router_rust/steps/precall.py

@@ -63,7 +63,7 @@ def count_tokens(
 ) -> int:
     """``_count_pre_call_check_tokens``: messages, or Responses input normalized to messages, plus
     the tools and Anthropic ``system`` block the message list never carries. Embedding media is left
-    out of the input when every item of ``deployments`` declares image embedding input. An
+    out of the input when an item of ``deployments`` declares image embedding input. An
     embedding batch (a list of texts, token IDs or token ID lists) counts as its longest input, and
     a text of at most ``fits_within`` UTF-8 bytes counts as its byte length without the tokenizer:
     a token is never shorter than a byte, so it fits any window of that size."""

model_prices_and_context_window.json

@@ -25620,7 +25620,8 @@
         "litellm_provider": "cohere",
         "max_input_tokens": 512,
         "mode": "embedding",
-        "output_cost_per_token": 0.0
+        "output_cost_per_token": 0.0,
+        "supports_embedding_image_input": true
     },
     "embed-english-v3.0": {
         "input_cost_per_image": 0.0001,

tests/unit/test_router/core/test_construction_precall_characterization_precall.py

@@ -18,7 +18,6 @@ import pytest
 
 import litellm
 from litellm.router_core import RouterCore
-from litellm.router_python_host import Endpoint
 from litellm.types.router import RouterRateLimitError
 from tests.unit.test_router.core.construction_precall_helpers import CallbackHost, restoring_model_cost, windowed
 from tests.unit.test_router.core.scripted_host import ScriptedHost, deployment, rust_core
@@ -56,26 +55,6 @@ UNSUPPORTED_MESSAGE: Final = (
 )
 
 
-class TokenizerHost(CallbackHost):
-    """Records each text the pre-call checks hand the tokenizer."""
-
-    __slots__ = ("tokenized",)
-
-    def __init__(self) -> None:
-        super().__init__()
-        self.tokenized: Final[list[object]] = []
-
-    def invoke(self, op: str) -> Endpoint:
-        if op != "token_counter":
-            return super().invoke(op)
-
-        def token_counter(**kwargs: object) -> int:  # kwargs-ok: litellm.token_counter's arguments
-            self.tokenized.append(kwargs.get("text"))
-            return litellm.token_counter(**kwargs)  # pyright: ignore[reportArgumentType]
-
-        return token_counter
-
-
 @pytest.fixture
 def host() -> ScriptedHost:
     return CallbackHost()
@@ -305,6 +284,7 @@ async def test_an_embedding_image_stays_out_of_the_context_window_on_a_backend_t
     "backend",
     [
         "bedrock/us.twelvelabs.marengo-embed-2-7-v1:0",
+        "cohere/embed-english-light-v3.0",
         "oci/cohere.embed-english-image-v3.0",
         "empiriolabs/tongyi-embedding-vision-plus",
         "voyage/voyage-multimodal-3.5",
@@ -324,19 +304,27 @@ async def test_an_embedding_image_stays_out_of_the_context_window_on_every_backe
 
 
 @pytest.mark.asyncio
-async def test_an_embedding_batch_tokenizes_only_the_inputs_longer_than_the_smallest_window_in_bytes() -> None:
-    host: Final = TokenizerHost()
-    core: Final = rust_core(
-        host, model_list=[windowed("g", "small", 50), windowed("g", "big", 10000)], enable_pre_call_checks=True
+async def test_an_embedding_image_stays_out_of_the_context_window_when_one_backend_in_the_group_embeds_it(
+    make_core: MakeCore,
+) -> None:
+    assert litellm.token_counter(text=EMBEDDING_IMAGE) > 100
+    core: Final = make_core(
+        model_list=[
+            deployment("g", "image", model="vertex_ai/gemini-embedding-2", model_info={"max_input_tokens": 100}),
+            deployment("g", "text", model="openai/text-embedding-3-small", model_info={"max_input_tokens": 100}),
+        ],
+        enable_pre_call_checks=True,
     )
-    at_the_window: Final = "word " * 10
-    over_it_in_bytes_only: Final = "word " * 20
-    assert len(at_the_window.encode()) == 50
+    assert await _picked(core, input=["a red car", EMBEDDING_IMAGE]) == {"image", "text"}
 
-    picked: Final = await _picked(core, input=[at_the_window] * 2048 + [over_it_in_bytes_only])
 
-    assert picked == {"small", "big"}
-    assert set(host.tokenized) == {over_it_in_bytes_only}
+@pytest.mark.asyncio
+async def test_an_embedding_input_over_the_smaller_window_in_bytes_but_not_in_tokens_fits_it(
+    make_core: MakeCore,
+) -> None:
+    short_words: Final = "word " * 8
+    assert litellm.token_counter(text=short_words) <= 10 < len(short_words.encode())
+    assert await _picked(_small_and_big(make_core, small=10), input=["hi", short_words]) == {"small", "big"}
 
 
 @pytest.mark.asyncio