mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-09 03:18:44 +00:00
391 lines
12 KiB
Python
391 lines
12 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from types import SimpleNamespace
|
|
from typing import Final
|
|
|
|
import pytest
|
|
|
|
from litellm.exceptions import APIError, AuthenticationError, InternalServerError, RateLimitError
|
|
from litellm.rust_bridge import bindings, runtime
|
|
|
|
|
|
class RustBridgeDeclined(Exception):
|
|
pass
|
|
|
|
|
|
class RustUpstreamError(Exception):
|
|
pass
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def native_exceptions(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setattr(
|
|
bindings,
|
|
"get_native_bridge",
|
|
lambda: SimpleNamespace(
|
|
RustBridgeDeclined=RustBridgeDeclined,
|
|
RustUpstreamError=RustUpstreamError,
|
|
),
|
|
)
|
|
|
|
|
|
def context() -> runtime.BridgeErrorContext:
|
|
return runtime.BridgeErrorContext(provider="anthropic", model="model")
|
|
|
|
|
|
def enabled() -> bool:
|
|
return True
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class FallbackCase:
|
|
process_enabled: bool | None = None
|
|
eligible: bool = True
|
|
binding_available: bool = True
|
|
declined: bool = False
|
|
expected_events: tuple[str, ...] = ()
|
|
|
|
|
|
FALLBACK_CASES: Final = (
|
|
pytest.param(
|
|
FallbackCase(process_enabled=False, expected_events=("python",)),
|
|
id="process-disabled",
|
|
),
|
|
pytest.param(
|
|
FallbackCase(eligible=False, expected_events=("python",)),
|
|
id="request-ineligible",
|
|
),
|
|
pytest.param(
|
|
FallbackCase(binding_available=False, expected_events=("load", "python")),
|
|
id="bridge-unavailable",
|
|
),
|
|
pytest.param(
|
|
FallbackCase(declined=True, expected_events=("load", "prepare", "rust", "python")),
|
|
id="bridge-declined",
|
|
),
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("case", FALLBACK_CASES)
|
|
def test_invoke_falls_back_only_before_provider_success(case: FallbackCase) -> None:
|
|
events: list[str] = []
|
|
|
|
def load() -> object | None:
|
|
events.append("load")
|
|
return object() if case.binding_available else None
|
|
|
|
def call(_binding: object, _request: object) -> int:
|
|
events.append("rust")
|
|
if case.declined:
|
|
raise RustBridgeDeclined("unsupported")
|
|
return 3
|
|
|
|
bridge: Final = runtime.EndpointBinding(
|
|
route="messages", load=load, enabled=lambda: case.process_enabled is not False
|
|
)
|
|
result: Final = bridge.invoke(
|
|
prepare=lambda: events.append("prepare"),
|
|
call=call,
|
|
fallback=lambda: events.append("python") or "fallback",
|
|
adapt=str,
|
|
error_context=context(),
|
|
eligible=case.eligible,
|
|
)
|
|
|
|
assert result == "fallback"
|
|
assert tuple(events) == case.expected_events
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("case", FALLBACK_CASES)
|
|
async def test_ainvoke_matches_sync_fallback_contract(case: FallbackCase) -> None:
|
|
events: list[str] = []
|
|
|
|
def load() -> object | None:
|
|
events.append("load")
|
|
return object() if case.binding_available else None
|
|
|
|
async def call(_binding: object, _request: object) -> int:
|
|
events.append("rust")
|
|
if case.declined:
|
|
raise RustBridgeDeclined("unsupported")
|
|
return 3
|
|
|
|
async def fallback() -> str:
|
|
events.append("python")
|
|
return "fallback"
|
|
|
|
bridge: Final = runtime.EndpointBinding(
|
|
route="messages", load=load, enabled=lambda: case.process_enabled is not False
|
|
)
|
|
result: Final = await bridge.ainvoke(
|
|
prepare=lambda: events.append("prepare"),
|
|
call=call,
|
|
fallback=fallback,
|
|
adapt=str,
|
|
error_context=context(),
|
|
eligible=case.eligible,
|
|
)
|
|
|
|
assert result == "fallback"
|
|
assert tuple(events) == case.expected_events
|
|
|
|
|
|
def test_invoke_adapts_native_success_without_fallback() -> None:
|
|
bridge: Final = runtime.EndpointBinding(route="messages", load=object, enabled=enabled)
|
|
|
|
result: Final = bridge.invoke(
|
|
prepare=lambda: 3,
|
|
call=lambda _binding, request: request * 2,
|
|
fallback=lambda: pytest.fail("fallback must not run"),
|
|
adapt=lambda value: f"adapted-{value}",
|
|
error_context=context(),
|
|
)
|
|
|
|
assert result == "adapted-6"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_ainvoke_adapts_native_success_without_fallback() -> None:
|
|
async def call(_binding: object, request: int) -> int:
|
|
return request * 2
|
|
|
|
async def fallback() -> str:
|
|
pytest.fail("fallback must not run")
|
|
|
|
bridge: Final = runtime.EndpointBinding(route="messages", load=object, enabled=enabled)
|
|
result: Final = await bridge.ainvoke(
|
|
prepare=lambda: 3,
|
|
call=call,
|
|
fallback=fallback,
|
|
adapt=lambda value: f"adapted-{value}",
|
|
error_context=context(),
|
|
)
|
|
|
|
assert result == "adapted-6"
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("error", "expected_type", "expected_status", "expected_message"),
|
|
(
|
|
pytest.param(RustUpstreamError(401, "unauthorized"), AuthenticationError, 401, "unauthorized", id="auth"),
|
|
pytest.param(RustUpstreamError(429, "rate limited"), RateLimitError, 429, "rate limited", id="rate-limit"),
|
|
pytest.param(RustUpstreamError(500, "failed"), InternalServerError, 500, "failed", id="server-error"),
|
|
pytest.param(RustUpstreamError(0, "connection reset"), APIError, 500, "connection reset", id="transport"),
|
|
pytest.param(RustUpstreamError(403, "forbidden"), APIError, 403, "forbidden", id="other-status"),
|
|
),
|
|
)
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("asynchronous", (False, True))
|
|
async def test_upstream_failure_maps_to_api_error_without_fallback(
|
|
asynchronous: bool,
|
|
error: RustUpstreamError,
|
|
expected_type: type[BaseException],
|
|
expected_status: int,
|
|
expected_message: str,
|
|
) -> None:
|
|
def fail(_binding: object, _request: object) -> object:
|
|
raise error
|
|
|
|
async def afail(binding: object, request: object) -> object:
|
|
return fail(binding, request)
|
|
|
|
async def fallback() -> str:
|
|
pytest.fail("fallback must not run")
|
|
|
|
bridge: Final = runtime.EndpointBinding(route="messages", load=object, enabled=enabled)
|
|
|
|
async def invoke() -> None:
|
|
if asynchronous:
|
|
await bridge.ainvoke(
|
|
prepare=lambda: None, call=afail, fallback=fallback, adapt=str, error_context=context()
|
|
)
|
|
else:
|
|
bridge.invoke(
|
|
prepare=lambda: None,
|
|
call=fail,
|
|
fallback=lambda: pytest.fail("fallback must not run"),
|
|
adapt=str,
|
|
error_context=context(),
|
|
)
|
|
|
|
with pytest.raises(expected_type, match=expected_message) as caught:
|
|
await invoke()
|
|
|
|
assert type(caught.value) is expected_type
|
|
assert caught.value.status_code == expected_status
|
|
assert caught.value.llm_provider == "anthropic"
|
|
assert caught.value.model == "model"
|
|
assert caught.value.__cause__ is error
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_upstream_failure_maps_to_api_error_without_fallback() -> None:
|
|
async def fail(_binding: object, _request: object) -> object:
|
|
raise RustUpstreamError(503, "overloaded")
|
|
|
|
async def fallback() -> object:
|
|
pytest.fail("fallback must not run")
|
|
|
|
bridge: Final = runtime.EndpointBinding(route="messages", load=object, enabled=enabled)
|
|
|
|
with pytest.raises(APIError, match="overloaded") as caught:
|
|
await bridge.ainvoke(prepare=lambda: None, call=fail, fallback=fallback, adapt=str, error_context=context())
|
|
|
|
assert caught.value.status_code == 503
|
|
|
|
|
|
def test_unknown_failure_is_preserved_without_fallback() -> None:
|
|
error: Final = RuntimeError("unknown")
|
|
bridge: Final = runtime.EndpointBinding(route="messages", load=object, enabled=enabled)
|
|
|
|
with pytest.raises(RuntimeError, match="unknown") as caught:
|
|
bridge.invoke(
|
|
prepare=lambda: None,
|
|
call=lambda _binding, _request: (_ for _ in ()).throw(error),
|
|
fallback=lambda: pytest.fail("fallback must not run"),
|
|
adapt=str,
|
|
error_context=context(),
|
|
)
|
|
|
|
assert caught.value is error
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("process_enabled", "binding_available", "declined", "expected_message"),
|
|
(
|
|
pytest.param(False, True, False, "native messages endpoint is disabled", id="disabled"),
|
|
pytest.param(None, False, False, "native messages endpoint is unavailable", id="unavailable"),
|
|
pytest.param(
|
|
None,
|
|
True,
|
|
True,
|
|
"native messages endpoint declined the request: unsupported",
|
|
id="declined",
|
|
),
|
|
),
|
|
)
|
|
def test_require_explains_why_rust_did_not_handle_request(
|
|
process_enabled: bool | None,
|
|
binding_available: bool,
|
|
declined: bool,
|
|
expected_message: str,
|
|
) -> None:
|
|
def call(_binding: object, _request: object) -> object:
|
|
if declined:
|
|
raise RustBridgeDeclined("unsupported")
|
|
return object()
|
|
|
|
bridge: Final = runtime.EndpointBinding(
|
|
route="messages",
|
|
load=object if binding_available else lambda: None,
|
|
enabled=lambda: process_enabled is not False,
|
|
)
|
|
|
|
with pytest.raises(RuntimeError, match=f"^{expected_message}$"):
|
|
bridge.require(
|
|
prepare=lambda: None,
|
|
call=call,
|
|
adapt=str,
|
|
error_context=context(),
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("state", "expected", "expected_events"),
|
|
(
|
|
pytest.param("disabled", False, (), id="disabled"),
|
|
pytest.param("ineligible", False, (), id="ineligible"),
|
|
pytest.param("unavailable", False, ("load",), id="unavailable"),
|
|
pytest.param("available", True, ("load",), id="available"),
|
|
),
|
|
)
|
|
def test_can_attempt_only_enabled_available_requests(
|
|
state: str,
|
|
expected: bool,
|
|
expected_events: tuple[str, ...],
|
|
) -> None:
|
|
events: list[str] = []
|
|
|
|
def load() -> object | None:
|
|
events.append("load")
|
|
return None if state == "unavailable" else object()
|
|
|
|
bridge: Final = runtime.EndpointBinding(route="messages", load=load, enabled=lambda: state != "disabled")
|
|
|
|
assert (
|
|
bridge.can_attempt(
|
|
eligible=state != "ineligible",
|
|
)
|
|
is expected
|
|
)
|
|
assert tuple(events) == expected_events
|
|
|
|
|
|
def test_native_endpoint_applies_partial_overrides_and_reset(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
def native_sync() -> str:
|
|
return "native"
|
|
|
|
async def native_async() -> str:
|
|
return "native async"
|
|
|
|
def replacement_sync() -> str:
|
|
return "replacement"
|
|
|
|
monkeypatch.setattr(
|
|
bindings,
|
|
"get_native_bridge",
|
|
lambda: SimpleNamespace(chat_completions=native_sync, achat_completions=native_async),
|
|
)
|
|
endpoint: Final[runtime.EndpointDispatch[object, object]] = runtime.EndpointDispatch.native(
|
|
route="test",
|
|
sync=lambda native: native.chat_completions,
|
|
asynchronous=lambda native: native.achat_completions,
|
|
enabled=enabled,
|
|
)
|
|
|
|
assert endpoint.sync.load() is native_sync
|
|
assert endpoint.asynchronous.load() is native_async
|
|
endpoint.override(sync=replacement_sync)
|
|
assert endpoint.sync.load() is replacement_sync
|
|
assert endpoint.asynchronous.load() is native_async
|
|
endpoint.override(asynchronous=None)
|
|
assert endpoint.sync.load() is replacement_sync
|
|
assert endpoint.asynchronous.load() is None
|
|
endpoint.reset()
|
|
assert endpoint.sync.load() is native_sync
|
|
assert endpoint.asynchronous.load() is native_async
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("asynchronous", (False, True))
|
|
async def test_response_adaptation_failure_never_authorizes_fallback(asynchronous: bool) -> None:
|
|
def adapt(value: str) -> str:
|
|
assert value == "provider response"
|
|
raise RustBridgeDeclined("adapter failed after provider response")
|
|
|
|
async def native(binding: object, request: object) -> str:
|
|
return "provider response"
|
|
|
|
async def fallback() -> str:
|
|
pytest.fail("a received response must not be retried")
|
|
|
|
bridge = runtime.EndpointBinding(route="messages", load=object, enabled=enabled)
|
|
|
|
async def invoke() -> None:
|
|
if asynchronous:
|
|
await bridge.ainvoke(
|
|
prepare=lambda: None, call=native, fallback=fallback, adapt=adapt, error_context=context()
|
|
)
|
|
else:
|
|
bridge.invoke(
|
|
prepare=lambda: None,
|
|
call=lambda binding, request: "provider response",
|
|
fallback=lambda: pytest.fail("a received response must not be retried"),
|
|
adapt=adapt,
|
|
error_context=context(),
|
|
)
|
|
|
|
with pytest.raises(RustBridgeDeclined, match="adapter failed"):
|
|
await invoke()
|