litellm/tests/test_litellm_rust/messages/test_dispatch.py
Yujong Lee ca21e2572f
Some checks failed
LiteLLM Rust / rust-lint (push) Has been cancelled
LiteLLM Rust / rust-test (push) Has been cancelled
LiteLLM Rust / rust-wheel (push) Has been cancelled
refactor(rust-bridge): decide native routing in the catalog, drop the bypass predicates
2026-10-19 09:38:54 -07:00

427 lines
17 KiB
Python

import inspect
from collections.abc import Awaitable, Callable, Mapping
from typing import Final, cast # noqa: TID251 # narrows legacy callable signatures for inspect
import pytest
import litellm
from litellm.llms.anthropic.pass_through.messages import handler as python_messages
from litellm.messages import dispatch
from litellm.messages.dispatch import (
_ADISPATCH, # pyright: ignore[reportPrivateUsage] # tests configured dispatch
_DISPATCH, # pyright: ignore[reportPrivateUsage] # tests configured dispatch
)
from litellm.proxy import proxy_server
from litellm.rust_bridge.bindings import NativeBinding
from litellm.rust_bridge.catalog import Feature, Route, RouteRule, Rules, Target
from litellm.rust_bridge.configuration import Rollout
from litellm.rust_bridge.messages.entrypoints import NATIVE_AMESSAGES, NATIVE_MESSAGES, NativeAmessages, NativeMessages
from litellm.rust_bridge.public_call import NativeCall
from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse
MESSAGES: Final = [{"role": "user", "content": "hi"}]
PYTHON_RULES: Final[Rules] = ()
RUST_RULES: Final[Rules] = (RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED),)
def messages_binding(native: NativeMessages | None) -> NativeBinding[NativeMessages]:
binding: Final[NativeBinding[NativeMessages]] = NativeBinding("anthropic_messages_handler", validate=lambda _: None)
binding.override(native)
return binding
def amessages_binding(native: NativeAmessages | None) -> NativeBinding[NativeAmessages]:
binding: Final[NativeBinding[NativeAmessages]] = NativeBinding("anthropic_messages", validate=lambda _: None)
binding.override(native)
return binding
def response(model: str = "claude-sonnet-4-5") -> AnthropicMessagesResponse:
return AnthropicMessagesResponse(id="msg_test", type="message", role="assistant", model=model, content=[])
def test_public_signature_is_the_legacy_signature() -> None:
public_messages: Final = cast(Callable[..., object], litellm.anthropic_messages_handler)
legacy_messages: Final = cast(Callable[..., object], python_messages.anthropic_messages_handler)
public_amessages: Final = cast(Callable[..., object], litellm.anthropic_messages)
legacy_amessages: Final = cast(Callable[..., object], python_messages.anthropic_messages)
assert inspect.signature(public_messages) == inspect.signature(legacy_messages)
assert inspect.signature(public_amessages) == inspect.signature(legacy_amessages)
def test_python_route_forwards_original_call_shape() -> None:
metadata: Final = {"user_id": "u"}
args: Final[tuple[object, ...]] = (16, MESSAGES, "claude-sonnet-4-5")
kwargs: Final[Mapping[str, object]] = {"temperature": 0.1, "litellm_metadata": metadata}
captured: Final[list[tuple[tuple[object, ...], Mapping[str, object]]]] = []
expected: Final = response()
def python(*call_args: object, **call_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: records call shape
captured.append((call_args, call_kwargs))
return expected
def native(
request: NativeCall,
) -> AnthropicMessagesResponse:
pytest.fail("Python-only dispatch must not call native")
result: Final = _DISPATCH.run(
args,
kwargs,
python=python,
binding=messages_binding(native),
native=lambda hook, request, call_args, call_kwargs: hook(request),
rules=PYTHON_RULES,
)
assert result is expected
call_args, call_kwargs = captured[0]
assert call_args == args
assert call_args[1] is MESSAGES
assert call_kwargs == kwargs
assert call_kwargs["litellm_metadata"] is metadata
assert kwargs == {"temperature": 0.1, "litellm_metadata": metadata}
@pytest.mark.asyncio
async def test_async_python_route_forwards_original_call_shape() -> None:
metadata: Final = {"user_id": "u"}
args: Final[tuple[object, ...]] = (16, MESSAGES, "claude-sonnet-4-5")
kwargs: Final[Mapping[str, object]] = {"temperature": 0.1, "litellm_metadata": metadata}
captured: Final[list[tuple[tuple[object, ...], Mapping[str, object]]]] = []
expected: Final = response()
async def python(
*call_args: object,
**call_kwargs: object, # kwargs-ok: records call shape
) -> AnthropicMessagesResponse:
captured.append((call_args, call_kwargs))
return expected
async def native(
request: NativeCall,
) -> AnthropicMessagesResponse:
pytest.fail("Python-only dispatch must not call native")
result: Final = await _ADISPATCH.arun(
args,
kwargs,
python=python,
binding=amessages_binding(native),
native=lambda hook, request, call_args, call_kwargs: hook(request),
rules=PYTHON_RULES,
)
assert result is expected
call_args, call_kwargs = captured[0]
assert call_args == args
assert call_args[1] is MESSAGES
assert call_kwargs == kwargs
assert call_kwargs["litellm_metadata"] is metadata
assert kwargs == {"temperature": 0.1, "litellm_metadata": metadata}
def test_native_receives_normalized_request_and_original_call_shape() -> None:
metadata: Final = {"user_id": "u"}
args: Final[tuple[object, ...]] = (16, MESSAGES, "anthropic/claude-sonnet-4-5")
kwargs: Final[Mapping[str, object]] = {
"stream": True,
"api_key": "sk-test",
"api_base": "https://example.invalid",
"custom_llm_provider": "anthropic",
"litellm_metadata": metadata,
}
captured: Final[list[tuple[NativeCall, tuple[object, ...], Mapping[str, object]]]] = []
expected: Final = response("anthropic/claude-sonnet-4-5")
def python(*call_args: object, **call_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: rejected fallback
pytest.fail("Required Rust dispatch must not call Python")
def native(
request: NativeCall,
) -> AnthropicMessagesResponse:
args: Final = request.args
kwargs: Final = request.kwargs
captured.append((request, args, kwargs))
return expected
result: Final = _DISPATCH.run(
args,
kwargs,
python=python,
binding=messages_binding(native),
native=lambda hook, request, call_args, call_kwargs: hook(request),
rules=RUST_RULES,
)
assert result is expected
request, call_args, call_kwargs = captured[0]
assert request.resolved["model"] == "anthropic/claude-sonnet-4-5"
assert request.resolved["messages"] is MESSAGES
assert request.resolved["max_tokens"] == 16
assert request.resolved["stream"] is True
assert request.resolved["api_key"] == "sk-test"
assert request.resolved["api_base"] == "https://example.invalid"
assert request.resolved["custom_llm_provider"] == "anthropic"
assert request.kwargs == kwargs
assert request.kwargs["litellm_metadata"] is metadata
assert call_args == args
assert call_args[1] is MESSAGES
assert call_kwargs == kwargs
assert call_kwargs["litellm_metadata"] is metadata
@pytest.mark.parametrize("stream, expected_timeout", ((True, 0.5), (False, 30.0)))
def test_native_streams_are_bounded_by_stream_timeout(stream: bool, expected_timeout: float) -> None:
captured: Final[list[NativeCall]] = []
def python(*call_args: object, **call_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: rejected fallback
pytest.fail("Required Rust dispatch must not call Python")
def native(request: NativeCall) -> AnthropicMessagesResponse:
captured.append(request)
return response()
_ = _DISPATCH.run(
(16, MESSAGES, "anthropic/claude-sonnet-4-5"),
{"stream": stream, "timeout": 30.0, "stream_timeout": 0.5},
python=python,
binding=messages_binding(native),
native=lambda hook, request, call_args, call_kwargs: hook(request),
rules=RUST_RULES,
)
assert captured[0].resolved["timeout"] == expected_timeout
@pytest.mark.parametrize(
("model", "kwargs", "passes_through"),
(
("vertex_ai/claude-sonnet-4-5@20250929", {}, True),
("claude-sonnet-4-5@20250929", {"custom_llm_provider": "vertex_ai"}, True),
("deepseek/deepseek-v4-flash", {}, True),
("openai/gpt-5", {"model_info": {"supported_endpoints": ["/v1/messages"]}}, True),
("vertex_ai/gemini-2.5-pro", {}, False),
("gemini-2.5-pro", {"custom_llm_provider": "vertex_ai"}, False),
("openai/gpt-5", {}, False),
("groq/llama-3.3-70b-versatile", {}, False),
("unknown-provider/unknown-model", {}, False),
),
)
def test_only_requests_that_pass_through_to_a_messages_api_reach_a_messages_target_rule(
model: str, kwargs: Mapping[str, object], passes_through: bool
) -> None:
python_result: Final = response(model)
native_result: Final = response(model)
def python(*call_args: object, **call_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: fallback
return python_result
result: Final = _DISPATCH.run(
(16, MESSAGES, model),
kwargs,
python=python,
binding=messages_binding(lambda request: native_result),
native=lambda hook, request, call_args, call_kwargs: hook(request),
rules=(RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED, targets=frozenset({Target.MESSAGES})),),
)
assert result is (native_result if passes_through else python_result)
DANGEROUS_TOOL_USE: Final = [{"type": "dangerous_tool_use", "classifier_context": {"permission_mode": "auto"}}]
CLASSIFIER_SETTINGS: Final = {"safeguards_classifier_model": "gpt-oss-120b"}
@pytest.mark.parametrize(
("model", "kwargs", "general_settings", "serves_python"),
(
("deepseek/deepseek-v4-flash", {"safeguards": DANGEROUS_TOOL_USE}, CLASSIFIER_SETTINGS, True),
(
"deepseek-v4-flash",
{"safeguards": DANGEROUS_TOOL_USE, "custom_llm_provider": "deepseek"},
CLASSIFIER_SETTINGS,
True,
),
("deepseek/deepseek-v4-flash", {"safeguards": DANGEROUS_TOOL_USE}, {}, False),
("deepseek/deepseek-v4-flash", {}, CLASSIFIER_SETTINGS, False),
("deepseek/deepseek-v4-flash", {"safeguards": [{"type": "other"}]}, CLASSIFIER_SETTINGS, False),
("anthropic/claude-sonnet-4-5", {"safeguards": DANGEROUS_TOOL_USE}, CLASSIFIER_SETTINGS, False),
),
)
def test_proxy_classified_safeguards_reach_the_feature_rule(
monkeypatch: pytest.MonkeyPatch,
model: str,
kwargs: Mapping[str, object],
general_settings: Mapping[str, object],
serves_python: bool,
) -> None:
monkeypatch.setattr(proxy_server, "general_settings", general_settings)
python_result: Final = response(model)
native_result: Final = response(model)
def python(*call_args: object, **call_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: fallback
return python_result
result: Final = _DISPATCH.run(
(16, MESSAGES, model),
kwargs,
python=python,
binding=messages_binding(lambda request: native_result),
native=lambda hook, request, call_args, call_kwargs: hook(request),
rules=(
RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY, features=frozenset({Feature.SAFEGUARDS_CLASSIFIER})),
RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED),
),
)
assert result is (python_result if serves_python else native_result)
@pytest.mark.parametrize(
("args", "kwargs"),
(
((16, MESSAGES, "claude-sonnet-4-5"), {"model": "duplicate"}),
((), {}),
),
)
def test_binding_errors_delegate_to_python(args: tuple[object, ...], kwargs: Mapping[str, object]) -> None:
captured: Final[list[tuple[tuple[object, ...], Mapping[str, object]]]] = []
expected: Final = response()
def python(*call_args: object, **call_kwargs: object) -> AnthropicMessagesResponse:
captured.append((call_args, call_kwargs))
return expected
def native(
request: NativeCall,
) -> AnthropicMessagesResponse:
pytest.fail("Binding failures must be delegated to Python")
result: Final = _DISPATCH.run(
args,
kwargs,
python=python,
binding=messages_binding(native),
native=lambda hook, request, call_args, call_kwargs: hook(request),
rules=RUST_RULES,
)
assert result is expected
assert captured == [(args, kwargs)]
def test_anthropic_create_routes_through_dispatch() -> None:
captured: Final[list[NativeCall]] = []
expected: Final = response()
def native(
request: NativeCall,
) -> AnthropicMessagesResponse:
captured.append(request)
return expected
NATIVE_MESSAGES.override(native)
public_create: Final = cast(Callable[..., AnthropicMessagesResponse], litellm.anthropic.create)
try:
result: Final = public_create(max_tokens=16, messages=MESSAGES, model="claude-sonnet-4-5")
finally:
NATIVE_MESSAGES.reset()
assert result is expected
assert [request.resolved["model"] for request in captured] == ["claude-sonnet-4-5"]
@pytest.mark.asyncio
async def test_anthropic_acreate_routes_through_dispatch() -> None:
captured: Final[list[NativeCall]] = []
expected: Final = response()
async def native(
request: NativeCall,
) -> AnthropicMessagesResponse:
captured.append(request)
return expected
NATIVE_AMESSAGES.override(native)
public_acreate: Final = cast(Callable[..., Awaitable[AnthropicMessagesResponse]], litellm.anthropic.acreate)
try:
result: Final = await public_acreate(max_tokens=16, messages=MESSAGES, model="claude-sonnet-4-5")
finally:
NATIVE_AMESSAGES.reset()
assert result is expected
assert [request.resolved["model"] for request in captured] == ["claude-sonnet-4-5"]
@pytest.mark.asyncio
async def test_truthy_is_async_hands_off_to_the_async_entry_like_true() -> None:
captured: Final[list[NativeCall]] = []
expected: Final = response()
async def native(request: NativeCall) -> AnthropicMessagesResponse:
captured.append(request)
return expected
def sync_native(request: NativeCall) -> AnthropicMessagesResponse:
pytest.fail("a truthy is_async must hand off to the async entry, as the Python handler does")
NATIVE_AMESSAGES.override(native)
NATIVE_MESSAGES.override(sync_native)
public_create: Final = cast(Callable[..., object], litellm.anthropic.create)
try:
pending: Final = public_create(max_tokens=16, messages=MESSAGES, model="claude-sonnet-4-5", is_async=1)
assert inspect.isawaitable(pending)
assert await pending is expected
reference: Final = public_create(max_tokens=16, messages=MESSAGES, model="claude-sonnet-4-5", is_async=True)
assert inspect.isawaitable(reference)
assert await reference is expected
finally:
NATIVE_AMESSAGES.reset()
NATIVE_MESSAGES.reset()
assert [dict(request.resolved) for request in captured] == [dict(captured[1].resolved)] * 2
def test_sync_messages_request_projects_public_arguments() -> None:
rules: Final[Rules] = (RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED),)
expected: Final = AnthropicMessagesResponse(model="claude-test")
def native(request: NativeCall) -> AnthropicMessagesResponse:
assert request.resolved["model"] == "claude-test"
assert request.resolved["messages"] == MESSAGES
assert request.resolved["max_tokens"] == 10
assert request.resolved["custom_llm_provider"] == "anthropic"
return expected
binding: Final[NativeBinding[NativeMessages]] = NativeBinding("messages", validate=lambda _: None)
binding.override(native)
response: Final = dispatch._DISPATCH.run( # pyright: ignore[reportPrivateUsage] # test an explicit route decision
(),
{
"model": "claude-test",
"messages": MESSAGES,
"max_tokens": 10,
"custom_llm_provider": "anthropic",
},
python=lambda *args, **kwargs: pytest.fail("required native route must handle this call"),
binding=binding,
native=lambda hook, request, args, kwargs: hook(request),
rules=rules,
)
assert response is expected
def test_messages_binding_error_delegates_unchanged_to_python() -> None:
rules: Final[Rules] = (RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED),)
expected: Final = AnthropicMessagesResponse(model="claude-test")
def python(*args: object, **kwargs: object) -> AnthropicMessagesResponse:
return expected
def native(request: NativeCall) -> AnthropicMessagesResponse:
pytest.fail("a call without max_tokens cannot project a request and must stay on Python")
binding: Final[NativeBinding[NativeMessages]] = NativeBinding("messages", validate=lambda _: None)
binding.override(native)
response: Final = dispatch._DISPATCH.run( # pyright: ignore[reportPrivateUsage] # test an explicit route decision
(),
{"model": "claude-test", "messages": MESSAGES, "custom_llm_provider": "anthropic"},
python=python,
binding=binding,
native=lambda hook, request, args, kwargs: hook(request),
rules=rules,
)
assert response is expected