diff --git a/litellm/router_config.py b/litellm/router_config.py index 36154af296c..1d38fe84d96 100644 --- a/litellm/router_config.py +++ b/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", diff --git a/litellm/router_python_core.py b/litellm/router_python_core.py index 8ca2afb9a99..7b4833401e5 100644 --- a/litellm/router_python_core.py +++ b/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: diff --git a/tests/unit/test_router/core/scripted_host.py b/tests/unit/test_router/core/scripted_host.py index f673ab98610..2bb934a3f58 100644 --- a/tests/unit/test_router/core/scripted_host.py +++ b/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, diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index a21777fbee6..eebb8f8bcba 100644 --- a/tests/unit/test_router/test_router.py +++ b/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