fix(router): count long embedding batches in slices so chat counts do not wait behind them

A text batch is now tokenized in slices of 65,536 characters, each its own offloaded call, so the FIFO count limiter can serve a chat count between slices. The no-OpenAI-tokenizer rule for text lists moved into the pre-call check, so RoutingContext.count_pre_call_check_tokens counts a text list by its longest text.

litellm/router_rust/steps/precall.py

@@ -13,7 +13,9 @@ deployment records rather than a step, and those need not be in the Rust registr
 """
 
 from collections.abc import Callable, Mapping, Sequence
+from dataclasses import dataclass
 from functools import partial
+from itertools import accumulate, groupby
 from types import MappingProxyType
 from typing import TYPE_CHECKING, Final, Protocol
 
@@ -38,6 +40,17 @@ if TYPE_CHECKING:
 
 # The chat params a deployment must take when the request sends them (`_PRE_CALL_CHECK_FILTERED_PARAMS`).
 _FILTERED_PARAMS: Final = frozenset({"response_format"})
+_CHARS_PER_OFFLOADED_COUNT: Final = 1 << 16
+
+
+@dataclass(frozen=True, slots=True)
+class _TextBatch:
+    texts: tuple[str, ...]
+
+
+@dataclass(frozen=True, slots=True)
+class _TokenIdBatch:
+    longest: int
 
 
 class PreCallRegistry(Protocol):
@@ -81,10 +94,14 @@ def count_tokens(
             LiteLLMCompletionResponsesConfig,
         )
 
-        text_input, _ = split_embedding_input_tokens(input, partial(deployment_backend_models, deployments))
-        embedding_tokens: Final = _longest_embedding_input_tokens(host, text_input, tokenizer_model, fits_within)
-        if embedding_tokens is not None:
-            return embedding_tokens
+        text_input: Final = _text_input(input, deployments)
+        match _embedding_batch(text_input):
+            case _TextBatch(texts):
+                return _longest_text_tokens(counter, tokenizer_model, texts, fits_within)
+            case _TokenIdBatch(longest):
+                return _token_id_count(host, longest)
+            case None:
+                pass
         input_messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages(  # pyright: ignore[reportUnknownMemberType]  # its input type is untyped
             input=text_input,  # pyright: ignore[reportArgumentType]  # str | list, as the request carries it
             responses_api_request={"instructions": instructions} if instructions is not None else {},
@@ -94,30 +111,39 @@ def count_tokens(
     raise ValueError("Either messages or input must be provided to count tokens")
 
 
-def _longest_embedding_input_tokens(
-    host: "PythonHost", value: object, tokenizer_model: str, fits_within: int
-) -> int | None:
-    items: Final = _items(value) if _is_list(value) else ()
+def _text_input(input: object, deployments: Sequence[Mapping[str, object]]) -> object:  # the request field's name
+    text_input, _ = split_embedding_input_tokens(input, partial(deployment_backend_models, deployments))
+    return text_input
+
+
+def _embedding_batch(text_input: object) -> _TextBatch | _TokenIdBatch | None:
+    items: Final = _items(text_input) if _is_list(text_input) else ()
     if not items:
         return None
     texts: Final = tuple(item for item in items if isinstance(item, str))
     if len(texts) == len(items):
-        return _longest_text_tokens(host, texts, tokenizer_model, fits_within)
+        return _TextBatch(texts)
     if all(_is_token_id(item) for item in items):
-        return _token_id_count(host, len(items))
+        return _TokenIdBatch(len(items))
     if all(_is_token_id_list(item) for item in items):
-        return _token_id_count(host, max(len(_items(item)) for item in items))
+        return _TokenIdBatch(max(len(_items(item)) for item in items))
     return None
 
 
-def _longest_text_tokens(host: "PythonHost", texts: Sequence[str], tokenizer_model: str, fits_within: int) -> int:
-    if not tokenizer_model:
-        # the default tokenizer overcounts non-English text, so only an OpenAI tokenizer may refuse a batch
-        return 0
-    counter: Final = host.invoke("token_counter")  # pyright: ignore[reportUnknownMemberType]  # see count_tokens
+def _longest_text_tokens(
+    counter: Callable[..., int], tokenizer_model: str, texts: Sequence[str], fits_within: int
+) -> int:
     return max(_text_tokens(counter, tokenizer_model, text, fits_within) for text in texts)
 
 
+def _offload_slices(texts: Sequence[str]) -> tuple[tuple[str, ...], ...]:
+    ends: Final = accumulate(len(text) for text in texts)
+    slices: Final = groupby(
+        zip(ends, texts, strict=True), key=lambda end_and_text: (end_and_text[0] - 1) // _CHARS_PER_OFFLOADED_COUNT
+    )
+    return tuple(tuple(text for _, text in texts_in_slice) for _, texts_in_slice in slices)
+
+
 def _text_tokens(counter: Callable[..., int], tokenizer_model: str, text: str, fits_within: int) -> int:
     size: Final = len(text.encode(errors="surrogatepass"))
     return size if size <= fits_within else counter(model=tokenizer_model, text=text)
@@ -294,17 +320,31 @@ class PreCallFacts:
     def _fits_within(self) -> int:
         return min((limit for limit in self._windows.values() if limit is not None), default=0)
 
+    def _text_batch(self) -> _TextBatch | None:
+        if self._messages is not None or self._input is None:
+            return None
+        batch: Final = _embedding_batch(_text_input(self._input, self._candidate_records()))
+        return batch if isinstance(batch, _TextBatch) else None
+
     def _count(self) -> int | None:
         try:
-            return count_tokens(
-                self._host,
-                self._messages,
-                self._input,
-                self._request,
-                self._tokenizer_model(),
-                self._candidate_records(),
-                self._fits_within(),
-            )
+            tokenizer_model: Final = self._tokenizer_model()
+            match self._text_batch():
+                case _TextBatch() if not tokenizer_model:
+                    return 0
+                case _TextBatch(texts):
+                    counter: Final = self._host.invoke("token_counter")  # pyright: ignore[reportUnknownMemberType]  # see count_tokens
+                    return _longest_text_tokens(counter, tokenizer_model, texts, self._fits_within())
+                case None:
+                    return count_tokens(
+                        self._host,
+                        self._messages,
+                        self._input,
+                        self._request,
+                        tokenizer_model,
+                        self._candidate_records(),
+                        self._fits_within(),
+                    )
         except Exception as e:  # noqa: BLE001  # best-effort: an uncountable prompt must not fail the request
             verbose_router_logger.error(
                 "litellm.router.py::_acount_pre_call_check_tokens: failed to count tokens. Got - %s", e
@@ -313,15 +353,30 @@ class PreCallFacts:
 
     async def _count_async(self) -> int | None:
         try:
-            return await offload_token_count(count_tokens)(
-                self._host,
-                self._messages,
-                self._input,
-                self._request,
-                self._tokenizer_model(),
-                self._candidate_records(),
-                self._fits_within(),
-            )
+            tokenizer_model: Final = self._tokenizer_model()
+            match self._text_batch():
+                case _TextBatch() if not tokenizer_model:
+                    return 0
+                case _TextBatch(texts):
+                    counter: Final = self._host.invoke("token_counter")  # pyright: ignore[reportUnknownMemberType]  # see count_tokens
+                    fits_within: Final = self._fits_within()
+                    count_slice: Final = offload_token_count(_longest_text_tokens)
+                    return max(
+                        [
+                            await count_slice(counter, tokenizer_model, texts_in_slice, fits_within)
+                            for texts_in_slice in _offload_slices(texts)
+                        ]
+                    )
+                case None:
+                    return await offload_token_count(count_tokens)(
+                        self._host,
+                        self._messages,
+                        self._input,
+                        self._request,
+                        tokenizer_model,
+                        self._candidate_records(),
+                        self._fits_within(),
+                    )
         except Exception as e:  # noqa: BLE001  # best-effort: an uncountable prompt must not fail the request
             verbose_router_logger.error(
                 "litellm.router.py::_acount_pre_call_check_tokens: failed to count tokens. Got - %s", e

tests/unit/test_router/test_router.py

@@ -4376,6 +4376,51 @@ async def test_aembedding_skips_the_context_window_check_when_token_counting_is_
     assert response.data
 
 
+@pytest.mark.asyncio
+async def test_a_chat_pre_call_count_does_not_wait_behind_long_embedding_batches():
+    from litellm.constants import TOKEN_COUNTER_MAX_CONCURRENT_COUNTS
+
+    window: Final = {"max_input_tokens": 10000}
+    router: Final = litellm.Router(
+        model_list=[
+            {
+                "model_name": "embed",
+                "litellm_params": {
+                    "model": "openai/text-embedding-3-small",
+                    "api_key": "sk-test",
+                    "mock_response": [0.1],
+                },
+                "model_info": window,
+            },
+            {
+                "model_name": "chat",
+                "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test", "mock_response": "hi"},
+                "model_info": window,
+            },
+        ],
+        enable_pre_call_checks=True,
+    )
+    batch: Final = ["word " * 2100] * 76
+    batches: Final = tuple(
+        asyncio.create_task(router.aembedding(model="embed", input=batch))
+        for _ in range(TOKEN_COUNTER_MAX_CONCURRENT_COUNTS)
+    )
+    await asyncio.sleep(0.05)
+    chat: Final = asyncio.create_task(router.acompletion(model="chat", messages=[{"role": "user", "content": "hi"}]))
+
+    first_done, _ = await asyncio.wait((chat, *batches), return_when=asyncio.FIRST_COMPLETED)
+    await asyncio.gather(chat, *batches)
+
+    assert first_done == {chat}
+
+
+def test_a_routing_plugin_counts_an_embedding_text_batch_by_its_longest_text():
+    long_text: Final = "word " * 4000
+    context: Final = RouterContext(_embedding_router(max_input_tokens=10))
+
+    assert context.count_pre_call_check_tokens(None, ["one", long_text]) == litellm.token_counter(text=long_text)
+
+
 @pytest.mark.asyncio
 @pytest.mark.usefixtures("local_model_cost_map", "isolate_router_model_cost_state")
 @pytest.mark.parametrize(