Merge bc1ba3b237 into aa7917a3d2

litellm/router_config.py

@@ -45,6 +45,7 @@ class RouterDynamicConfig:
     context_window_fallbacks: Sequence[object] | None
     content_policy_fallbacks: Sequence[object] | None
     drop_params: bool | None
+    disable_token_counter: bool
     provider_list: Sequence[str]
     expose_router_debug_in_errors: bool
     redact_user_api_key_info: bool | None
@@ -102,6 +103,7 @@ def current_dynamic_config() -> RouterDynamicConfig:
         context_window_fallbacks=litellm.context_window_fallbacks,
         content_policy_fallbacks=litellm.content_policy_fallbacks,
         drop_params=litellm.drop_params,
+        disable_token_counter=litellm.disable_token_counter,
         provider_list=litellm.provider_list,
         expose_router_debug_in_errors=litellm.expose_router_debug_in_errors,
         redact_user_api_key_info=litellm.redact_user_api_key_info,
@@ -130,6 +132,7 @@ _LIVE_SOURCES: Final = MappingProxyType(
         "context_window_fallbacks": "context_window_fallbacks",
         "content_policy_fallbacks": "content_policy_fallbacks",
         "drop_params": "drop_params",
+        "disable_token_counter": "disable_token_counter",
         "provider_list": "provider_list",
         "expose_router_debug_in_errors": "expose_router_debug_in_errors",
         "redact_user_api_key_info": "redact_user_api_key_info",

litellm/router_python_core.py

@@ -698,6 +698,7 @@ _MODEL_INFO_ADAPTER: Final = TypeAdapter(Mapping[str, object])
 _CALL_KWARGS_ADAPTER: Final = TypeAdapter(Mapping[str, object])
 _LITELLM_PARAMS_ADAPTER: Final = TypeAdapter(Mapping[str, object])
 _SILENT_MODEL_ADAPTER: Final = TypeAdapter(str | list[str])
+_EMBEDDING_TOKEN_ID_BATCH_ADAPTER: Final = TypeAdapter(list[list[int]])
 _RESOLVED_RETRY_POLICY_ADAPTER: Final = TypeAdapter(RetryPolicy | None)
 _ROUTING_KWARGS_ADAPTER: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(Mapping[str, object] | None)
 _FALLBACK_HOP_ADAPTER: Final = TypeAdapter(Mapping[str, object])
@@ -15113,6 +15114,9 @@ class PythonCore(DeprecatedRouterMembers):
         Prompt content the message list never carries is read from `request_kwargs`:
         `tools` (Chat Completions, Responses and Anthropic Messages shapes) and the
         Anthropic Messages top-level `system` block.
+
+        An embedding batch (a list of texts, token IDs or token ID lists) counts as its
+        longest input. Media the deployments embed as media is left out of that count.
         """
         from litellm.llms.anthropic.messages.native.utils import (
             anthropic_system_to_openai_message,
@@ -15140,6 +15144,11 @@ class PythonCore(DeprecatedRouterMembers):
             )
 
             text_input, _ = split_embedding_input_tokens(input, partial(deployment_backend_models, deployments))
+            if (
+                isinstance(text_input, list)
+                and (embedding_tokens := self._longest_embedding_input_tokens(text_input, tokenizer_model)) is not None
+            ):
+                return embedding_tokens
             typed_input: Final = cast(  # cast-ok: str | list matches transform input
                 str | ResponseInputParam, text_input
             )
@@ -15155,6 +15164,22 @@ class PythonCore(DeprecatedRouterMembers):
             )
         raise ValueError("Either messages or input must be provided to count tokens")
 
+    def _longest_embedding_input_tokens(self, items: Sequence[object], tokenizer_model: str) -> int | None:
+        """An embedding model's context window bounds each input on its own, so a batch is never summed."""
+        if not items:
+            return None
+        if all(isinstance(item, str) for item in items):
+            count_tokens: Final = self._host.invoke("token_counter")
+            return max(count_tokens(model=tokenizer_model, text=item) for item in items)
+        counting_disabled: Final = self._host.dynamic().disable_token_counter
+        if all(isinstance(item, int) and not isinstance(item, bool) for item in items):
+            return 0 if counting_disabled else len(items)
+        try:
+            token_id_batch: Final = _EMBEDDING_TOKEN_ID_BATCH_ADAPTER.validate_python(items, strict=True)
+        except ValidationError:
+            return None
+        return 0 if counting_disabled else max(len(token_ids) for token_ids in token_id_batch)
+
     def _deployment_max_input_tokens(self, model: str, deployment: Mapping[str, object]) -> int | None:
         """The deployment's declared context window, or None when it declares none or cannot be resolved."""
         try:

tests/unit/test_router/core/scripted_host.py

@@ -167,6 +167,7 @@ def dynamic_config(**changes: object) -> RouterDynamicConfig:  # kwargs-ok: Rout
         context_window_fallbacks=None,
         content_policy_fallbacks=None,
         drop_params=None,
+        disable_token_counter=False,
         provider_list=litellm.provider_list,
         expose_router_debug_in_errors=False,
         redact_user_api_key_info=None,

tests/unit/test_router/test_router.py

@@ -8519,6 +8519,35 @@ def test_context_window_pre_call_check_counts_an_embedding_image_as_text_on_a_te
         )
 
 
