mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-09 08:18:12 +00:00
fix(router): count embedding batches per input in pre-call checks (#49851)
* fix(router): count embedding batches per input in pre-call checks Embedding input lists went through the Responses input transform, which calls .get() on each item, so every list-input /embeddings call logged "failed to count tokens" and skipped the context window check. Lists of texts, token IDs or token ID lists now count as their longest input, since an embedding model's context window applies to each input on its own * fix(router): keep empty and boolean input lists on the Responses count path An empty list carries no embedding evidence, and counting it as a zero token batch skipped the instructions and tools of an empty Responses input. Boolean lists are not token IDs * fix(router): honor disable_token_counter for token ID embedding input Text inputs already counted as zero through token_counter when token counting is disabled, while token ID inputs were still measured by length and could be rejected. The router reads the flag through its dynamic config like its other process settings * refactor(router): type the embedding count helper with a read-only sequence and final locals * fix(router): count embedding batches per input on the media-free text projection The per-input count now runs on the text projection from the embedding media split, so data URIs a deployment embeds as media keep the media path and only text items are checked against the context window. A caption next to an image is checked on its own, and a text-only deployment still counts an image data URI as text * fix(router): keep embedding images out of the pre-call count only on image backends, and skip tokenizing inputs that fit by bytes Media leaves the embedding count only when every deployment in the group declares image embedding input, so a text-only backend still counts a data URI as text and an image backend outside the TPM media list no longer gets a false 400. A text no longer than the smallest window in UTF-8 bytes counts as its byte length without the tokenizer, since a token is never shorter than a byte. * 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. * fix(router): check embedding text batches only with an exact OpenAI tokenizer litellm's default tokenizer counts non-English text several times higher than non-OpenAI embedding providers do, so counting batches with it refused inputs the provider accepts. Text batches on those groups now stay unchecked as at the merge base, without the error log. Main's per-provider media split and the cost map are restored, so strings and embedding media keep main's path. * 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. * fix(router): read embedding media off the event loop and wait on the count limiter in the fairness test --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
m.torres
GitHub
parent
1f15088fc9
commit
150f4b706b
6 changed files with 429 additions and 36 deletions
|
|
@ -202,7 +202,7 @@ def _tpm_embedding_leaf(value: object) -> _TPMMediaProjection:
|
|||
part: Final = as_str_mapping(value)
|
||||
if part is not None and part.get("type") == "file":
|
||||
return _tpm_embedding_file_block(part)
|
||||
if not isinstance(value, str) or not value.startswith("data:"):
|
||||
if not isinstance(value, str) or is_plain_embedding_text(value):
|
||||
return _tpm_media_leaf(value)
|
||||
if not _is_base64_data_uri(value):
|
||||
return _TPMMediaProjection(value)
|
||||
|
|
@ -210,6 +210,10 @@ def _tpm_embedding_leaf(value: object) -> _TPMMediaProjection:
|
|||
return _TPMMediaProjection("" if media else value, media)
|
||||
|
||||
|
||||
def is_plain_embedding_text(value: object) -> bool:
|
||||
return isinstance(value, str) and not value.startswith("data:")
|
||||
|
||||
|
||||
def _is_json_array(value: object) -> bool:
|
||||
return isinstance(value, (list, tuple))
|
||||
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
@ -21,6 +23,7 @@ from litellm._logging import verbose_router_logger
|
|||
from litellm.litellm_core_utils.core_helpers import as_str_mapping
|
||||
from litellm.litellm_core_utils.deployment_pricing import deployment_backend_models
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import declared_authenticating_provider
|
||||
from litellm.litellm_core_utils.prompt_templates.tpm_media_utils import is_plain_embedding_text
|
||||
from litellm.litellm_core_utils.token_counter import (
|
||||
offload_token_count,
|
||||
served_openai_models,
|
||||
|
|
@ -38,6 +41,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):
|
||||
|
|
@ -59,35 +73,111 @@ def count_tokens(
|
|||
request_kwargs: Mapping[str, object] | None,
|
||||
tokenizer_model: str,
|
||||
deployments: Sequence[Mapping[str, object]] = (),
|
||||
fits_within: int = 0,
|
||||
) -> 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 an
|
||||
item of ``deployments`` reads as media is left out of the input: it never counts as text."""
|
||||
from litellm.llms.anthropic.messages.native.utils import anthropic_system_to_openai_message
|
||||
|
||||
extras: Final[Mapping[str, object]] = request_kwargs if request_kwargs is not None else MappingProxyType({})
|
||||
raw_instructions: Final = extras.get("instructions")
|
||||
instructions: Final = raw_instructions if isinstance(raw_instructions, str) else None
|
||||
raw_tools: Final = extras.get("tools")
|
||||
tools: Final = raw_tools if _is_list(raw_tools) and raw_tools else None
|
||||
if messages is None and input is not None:
|
||||
return _input_tokens(host, _text_input(input, deployments), request_kwargs, tokenizer_model, fits_within)
|
||||
if messages is None:
|
||||
raise ValueError("Either messages or input must be provided to count tokens")
|
||||
extras: Final = _extras(request_kwargs)
|
||||
system_message: Final = anthropic_system_to_openai_message(extras.get("system"))
|
||||
counter: Final = host.invoke("token_counter") # pyright: ignore[reportUnknownMemberType] # the host's overloads name litellm's untyped returns
|
||||
if messages is not None:
|
||||
counted_messages: Final = (system_message, *_items(messages)) if system_message is not None else messages
|
||||
return counter(model=tokenizer_model, messages=counted_messages, tools=tools)
|
||||
if input is not None:
|
||||
from litellm.responses.litellm_completion_transformation.transformation import (
|
||||
LiteLLMCompletionResponsesConfig,
|
||||
)
|
||||
counted_messages: Final = (system_message, *_items(messages)) if system_message is not None else messages
|
||||
return counter(model=tokenizer_model, messages=counted_messages, tools=_tools(extras))
|
||||
|
||||
text_input, _ = split_embedding_input_tokens(input, partial(deployment_backend_models, deployments))
|
||||
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 {},
|
||||
only_since_compaction=True,
|
||||
)
|
||||
return counter(model=tokenizer_model, messages=input_messages, tools=tools)
|
||||
raise ValueError("Either messages or input must be provided to count tokens")
|
||||
|
||||
def _input_tokens(
|
||||
host: "PythonHost",
|
||||
text_input: object,
|
||||
request_kwargs: Mapping[str, object] | None,
|
||||
tokenizer_model: str,
|
||||
fits_within: int,
|
||||
) -> int:
|
||||
from litellm.responses.litellm_completion_transformation.transformation import (
|
||||
LiteLLMCompletionResponsesConfig,
|
||||
)
|
||||
|
||||
counter: Final = host.invoke("token_counter") # pyright: ignore[reportUnknownMemberType] # see count_tokens
|
||||
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
|
||||
extras: Final = _extras(request_kwargs)
|
||||
raw_instructions: Final = extras.get("instructions")
|
||||
instructions: Final = raw_instructions if isinstance(raw_instructions, str) else None
|
||||
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 {},
|
||||
only_since_compaction=True,
|
||||
)
|
||||
return counter(model=tokenizer_model, messages=input_messages, tools=_tools(extras))
|
||||
|
||||
|
||||
def _extras(request_kwargs: Mapping[str, object] | None) -> Mapping[str, object]:
|
||||
return request_kwargs if request_kwargs is not None else MappingProxyType({})
|
||||
|
||||
|
||||
def _tools(extras: Mapping[str, object]) -> object:
|
||||
raw_tools: Final = extras.get("tools")
|
||||
return raw_tools if _is_list(raw_tools) and raw_tools else None
|
||||
|
||||
|
||||
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 _TextBatch(texts)
|
||||
if all(_is_token_id(item) for item in items):
|
||||
return _TokenIdBatch(len(items))
|
||||
if all(_is_token_id_list(item) for item in items):
|
||||
return _TokenIdBatch(max(len(_items(item)) for item in items))
|
||||
return None
|
||||
|
||||
|
||||
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)
|
||||
|
||||
|
||||
def _is_token_id(value: object) -> bool:
|
||||
return isinstance(value, int) and not isinstance(value, bool)
|
||||
|
||||
|
||||
def _is_token_id_list(value: object) -> bool:
|
||||
return _is_list(value) and all(_is_token_id(token) for token in _items(value))
|
||||
|
||||
|
||||
def _token_id_count(host: "PythonHost", count: int) -> int:
|
||||
return 0 if host.dynamic().disable_token_counter else count
|
||||
|
||||
|
||||
def _max_input_tokens(registry: PreCallRegistry, model: str, deployment: Mapping[str, object]) -> int | None:
|
||||
|
|
@ -246,16 +336,31 @@ class PreCallFacts:
|
|||
def _candidate_records(self) -> Sequence[Mapping[str, object]]:
|
||||
return tuple(self._record(id) for id in self._candidates)
|
||||
|
||||
def _fits_within(self) -> int:
|
||||
return min((limit for limit in self._windows.values() if limit is not None), default=0)
|
||||
|
||||
def _input_is_a_list(self) -> bool:
|
||||
return self._messages is None and _is_list(self._input)
|
||||
|
||||
def _count(self) -> int | None:
|
||||
try:
|
||||
return count_tokens(
|
||||
self._host,
|
||||
self._messages,
|
||||
self._input,
|
||||
self._request,
|
||||
self._tokenizer_model(),
|
||||
self._candidate_records(),
|
||||
)
|
||||
tokenizer_model: Final = self._tokenizer_model()
|
||||
if not self._input_is_a_list():
|
||||
return count_tokens(
|
||||
self._host,
|
||||
self._messages,
|
||||
self._input,
|
||||
self._request,
|
||||
tokenizer_model,
|
||||
self._candidate_records(),
|
||||
self._fits_within(),
|
||||
)
|
||||
text_input: Final = _text_input(self._input, self._candidate_records())
|
||||
match _embedding_batch(text_input):
|
||||
case _TextBatch() if not tokenizer_model:
|
||||
return 0
|
||||
case _:
|
||||
return _input_tokens(self._host, text_input, self._request, tokenizer_model, 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
|
||||
|
|
@ -264,14 +369,39 @@ 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(),
|
||||
tokenizer_model: Final = self._tokenizer_model()
|
||||
if not self._input_is_a_list():
|
||||
return await offload_token_count(count_tokens)(
|
||||
self._host,
|
||||
self._messages,
|
||||
self._input,
|
||||
self._request,
|
||||
tokenizer_model,
|
||||
self._candidate_records(),
|
||||
self._fits_within(),
|
||||
)
|
||||
text_input: Final = (
|
||||
self._input
|
||||
if all(map(is_plain_embedding_text, _items(self._input)))
|
||||
else await offload_token_count(_text_input)(self._input, self._candidate_records())
|
||||
)
|
||||
match _embedding_batch(text_input):
|
||||
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 _:
|
||||
return await offload_token_count(_input_tokens)(
|
||||
self._host, text_input, self._request, tokenizer_model, 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
|
||||
|
|
|
|||
|
|
@ -166,6 +166,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,
|
||||
|
|
|
|||
|
|
@ -27,6 +27,7 @@ PICKS: Final = 60 # three-way shuffles miss a deployment in 20 picks about once
|
|||
SHORT: Final = [{"role": "user", "content": "hi"}]
|
||||
LONG_TEXT: Final = " ".join(["word"] * 100)
|
||||
LONG: Final = [{"role": "user", "content": LONG_TEXT}]
|
||||
OPENAI_EMBEDDER: Final = "text-embedding-3-small"
|
||||
EMBEDDING_IMAGE: Final = f"data:image/png;base64,{base64.b64encode(bytes(range(256)) * 1024).decode()}"
|
||||
IMAGE: Final = [
|
||||
{
|
||||
|
|
@ -86,6 +87,14 @@ def _small_and_big(
|
|||
)
|
||||
|
||||
|
||||
def _small_and_big_openai_embedders(make_core: MakeCore, small: int) -> RouterCore:
|
||||
counted_as: Final = {"model_info": {"base_model": OPENAI_EMBEDDER}}
|
||||
return make_core(
|
||||
model_list=[windowed("g", "small", small, **counted_as), windowed("g", "big", 10000, **counted_as)],
|
||||
enable_pre_call_checks=True,
|
||||
)
|
||||
|
||||
|
||||
def _json_mode_group(make_core: MakeCore, **arguments: object) -> RouterCore: # kwargs-ok: constructor arguments
|
||||
return make_core(
|
||||
model_list=[deployment("g", "old", model=NO_JSON_MODE_MODEL), deployment("g", "new", model=JSON_MODE_MODEL)],
|
||||
|
|
@ -279,6 +288,43 @@ async def test_an_embedding_image_stays_out_of_the_context_window_on_a_backend_t
|
|||
assert await _picked(core, times=2, input=EMBEDDING_IMAGE) == {"image"}
|
||||
|
||||
|
||||
@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(model=OPENAI_EMBEDDER, text=short_words) <= 10 < len(short_words.encode())
|
||||
assert await _picked(_small_and_big_openai_embedders(make_core, small=10), input=["hi", short_words]) == {
|
||||
"small",
|
||||
"big",
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_embedding_input_that_fits_the_window_in_characters_but_not_in_tokens_is_counted(
|
||||
make_core: MakeCore,
|
||||
) -> None:
|
||||
dense: Final = "\ua66e" * 40
|
||||
assert len(dense) <= 50 < litellm.token_counter(model=OPENAI_EMBEDDER, text=dense)
|
||||
assert await _picked(_small_and_big_openai_embedders(make_core, small=50), times=10, input=["hi", dense]) == {"big"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_embedding_batch_is_never_refused_on_the_default_tokenizers_count(make_core: MakeCore) -> None:
|
||||
dense: Final = "\ua66e" * 40
|
||||
assert litellm.token_counter(text=dense) > 50
|
||||
assert await _picked(_small_and_big(make_core, small=50), input=["hi", dense]) == {"small", "big"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_embedding_input_with_a_lone_surrogate_is_counted(make_core: MakeCore) -> None:
|
||||
unpaired: Final = "\ud800" + "word " * 20
|
||||
assert litellm.token_counter(model=OPENAI_EMBEDDER, text=unpaired) > 10
|
||||
assert await _picked(_small_and_big_openai_embedders(make_core, small=10), times=10, input=["hi", unpaired]) == {
|
||||
"big"
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_embedding_image_counts_as_text_on_a_text_only_backend(make_core: MakeCore) -> None:
|
||||
core: Final = make_core(
|
||||
|
|
|
|||
|
|
@ -4396,6 +4396,48 @@ def test_context_window_pre_call_check_counts_an_embedding_image_as_text_on_a_te
|
|||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("backend", "fits"),
|
||||
[
|
||||
pytest.param("vertex_ai/gemini-embedding-2", True, id="image-backend-without-an-openai-tokenizer"),
|
||||
pytest.param("text-embedding-3-small", False, id="text-only-backend-counts-the-image"),
|
||||
],
|
||||
)
|
||||
def test_context_window_pre_call_check_on_a_caption_next_to_an_embedding_image(backend: str, fits: bool, caplog):
|
||||
caption: Final = "word " * 40
|
||||
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
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_async_pre_call_check_reads_embedding_media_off_the_event_loop():
|
||||
from tests.unit.litellm_core_utils.event_loop_lag import assert_loop_stayed_free, timed_with_loop_lags
|
||||
|
||||
router: Final = _embedding_context_router("vertex_ai/gemini-embedding-2", window=20)
|
||||
caption_and_images: Final = ["a red square", *[_EMBEDDING_IMAGE_INPUT] * 400]
|
||||
|
||||
deployment, took, lags = await timed_with_loop_lags(
|
||||
lambda: router.select("embedder", {}, input=caption_and_images)
|
||||
)
|
||||
|
||||
assert deployment["litellm_params"]["model"] == "vertex_ai/gemini-embedding-2"
|
||||
assert_loop_stayed_free(took, lags)
|
||||
|
||||
|
||||
async def test_a_wildcard_route_checks_the_context_window_with_the_requested_openai_models_tokenizer():
|
||||
limit: Final = (
|
||||
litellm.token_counter(model="openai/gpt-6.1-sol", messages=_MULTILINGUAL_PROMPT)
|
||||
|
|
@ -4547,6 +4589,173 @@ 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
|
||||
|
||||
|
||||
async def _until_every_count_slot_is_taken() -> None:
|
||||
from litellm.constants import TOKEN_COUNTER_MAX_CONCURRENT_COUNTS
|
||||
from litellm.litellm_core_utils.token_counter import _count_offload_limiter_for_this_loop
|
||||
|
||||
while _count_offload_limiter_for_this_loop().borrowed_tokens < TOKEN_COUNTER_MAX_CONCURRENT_COUNTS:
|
||||
await asyncio.sleep(0)
|
||||
|
||||
|
||||
@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.wait_for(_until_every_count_slot_is_taken(), timeout=30)
|
||||
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(
|
||||
"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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue