mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
c76882b51b
* fix(auto-router): stop the embedding model's context window from failing long requests The auto-router embeds the last user message to pick a model and sent it to the embedding model unbounded. Embedding models carry 512 to 8k token windows while the chat models they route to carry 200k+, so any prompt over the encoder's window failed at the routing step with a 400 the destination model would never have raised. Cut every doc to a character cap inside LiteLLMRouterEncoder, which is the one choke point the auto-router, complexity-router, semantic guard and MCP tool filter all share. Default 2000 chars, roughly 500 tokens, which fits even a 512-token self-hosted encoder, overridable per deployment with auto_router_max_input_chars and globally with DEFAULT_MAX_EMBEDDING_INPUT_CHARS. Truncation alone cannot cover provider-side batch and byte limits, so any failure of the route call now falls back to the auto-router's default model instead of propagating. That path also fixes two latent bugs: a no-match left the auto-router alias in place as the model name, which fails downstream with "Unmapped LLM provider" rather than reaching default_model, and an empty route list raised IndexError. Fixes #17869 Fixes #20277 * fix(auto-router): make the embedding input cap opt-in so guards still see whole prompts Defaulting the cap inside the shared encoder truncated every consumer, not just the auto-router. The semantic guard builds the same encoder, so its pre-call check would have classified only the first 2000 characters while the full message still reached the model, which a benign opener in front of an injection payload walks straight past. The MCP tool filter and complexity router were silently narrowed the same way. The encoder now defaults to sending docs whole and cuts only when a caller passes max_input_chars. The auto-router is the only caller that does, so guard, MCP filter and complexity-router behaviour is unchanged from before this branch. DEFAULT_MAX_EMBEDDING_INPUT_CHARS becomes DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS, since it is now specific to the auto-router, and drops its env override: the per-deployment auto_router_max_input_chars already covers it, and every env var in constants.py has to be documented, which is what broke the documentation and code-quality checks. Also drops the added comments and the redundant type: ignore that review flagged. * test(auto-router): cover the max_input_chars wiring from litellm_params Nothing asserted that auto_router_max_input_chars on the deployment reaches the AutoRouter that embeds prompts. Dropping the wiring left every test green while the cap silently reverted to the default, so an operator with a 512-token embedding model could not lower it and every long prompt would fall back to the default model instead of being routed. * test(auto-router): cover the populated route-choice list branch The route layer can hand back a list, and picking its first element is where the IndexError lived: the empty case was covered but the populated one was not, so the branch that reads route_choice[0].name could be deleted with every test still green.
134 lines
4.8 KiB
Python
134 lines
4.8 KiB
Python
"""Tests for litellm/router_strategy/auto_router/litellm_encoder.py"""
|
|
|
|
import os
|
|
import sys
|
|
from typing import Any, Final
|
|
|
|
import pytest
|
|
|
|
sys.path.insert(0, os.path.abspath("../../.."))
|
|
|
|
import litellm
|
|
from litellm.constants import DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS
|
|
from litellm.router_strategy.auto_router.litellm_encoder import LiteLLMRouterEncoder
|
|
|
|
|
|
class RecordingRouter:
|
|
"""Stand-in for the LiteLLM Router that records what reached the embedding call.
|
|
|
|
Injected through the encoder's constructor so the assertion is on the real code path
|
|
rather than on a patched attribute.
|
|
"""
|
|
|
|
def __init__(self) -> None:
|
|
self.embedded_inputs: list[list[str]] = []
|
|
|
|
def _response(self, input: list[str]) -> litellm.EmbeddingResponse:
|
|
self.embedded_inputs.append(list(input))
|
|
return litellm.EmbeddingResponse(
|
|
data=[{"embedding": [0.1, 0.2], "index": i, "object": "embedding"} for i in range(len(input))]
|
|
)
|
|
|
|
def embedding(self, input: list[str], model: str, **kwargs: Any) -> litellm.EmbeddingResponse:
|
|
return self._response(input)
|
|
|
|
async def aembedding(self, input: list[str], model: str, **kwargs: Any) -> litellm.EmbeddingResponse:
|
|
return self._response(input)
|
|
|
|
|
|
def _encoder(router: RecordingRouter, **kwargs: Any) -> LiteLLMRouterEncoder:
|
|
return LiteLLMRouterEncoder(
|
|
litellm_router_instance=router, # pyright: ignore[reportArgumentType] # test double stands in for Router
|
|
model_name="text-embedding-3-small",
|
|
**kwargs,
|
|
)
|
|
|
|
|
|
class TestSendsDocsWholeUnlessAskedNotTo:
|
|
"""Cutting is opt-in: a caller that must see the whole text is never cut behind its back.
|
|
|
|
The semantic guard and the MCP tool filter build this encoder without a limit. If a default
|
|
limit ever creeps back in, a prompt-injection payload placed after a benign opener would be
|
|
invisible to the guard while the full message still reaches the model.
|
|
"""
|
|
|
|
def test_should_send_a_long_doc_whole_when_no_cap_is_configured(self):
|
|
router: Final = RecordingRouter()
|
|
long_doc: Final = "x" * 50_000
|
|
|
|
_encoder(router).encode_queries([long_doc])
|
|
|
|
assert router.embedded_inputs == [[long_doc]]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_should_send_a_long_doc_whole_when_no_cap_is_configured_async(self):
|
|
router: Final = RecordingRouter()
|
|
long_doc: Final = "x" * 50_000
|
|
|
|
await _encoder(router).aencode_queries([long_doc])
|
|
|
|
assert router.embedded_inputs == [[long_doc]]
|
|
|
|
def test_should_send_a_long_doc_whole_when_the_cap_is_not_positive(self):
|
|
router: Final = RecordingRouter()
|
|
long_doc: Final = "x" * 50_000
|
|
|
|
_encoder(router, max_input_chars=0).encode_queries([long_doc])
|
|
|
|
assert router.embedded_inputs == [[long_doc]]
|
|
|
|
|
|
class TestEmbeddingInputCap:
|
|
"""With a cap configured, an embedding model's context window cannot fail the caller's request."""
|
|
|
|
def test_should_cut_a_long_doc_to_the_cap_on_sync_encode(self):
|
|
router: Final = RecordingRouter()
|
|
|
|
_encoder(router, max_input_chars=DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS).encode_queries(["x" * 50_000])
|
|
|
|
assert router.embedded_inputs == [["x" * DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS]]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_should_cut_a_long_doc_to_the_cap_on_async_encode(self):
|
|
router: Final = RecordingRouter()
|
|
|
|
await _encoder(router, max_input_chars=DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS).aencode_queries(["x" * 50_000])
|
|
|
|
assert router.embedded_inputs == [["x" * DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS]]
|
|
|
|
def test_should_cut_documents_as_well_as_queries(self):
|
|
router: Final = RecordingRouter()
|
|
|
|
_encoder(router, max_input_chars=100).encode_documents(["y" * 9_000])
|
|
|
|
assert router.embedded_inputs == [["y" * 100]]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_should_cut_documents_as_well_as_queries_async(self):
|
|
router: Final = RecordingRouter()
|
|
|
|
await _encoder(router, max_input_chars=100).aencode_documents(["y" * 9_000])
|
|
|
|
assert router.embedded_inputs == [["y" * 100]]
|
|
|
|
def test_should_leave_docs_within_the_cap_untouched(self):
|
|
router: Final = RecordingRouter()
|
|
docs: Final = ["a short prompt", "b" * 100]
|
|
|
|
_encoder(router, max_input_chars=100).encode_queries(docs)
|
|
|
|
assert router.embedded_inputs == [docs]
|
|
|
|
def test_should_cut_each_doc_in_a_batch_independently(self):
|
|
router: Final = RecordingRouter()
|
|
|
|
_encoder(router, max_input_chars=100).encode_queries(["short", "z" * 5_000])
|
|
|
|
assert router.embedded_inputs == [["short", "z" * 100]]
|
|
|
|
def test_should_keep_the_head_of_the_doc(self):
|
|
router: Final = RecordingRouter()
|
|
|
|
_encoder(router, max_input_chars=20).encode_queries(["route me please, then a huge pasted file"])
|
|
|
|
assert router.embedded_inputs == [["route me please, the"]]
|