+
+@pytest.mark.parametrize(
+    ("backend", "caption", "fits"),
+    [
+        pytest.param("vertex_ai/gemini-embedding-2", "a red car", True, id="image-backend-short-caption"),
+        pytest.param("vertex_ai/gemini-embedding-2", "word " * 40, False, id="image-backend-oversized-caption"),
+        pytest.param("text-embedding-3-small", "a red car", False, id="text-only-backend-counts-the-image"),
+    ],
+)
+def test_context_window_pre_call_check_counts_each_caption_next_to_an_embedding_image(
+    backend: str, caption: str, fits: bool, caplog
+):
+    router: Final = _embedding_context_router(backend, window=20)
+    deployments: Final = router.registry().deployments("embedder")
+
+    with caplog.at_level(logging.ERROR, logger="LiteLLM Router"):
+        if fits:
+            kept: Final = router._pre_call_checks(
+                model="embedder", healthy_deployments=deployments, input=[caption, _EMBEDDING_IMAGE_INPUT]
+            )
+            assert len(kept) == 1
+        else:
+            with pytest.raises(litellm.ContextWindowExceededError):
+                router._pre_call_checks(
+                    model="embedder", healthy_deployments=deployments, input=[caption, _EMBEDDING_IMAGE_INPUT]
+                )
+
+    assert "failed to count tokens" not in caplog.text
+
 async def test_acount_pre_call_check_tokens_counts_a_group_of_one_openai_model_with_that_models_tokenizer():
     router: Final = _context_limited_router("openai/gpt-6.1-sol")
 
@@ -9191,6 +9220,120 @@ async def test_aresponses_enforces_context_window_pre_call_check():
         )
 
 
+def _embedding_router(max_input_tokens: int) -> litellm.Router:
+    return litellm.Router(
+        model_list=[
+            {
+                "model_name": "embed",
+                "litellm_params": {
+                    "model": "openai/text-embedding-3-small",
+                    "api_key": "sk-test",
+                    "mock_response": [0.1, 0.2],
+                },
+                "model_info": {"max_input_tokens": max_input_tokens},
+            }
+        ],
+        enable_pre_call_checks=True,
+    )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "embedding_input",
+    [
+        pytest.param(["one two three four five six"] * 3, id="texts"),
+        pytest.param([list(range(6))] * 3, id="token-id-lists"),
+        pytest.param(list(range(10)), id="token-ids-at-the-limit"),
+    ],
+)
+async def test_aembedding_accepts_a_batch_whose_inputs_each_fit_the_context_window(embedding_input, caplog):
+    router: Final = _embedding_router(max_input_tokens=10)
+
+    with caplog.at_level(logging.ERROR, logger="LiteLLM Router"):
+        response: Final = await router.aembedding(model="embed", input=embedding_input)
+
+    assert response.data
+    assert "failed to count tokens" not in caplog.text
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "embedding_input",
+    [
+        pytest.param(["one", " ".join(["word"] * 20)], id="texts"),
+        pytest.param(list(range(11)), id="token-ids"),
+        pytest.param([[1], list(range(11))], id="token-id-lists"),
+    ],
+)
+async def test_aembedding_rejects_a_batch_with_one_input_over_the_context_window(embedding_input):
+    router: Final = _embedding_router(max_input_tokens=10)
+
+    with pytest.raises(litellm.ContextWindowExceededError):
+        await router.aembedding(model="embed", input=embedding_input)
+
+
+def test_embedding_rejects_a_batch_with_one_input_over_the_context_window():
+    router: Final = _embedding_router(max_input_tokens=10)
+
+    with pytest.raises(litellm.ContextWindowExceededError):
+        router.embedding(model="embed", input=["one", " ".join(["word"] * 20)])
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "embedding_input",
+    [
+        pytest.param(["one", " ".join(["word"] * 20)], id="texts"),
+        pytest.param(list(range(11)), id="token-ids"),
+        pytest.param([[1], list(range(11))], id="token-id-lists"),
+    ],
+)
+async def test_aembedding_skips_the_context_window_check_when_token_counting_is_disabled(embedding_input, monkeypatch):
+    monkeypatch.setattr(litellm, "disable_token_counter", True)
+    router: Final = _embedding_router(max_input_tokens=10)
+
+    response: Final = await router.aembedding(model="embed", input=embedding_input)
+
+    assert response.data
+
+
+@pytest.mark.asyncio
+@pytest.mark.usefixtures("local_model_cost_map", "isolate_router_model_cost_state")
+@pytest.mark.parametrize(
+    "prompt_kwargs",
+    [
+        pytest.param({"instructions": "you are a helpful assistant. " * 20}, id="instructions"),
+        pytest.param(
+            {
+                "tools": [
+                    {
+                        "type": "function",
+                        "name": "lookup",
+                        "description": "look up a record by its identifier. " * 20,
+                        "parameters": {"type": "object", "properties": {"q": {"type": "string"}}},
+                    }
+                ]
+            },
+            id="tools",
+        ),
+    ],
+)
+async def test_aresponses_with_empty_input_still_counts_its_other_prompt_content(prompt_kwargs):
+    router: Final = litellm.Router(
+        model_list=[
+            {
+                "model_name": "small-ctx",
+                "litellm_params": {"model": "gpt-3.5-turbo", "mock_response": "hi"},
+                "model_info": {"max_input_tokens": 50},
+            }
+        ],
+        enable_pre_call_checks=True,
+    )
+
+    with pytest.raises(litellm.ContextWindowExceededError):
+        await router.aresponses(model="small-ctx", input=[], **prompt_kwargs)
+
+
 def test_get_deployment_model_info_base_model_flow():
     """Test that get_deployment_model_info correctly handles the base model flow"""
     from unittest.mock import patch