mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-25 05:54:53 +00:00
331 lines
12 KiB
Python
331 lines
12 KiB
Python
from collections.abc import AsyncGenerator, Awaitable, Callable, Iterator, Mapping
|
|
from dataclasses import dataclass
|
|
from typing import Final, NoReturn, TypeAlias
|
|
|
|
import pytest
|
|
|
|
from litellm.rust_bridge.bindings import NativeBinding
|
|
from litellm.rust_bridge.catalog import Feature, Route, RouteContext, RouteRule, Rules, Target
|
|
from litellm.rust_bridge.configuration import Rollout
|
|
from litellm.rust_bridge.dispatch import PublicDispatch
|
|
from litellm.rust_bridge.runtime import NO_PYTHON, NoPythonImplementationError
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class Request:
|
|
model: str
|
|
|
|
|
|
def binding() -> NativeBinding[object]:
|
|
bound: Final[NativeBinding[object]] = NativeBinding("unused", validate=lambda value: value)
|
|
bound.override(None)
|
|
return bound
|
|
|
|
|
|
@pytest.mark.parametrize("rules", ((), (RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED),)))
|
|
def test_route_without_rules_forwards_before_request_projection(rules: Rules) -> None:
|
|
stream: Final[Iterator[int]] = iter((1, 2))
|
|
|
|
def reject_request(args: tuple[object, ...], kwargs: Mapping[str, object]) -> Request:
|
|
pytest.fail("Python-only routes must not project the request")
|
|
|
|
dispatch: Final = PublicDispatch(
|
|
route=Route.CHAT_COMPLETIONS, request=reject_request, context=lambda _: RouteContext(Route.CHAT_COMPLETIONS)
|
|
)
|
|
result: Final = dispatch.run(
|
|
("model",),
|
|
{"stream": True},
|
|
python=lambda *args, **kwargs: stream,
|
|
binding=binding(),
|
|
native=lambda hook, request, args, kwargs: pytest.fail("Python-only routes must not call native"),
|
|
rules=rules,
|
|
)
|
|
assert result is stream
|
|
|
|
|
|
def test_unconditional_python_rule_prevents_later_rust_rule_projection() -> None:
|
|
rules: Final[Rules] = (
|
|
RouteRule(Route.CHAT_COMPLETIONS, Rollout.PYTHON_ONLY),
|
|
RouteRule(Route.CHAT_COMPLETIONS, Rollout.RUST_REQUIRED),
|
|
)
|
|
|
|
def reject_request(args: tuple[object, ...], kwargs: Mapping[str, object]) -> Request:
|
|
pytest.fail("First-match Python rule must prevent request projection")
|
|
|
|
dispatch: Final = PublicDispatch(
|
|
route=Route.CHAT_COMPLETIONS,
|
|
request=reject_request,
|
|
context=lambda _: RouteContext(Route.CHAT_COMPLETIONS),
|
|
)
|
|
expected: Final = object()
|
|
result: Final = dispatch.run(
|
|
("model",),
|
|
{},
|
|
python=lambda *args, **kwargs: expected,
|
|
binding=binding(),
|
|
native=lambda hook, request, args, kwargs: pytest.fail("First-match Python rule must prevent native"),
|
|
rules=rules,
|
|
)
|
|
assert result is expected
|
|
|
|
|
|
def test_native_stream_result_is_not_consumed_or_wrapped() -> None:
|
|
request: Final = Request(model="streaming-model")
|
|
stream: Final[Iterator[int]] = iter((1, 2))
|
|
rules: Final[Rules] = (
|
|
RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY),
|
|
RouteRule(Route.CHAT_COMPLETIONS, Rollout.RUST_REQUIRED),
|
|
)
|
|
dispatch: Final = PublicDispatch(
|
|
route=Route.CHAT_COMPLETIONS,
|
|
request=lambda args, kwargs: request,
|
|
context=lambda value: RouteContext(Route.CHAT_COMPLETIONS, model=value.model),
|
|
)
|
|
|
|
def native(request: Request, args: tuple[object, ...], kwargs: Mapping[str, object]) -> Iterator[int]:
|
|
return stream
|
|
|
|
native_binding: Final[
|
|
NativeBinding[Callable[[Request, tuple[object, ...], Mapping[str, object]], Iterator[int]]]
|
|
] = NativeBinding("stream", validate=lambda _: None)
|
|
native_binding.override(native)
|
|
result: Final = dispatch.run(
|
|
("streaming-model",),
|
|
{"stream": True},
|
|
python=lambda *args, **kwargs: pytest.fail("Required native stream dispatch must not call Python"),
|
|
binding=native_binding,
|
|
native=lambda hook, value, args, kwargs: hook(value, args, kwargs),
|
|
rules=rules,
|
|
)
|
|
assert result is stream
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("rules", ((), (RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED),)))
|
|
async def test_async_route_without_rules_preserves_async_iterator_result(rules: Rules) -> None:
|
|
async def chunks() -> AsyncGenerator[int, None]:
|
|
yield 1
|
|
|
|
stream: Final = chunks()
|
|
|
|
def reject_request(args: tuple[object, ...], kwargs: Mapping[str, object]) -> Request:
|
|
pytest.fail("Python-only routes must not project the request")
|
|
|
|
async def python(*args: object, **kwargs: object) -> AsyncGenerator[int, None]: # kwargs-ok: pass-through shape
|
|
return stream
|
|
|
|
dispatch: Final = PublicDispatch(
|
|
route=Route.RESPONSES, request=reject_request, context=lambda _: RouteContext(Route.RESPONSES)
|
|
)
|
|
result: Final = await dispatch.arun(
|
|
("model",),
|
|
{"stream": True},
|
|
python=python,
|
|
binding=binding(),
|
|
native=lambda hook, request, args, kwargs: pytest.fail("Python-only routes must not call native"),
|
|
rules=rules,
|
|
)
|
|
assert result is stream
|
|
await stream.aclose()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_dispatch_accepts_websocket_style_none_result() -> None:
|
|
request: Final = Request(model="realtime-model")
|
|
rules: Final[Rules] = (RouteRule(Route.RESPONSES, Rollout.RUST_REQUIRED),)
|
|
dispatch: Final = PublicDispatch(
|
|
route=Route.RESPONSES,
|
|
request=lambda args, kwargs: request,
|
|
context=lambda value: RouteContext(Route.RESPONSES, model=value.model),
|
|
)
|
|
|
|
async def python(*args: object, **kwargs: object) -> None: # kwargs-ok: public pass-through shape
|
|
pytest.fail("Required native WebSocket dispatch must not call Python")
|
|
|
|
async def native(request: Request, args: tuple[object, ...], kwargs: Mapping[str, object]) -> None:
|
|
return None
|
|
|
|
native_binding: Final[
|
|
NativeBinding[Callable[[Request, tuple[object, ...], Mapping[str, object]], Awaitable[None]]]
|
|
] = NativeBinding("websocket", validate=lambda _: None)
|
|
native_binding.override(native)
|
|
|
|
result: Final = await dispatch.arun(
|
|
("realtime-model",),
|
|
{},
|
|
python=python,
|
|
binding=native_binding,
|
|
native=lambda hook, value, args, kwargs: hook(value, args, kwargs),
|
|
rules=rules,
|
|
)
|
|
assert result is None
|
|
|
|
|
|
def test_rules_for_other_routes_and_constrained_python_rules_skip_projection() -> None:
|
|
rules: Final[Rules] = (
|
|
RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED),
|
|
RouteRule(Route.OCR, Rollout.PYTHON_ONLY, providers=frozenset({"mistral"})),
|
|
)
|
|
|
|
def reject_request(args: tuple[object, ...], kwargs: Mapping[str, object]) -> Request:
|
|
pytest.fail("Rules that cannot select Rust must not project the request")
|
|
|
|
dispatch: Final = PublicDispatch(route=Route.OCR, request=reject_request, context=lambda _: RouteContext(Route.OCR))
|
|
expected: Final = object()
|
|
result: Final = dispatch.run(
|
|
("model",),
|
|
{},
|
|
python=lambda *args, **kwargs: expected,
|
|
binding=binding(),
|
|
native=lambda hook, request, args, kwargs: pytest.fail("Rules that cannot select Rust must not call native"),
|
|
rules=rules,
|
|
)
|
|
assert result is expected
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"constrained",
|
|
(
|
|
RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY, providers=frozenset({"other"})),
|
|
RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY, models=frozenset({"other"})),
|
|
RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY, targets=frozenset({Target.CHAT_COMPLETIONS})),
|
|
RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY, features=frozenset({Feature.SAFEGUARDS_CLASSIFIER})),
|
|
),
|
|
ids=("providers", "models", "targets", "features"),
|
|
)
|
|
def test_constrained_python_rule_still_projects_for_a_later_rust_rule(constrained: RouteRule) -> None:
|
|
expected: Final = object()
|
|
native_binding: Final[NativeBinding[Callable[[Request], object]]] = NativeBinding("x", validate=lambda _: None)
|
|
native_binding.override(lambda request: expected)
|
|
dispatch: Final = PublicDispatch(
|
|
route=Route.MESSAGES,
|
|
request=lambda args, kwargs: Request(model="model"),
|
|
context=lambda value: RouteContext(Route.MESSAGES, model=value.model, target=Target.MESSAGES),
|
|
)
|
|
|
|
result: Final = dispatch.run(
|
|
("model",),
|
|
{},
|
|
python=lambda *args, **kwargs: pytest.fail("a constrained Python rule must not hide the Rust rule after it"),
|
|
binding=native_binding,
|
|
native=lambda hook, value, args, kwargs: hook(value),
|
|
rules=(constrained, RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED)),
|
|
)
|
|
|
|
assert result is expected
|
|
|
|
|
|
NativeRoute: TypeAlias = Callable[[Request, tuple[object, ...], Mapping[str, object]], object]
|
|
|
|
|
|
def no_python_binding(result: object) -> NativeBinding[NativeRoute]:
|
|
bound: Final[NativeBinding[NativeRoute]] = NativeBinding("no_python", validate=lambda _: None)
|
|
bound.override(lambda request, args, kwargs: (result, request, args, dict(kwargs)))
|
|
return bound
|
|
|
|
|
|
async def dispatch_without_python(
|
|
dispatch: PublicDispatch[Request], bound: NativeBinding[NativeRoute], rules: Rules, *, asynchronous: bool
|
|
) -> object:
|
|
if not asynchronous:
|
|
return dispatch.run(
|
|
("model",),
|
|
{"page": 1},
|
|
python=NO_PYTHON,
|
|
binding=bound,
|
|
native=lambda hook, value, args, kwargs: hook(value, args, kwargs),
|
|
rules=rules,
|
|
)
|
|
|
|
async def native(
|
|
hook: NativeRoute, value: Request, args: tuple[object, ...], kwargs: Mapping[str, object]
|
|
) -> object:
|
|
return hook(value, args, kwargs)
|
|
|
|
return await dispatch.arun(("model",), {"page": 1}, python=NO_PYTHON, binding=bound, native=native, rules=rules)
|
|
|
|
|
|
def failing_context(request: Request) -> RouteContext:
|
|
raise RuntimeError("detector exploded")
|
|
|
|
|
|
def failing_dispatch() -> PublicDispatch[Request]:
|
|
return PublicDispatch(
|
|
route=Route.CHAT_COMPLETIONS, request=lambda args, kwargs: Request(model="model"), context=failing_context
|
|
)
|
|
|
|
|
|
def reject_native(hook: object, request: Request, args: tuple[object, ...], kwargs: Mapping[str, object]) -> NoReturn:
|
|
pytest.fail("a failed context build must not reach native")
|
|
|
|
|
|
def test_a_detector_error_keeps_the_call_on_python() -> None:
|
|
expected: Final = object()
|
|
result: Final = failing_dispatch().run(
|
|
("model",),
|
|
{},
|
|
python=lambda *args, **kwargs: expected,
|
|
binding=binding(),
|
|
native=reject_native,
|
|
rules=(RouteRule(Route.CHAT_COMPLETIONS, Rollout.RUST_REQUIRED),),
|
|
)
|
|
assert result is expected
|
|
|
|
|
|
async def test_a_detector_error_keeps_the_async_call_on_python() -> None:
|
|
expected: Final = object()
|
|
|
|
async def python(*args: object, **kwargs: object) -> object: # kwargs-ok: stub
|
|
return expected
|
|
|
|
result: Final = await failing_dispatch().arun(
|
|
("model",),
|
|
{},
|
|
python=python,
|
|
binding=binding(),
|
|
native=reject_native,
|
|
rules=(RouteRule(Route.CHAT_COMPLETIONS, Rollout.RUST_REQUIRED),),
|
|
)
|
|
assert result is expected
|
|
|
|
|
|
def ocr_dispatch(request: Request | None) -> PublicDispatch[Request]:
|
|
return PublicDispatch(
|
|
route=Route.OCR,
|
|
request=lambda args, kwargs: request,
|
|
context=lambda value: RouteContext(Route.OCR, model=value.model),
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("asynchronous", (False, True))
|
|
async def test_dispatch_without_python_hands_every_call_to_native(asynchronous: bool) -> None:
|
|
request: Final = Request(model="model")
|
|
result: Final = object()
|
|
dispatched: Final = await dispatch_without_python(
|
|
ocr_dispatch(request),
|
|
no_python_binding(result),
|
|
(RouteRule(Route.OCR, Rollout.RUST_REQUIRED),),
|
|
asynchronous=asynchronous,
|
|
)
|
|
|
|
assert dispatched == (result, request, ("model",), {"page": 1})
|
|
|
|
|
|
@pytest.mark.parametrize("asynchronous", (False, True))
|
|
@pytest.mark.parametrize(
|
|
("request_value", "rules", "reason"),
|
|
(
|
|
(Request(model="model"), (), "must resolve to RUST_REQUIRED"),
|
|
(Request(model="model"), (RouteRule(Route.OCR, Rollout.RUST_OPT_OUT),), "must resolve to RUST_REQUIRED"),
|
|
(None, (RouteRule(Route.OCR, Rollout.RUST_REQUIRED),), "must project to a native request"),
|
|
),
|
|
ids=("no-rule", "opt-out-rule", "unprojectable-call"),
|
|
)
|
|
async def test_dispatch_without_python_never_falls_back(
|
|
asynchronous: bool, request_value: Request | None, rules: Rules, reason: str
|
|
) -> None:
|
|
bound: Final[NativeBinding[NativeRoute]] = NativeBinding("no_python", validate=lambda _: None)
|
|
bound.override(lambda request, args, kwargs: pytest.fail("a misdeclared route must not reach native"))
|
|
|
|
with pytest.raises(NoPythonImplementationError, match=f"ocr has no Python implementation, so .*{reason}"):
|
|
await dispatch_without_python(ocr_dispatch(request_value), bound, rules, asynchronous=asynchronous)
|