mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-09 22:31:41 +00:00
6a0d03914c
* test: drop the cwd-relative sys.path.insert calls from the test suite
TQ003 stands at 1,077 across 1,058 files, and 1,015 of them are the same shape:
sys.path.insert(0, os.path.abspath("../..")) and its deeper siblings. The
argument resolves against the working directory rather than the file, so from
the repo root, where every job runs pytest, it inserts the directory two levels
above the checkout. It has never pointed at litellm. The package is installed
into the environment anyway, which is what actually makes the import work, and
what the rule's message has said all along.
Removing them leaves 1,634 imports of sys and os with no remaining reference,
and those go too, except where another test module imports the name back out of
the file. The rest of TQ003 is 62 call sites that resolve against __file__ or a
variable, which are a different question and are left alone.
Collection is identical either way: 45,871 tests and the same 51 pre-existing
collection errors before and after, and ruff reports no new undefined name.
* test: drop the duplicate imports the sys.path sweep exposed to F811
* test(pre-call-utils): restore the os import the new bedrock tests need
131 lines
4.8 KiB
Python
131 lines
4.8 KiB
Python
"""Tests for litellm/router_strategy/auto_router/litellm_encoder.py"""
|
|
|
|
from typing import Any, Final
|
|
|
|
import pytest
|
|
|
|
|
|
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"]]
|