mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-24 05:45:11 +00:00
refactor(rust-bridge): decide native routing in the catalog, drop the bypass predicates
This commit is contained in:
parent
2dfeb80e04
commit
aaf491e80e
23 changed files with 862 additions and 888 deletions
7
.github/workflows/test-unit.yml
vendored
7
.github/workflows/test-unit.yml
vendored
|
|
@ -430,6 +430,13 @@ jobs:
|
|||
|
||||
- shard: rust-bridge-harness
|
||||
test-path: >-
|
||||
tests/test_litellm_rust/test_dispatch.py
|
||||
tests/test_litellm_rust/test_public_entries.py
|
||||
tests/test_litellm_rust/chat_completions/test_dispatch.py
|
||||
tests/test_litellm_rust/embeddings/test_dispatch.py
|
||||
tests/test_litellm_rust/messages/test_dispatch.py
|
||||
tests/test_litellm_rust/ocr/test_dispatch.py
|
||||
tests/test_litellm_rust/responses/test_dispatch.py
|
||||
tests/unit/compression
|
||||
tests/unit/harness
|
||||
tests/unit/rust_bridge
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ from litellm.rust_bridge.chat_completions.entrypoints import (
|
|||
)
|
||||
from litellm.rust_bridge.dispatch import PublicDispatch
|
||||
from litellm.rust_bridge.public_call import (
|
||||
AsyncFlag,
|
||||
NativeCall,
|
||||
bind,
|
||||
native_call,
|
||||
|
|
@ -44,6 +45,7 @@ def _python_acompletion() -> PythonAcompletion:
|
|||
|
||||
_PYTHON_COMPLETION: Final = _python_completion()
|
||||
_COMPLETION: Final = signature(_PYTHON_COMPLETION)
|
||||
_FLAG: Final = AsyncFlag("acompletion", sync=_COMPLETION)
|
||||
_PYTHON_ACOMPLETION: Final = _python_acompletion()
|
||||
_ACOMPLETION: Final = signature(_PYTHON_ACOMPLETION)
|
||||
|
||||
|
|
@ -73,7 +75,6 @@ _DISPATCH: Final = PublicDispatch(
|
|||
route=Route.CHAT_COMPLETIONS,
|
||||
request=lambda args, kwargs: _public_request(_COMPLETION, args, kwargs),
|
||||
context=_context,
|
||||
bypass=lambda request: request.kwargs.get("acompletion") is True,
|
||||
)
|
||||
|
||||
_ADISPATCH: Final = PublicDispatch(
|
||||
|
|
@ -88,9 +89,11 @@ def completion(
|
|||
**kwargs: object, # kwargs-ok: preserve the public chat completions call shape
|
||||
) -> ChatResult | Coroutine[object, object, ChatResult]:
|
||||
python: Final = _PYTHON_COMPLETION
|
||||
if _FLAG.requested(kwargs):
|
||||
return acompletion(**_FLAG.keywords(args, kwargs))
|
||||
return _DISPATCH.run(
|
||||
args,
|
||||
kwargs,
|
||||
_FLAG.dropped(kwargs),
|
||||
python=python,
|
||||
binding=NATIVE_COMPLETION,
|
||||
native=native_call_hook,
|
||||
|
|
@ -101,7 +104,7 @@ async def acompletion(*args: object, **kwargs: object) -> ChatResult: # kwargs-
|
|||
python: Final = _PYTHON_ACOMPLETION
|
||||
return await _ADISPATCH.arun(
|
||||
args,
|
||||
kwargs,
|
||||
_FLAG.dropped(kwargs),
|
||||
python=python,
|
||||
binding=NATIVE_ACOMPLETION,
|
||||
native=native_call_hook,
|
||||
|
|
|
|||
|
|
@ -9,7 +9,15 @@ from litellm.rust_bridge.embeddings.entrypoints import (
|
|||
NATIVE_AEMBEDDING,
|
||||
NATIVE_EMBEDDING,
|
||||
)
|
||||
from litellm.rust_bridge.public_call import NativeCall, bind, native_call, native_call_hook, optional_str, signature
|
||||
from litellm.rust_bridge.public_call import (
|
||||
AsyncFlag,
|
||||
NativeCall,
|
||||
bind,
|
||||
native_call,
|
||||
native_call_hook,
|
||||
optional_str,
|
||||
signature,
|
||||
)
|
||||
from litellm.types.utils import EmbeddingResponse
|
||||
|
||||
__all__ = ("aembedding", "embedding")
|
||||
|
|
@ -24,6 +32,7 @@ _PYTHON_AEMBEDDING: Final = cast( # cast-ok: [LIT006] preserve the legacy publi
|
|||
PythonAembedding, main.aembedding
|
||||
)
|
||||
_EMBEDDING_SIGNATURE: Final = signature(_PYTHON_EMBEDDING)
|
||||
_FLAG: Final = AsyncFlag("aembedding", sync=_EMBEDDING_SIGNATURE)
|
||||
|
||||
|
||||
def _public_request(
|
||||
|
|
@ -50,7 +59,6 @@ _DISPATCH: Final = PublicDispatch(
|
|||
route=Route.EMBEDDINGS,
|
||||
request=lambda args, kwargs: _public_request(_EMBEDDING_SIGNATURE, args, kwargs),
|
||||
context=_context,
|
||||
bypass=lambda request: request.kwargs.get("aembedding") is True,
|
||||
)
|
||||
|
||||
_ADISPATCH: Final = PublicDispatch(
|
||||
|
|
@ -64,9 +72,11 @@ def embedding(
|
|||
*args: object,
|
||||
**kwargs: object, # kwargs-ok: preserve the public embedding call shape
|
||||
) -> EmbeddingResponse | Coroutine[object, object, EmbeddingResponse]:
|
||||
if _FLAG.requested(kwargs):
|
||||
return aembedding(**_FLAG.keywords(args, kwargs))
|
||||
return _DISPATCH.run(
|
||||
args,
|
||||
kwargs,
|
||||
_FLAG.dropped(kwargs),
|
||||
python=_PYTHON_EMBEDDING,
|
||||
binding=NATIVE_EMBEDDING,
|
||||
native=native_call_hook,
|
||||
|
|
@ -76,7 +86,7 @@ def embedding(
|
|||
async def aembedding(*args: object, **kwargs: object) -> EmbeddingResponse: # kwargs-ok: preserve the public call shape
|
||||
return await _ADISPATCH.arun(
|
||||
args,
|
||||
kwargs,
|
||||
_FLAG.dropped(kwargs),
|
||||
python=_PYTHON_AEMBEDDING,
|
||||
binding=NATIVE_AEMBEDDING,
|
||||
native=native_call_hook,
|
||||
|
|
|
|||
|
|
@ -1,18 +1,20 @@
|
|||
import inspect
|
||||
from collections.abc import AsyncIterator, Awaitable, Callable, Coroutine, Iterator, Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import Final, TypeAlias, cast # noqa: TID251 # native binding selects a sync result or an async awaitable
|
||||
|
||||
from litellm.exceptions import BadRequestError
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
|
||||
from litellm.llms.anthropic.pass_through.messages import handler as main
|
||||
from litellm.llms.anthropic.pass_through.safeguards import native_route_classifies
|
||||
from litellm.rust_bridge.catalog import Route, RouteContext
|
||||
from litellm.rust_bridge.catalog import Feature, Route, RouteContext, Target
|
||||
from litellm.rust_bridge.dispatch import PublicDispatch
|
||||
from litellm.rust_bridge.messages.entrypoints import (
|
||||
NATIVE_AMESSAGES,
|
||||
NATIVE_MESSAGES,
|
||||
)
|
||||
from litellm.rust_bridge.public_call import (
|
||||
AsyncFlag,
|
||||
NativeCall,
|
||||
bind,
|
||||
native_call,
|
||||
|
|
@ -46,6 +48,7 @@ def _python_amessages() -> PythonAmessages:
|
|||
|
||||
_PYTHON_MESSAGES: Final = _python_messages()
|
||||
_MESSAGES: Final = signature(_PYTHON_MESSAGES)
|
||||
_FLAG: Final = AsyncFlag("is_async", sync=_MESSAGES, is_set=bool)
|
||||
_PYTHON_AMESSAGES: Final = _python_amessages()
|
||||
_AMESSAGES: Final = signature(_PYTHON_AMESSAGES)
|
||||
|
||||
|
|
@ -71,30 +74,49 @@ def _public_request(
|
|||
return native_call(legacy, args, _native_kwargs(fields, kwargs))
|
||||
|
||||
|
||||
def _resolved_provider(request: NativeCall) -> str | None:
|
||||
try:
|
||||
return get_llm_provider(
|
||||
str(request.resolved["model"]), optional_str(request.resolved.get("custom_llm_provider"))
|
||||
)[1]
|
||||
except BadRequestError:
|
||||
return optional_str(request.resolved.get("custom_llm_provider"))
|
||||
|
||||
|
||||
def _python_serves(request: NativeCall) -> bool:
|
||||
"""Vertex AI serves only Claude through the Anthropic Messages config, as
|
||||
``get_provider_anthropic_messages_config`` does; every other Vertex model keeps Python's
|
||||
translation through Chat Completions. A request Python would answer with classifier
|
||||
``safeguard_results`` stays on Python too."""
|
||||
def _resolved(request: NativeCall) -> tuple[str, str | None]:
|
||||
model: Final = str(request.resolved["model"])
|
||||
vertex_without_claude: Final = _resolved_provider(request) == "vertex_ai" and "claude" not in model.lower()
|
||||
return vertex_without_claude or native_route_classifies(model, request.kwargs.get("safeguards"))
|
||||
provider: Final = optional_str(request.resolved.get("custom_llm_provider"))
|
||||
try:
|
||||
resolved_model, resolved_provider, _, _ = get_llm_provider(model, provider)
|
||||
except BadRequestError:
|
||||
return model, provider
|
||||
return resolved_model, resolved_provider
|
||||
|
||||
|
||||
def _target(request: NativeCall, model: str, provider: str | None) -> Target:
|
||||
requested_model: Final = str(request.resolved["model"])
|
||||
model_info: Final = request.resolved.get("model_info")
|
||||
api_base: Final = request.resolved.get("api_base")
|
||||
native_config: Final = main._native_messages_config( # pyright: ignore[reportPrivateUsage] # the handler's own passthrough lookup
|
||||
model, provider, model_info, api_base if isinstance(api_base, str) else None
|
||||
)
|
||||
if native_config is not None:
|
||||
return Target.MESSAGES
|
||||
compacting: Final = request.resolved.get("compaction") is not None
|
||||
if not compacting and main._should_route_to_responses_api(provider, requested_model, model, model_info): # pyright: ignore[reportPrivateUsage] # the handler's own translation choice
|
||||
return Target.RESPONSES
|
||||
return Target.CHAT_COMPLETIONS
|
||||
|
||||
|
||||
_FEATURES: Final[Mapping[Feature, Callable[[NativeCall], bool]]] = MappingProxyType(
|
||||
{
|
||||
Feature.SAFEGUARDS_CLASSIFIER: lambda request: native_route_classifies(
|
||||
str(request.resolved["model"]), request.resolved.get("safeguards")
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _context(request: NativeCall) -> RouteContext:
|
||||
requested_model: Final = str(request.resolved["model"])
|
||||
model, provider = _resolved(request)
|
||||
return RouteContext(
|
||||
Route.MESSAGES,
|
||||
provider=_resolved_provider(request),
|
||||
model=str(request.resolved["model"]),
|
||||
provider=provider,
|
||||
model=requested_model,
|
||||
target=_target(request, model, provider),
|
||||
features=frozenset(feature for feature, detected in _FEATURES.items() if detected(request)),
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -102,14 +124,12 @@ _DISPATCH: Final = PublicDispatch(
|
|||
route=Route.MESSAGES,
|
||||
request=lambda args, kwargs: _public_request(_MESSAGES, args, kwargs),
|
||||
context=_context,
|
||||
bypass=lambda request: request.kwargs.get("is_async") is True or _python_serves(request),
|
||||
)
|
||||
|
||||
_ADISPATCH: Final = PublicDispatch(
|
||||
route=Route.MESSAGES,
|
||||
request=lambda args, kwargs: _public_request(_AMESSAGES, args, kwargs),
|
||||
context=_context,
|
||||
bypass=_python_serves,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -118,9 +138,11 @@ def anthropic_messages_handler(
|
|||
**kwargs: object, # kwargs-ok: preserve the public Anthropic Messages call shape
|
||||
) -> MessagesResult | Coroutine[object, object, MessagesResult]:
|
||||
python: Final = _PYTHON_MESSAGES
|
||||
if _FLAG.requested(kwargs):
|
||||
return anthropic_messages(**_FLAG.keywords(args, kwargs))
|
||||
return _DISPATCH.run(
|
||||
args,
|
||||
kwargs,
|
||||
_FLAG.dropped(kwargs),
|
||||
python=python,
|
||||
binding=NATIVE_MESSAGES,
|
||||
native=native_call_hook,
|
||||
|
|
@ -131,7 +153,7 @@ async def anthropic_messages(*args: object, **kwargs: object) -> MessagesResult:
|
|||
python: Final = _PYTHON_AMESSAGES
|
||||
return await _ADISPATCH.arun(
|
||||
args,
|
||||
kwargs,
|
||||
_FLAG.dropped(kwargs),
|
||||
python=python,
|
||||
binding=NATIVE_AMESSAGES,
|
||||
native=native_call_hook,
|
||||
|
|
|
|||
|
|
@ -9,7 +9,14 @@ from litellm.rust_bridge import runtime
|
|||
from litellm.rust_bridge.catalog import Route, RouteContext
|
||||
from litellm.rust_bridge.dispatch import PublicDispatch
|
||||
from litellm.rust_bridge.ocr.entrypoints import NATIVE_AOCR, NATIVE_OCR
|
||||
from litellm.rust_bridge.public_call import NativeCall, native_call, native_call_hook, optional_str, signature
|
||||
from litellm.rust_bridge.public_call import (
|
||||
AsyncFlag,
|
||||
NativeCall,
|
||||
native_call,
|
||||
native_call_hook,
|
||||
optional_str,
|
||||
signature,
|
||||
)
|
||||
|
||||
__all__ = ("aocr", "ocr")
|
||||
|
||||
|
|
@ -39,6 +46,7 @@ def _bind_request(
|
|||
|
||||
|
||||
_OCR: Final = signature(_bind_request)
|
||||
_FLAG: Final = AsyncFlag("aocr", sync=_OCR)
|
||||
|
||||
|
||||
def _public_request(name: str, args: tuple[object, ...], kwargs: Mapping[str, object]) -> NativeCall:
|
||||
|
|
@ -59,7 +67,6 @@ _DISPATCH: Final = PublicDispatch(
|
|||
route=Route.OCR,
|
||||
request=lambda args, kwargs: _public_request("ocr", args, kwargs),
|
||||
context=_context,
|
||||
bypass=lambda request: request.kwargs.get("aocr") is True,
|
||||
)
|
||||
|
||||
_ADISPATCH: Final = PublicDispatch(
|
||||
|
|
@ -73,9 +80,11 @@ def ocr(
|
|||
*args: object,
|
||||
**kwargs: object, # kwargs-ok: preserve the public OCR call shape
|
||||
) -> OCRResponse | Coroutine[object, object, OCRResponse]:
|
||||
if _FLAG.requested(kwargs):
|
||||
return aocr(**_FLAG.keywords(args, kwargs))
|
||||
return _DISPATCH.run(
|
||||
args,
|
||||
kwargs,
|
||||
_FLAG.dropped(kwargs),
|
||||
python=runtime.NO_PYTHON,
|
||||
binding=NATIVE_OCR,
|
||||
native=native_call_hook,
|
||||
|
|
@ -85,7 +94,7 @@ def ocr(
|
|||
async def aocr(*args: object, **kwargs: object) -> OCRResponse: # kwargs-ok: preserve the public OCR call shape
|
||||
return await _ADISPATCH.arun(
|
||||
args,
|
||||
kwargs,
|
||||
_FLAG.dropped(kwargs),
|
||||
python=runtime.NO_PYTHON,
|
||||
binding=NATIVE_AOCR,
|
||||
native=native_call_hook,
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterat
|
|||
from litellm.rust_bridge.catalog import Route, RouteContext
|
||||
from litellm.rust_bridge.dispatch import PublicDispatch
|
||||
from litellm.rust_bridge.public_call import (
|
||||
AsyncFlag,
|
||||
NativeCall,
|
||||
bind,
|
||||
native_call,
|
||||
|
|
@ -43,6 +44,7 @@ def _python_aresponses() -> PythonAresponses:
|
|||
|
||||
_PYTHON_RESPONSES: Final = _python_responses()
|
||||
_RESPONSES: Final = signature(_PYTHON_RESPONSES)
|
||||
_FLAG: Final = AsyncFlag("aresponses", sync=_RESPONSES)
|
||||
_PYTHON_ARESPONSES: Final = _python_aresponses()
|
||||
_ARESPONSES: Final = signature(_PYTHON_ARESPONSES)
|
||||
|
||||
|
|
@ -71,7 +73,6 @@ _DISPATCH: Final = PublicDispatch(
|
|||
route=Route.RESPONSES,
|
||||
request=lambda args, kwargs: _public_request(_RESPONSES, args, kwargs),
|
||||
context=_context,
|
||||
bypass=lambda request: request.kwargs.get("aresponses") is True,
|
||||
)
|
||||
|
||||
_ADISPATCH: Final = PublicDispatch(
|
||||
|
|
@ -86,9 +87,11 @@ def responses(
|
|||
**kwargs: object, # kwargs-ok: preserve the public Responses call shape
|
||||
) -> ResponsesResult | Coroutine[object, object, ResponsesResult]:
|
||||
python: Final = _PYTHON_RESPONSES
|
||||
if _FLAG.requested(kwargs):
|
||||
return aresponses(**_FLAG.keywords(args, kwargs))
|
||||
return _DISPATCH.run(
|
||||
args,
|
||||
kwargs,
|
||||
_FLAG.dropped(kwargs),
|
||||
python=python,
|
||||
binding=NATIVE_RESPONSES,
|
||||
native=native_call_hook,
|
||||
|
|
@ -99,7 +102,7 @@ async def aresponses(*args: object, **kwargs: object) -> ResponsesResult: # kwa
|
|||
python: Final = _PYTHON_ARESPONSES
|
||||
return await _ADISPATCH.arun(
|
||||
args,
|
||||
kwargs,
|
||||
_FLAG.dropped(kwargs),
|
||||
python=python,
|
||||
binding=NATIVE_ARESPONSES,
|
||||
native=native_call_hook,
|
||||
|
|
|
|||
|
|
@ -13,3 +13,9 @@ Await each selected suspension in the caller's task so hooks retain thread, loop
|
|||
Native admission may select legacy before execution starts. A failure after execution starts is terminal and never authorizes fallback. Creating an unstarted native coroutine must not acquire clients or credentials or capture execution context
|
||||
|
||||
Verify cross-language behavior using a fresh installed extension, including unstarted-call release, setup ordering, exception identity, sync and async streams, cancellation and GC. Keep `_native.pyi` accurate about Future-returning and coroutine-returning APIs
|
||||
|
||||
A route with no Python implementation passes `runtime.NO_PYTHON` from both public entries, and every catalog rule for it must be `RUST_REQUIRED`. Dispatch raises `NoPythonImplementationError` for any other rollout instead of falling back. `tests/test_litellm_rust/test_public_entries.py` marks such a route with `python=None`, so dropping a route's Python changes the entries, its rules and that one table field together
|
||||
|
||||
A route's context, its target as well as its features, is built from `NativeCall.resolved`, so an argument counts the same whether it came by position, by keyword or as a signature default
|
||||
|
||||
Every public route serves `a<entry>(...)`, `<entry>(..., a<entry>=True)` and `a<entry>(..., a<entry>=True)` the same way. Each hands native one identical `NativeCall` with no async flag, on the async binding. `test_public_entries.py` holds this for every route in its table, so a new public route gets a row there
|
||||
|
|
|
|||
|
|
@ -25,11 +25,26 @@ class Route(str, Enum):
|
|||
TOKENIZER = "tokenizer"
|
||||
|
||||
|
||||
class Target(str, Enum):
|
||||
CHAT_COMPLETIONS = "chat_completions"
|
||||
MESSAGES = "messages"
|
||||
RESPONSES = "responses"
|
||||
|
||||
|
||||
class Feature(str, Enum):
|
||||
SAFEGUARDS_CLASSIFIER = "safeguards_classifier"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class RouteContext:
|
||||
route: Route
|
||||
provider: str | None = None
|
||||
model: str | None = None
|
||||
target: Target | None = None
|
||||
features: frozenset[Feature] = frozenset()
|
||||
|
||||
def axes(self) -> tuple[frozenset[object], ...]:
|
||||
return (frozenset({self.provider}), frozenset({self.model}), frozenset({self.target}), self.features)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
@ -38,12 +53,16 @@ class RouteRule:
|
|||
rollout: Rollout
|
||||
providers: frozenset[str] | None = None
|
||||
models: frozenset[str] | None = None
|
||||
targets: frozenset[Target] | None = None
|
||||
features: frozenset[Feature] | None = None
|
||||
|
||||
def constraints(self) -> tuple[frozenset[object] | None, ...]:
|
||||
return (self.providers, self.models, self.targets, self.features)
|
||||
|
||||
def matches(self, context: RouteContext) -> bool:
|
||||
return (
|
||||
context.route is self.route
|
||||
and (self.providers is None or context.provider in self.providers)
|
||||
and (self.models is None or context.model in self.models)
|
||||
return context.route is self.route and all(
|
||||
allowed is None or not allowed.isdisjoint(values)
|
||||
for allowed, values in zip(self.constraints(), context.axes(), strict=True)
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -53,7 +72,9 @@ RULES: Final[Rules] = (
|
|||
RouteRule(Route.CHAT_COMPLETIONS, Rollout.PYTHON_ONLY),
|
||||
RouteRule(Route.EMBEDDINGS, Rollout.PYTHON_ONLY),
|
||||
RouteRule(Route.OCR, Rollout.RUST_REQUIRED),
|
||||
RouteRule(Route.MESSAGES, Rollout.RUST_OPT_IN, providers=frozenset({"anthropic", "deepseek", "vertex_ai"})),
|
||||
RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY, targets=frozenset({Target.CHAT_COMPLETIONS, Target.RESPONSES})),
|
||||
RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY, features=frozenset({Feature.SAFEGUARDS_CLASSIFIER})),
|
||||
RouteRule(Route.MESSAGES, Rollout.RUST_OPT_IN, providers=frozenset({"anthropic"})),
|
||||
RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY),
|
||||
RouteRule(Route.RESPONSES, Rollout.PYTHON_ONLY),
|
||||
RouteRule(Route.TOKEN_COUNTER, Rollout.PYTHON_ONLY),
|
||||
|
|
@ -69,3 +90,9 @@ def rollout(context: RouteContext, rules: Rules | None = None) -> Rollout:
|
|||
|
||||
def decision(context: RouteContext, rules: Rules | None = None) -> Decision:
|
||||
return _decision(rollout(context, rules))
|
||||
|
||||
|
||||
def may_leave_python(route: Route, rules: Rules) -> bool:
|
||||
own: Final = tuple(rule for rule in rules if rule.route is route)
|
||||
last: Final = next((i for i, rule in enumerate(own) if all(c is None for c in rule.constraints())), len(own))
|
||||
return any(_decision(rule.rollout) is not Decision.PYTHON for rule in own[: last + 1])
|
||||
|
|
|
|||
|
|
@ -7,8 +7,6 @@ from typing import Final, Generic, TypeAlias, TypeVar
|
|||
from litellm.rust_bridge import catalog, runtime
|
||||
from litellm.rust_bridge.bindings import NativeBinding
|
||||
from litellm.rust_bridge.catalog import Route, RouteContext, Rules
|
||||
from litellm.rust_bridge.configuration import Decision
|
||||
from litellm.rust_bridge.configuration import decision as rollout_decision
|
||||
|
||||
RequestT: Final = TypeVar("RequestT")
|
||||
NativeT: Final = TypeVar("NativeT")
|
||||
|
|
@ -31,18 +29,6 @@ class PublicDispatch(Generic[RequestT]):
|
|||
route: Route
|
||||
request: Callable[[tuple[object, ...], Mapping[str, object]], RequestT | None]
|
||||
context: Callable[[RequestT], RouteContext]
|
||||
bypass: Callable[[RequestT], bool] | None = None
|
||||
|
||||
def _requires_projection(self, rules: Rules) -> bool:
|
||||
for rule in rules:
|
||||
if rule.route is not self.route:
|
||||
continue
|
||||
if rule.providers is not None or rule.models is not None:
|
||||
if rollout_decision(rule.rollout) is not Decision.PYTHON:
|
||||
return True
|
||||
continue
|
||||
return rollout_decision(rule.rollout) is not Decision.PYTHON
|
||||
return False
|
||||
|
||||
def _native_request(self, args: tuple[object, ...], kwargs: Mapping[str, object]) -> RequestT:
|
||||
request: Final = self.request(args, kwargs)
|
||||
|
|
@ -50,11 +36,6 @@ class PublicDispatch(Generic[RequestT]):
|
|||
raise runtime.NoPythonImplementationError(
|
||||
f"{self.route.value} has no Python implementation, so every call must project to a native request"
|
||||
)
|
||||
if self.bypass is not None and self.bypass(request):
|
||||
raise runtime.NoPythonImplementationError(
|
||||
f"{self.route.value} has no Python implementation, so a call its bypass predicate matches cannot "
|
||||
"be served"
|
||||
)
|
||||
return request
|
||||
|
||||
def run(
|
||||
|
|
@ -77,10 +58,10 @@ class PublicDispatch(Generic[RequestT]):
|
|||
python=python,
|
||||
rules=selected_rules,
|
||||
)
|
||||
if not self._requires_projection(selected_rules):
|
||||
if not catalog.may_leave_python(self.route, selected_rules):
|
||||
return python(*args, **kwargs)
|
||||
request: Final = self.request(args, kwargs)
|
||||
if request is None or (self.bypass is not None and self.bypass(request)):
|
||||
if request is None:
|
||||
return python(*args, **kwargs)
|
||||
return runtime.run(
|
||||
self.context(request),
|
||||
|
|
@ -110,10 +91,10 @@ class PublicDispatch(Generic[RequestT]):
|
|||
python=python,
|
||||
rules=selected_rules,
|
||||
)
|
||||
if not self._requires_projection(selected_rules):
|
||||
if not catalog.may_leave_python(self.route, selected_rules):
|
||||
return await python(*args, **kwargs)
|
||||
request: Final = self.request(args, kwargs)
|
||||
if request is None or (self.bypass is not None and self.bypass(request)):
|
||||
if request is None:
|
||||
return await python(*args, **kwargs)
|
||||
return await runtime.arun(
|
||||
self.context(request),
|
||||
|
|
|
|||
|
|
@ -113,6 +113,33 @@ def _without_bags(legacy: inspect.Signature, named: Mapping[str, object]) -> Map
|
|||
return MappingProxyType({name: value for name, value in named.items() if legacy.parameters[name].kind not in _BAGS})
|
||||
|
||||
|
||||
def _is_true(value: object) -> bool:
|
||||
return value is True
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class AsyncFlag:
|
||||
"""The keyword that makes a sync entry hand off to its async entry, as in ``completion(..., acompletion=True)``.
|
||||
|
||||
``is_set`` is the route's own reading of the flag. ``@client`` hands off on ``is True``, the Messages handler on truthiness.
|
||||
"""
|
||||
|
||||
name: str
|
||||
sync: inspect.Signature
|
||||
is_set: Callable[[object], bool] = _is_true
|
||||
|
||||
def requested(self, kwargs: Mapping[str, object]) -> bool:
|
||||
return self.is_set(kwargs.get(self.name))
|
||||
|
||||
def dropped(self, kwargs: Mapping[str, object]) -> Mapping[str, object]:
|
||||
return MappingProxyType({name: value for name, value in kwargs.items() if name != self.name})
|
||||
|
||||
def keywords(self, args: tuple[object, ...], kwargs: Mapping[str, object]) -> Mapping[str, object]:
|
||||
"""Name positionals by the sync signature, since the async entry may order its parameters differently."""
|
||||
named: Final = _without_bags(self.sync, self.sync.bind_partial(*args, **kwargs).arguments)
|
||||
return self.dropped({**named, **kwargs})
|
||||
|
||||
|
||||
def native_call(legacy: inspect.Signature, args: tuple[object, ...], kwargs: Mapping[str, object]) -> NativeCall:
|
||||
positional: Final = legacy.bind_partial(*args)
|
||||
positional.apply_defaults()
|
||||
|
|
|
|||
|
|
@ -1,20 +1,24 @@
|
|||
# Invariants
|
||||
|
||||
- This directory tests the Rust path and nothing else. Treat the Python implementation of every route as already removed. Never call it on purpose, never run one request through both engines, and never compare Python and Rust outputs, headers, cache keys or logs
|
||||
- A test that would still pass with Python serving the call does not belong here. Every test asserts the call was served natively with `support.routing.assert_served_natively`
|
||||
- Outside the dispatch tests below, this directory tests the Rust path and nothing else. Treat the Python implementation of every route as already removed. Never call it on purpose, never run one request through both engines, and never compare Python and Rust outputs, headers, cache keys or logs
|
||||
- A test that would still pass with Python serving the call does not belong here. Every test except the dispatch tests asserts the call was served natively with `support.routing.assert_served_natively`
|
||||
- Dispatch tests check which path a public call takes under given rules. They are each route folder's `test_dispatch.py`, plus `test_dispatch.py` and `test_public_entries.py` at the top of this directory
|
||||
- They pass `rules=` to a dispatch directly or wrap the whole test in one `support.routing.routed(...)` block, and they stub the native bindings
|
||||
- They may call Python, since the choice between the two paths is what they test
|
||||
- They live here until they can run in `tests/unit/` without the compiled extension
|
||||
- Select the Rust path with `support.routing.native_route(...)`, which prepends `RUST_REQUIRED` rules to `catalog.RULES` and restores them on exit. Each route folder's `conftest.py` applies it to every test through an autouse fixture. A test outside a route folder wraps its whole body in one `native_route(...)` block, and a child interpreter script assigns `catalog.RULES` once at its top. Never select a backend by reading or setting `LITELLM_RUST`, calling `litellm.rust(...)` or rebinding `catalog.RULES` any other way. The rollout in `litellm/rust_bridge/catalog.py` is a product decision, and a test here must pass no matter what it says, including for a route marked `PYTHON_ONLY`
|
||||
- `catalog.RULES` is static at runtime and `native_route(...)` swaps a module global, so open it once per test and never per call. Two overlapping blocks, such as one per task inside an `asyncio.gather`, restore out of order and leave the wrong rules behind for later tests
|
||||
- A test whose subject is the rollout switch itself (env, `litellm.rust(...)`, catalog rules) is a unit test of `litellm/rust_bridge/configuration.py` or `catalog.py` and lives in `tests/unit/rust_bridge/`
|
||||
- No test here depends on the rollout switch. Only no rule, `PYTHON_ONLY` and `RUST_REQUIRED` pick a path without it, so those are the only rollouts a test here uses, and Rust with Python fallback is not reachable from this directory. What `LITELLM_RUST` and `litellm.rust(...)` mean is tested in `tests/unit/rust_bridge/test_configuration.py`, mostly through `decide`, which takes both as plain values
|
||||
- A test lives here only when it cannot be written in Rust. That means it needs the Python public API, Python-owned state such as `litellm.cache` and callbacks, or the process boundary (fork, threads, event loops). Anything checkable from the Rust side goes in the crate's own tests
|
||||
- The wire is the local recording server. No test here reaches a real provider
|
||||
- Never parametrize one test across routes. Each route folder owns its tests, its `server` fixture, its request data and its `call`
|
||||
- Never parametrize one test across routes. Each route folder owns its tests, its `server` fixture, its request data and its `call`. The one exception is `test_public_entries.py`, whose subject is that every route serves its public call forms the same way
|
||||
- Call `_native` directly only for a case the public API cannot reach
|
||||
|
||||
# Layout
|
||||
|
||||
Each route folder (`ocr/`, `messages/`, `responses/`, `chat_completions/`) uses the same file names, one concern per file
|
||||
|
||||
- `test_dispatch.py` the call reaches the native route
|
||||
- `test_dispatch.py` the call reaches the native route, and which path the rules pick for each public call
|
||||
- `test_requests.py` Python inputs that reach the wire
|
||||
- `test_failures.py` errors raised to the caller and rejection before provider IO
|
||||
- `test_settings.py` globals, env, named credentials and their precedence
|
||||
|
|
@ -24,4 +28,6 @@ Each route folder (`ocr/`, `messages/`, `responses/`, `chat_completions/`) uses
|
|||
- `test_lifecycle.py` cancellation, concurrency, context and nested calls
|
||||
- `test_cache.py` the configured cache seen from this route
|
||||
|
||||
`embeddings/` holds only `test_dispatch.py` so far. At the top, `test_dispatch.py` tests `PublicDispatch` with stub routes and `test_public_entries.py` runs every route's public call forms through one table
|
||||
|
||||
Route data and the route's `call` live in `<route>/data.py`, and `<route>/conftest.py` owns the `server` fixture and the autouse `native` fixture. `support/` holds only route-neutral helpers
|
||||
|
|
|
|||
|
|
@ -9,7 +9,6 @@ from litellm.chat_completions.dispatch import (
|
|||
_ADISPATCH, # pyright: ignore[reportPrivateUsage] # tests configured dispatch
|
||||
_DISPATCH, # pyright: ignore[reportPrivateUsage] # tests configured dispatch
|
||||
)
|
||||
from litellm.rust_bridge import catalog
|
||||
from litellm.rust_bridge.bindings import NativeBinding
|
||||
from litellm.rust_bridge.catalog import Route, RouteRule, Rules
|
||||
from litellm.rust_bridge.chat_completions.entrypoints import (
|
||||
|
|
@ -147,29 +146,6 @@ def test_native_receives_bound_request_and_original_call_shape() -> None:
|
|||
assert call_kwargs["metadata"] is metadata
|
||||
|
||||
|
||||
def test_internal_async_marker_bypasses_native() -> None:
|
||||
response: Final = ModelResponse()
|
||||
called: Final[list[bool]] = []
|
||||
|
||||
def python(*call_args: object, **call_kwargs: object) -> ModelResponse: # kwargs-ok: records public call shape
|
||||
called.append(True)
|
||||
return response
|
||||
|
||||
def native(request: NativeCall) -> ModelResponse:
|
||||
pytest.fail("acompletion's inner completion call must stay on Python")
|
||||
|
||||
result: Final = _DISPATCH.run(
|
||||
("gpt-4o", MESSAGES),
|
||||
{"acompletion": True},
|
||||
python=python,
|
||||
binding=completion_binding(native),
|
||||
native=native_call_hook,
|
||||
rules=RUST_RULES,
|
||||
)
|
||||
assert result is response
|
||||
assert called == [True]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("args", "kwargs"),
|
||||
(
|
||||
|
|
@ -202,56 +178,6 @@ def test_binding_errors_delegate_to_python(args: tuple[object, ...], kwargs: Map
|
|||
assert captured == [(args, kwargs)]
|
||||
|
||||
|
||||
def test_public_completion_routes_through_dispatch(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
captured: Final[list[NativeCall]] = []
|
||||
expected: Final = ModelResponse()
|
||||
|
||||
def native(request: NativeCall) -> ModelResponse:
|
||||
captured.append(request)
|
||||
return expected
|
||||
|
||||
NATIVE_COMPLETION.override(native)
|
||||
monkeypatch.setattr(catalog, "RULES", RUST_RULES)
|
||||
public_completion: Final = cast(Callable[..., ModelResponse], litellm.completion)
|
||||
try:
|
||||
result: Final = public_completion(model="gpt-4o", messages=MESSAGES)
|
||||
finally:
|
||||
NATIVE_COMPLETION.reset()
|
||||
assert result is expected
|
||||
assert [request.resolved["model"] for request in captured] == ["gpt-4o"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_public_acompletion_routes_through_dispatch(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
captured: Final[list[NativeCall]] = []
|
||||
expected: Final = ModelResponse()
|
||||
|
||||
async def native(request: NativeCall) -> ModelResponse:
|
||||
captured.append(request)
|
||||
return expected
|
||||
|
||||
NATIVE_ACOMPLETION.override(native)
|
||||
monkeypatch.setattr(catalog, "RULES", RUST_RULES)
|
||||
public_acompletion: Final = cast(Callable[..., Awaitable[ModelResponse]], litellm.acompletion)
|
||||
try:
|
||||
result: Final = await public_acompletion(model="gpt-4o", messages=MESSAGES)
|
||||
finally:
|
||||
NATIVE_ACOMPLETION.reset()
|
||||
assert result is expected
|
||||
assert [request.resolved["model"] for request in captured] == ["gpt-4o"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_public_completion_calls_keep_the_python_result() -> None:
|
||||
sync_response: Final = litellm.completion(model="openai/test-model", messages=MESSAGES, mock_response="ok")
|
||||
async_response: Final = await litellm.acompletion(model="openai/test-model", messages=MESSAGES, mock_response="ok")
|
||||
|
||||
assert isinstance(sync_response, ModelResponse)
|
||||
assert isinstance(async_response, ModelResponse)
|
||||
assert sync_response.choices[0].message.content == "ok"
|
||||
assert async_response.choices[0].message.content == "ok"
|
||||
|
||||
|
||||
def test_sync_completion_request_projects_public_arguments() -> None:
|
||||
rules: Final[Rules] = (RouteRule(Route.CHAT_COMPLETIONS, Rollout.RUST_REQUIRED),)
|
||||
expected: Final = ModelResponse()
|
||||
|
|
@ -277,64 +203,46 @@ def test_sync_completion_request_projects_public_arguments() -> None:
|
|||
assert response is expected
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_completion_falls_back_after_native_declines() -> None:
|
||||
from litellm.rust_bridge.bindings import native_exception_types
|
||||
|
||||
native_types: Final = native_exception_types()
|
||||
if native_types is None:
|
||||
pytest.skip("native bridge is unavailable")
|
||||
declined, _ = native_types
|
||||
expected: Final = ModelResponse()
|
||||
rules: Final[Rules] = (RouteRule(Route.CHAT_COMPLETIONS, Rollout.RUST_OPT_OUT),)
|
||||
|
||||
async def native(request: NativeCall) -> ModelResponse:
|
||||
raise declined("unsupported")
|
||||
|
||||
async def python(*args: object, **kwargs: object) -> ModelResponse:
|
||||
return expected
|
||||
|
||||
binding: Final[NativeBinding[NativeAcompletion]] = NativeBinding("acompletion", validate=lambda _: None)
|
||||
binding.override(native)
|
||||
response: Final = await dispatch._ADISPATCH.arun( # pyright: ignore[reportPrivateUsage] # test an explicit route decision
|
||||
("test-model", MESSAGES),
|
||||
{},
|
||||
python=python,
|
||||
binding=binding,
|
||||
native=native_call_hook,
|
||||
rules=rules,
|
||||
)
|
||||
|
||||
assert response is expected
|
||||
|
||||
|
||||
def test_internal_acompletion_marker_bypasses_native() -> None:
|
||||
rules: Final[Rules] = (RouteRule(Route.CHAT_COMPLETIONS, Rollout.RUST_REQUIRED),)
|
||||
expected: Final = ModelResponse()
|
||||
|
||||
def python(*args: object, **kwargs: object) -> ModelResponse:
|
||||
return expected
|
||||
|
||||
def native(request: NativeCall) -> ModelResponse:
|
||||
pytest.fail("acompletion's inner completion call must stay on Python")
|
||||
|
||||
binding: Final[NativeBinding[NativeCompletion]] = NativeBinding("completion", validate=lambda _: None)
|
||||
binding.override(native)
|
||||
response: Final = dispatch._DISPATCH.run( # pyright: ignore[reportPrivateUsage] # test an explicit route decision
|
||||
("test-model", MESSAGES),
|
||||
{"custom_llm_provider": "openai", "acompletion": True},
|
||||
python=python,
|
||||
binding=binding,
|
||||
native=native_call_hook,
|
||||
rules=rules,
|
||||
)
|
||||
|
||||
assert response is expected
|
||||
|
||||
|
||||
def test_positional_parameters_remain_available_to_native_projection() -> None:
|
||||
request: Final = _DISPATCH.request(("anthropic/test-model", MESSAGES, 12.0, 0.25), {})
|
||||
assert request is not None
|
||||
assert request.resolved["timeout"] == 12.0
|
||||
assert request.resolved["temperature"] == 0.25
|
||||
assert request.resolved["messages"] is MESSAGES
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_acompletion_marker_binds_positionals_by_the_sync_signature() -> None:
|
||||
captured: Final[list[NativeCall]] = []
|
||||
|
||||
async def native(call: NativeCall) -> ModelResponse:
|
||||
captured.append(call)
|
||||
return ModelResponse()
|
||||
|
||||
NATIVE_ACOMPLETION.override(native)
|
||||
public_completion: Final = cast(Callable[..., Awaitable[ModelResponse]], litellm.completion)
|
||||
try:
|
||||
_ = await public_completion("gpt-4o", MESSAGES, 30.0, acompletion=True)
|
||||
finally:
|
||||
NATIVE_ACOMPLETION.reset()
|
||||
assert [(request.resolved["timeout"], request.resolved["functions"]) for request in captured] == [(30.0, None)]
|
||||
|
||||
|
||||
def test_acompletion_flag_hands_off_only_when_true() -> None:
|
||||
expected: Final = ModelResponse()
|
||||
|
||||
def native(call: NativeCall) -> ModelResponse:
|
||||
return expected
|
||||
|
||||
async def anative(call: NativeCall) -> ModelResponse:
|
||||
pytest.fail("completion(..., acompletion=1) stays sync, since the Python @client only hands off on is True")
|
||||
|
||||
NATIVE_COMPLETION.override(native)
|
||||
NATIVE_ACOMPLETION.override(anative)
|
||||
public_completion: Final = cast(Callable[..., object], litellm.completion)
|
||||
try:
|
||||
result: Final = public_completion("gpt-4o", MESSAGES, acompletion=1)
|
||||
finally:
|
||||
NATIVE_COMPLETION.reset()
|
||||
NATIVE_ACOMPLETION.reset()
|
||||
assert result is expected
|
||||
|
|
|
|||
|
|
@ -1,12 +1,10 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Awaitable, Callable
|
||||
from collections.abc import Callable
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
import litellm
|
||||
from litellm.embeddings import dispatch
|
||||
from litellm.rust_bridge.bindings import NativeBinding
|
||||
from litellm.rust_bridge.catalog import Route, RouteRule, Rules
|
||||
|
|
@ -15,19 +13,6 @@ from litellm.rust_bridge.public_call import NativeCall, native_call_hook
|
|||
from litellm.types.utils import EmbeddingResponse
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_public_embedding_calls_keep_the_python_result() -> None:
|
||||
vector: Final = [0.1, 0.2]
|
||||
|
||||
sync_response: Final = litellm.embedding(model="openai/test-model", input="hello", mock_response=vector)
|
||||
async_response: Final = await litellm.aembedding(model="openai/test-model", input="hello", mock_response=vector)
|
||||
|
||||
assert isinstance(sync_response, EmbeddingResponse)
|
||||
rows: Final = TypeAdapter(list[dict[str, object]])
|
||||
assert rows.validate_python(sync_response.model_dump()["data"])[0]["embedding"] == vector
|
||||
assert rows.validate_python(async_response.model_dump()["data"])[0]["embedding"] == vector
|
||||
|
||||
|
||||
def test_sync_embedding_request_projects_public_arguments() -> None:
|
||||
rules: Final[Rules] = (RouteRule(Route.EMBEDDINGS, Rollout.RUST_REQUIRED),)
|
||||
expected: Final = EmbeddingResponse(model="test-model", data=[])
|
||||
|
|
@ -52,36 +37,3 @@ def test_sync_embedding_request_projects_public_arguments() -> None:
|
|||
)
|
||||
|
||||
assert response is expected
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_embedding_falls_back_after_native_declines() -> None:
|
||||
from litellm.rust_bridge.bindings import native_exception_types
|
||||
|
||||
native_types: Final = native_exception_types()
|
||||
if native_types is None:
|
||||
pytest.skip("native bridge is unavailable")
|
||||
declined, _ = native_types
|
||||
expected: Final = EmbeddingResponse(model="test-model", data=[])
|
||||
rules: Final[Rules] = (RouteRule(Route.EMBEDDINGS, Rollout.RUST_OPT_OUT),)
|
||||
|
||||
async def native(request: NativeCall) -> EmbeddingResponse:
|
||||
raise declined("unsupported")
|
||||
|
||||
async def python(*args: object, **kwargs: object) -> EmbeddingResponse:
|
||||
return expected
|
||||
|
||||
binding: Final[NativeBinding[Callable[[NativeCall], Awaitable[EmbeddingResponse]]]] = NativeBinding(
|
||||
"aembedding", validate=lambda _: None
|
||||
)
|
||||
binding.override(native)
|
||||
response: Final = await dispatch._ADISPATCH.arun( # pyright: ignore[reportPrivateUsage] # test an explicit route decision
|
||||
("test-model", "hello"),
|
||||
{},
|
||||
python=python,
|
||||
binding=binding,
|
||||
native=native_call_hook,
|
||||
rules=rules,
|
||||
)
|
||||
|
||||
assert response is expected
|
||||
|
|
|
|||
|
|
@ -3,7 +3,6 @@ from collections.abc import Awaitable, Callable, Mapping
|
|||
from typing import Final, cast # noqa: TID251 # narrows legacy callable signatures for inspect
|
||||
|
||||
import pytest
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
import litellm
|
||||
from litellm.llms.anthropic.pass_through.messages import handler as python_messages
|
||||
|
|
@ -13,17 +12,14 @@ from litellm.messages.dispatch import (
|
|||
_DISPATCH, # pyright: ignore[reportPrivateUsage] # tests configured dispatch
|
||||
)
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.rust_bridge import catalog
|
||||
from litellm.rust_bridge.bindings import NativeBinding
|
||||
from litellm.rust_bridge.catalog import Route, RouteRule, Rules
|
||||
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"}]
|
||||
DANGEROUS_TOOL_USE: Final = [{"type": "dangerous_tool_use", "classifier_context": {"permission_mode": "auto"}}]
|
||||
CLASSIFIER_SETTINGS: Final = {"safeguards_classifier_model": "gpt-oss-120b"}
|
||||
PYTHON_RULES: Final[Rules] = ()
|
||||
RUST_RULES: Final[Rules] = (RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED),)
|
||||
|
||||
|
|
@ -194,81 +190,44 @@ def test_native_streams_are_bounded_by_stream_timeout(stream: bool, expected_tim
|
|||
assert captured[0].resolved["timeout"] == expected_timeout
|
||||
|
||||
|
||||
def test_internal_async_marker_bypasses_native() -> None:
|
||||
args: Final[tuple[object, ...]] = (16, MESSAGES, "claude-sonnet-4-5")
|
||||
kwargs: Final[Mapping[str, object]] = {"is_async": True}
|
||||
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("The async handler's inner sync call must stay on 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)]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("model", "custom_llm_provider"),
|
||||
[
|
||||
("vertex_ai/gemini-2.5-pro", None),
|
||||
("gemini-2.5-pro", "vertex_ai"),
|
||||
],
|
||||
("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_vertex_models_other_than_claude_stay_on_python(model: str, custom_llm_provider: str | None) -> None:
|
||||
args: Final[tuple[object, ...]] = (16, MESSAGES, model)
|
||||
kwargs: Final[Mapping[str, object]] = {"custom_llm_provider": custom_llm_provider}
|
||||
expected: Final = response(model)
|
||||
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: records call shape
|
||||
return expected
|
||||
|
||||
def native(request: NativeCall) -> AnthropicMessagesResponse:
|
||||
pytest.fail("Rust serves only Claude on Vertex AI")
|
||||
def python(*call_args: object, **call_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: fallback
|
||||
return python_result
|
||||
|
||||
result: Final = _DISPATCH.run(
|
||||
args,
|
||||
(16, MESSAGES, model),
|
||||
kwargs,
|
||||
python=python,
|
||||
binding=messages_binding(native),
|
||||
binding=messages_binding(lambda request: native_result),
|
||||
native=lambda hook, request, call_args, call_kwargs: hook(request),
|
||||
rules=RUST_RULES,
|
||||
rules=(RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED, targets=frozenset({Target.MESSAGES})),),
|
||||
)
|
||||
assert result is expected
|
||||
assert result is (native_result if passes_through else python_result)
|
||||
|
||||
|
||||
def test_vertex_claude_reaches_native() -> None:
|
||||
args: Final[tuple[object, ...]] = (16, MESSAGES, "vertex_ai/claude-sonnet-4-5@20250929")
|
||||
expected: Final = response("vertex_ai/claude-sonnet-4-5@20250929")
|
||||
|
||||
def python(*call_args: object, **call_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: records call shape
|
||||
pytest.fail("Claude on Vertex AI is served natively")
|
||||
|
||||
result: Final = _DISPATCH.run(
|
||||
args,
|
||||
{},
|
||||
python=python,
|
||||
binding=messages_binding(lambda request: expected),
|
||||
native=lambda hook, request, call_args, call_kwargs: hook(request),
|
||||
rules=RUST_RULES,
|
||||
)
|
||||
assert result is expected
|
||||
DANGEROUS_TOOL_USE: Final = [{"type": "dangerous_tool_use", "classifier_context": {"permission_mode": "auto"}}]
|
||||
CLASSIFIER_SETTINGS: Final = {"safeguards_classifier_model": "gpt-oss-120b"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("model", "kwargs", "general_settings", "serves_python"),
|
||||
(
|
||||
|
|
@ -285,7 +244,7 @@ def test_vertex_claude_reaches_native() -> None:
|
|||
("anthropic/claude-sonnet-4-5", {"safeguards": DANGEROUS_TOOL_USE}, CLASSIFIER_SETTINGS, False),
|
||||
),
|
||||
)
|
||||
async def test_python_keeps_the_requests_it_answers_with_safeguard_results(
|
||||
def test_proxy_classified_safeguards_reach_the_feature_rule(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
model: str,
|
||||
kwargs: Mapping[str, object],
|
||||
|
|
@ -296,19 +255,19 @@ async def test_python_keeps_the_requests_it_answers_with_safeguard_results(
|
|||
python_result: Final = response(model)
|
||||
native_result: Final = response(model)
|
||||
|
||||
async def python(*call_args: object, **call_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: fallback
|
||||
def python(*call_args: object, **call_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: fallback
|
||||
return python_result
|
||||
|
||||
async def native(request: NativeCall) -> AnthropicMessagesResponse:
|
||||
return native_result
|
||||
|
||||
result: Final = await _ADISPATCH.arun(
|
||||
result: Final = _DISPATCH.run(
|
||||
(16, MESSAGES, model),
|
||||
kwargs,
|
||||
python=python,
|
||||
binding=amessages_binding(native),
|
||||
binding=messages_binding(lambda request: native_result),
|
||||
native=lambda hook, request, call_args, call_kwargs: hook(request),
|
||||
rules=RUST_RULES,
|
||||
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)
|
||||
|
||||
|
|
@ -345,7 +304,7 @@ def test_binding_errors_delegate_to_python(args: tuple[object, ...], kwargs: Map
|
|||
assert captured == [(args, kwargs)]
|
||||
|
||||
|
||||
def test_anthropic_create_routes_through_dispatch(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
def test_anthropic_create_routes_through_dispatch() -> None:
|
||||
captured: Final[list[NativeCall]] = []
|
||||
expected: Final = response()
|
||||
|
||||
|
|
@ -356,7 +315,6 @@ def test_anthropic_create_routes_through_dispatch(monkeypatch: pytest.MonkeyPatc
|
|||
return expected
|
||||
|
||||
NATIVE_MESSAGES.override(native)
|
||||
monkeypatch.setattr(catalog, "RULES", RUST_RULES)
|
||||
public_create: Final = cast(Callable[..., AnthropicMessagesResponse], litellm.anthropic.create)
|
||||
try:
|
||||
result: Final = public_create(max_tokens=16, messages=MESSAGES, model="claude-sonnet-4-5")
|
||||
|
|
@ -367,7 +325,7 @@ def test_anthropic_create_routes_through_dispatch(monkeypatch: pytest.MonkeyPatc
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_anthropic_acreate_routes_through_dispatch(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
async def test_anthropic_acreate_routes_through_dispatch() -> None:
|
||||
captured: Final[list[NativeCall]] = []
|
||||
expected: Final = response()
|
||||
|
||||
|
|
@ -378,7 +336,6 @@ async def test_anthropic_acreate_routes_through_dispatch(monkeypatch: pytest.Mon
|
|||
return expected
|
||||
|
||||
NATIVE_AMESSAGES.override(native)
|
||||
monkeypatch.setattr(catalog, "RULES", RUST_RULES)
|
||||
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")
|
||||
|
|
@ -389,14 +346,31 @@ async def test_anthropic_acreate_routes_through_dispatch(monkeypatch: pytest.Mon
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_public_anthropic_messages_keeps_the_python_result() -> None:
|
||||
response: Final = await litellm.anthropic_messages(
|
||||
model="anthropic/claude-sonnet-4-5", messages=MESSAGES, max_tokens=10, mock_response="ok"
|
||||
)
|
||||
async def test_truthy_is_async_hands_off_to_the_async_entry_like_true() -> None:
|
||||
captured: Final[list[NativeCall]] = []
|
||||
expected: Final = response()
|
||||
|
||||
assert isinstance(response, dict)
|
||||
content: Final = TypeAdapter(list[dict[str, object]]).validate_python(response.get("content", []))
|
||||
assert content[0]["text"] == "ok"
|
||||
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:
|
||||
|
|
@ -451,64 +425,3 @@ def test_messages_binding_error_delegates_unchanged_to_python() -> None:
|
|||
)
|
||||
|
||||
assert response is expected
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_messages_falls_back_after_native_declines() -> None:
|
||||
from litellm.rust_bridge.bindings import native_exception_types
|
||||
|
||||
native_types: Final = native_exception_types()
|
||||
if native_types is None:
|
||||
pytest.skip("native bridge is unavailable")
|
||||
declined, _ = native_types
|
||||
expected: Final = AnthropicMessagesResponse(model="claude-test")
|
||||
rules: Final[Rules] = (RouteRule(Route.MESSAGES, Rollout.RUST_OPT_OUT),)
|
||||
|
||||
async def native(request: NativeCall) -> AnthropicMessagesResponse:
|
||||
raise declined("unsupported")
|
||||
|
||||
async def python(*args: object, **kwargs: object) -> AnthropicMessagesResponse:
|
||||
return expected
|
||||
|
||||
binding: Final[NativeBinding[NativeAmessages]] = NativeBinding("amessages", validate=lambda _: None)
|
||||
binding.override(native)
|
||||
response: Final = await dispatch._ADISPATCH.arun( # pyright: ignore[reportPrivateUsage] # test an explicit route decision
|
||||
(),
|
||||
{"model": "claude-test", "messages": MESSAGES, "max_tokens": 10},
|
||||
python=python,
|
||||
binding=binding,
|
||||
native=lambda hook, request, args, kwargs: hook(request),
|
||||
rules=rules,
|
||||
)
|
||||
|
||||
assert response is expected
|
||||
|
||||
|
||||
def test_internal_is_async_marker_bypasses_native() -> 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("anthropic_messages' inner handler call 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,
|
||||
"max_tokens": 10,
|
||||
"custom_llm_provider": "anthropic",
|
||||
"is_async": True,
|
||||
},
|
||||
python=python,
|
||||
binding=binding,
|
||||
native=lambda hook, request, args, kwargs: hook(request),
|
||||
rules=rules,
|
||||
)
|
||||
|
||||
assert response is expected
|
||||
|
|
|
|||
|
|
@ -1,8 +1,21 @@
|
|||
from collections.abc import Mapping
|
||||
from typing import Final, NoReturn
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.llms.base_llm.ocr.transformation import OCRResponse
|
||||
from litellm.ocr.dispatch import (
|
||||
_ADISPATCH, # pyright: ignore[reportPrivateUsage] # tests configured dispatch
|
||||
_DISPATCH, # pyright: ignore[reportPrivateUsage] # tests configured dispatch
|
||||
)
|
||||
from litellm.rust_bridge import runtime
|
||||
from litellm.rust_bridge.bindings import NativeBinding
|
||||
from litellm.rust_bridge.catalog import Route, RouteRule, Rules
|
||||
from litellm.rust_bridge.configuration import Rollout
|
||||
from litellm.rust_bridge.ocr.entrypoints import NativeAocr, NativeOcr
|
||||
from litellm.rust_bridge.public_call import NativeCall
|
||||
from tests.test_litellm_rust.ocr.data import OCR_DOCUMENT, OCR_MODEL
|
||||
from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
|
||||
from tests.test_litellm_rust.support.routing import assert_served_natively
|
||||
|
|
@ -110,3 +123,214 @@ async def test_ocr_health_check_uses_the_native_route(recording_server: Recordin
|
|||
assert "error" not in response
|
||||
assert_served_natively(recording_server)
|
||||
assert recording_server.requests[0].body["document"]["image_url"].startswith("data:image/png;base64,")
|
||||
|
||||
|
||||
RUST_RULES: Final[Rules] = (RouteRule(Route.OCR, Rollout.RUST_REQUIRED),)
|
||||
|
||||
|
||||
def ocr_binding(native: NativeOcr | None) -> NativeBinding[NativeOcr]:
|
||||
binding: Final[NativeBinding[NativeOcr]] = NativeBinding("ocr", validate=lambda _: None)
|
||||
binding.override(native)
|
||||
return binding
|
||||
|
||||
|
||||
def aocr_binding(native: NativeAocr | None) -> NativeBinding[NativeAocr]:
|
||||
binding: Final[NativeBinding[NativeAocr]] = NativeBinding("aocr", validate=lambda _: None)
|
||||
binding.override(native)
|
||||
return binding
|
||||
|
||||
|
||||
def response(model: str = "mistral/mistral-ocr-latest") -> OCRResponse:
|
||||
return OCRResponse(pages=[], model=model)
|
||||
|
||||
|
||||
def test_native_receives_normalized_positional_request_and_original_call_shape() -> None:
|
||||
document: Final[Mapping[str, object]] = {"type": "file", "file": b"pdf"}
|
||||
timeout: Final = httpx.Timeout(30)
|
||||
extra_headers: Final[dict[str, object]] = {"x-test": "1"}
|
||||
pages: Final = [0, 2]
|
||||
args: Final[tuple[object, ...]] = ("mistral/mistral-ocr-latest", document)
|
||||
kwargs: Final[Mapping[str, object]] = {
|
||||
"api_key": "test-key",
|
||||
"api_base": "https://example.invalid",
|
||||
"timeout": timeout,
|
||||
"custom_llm_provider": "mistral",
|
||||
"extra_headers": extra_headers,
|
||||
"pages": pages,
|
||||
}
|
||||
captured: Final[list[tuple[NativeCall, tuple[object, ...], Mapping[str, object]]]] = []
|
||||
expected: Final = response()
|
||||
|
||||
def native(
|
||||
request: NativeCall,
|
||||
) -> OCRResponse:
|
||||
args: Final = request.args
|
||||
kwargs: Final = request.kwargs
|
||||
captured.append((request, args, kwargs))
|
||||
return expected
|
||||
|
||||
result: Final = _DISPATCH.run(
|
||||
args,
|
||||
kwargs,
|
||||
python=runtime.NO_PYTHON,
|
||||
binding=ocr_binding(native),
|
||||
native=lambda hook, request, call_args, call_kwargs: hook(request),
|
||||
rules=RUST_RULES,
|
||||
)
|
||||
|
||||
request, call_args, call_kwargs = captured[0]
|
||||
assert result is expected
|
||||
assert request.resolved["model"] == "mistral/mistral-ocr-latest"
|
||||
assert request.resolved["document"] is document
|
||||
assert request.resolved["api_key"] == "test-key"
|
||||
assert request.resolved["api_base"] == "https://example.invalid"
|
||||
assert request.resolved["timeout"] is timeout
|
||||
assert request.resolved["custom_llm_provider"] == "mistral"
|
||||
assert request.resolved["extra_headers"] is extra_headers
|
||||
assert request.kwargs == kwargs
|
||||
assert request.kwargs["pages"] is pages
|
||||
assert call_args is args
|
||||
assert call_kwargs is kwargs
|
||||
|
||||
|
||||
def test_native_preserves_keyword_model_and_document_in_original_call_shape() -> None:
|
||||
document: Final[Mapping[str, object]] = {
|
||||
"type": "document_url",
|
||||
"document_url": "https://example.invalid/document.pdf",
|
||||
}
|
||||
pages: Final = [1]
|
||||
args: Final[tuple[object, ...]] = ()
|
||||
kwargs: Final[Mapping[str, object]] = {
|
||||
"model": "mistral/mistral-ocr-latest",
|
||||
"document": document,
|
||||
"pages": pages,
|
||||
}
|
||||
captured: Final[list[tuple[NativeCall, tuple[object, ...], Mapping[str, object]]]] = []
|
||||
expected: Final = response()
|
||||
|
||||
def native(
|
||||
request: NativeCall,
|
||||
) -> OCRResponse:
|
||||
args: Final = request.args
|
||||
kwargs: Final = request.kwargs
|
||||
captured.append((request, args, kwargs))
|
||||
return expected
|
||||
|
||||
result: Final = _DISPATCH.run(
|
||||
args,
|
||||
kwargs,
|
||||
python=runtime.NO_PYTHON,
|
||||
binding=ocr_binding(native),
|
||||
native=lambda hook, request, call_args, call_kwargs: hook(request),
|
||||
rules=RUST_RULES,
|
||||
)
|
||||
|
||||
request, call_args, call_kwargs = captured[0]
|
||||
assert result is expected
|
||||
assert request.resolved["model"] == "mistral/mistral-ocr-latest"
|
||||
assert request.resolved["document"] is document
|
||||
assert request.kwargs == kwargs
|
||||
assert call_args is args
|
||||
assert call_kwargs is kwargs
|
||||
assert call_kwargs["model"] == "mistral/mistral-ocr-latest"
|
||||
assert call_kwargs["document"] is document
|
||||
|
||||
|
||||
def test_missing_native_binding_is_a_required_rust_error() -> None:
|
||||
args: Final[tuple[object, ...]] = ("mistral/mistral-ocr-latest", {"type": "file", "file": b"pdf"})
|
||||
|
||||
with pytest.raises(RuntimeError, match="Rust ocr bridge"):
|
||||
_DISPATCH.run(
|
||||
args,
|
||||
{},
|
||||
python=runtime.NO_PYTHON,
|
||||
binding=ocr_binding(None),
|
||||
native=lambda hook, request, call_args, call_kwargs: hook(request),
|
||||
rules=RUST_RULES,
|
||||
)
|
||||
|
||||
|
||||
def test_non_required_rule_cannot_be_served_without_python() -> None:
|
||||
args: Final[tuple[object, ...]] = ("mistral/mistral-ocr-latest", {"type": "file", "file": b"pdf"})
|
||||
|
||||
def native(
|
||||
request: NativeCall,
|
||||
) -> OCRResponse:
|
||||
return response()
|
||||
|
||||
with pytest.raises(runtime.NoPythonImplementationError, match="RUST_REQUIRED"):
|
||||
_DISPATCH.run(
|
||||
args,
|
||||
{},
|
||||
python=runtime.NO_PYTHON,
|
||||
binding=ocr_binding(native),
|
||||
native=lambda hook, request, call_args, call_kwargs: hook(request),
|
||||
rules=(RouteRule(Route.OCR, Rollout.PYTHON_ONLY),),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("args", "kwargs", "message"),
|
||||
(
|
||||
(
|
||||
("mistral/mistral-ocr-latest", {"type": "file", "file": b"pdf"}),
|
||||
{"model": "duplicate"},
|
||||
r"ocr\(\) got multiple values for argument 'model'",
|
||||
),
|
||||
(
|
||||
("mistral/mistral-ocr-latest",),
|
||||
{},
|
||||
r"ocr\(\) missing 1 required positional argument: 'document'",
|
||||
),
|
||||
),
|
||||
)
|
||||
def test_ocr_parser_errors_before_native(args: tuple[object, ...], kwargs: Mapping[str, object], message: str) -> None:
|
||||
def native(
|
||||
request: NativeCall,
|
||||
) -> OCRResponse:
|
||||
pytest.fail("OCR parser failures must not call native")
|
||||
|
||||
with pytest.raises(TypeError, match=message):
|
||||
_DISPATCH.run(
|
||||
args,
|
||||
kwargs,
|
||||
python=runtime.NO_PYTHON,
|
||||
binding=ocr_binding(native),
|
||||
native=lambda hook, request, call_args, call_kwargs: hook(request),
|
||||
rules=RUST_RULES,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("args", "kwargs", "message"),
|
||||
(
|
||||
(
|
||||
("mistral/mistral-ocr-latest", {"type": "file", "file": b"pdf"}),
|
||||
{"model": "duplicate"},
|
||||
r"aocr\(\) got multiple values for argument 'model'",
|
||||
),
|
||||
(
|
||||
("mistral/mistral-ocr-latest",),
|
||||
{},
|
||||
r"aocr\(\) missing 1 required positional argument: 'document'",
|
||||
),
|
||||
),
|
||||
)
|
||||
async def test_aocr_parser_errors_before_native(
|
||||
args: tuple[object, ...], kwargs: Mapping[str, object], message: str
|
||||
) -> None:
|
||||
async def native(
|
||||
request: NativeCall,
|
||||
) -> OCRResponse:
|
||||
pytest.fail("OCR parser failures must not call native")
|
||||
|
||||
with pytest.raises(TypeError, match=message):
|
||||
await _ADISPATCH.arun(
|
||||
args,
|
||||
kwargs,
|
||||
python=runtime.NO_PYTHON,
|
||||
binding=aocr_binding(native),
|
||||
native=lambda hook, request, call_args, call_kwargs: hook(request),
|
||||
rules=RUST_RULES,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1,300 +0,0 @@
|
|||
from collections.abc import Awaitable, Callable, Mapping
|
||||
from typing import Final, cast # noqa: TID251 # narrows legacy callable signatures for inspect
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.llms.base_llm.ocr.transformation import OCRResponse
|
||||
from litellm.ocr.dispatch import (
|
||||
_ADISPATCH, # pyright: ignore[reportPrivateUsage] # tests configured dispatch
|
||||
_DISPATCH, # pyright: ignore[reportPrivateUsage] # tests configured dispatch
|
||||
)
|
||||
from litellm.rust_bridge import catalog, runtime
|
||||
from litellm.rust_bridge.bindings import NativeBinding
|
||||
from litellm.rust_bridge.catalog import Route, RouteRule, Rules
|
||||
from litellm.rust_bridge.configuration import Rollout
|
||||
from litellm.rust_bridge.ocr.entrypoints import NATIVE_AOCR, NATIVE_OCR, NativeAocr, NativeOcr
|
||||
from litellm.rust_bridge.public_call import NativeCall
|
||||
|
||||
RUST_RULES: Final[Rules] = (RouteRule(Route.OCR, Rollout.RUST_REQUIRED),)
|
||||
|
||||
|
||||
def ocr_binding(native: NativeOcr | None) -> NativeBinding[NativeOcr]:
|
||||
binding: Final[NativeBinding[NativeOcr]] = NativeBinding("ocr", validate=lambda _: None)
|
||||
binding.override(native)
|
||||
return binding
|
||||
|
||||
|
||||
def aocr_binding(native: NativeAocr | None) -> NativeBinding[NativeAocr]:
|
||||
binding: Final[NativeBinding[NativeAocr]] = NativeBinding("aocr", validate=lambda _: None)
|
||||
binding.override(native)
|
||||
return binding
|
||||
|
||||
|
||||
def response(model: str = "mistral/mistral-ocr-latest") -> OCRResponse:
|
||||
return OCRResponse(pages=[], model=model)
|
||||
|
||||
|
||||
def test_native_receives_normalized_positional_request_and_original_call_shape() -> None:
|
||||
document: Final[Mapping[str, object]] = {"type": "file", "file": b"pdf"}
|
||||
timeout: Final = httpx.Timeout(30)
|
||||
extra_headers: Final[dict[str, object]] = {"x-test": "1"}
|
||||
pages: Final = [0, 2]
|
||||
args: Final[tuple[object, ...]] = ("mistral/mistral-ocr-latest", document)
|
||||
kwargs: Final[Mapping[str, object]] = {
|
||||
"api_key": "test-key",
|
||||
"api_base": "https://example.invalid",
|
||||
"timeout": timeout,
|
||||
"custom_llm_provider": "mistral",
|
||||
"extra_headers": extra_headers,
|
||||
"pages": pages,
|
||||
}
|
||||
captured: Final[list[tuple[NativeCall, tuple[object, ...], Mapping[str, object]]]] = []
|
||||
expected: Final = response()
|
||||
|
||||
def native(
|
||||
request: NativeCall,
|
||||
) -> OCRResponse:
|
||||
args: Final = request.args
|
||||
kwargs: Final = request.kwargs
|
||||
captured.append((request, args, kwargs))
|
||||
return expected
|
||||
|
||||
result: Final = _DISPATCH.run(
|
||||
args,
|
||||
kwargs,
|
||||
python=runtime.NO_PYTHON,
|
||||
binding=ocr_binding(native),
|
||||
native=lambda hook, request, call_args, call_kwargs: hook(request),
|
||||
rules=RUST_RULES,
|
||||
)
|
||||
|
||||
request, call_args, call_kwargs = captured[0]
|
||||
assert result is expected
|
||||
assert request.resolved["model"] == "mistral/mistral-ocr-latest"
|
||||
assert request.resolved["document"] is document
|
||||
assert request.resolved["api_key"] == "test-key"
|
||||
assert request.resolved["api_base"] == "https://example.invalid"
|
||||
assert request.resolved["timeout"] is timeout
|
||||
assert request.resolved["custom_llm_provider"] == "mistral"
|
||||
assert request.resolved["extra_headers"] is extra_headers
|
||||
assert request.kwargs == kwargs
|
||||
assert request.kwargs["pages"] is pages
|
||||
assert call_args is args
|
||||
assert call_kwargs is kwargs
|
||||
|
||||
|
||||
def test_native_preserves_keyword_model_and_document_in_original_call_shape() -> None:
|
||||
document: Final[Mapping[str, object]] = {
|
||||
"type": "document_url",
|
||||
"document_url": "https://example.invalid/document.pdf",
|
||||
}
|
||||
pages: Final = [1]
|
||||
args: Final[tuple[object, ...]] = ()
|
||||
kwargs: Final[Mapping[str, object]] = {
|
||||
"model": "mistral/mistral-ocr-latest",
|
||||
"document": document,
|
||||
"pages": pages,
|
||||
}
|
||||
captured: Final[list[tuple[NativeCall, tuple[object, ...], Mapping[str, object]]]] = []
|
||||
expected: Final = response()
|
||||
|
||||
def native(
|
||||
request: NativeCall,
|
||||
) -> OCRResponse:
|
||||
args: Final = request.args
|
||||
kwargs: Final = request.kwargs
|
||||
captured.append((request, args, kwargs))
|
||||
return expected
|
||||
|
||||
result: Final = _DISPATCH.run(
|
||||
args,
|
||||
kwargs,
|
||||
python=runtime.NO_PYTHON,
|
||||
binding=ocr_binding(native),
|
||||
native=lambda hook, request, call_args, call_kwargs: hook(request),
|
||||
rules=RUST_RULES,
|
||||
)
|
||||
|
||||
request, call_args, call_kwargs = captured[0]
|
||||
assert result is expected
|
||||
assert request.resolved["model"] == "mistral/mistral-ocr-latest"
|
||||
assert request.resolved["document"] is document
|
||||
assert request.kwargs == kwargs
|
||||
assert call_args is args
|
||||
assert call_kwargs is kwargs
|
||||
assert call_kwargs["model"] == "mistral/mistral-ocr-latest"
|
||||
assert call_kwargs["document"] is document
|
||||
|
||||
|
||||
def test_aocr_marker_cannot_be_served_without_python() -> None:
|
||||
document: Final[Mapping[str, object]] = {"type": "file", "file": b"pdf"}
|
||||
args: Final[tuple[object, ...]] = ("mistral/mistral-ocr-latest", document)
|
||||
kwargs: Final[Mapping[str, object]] = {"aocr": True}
|
||||
|
||||
def native(
|
||||
request: NativeCall,
|
||||
) -> OCRResponse:
|
||||
pytest.fail("the aocr bypass marker must not reach native")
|
||||
|
||||
with pytest.raises(runtime.NoPythonImplementationError, match="bypass"):
|
||||
_DISPATCH.run(
|
||||
args,
|
||||
kwargs,
|
||||
python=runtime.NO_PYTHON,
|
||||
binding=ocr_binding(native),
|
||||
native=lambda hook, request, call_args, call_kwargs: hook(request),
|
||||
rules=RUST_RULES,
|
||||
)
|
||||
|
||||
|
||||
def test_missing_native_binding_is_a_required_rust_error() -> None:
|
||||
args: Final[tuple[object, ...]] = ("mistral/mistral-ocr-latest", {"type": "file", "file": b"pdf"})
|
||||
|
||||
with pytest.raises(RuntimeError, match="Rust ocr bridge"):
|
||||
_DISPATCH.run(
|
||||
args,
|
||||
{},
|
||||
python=runtime.NO_PYTHON,
|
||||
binding=ocr_binding(None),
|
||||
native=lambda hook, request, call_args, call_kwargs: hook(request),
|
||||
rules=RUST_RULES,
|
||||
)
|
||||
|
||||
|
||||
def test_non_required_rule_cannot_be_served_without_python() -> None:
|
||||
args: Final[tuple[object, ...]] = ("mistral/mistral-ocr-latest", {"type": "file", "file": b"pdf"})
|
||||
|
||||
def native(
|
||||
request: NativeCall,
|
||||
) -> OCRResponse:
|
||||
return response()
|
||||
|
||||
with pytest.raises(runtime.NoPythonImplementationError, match="RUST_REQUIRED"):
|
||||
_DISPATCH.run(
|
||||
args,
|
||||
{},
|
||||
python=runtime.NO_PYTHON,
|
||||
binding=ocr_binding(native),
|
||||
native=lambda hook, request, call_args, call_kwargs: hook(request),
|
||||
rules=(RouteRule(Route.OCR, Rollout.PYTHON_ONLY),),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("args", "kwargs", "message"),
|
||||
(
|
||||
(
|
||||
("mistral/mistral-ocr-latest", {"type": "file", "file": b"pdf"}),
|
||||
{"model": "duplicate"},
|
||||
r"ocr\(\) got multiple values for argument 'model'",
|
||||
),
|
||||
(
|
||||
("mistral/mistral-ocr-latest",),
|
||||
{},
|
||||
r"ocr\(\) missing 1 required positional argument: 'document'",
|
||||
),
|
||||
),
|
||||
)
|
||||
def test_ocr_parser_errors_before_native(args: tuple[object, ...], kwargs: Mapping[str, object], message: str) -> None:
|
||||
def native(
|
||||
request: NativeCall,
|
||||
) -> OCRResponse:
|
||||
pytest.fail("OCR parser failures must not call native")
|
||||
|
||||
with pytest.raises(TypeError, match=message):
|
||||
_DISPATCH.run(
|
||||
args,
|
||||
kwargs,
|
||||
python=runtime.NO_PYTHON,
|
||||
binding=ocr_binding(native),
|
||||
native=lambda hook, request, call_args, call_kwargs: hook(request),
|
||||
rules=RUST_RULES,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("args", "kwargs", "message"),
|
||||
(
|
||||
(
|
||||
("mistral/mistral-ocr-latest", {"type": "file", "file": b"pdf"}),
|
||||
{"model": "duplicate"},
|
||||
r"aocr\(\) got multiple values for argument 'model'",
|
||||
),
|
||||
(
|
||||
("mistral/mistral-ocr-latest",),
|
||||
{},
|
||||
r"aocr\(\) missing 1 required positional argument: 'document'",
|
||||
),
|
||||
),
|
||||
)
|
||||
async def test_aocr_parser_errors_before_native(
|
||||
args: tuple[object, ...], kwargs: Mapping[str, object], message: str
|
||||
) -> None:
|
||||
async def native(
|
||||
request: NativeCall,
|
||||
) -> OCRResponse:
|
||||
pytest.fail("OCR parser failures must not call native")
|
||||
|
||||
with pytest.raises(TypeError, match=message):
|
||||
await _ADISPATCH.arun(
|
||||
args,
|
||||
kwargs,
|
||||
python=runtime.NO_PYTHON,
|
||||
binding=aocr_binding(native),
|
||||
native=lambda hook, request, call_args, call_kwargs: hook(request),
|
||||
rules=RUST_RULES,
|
||||
)
|
||||
|
||||
|
||||
def test_public_ocr_routes_through_dispatch(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
document: Final[Mapping[str, object]] = {
|
||||
"type": "document_url",
|
||||
"document_url": "https://example.invalid/document.pdf",
|
||||
}
|
||||
captured: Final[list[NativeCall]] = []
|
||||
expected: Final = response()
|
||||
|
||||
def native(
|
||||
request: NativeCall,
|
||||
) -> OCRResponse:
|
||||
captured.append(request)
|
||||
return expected
|
||||
|
||||
NATIVE_OCR.override(native)
|
||||
monkeypatch.setattr(catalog, "RULES", RUST_RULES)
|
||||
public_ocr: Final = cast(Callable[..., OCRResponse], litellm.ocr)
|
||||
try:
|
||||
result: Final = public_ocr(model="mistral/mistral-ocr-latest", document=document)
|
||||
finally:
|
||||
NATIVE_OCR.reset()
|
||||
assert result is expected
|
||||
assert [request.resolved["model"] for request in captured] == ["mistral/mistral-ocr-latest"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_public_aocr_routes_through_dispatch(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
document: Final[Mapping[str, object]] = {
|
||||
"type": "document_url",
|
||||
"document_url": "https://example.invalid/document.pdf",
|
||||
}
|
||||
captured: Final[list[NativeCall]] = []
|
||||
expected: Final = response()
|
||||
|
||||
async def native(
|
||||
request: NativeCall,
|
||||
) -> OCRResponse:
|
||||
captured.append(request)
|
||||
return expected
|
||||
|
||||
NATIVE_AOCR.override(native)
|
||||
monkeypatch.setattr(catalog, "RULES", RUST_RULES)
|
||||
public_aocr: Final = cast(Callable[..., Awaitable[OCRResponse]], litellm.aocr)
|
||||
try:
|
||||
result: Final = await public_aocr(model="mistral/mistral-ocr-latest", document=document)
|
||||
finally:
|
||||
NATIVE_AOCR.reset()
|
||||
assert result is expected
|
||||
assert [request.resolved["model"] for request in captured] == ["mistral/mistral-ocr-latest"]
|
||||
|
|
@ -1,5 +1,5 @@
|
|||
import inspect
|
||||
from collections.abc import Awaitable, Callable, Mapping
|
||||
from collections.abc import Callable, Mapping
|
||||
from typing import Final, cast # noqa: TID251 # narrows legacy callable signatures for inspect
|
||||
|
||||
import pytest
|
||||
|
|
@ -11,14 +11,11 @@ from litellm.responses.dispatch import (
|
|||
_ADISPATCH, # pyright: ignore[reportPrivateUsage] # tests configured dispatch
|
||||
_DISPATCH, # pyright: ignore[reportPrivateUsage] # tests configured dispatch
|
||||
)
|
||||
from litellm.rust_bridge import catalog
|
||||
from litellm.rust_bridge.bindings import NativeBinding
|
||||
from litellm.rust_bridge.catalog import Route, RouteRule
|
||||
from litellm.rust_bridge.configuration import Rollout
|
||||
from litellm.rust_bridge.public_call import NativeCall
|
||||
from litellm.rust_bridge.responses.entrypoints import (
|
||||
NATIVE_ARESPONSES,
|
||||
NATIVE_RESPONSES,
|
||||
NativeAresponses,
|
||||
NativeResponses,
|
||||
)
|
||||
|
|
@ -181,35 +178,6 @@ def test_native_receives_normalized_request_and_original_call_shape() -> None:
|
|||
assert call_kwargs["litellm_metadata"] is metadata
|
||||
|
||||
|
||||
def test_internal_async_marker_bypasses_native() -> None:
|
||||
args: Final[tuple[object, ...]] = (INPUT, "gpt-4o")
|
||||
kwargs: Final[Mapping[str, object]] = {"aresponses": True}
|
||||
captured: Final[list[tuple[tuple[object, ...], Mapping[str, object]]]] = []
|
||||
response: Final = _response()
|
||||
|
||||
def python(*call_args: object, **call_kwargs: object) -> ResponsesAPIResponse: # kwargs-ok: records call shape
|
||||
captured.append((call_args, call_kwargs))
|
||||
return response
|
||||
|
||||
def native(
|
||||
request: NativeCall,
|
||||
) -> ResponsesAPIResponse:
|
||||
pytest.fail("aresponses' inner responses call must stay on Python")
|
||||
|
||||
assert (
|
||||
_DISPATCH.run(
|
||||
args,
|
||||
kwargs,
|
||||
python=python,
|
||||
binding=responses_binding(native),
|
||||
native=lambda hook, request, call_args, call_kwargs: hook(request),
|
||||
rules=RUST_RULES,
|
||||
)
|
||||
is response
|
||||
)
|
||||
assert captured == [(args, kwargs)]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("args", "kwargs"),
|
||||
(
|
||||
|
|
@ -244,49 +212,6 @@ def test_binding_errors_delegate_unchanged_to_python(args: tuple[object, ...], k
|
|||
assert captured == [(args, kwargs)]
|
||||
|
||||
|
||||
def test_public_responses_routes_through_dispatch(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
captured: Final[list[NativeCall]] = []
|
||||
expected: Final = _response()
|
||||
|
||||
def native(
|
||||
request: NativeCall,
|
||||
) -> ResponsesAPIResponse:
|
||||
captured.append(request)
|
||||
return expected
|
||||
|
||||
NATIVE_RESPONSES.override(native)
|
||||
monkeypatch.setattr(catalog, "RULES", RUST_RULES)
|
||||
public_responses: Final = cast(Callable[..., ResponsesAPIResponse], litellm.responses)
|
||||
try:
|
||||
result: Final = public_responses(input=INPUT, model="gpt-4o")
|
||||
finally:
|
||||
NATIVE_RESPONSES.reset()
|
||||
assert result is expected
|
||||
assert [request.resolved["model"] for request in captured] == ["gpt-4o"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_public_aresponses_routes_through_dispatch(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
captured: Final[list[NativeCall]] = []
|
||||
expected: Final = _response()
|
||||
|
||||
async def native(
|
||||
request: NativeCall,
|
||||
) -> ResponsesAPIResponse:
|
||||
captured.append(request)
|
||||
return expected
|
||||
|
||||
NATIVE_ARESPONSES.override(native)
|
||||
monkeypatch.setattr(catalog, "RULES", RUST_RULES)
|
||||
public_aresponses: Final = cast(Callable[..., Awaitable[ResponsesAPIResponse]], litellm.aresponses)
|
||||
try:
|
||||
result: Final = await public_aresponses(input=INPUT, model="gpt-4o")
|
||||
finally:
|
||||
NATIVE_ARESPONSES.reset()
|
||||
assert result is expected
|
||||
assert [request.resolved["model"] for request in captured] == ["gpt-4o"]
|
||||
|
||||
|
||||
def test_responses_with_retries_uses_the_dispatch_entrypoint(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
calls: Final[list[Mapping[str, object]]] = []
|
||||
expected: Final = _response()
|
||||
|
|
|
|||
|
|
@ -11,7 +11,7 @@ from tests.test_litellm_rust.support.isolation import rebound
|
|||
from tests.test_litellm_rust.support.recording_server import RecordingServer
|
||||
from tests.test_litellm_rust.support.served import assert_served_natively
|
||||
|
||||
__all__ = ["CALL_ID", "assert_served_natively", "call_public", "native_route"]
|
||||
__all__ = ["CALL_ID", "assert_served_natively", "call_public", "native_route", "routed"]
|
||||
|
||||
CALL_ID: Final = "callback-contract"
|
||||
|
||||
|
|
@ -22,6 +22,12 @@ def native_route(*routes: Route) -> Generator[None]:
|
|||
yield
|
||||
|
||||
|
||||
@contextmanager
|
||||
def routed(*rules: RouteRule) -> Generator[None]:
|
||||
with rebound(catalog, "RULES", rules):
|
||||
yield
|
||||
|
||||
|
||||
async def call_public(
|
||||
create: Callable[..., object],
|
||||
acreate: Callable[..., Awaitable[object]],
|
||||
|
|
|
|||
|
|
@ -4,9 +4,8 @@ from typing import Final, TypeAlias
|
|||
|
||||
import pytest
|
||||
|
||||
from litellm.rust_bridge import configuration
|
||||
from litellm.rust_bridge.bindings import NativeBinding
|
||||
from litellm.rust_bridge.catalog import Route, RouteContext, RouteRule, Rules
|
||||
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
|
||||
|
|
@ -70,29 +69,6 @@ def test_unconditional_python_rule_prevents_later_rust_rule_projection() -> None
|
|||
assert result is expected
|
||||
|
||||
|
||||
def test_disabled_optional_rust_rule_forwards_before_projection() -> None:
|
||||
rules: Final[Rules] = (RouteRule(Route.OCR, Rollout.RUST_OPT_OUT),)
|
||||
|
||||
def reject_request(args: tuple[object, ...], kwargs: Mapping[str, object]) -> Request:
|
||||
pytest.fail("Disabled optional Rust must not project the request")
|
||||
|
||||
dispatch: Final = PublicDispatch(route=Route.OCR, request=reject_request, context=lambda _: RouteContext(Route.OCR))
|
||||
expected: Final = object()
|
||||
configuration.rust(False)
|
||||
try:
|
||||
result: Final = dispatch.run(
|
||||
("model",),
|
||||
{},
|
||||
python=lambda *args, **kwargs: expected,
|
||||
binding=binding(),
|
||||
native=lambda hook, request, args, kwargs: pytest.fail("Disabled optional Rust must not call native"),
|
||||
rules=rules,
|
||||
)
|
||||
finally:
|
||||
configuration.rust(None)
|
||||
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))
|
||||
|
|
@ -207,36 +183,42 @@ def test_rules_for_other_routes_and_constrained_python_rules_skip_projection() -
|
|||
assert result is expected
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_bypass_forwards_to_python_without_native() -> None:
|
||||
request: Final = Request(model="bypassed-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),
|
||||
bypass=lambda value: value.model == "bypassed-model",
|
||||
)
|
||||
@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()
|
||||
|
||||
async def python(*args: object, **kwargs: object) -> object: # kwargs-ok: public pass-through shape
|
||||
return expected
|
||||
|
||||
result: Final = await dispatch.arun(
|
||||
("bypassed-model",),
|
||||
{},
|
||||
python=python,
|
||||
binding=binding(),
|
||||
native=lambda hook, value, args, kwargs: pytest.fail("Bypassed requests must not call native"),
|
||||
rules=rules,
|
||||
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 native_route(result: object) -> NativeBinding[NativeRoute]:
|
||||
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
|
||||
|
|
@ -263,52 +245,43 @@ async def dispatch_without_python(
|
|||
return await dispatch.arun(("model",), {"page": 1}, python=NO_PYTHON, binding=bound, native=native, rules=rules)
|
||||
|
||||
|
||||
def ocr_dispatch(request: Request | None, *, bypass: bool = False) -> PublicDispatch[Request]:
|
||||
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),
|
||||
bypass=lambda _: bypass,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("asynchronous", (False, True))
|
||||
@pytest.mark.parametrize("switch", (None, False))
|
||||
async def test_dispatch_without_python_hands_every_call_to_native(asynchronous: bool, switch: bool | None) -> None:
|
||||
async def test_dispatch_without_python_hands_every_call_to_native(asynchronous: bool) -> None:
|
||||
request: Final = Request(model="model")
|
||||
result: Final = object()
|
||||
configuration.rust(switch)
|
||||
try:
|
||||
dispatched: Final = await dispatch_without_python(
|
||||
ocr_dispatch(request),
|
||||
native_route(result),
|
||||
(RouteRule(Route.OCR, Rollout.RUST_REQUIRED),),
|
||||
asynchronous=asynchronous,
|
||||
)
|
||||
finally:
|
||||
configuration.rust(None)
|
||||
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", "bypass", "rules", "reason"),
|
||||
("request_value", "rules", "reason"),
|
||||
(
|
||||
(Request(model="model"), False, (), "must resolve to RUST_REQUIRED"),
|
||||
(Request(model="model"), False, (RouteRule(Route.OCR, Rollout.RUST_OPT_OUT),), "must resolve to RUST_REQUIRED"),
|
||||
(None, False, (RouteRule(Route.OCR, Rollout.RUST_REQUIRED),), "must project to a native request"),
|
||||
(Request(model="model"), True, (RouteRule(Route.OCR, Rollout.RUST_REQUIRED),), "bypass predicate matches"),
|
||||
(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", "bypassed-call"),
|
||||
ids=("no-rule", "opt-out-rule", "unprojectable-call"),
|
||||
)
|
||||
async def test_dispatch_without_python_never_falls_back(
|
||||
asynchronous: bool, request_value: Request | None, bypass: bool, rules: Rules, reason: str
|
||||
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, bypass=bypass), bound, rules, asynchronous=asynchronous
|
||||
)
|
||||
await dispatch_without_python(ocr_dispatch(request_value), bound, rules, asynchronous=asynchronous)
|
||||
|
|
|
|||
287
tests/test_litellm_rust/test_public_entries.py
Normal file
287
tests/test_litellm_rust/test_public_entries.py
Normal file
|
|
@ -0,0 +1,287 @@
|
|||
import inspect
|
||||
from collections.abc import Awaitable, Callable, Iterator, Mapping
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum, auto
|
||||
from typing import Final, TypeAlias, cast # noqa: TID251 # public entries return union results
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.messages import dispatch as messages
|
||||
from litellm.rust_bridge.bindings import NativeBinding
|
||||
from litellm.rust_bridge.catalog import Route, RouteRule
|
||||
from litellm.rust_bridge.chat_completions.entrypoints import NATIVE_ACOMPLETION, NATIVE_COMPLETION
|
||||
from litellm.rust_bridge.configuration import Rollout
|
||||
from litellm.rust_bridge.embeddings.entrypoints import NATIVE_AEMBEDDING, NATIVE_EMBEDDING
|
||||
from litellm.rust_bridge.messages.entrypoints import NATIVE_AMESSAGES, NATIVE_MESSAGES
|
||||
from litellm.rust_bridge.ocr.entrypoints import NATIVE_AOCR, NATIVE_OCR
|
||||
from litellm.rust_bridge.public_call import NativeCall
|
||||
from litellm.rust_bridge.responses.entrypoints import NATIVE_ARESPONSES, NATIVE_RESPONSES
|
||||
from litellm.rust_bridge.runtime import NoPythonImplementationError
|
||||
from litellm.types.utils import ModelResponse
|
||||
from tests.test_litellm_rust.support.routing import routed
|
||||
|
||||
MESSAGES: Final = [{"role": "user", "content": "hi"}]
|
||||
|
||||
Served: TypeAlias = Callable[[object], bool]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Entry:
|
||||
"""``python`` is ``None`` for a route with no Python implementation, which only RUST_REQUIRED rules serve."""
|
||||
|
||||
route: Route
|
||||
sync: Callable[..., object]
|
||||
asynchronous: Callable[..., Awaitable[object]]
|
||||
flag: str
|
||||
sync_binding: NativeBinding[object]
|
||||
async_binding: NativeBinding[object]
|
||||
arguments: Mapping[str, object]
|
||||
positional: int
|
||||
python: Served | None
|
||||
|
||||
def args(self) -> tuple[object, ...]:
|
||||
return tuple(self.arguments.values())[: self.positional]
|
||||
|
||||
def kwargs(self, **extra: object) -> Mapping[str, object]: # kwargs-ok: forwards the async flag
|
||||
return {**dict(tuple(self.arguments.items())[self.positional :]), **extra}
|
||||
|
||||
|
||||
def _binding(value: object) -> NativeBinding[object]:
|
||||
return cast("NativeBinding[object]", value) # cast-ok: the table only overrides and resets bindings
|
||||
|
||||
|
||||
def _chat(result: object) -> bool:
|
||||
return isinstance(result, ModelResponse) and result.choices[0].message.content == "python" # pyright: ignore[reportAttributeAccessIssue] # non-streaming choice
|
||||
|
||||
|
||||
def _embedding(result: object) -> bool:
|
||||
return getattr(result, "data")[0]["embedding"] == [0.5] # noqa: B009 # union result has no common attribute
|
||||
|
||||
|
||||
def _responses(result: object) -> bool:
|
||||
return getattr(result, "output_text") == "python" # noqa: B009 # union result has no common attribute
|
||||
|
||||
|
||||
def _messages(result: object) -> bool:
|
||||
return isinstance(result, Mapping) and result["content"][0]["text"] == "python"
|
||||
|
||||
|
||||
ENTRIES: Final = (
|
||||
Entry(
|
||||
route=Route.CHAT_COMPLETIONS,
|
||||
sync=litellm.completion,
|
||||
asynchronous=litellm.acompletion,
|
||||
flag="acompletion",
|
||||
sync_binding=_binding(NATIVE_COMPLETION),
|
||||
async_binding=_binding(NATIVE_ACOMPLETION),
|
||||
arguments={"model": "openai/gpt-4o", "messages": MESSAGES, "mock_response": "python"},
|
||||
positional=2,
|
||||
python=_chat,
|
||||
),
|
||||
Entry(
|
||||
route=Route.EMBEDDINGS,
|
||||
sync=litellm.embedding,
|
||||
asynchronous=litellm.aembedding,
|
||||
flag="aembedding",
|
||||
sync_binding=_binding(NATIVE_EMBEDDING),
|
||||
async_binding=_binding(NATIVE_AEMBEDDING),
|
||||
arguments={"model": "openai/text-embedding-3-small", "input": ["hi"], "mock_response": [0.5]},
|
||||
positional=2,
|
||||
python=_embedding,
|
||||
),
|
||||
Entry(
|
||||
route=Route.RESPONSES,
|
||||
sync=litellm.responses,
|
||||
asynchronous=litellm.aresponses,
|
||||
flag="aresponses",
|
||||
sync_binding=_binding(NATIVE_RESPONSES),
|
||||
async_binding=_binding(NATIVE_ARESPONSES),
|
||||
arguments={"input": "hi", "model": "openai/gpt-4o", "mock_response": "python"},
|
||||
positional=2,
|
||||
python=_responses,
|
||||
),
|
||||
Entry(
|
||||
route=Route.MESSAGES,
|
||||
sync=messages.anthropic_messages_handler,
|
||||
asynchronous=messages.anthropic_messages,
|
||||
flag="is_async",
|
||||
sync_binding=_binding(NATIVE_MESSAGES),
|
||||
async_binding=_binding(NATIVE_AMESSAGES),
|
||||
arguments={
|
||||
"max_tokens": 16,
|
||||
"messages": MESSAGES,
|
||||
"model": "anthropic/claude-sonnet-4-5",
|
||||
"mock_response": "python",
|
||||
},
|
||||
positional=0,
|
||||
python=_messages,
|
||||
),
|
||||
Entry(
|
||||
route=Route.OCR,
|
||||
sync=litellm.ocr,
|
||||
asynchronous=litellm.aocr,
|
||||
flag="aocr",
|
||||
sync_binding=_binding(NATIVE_OCR),
|
||||
async_binding=_binding(NATIVE_AOCR),
|
||||
arguments={"model": "mistral/mistral-ocr-latest", "document": {"type": "file", "file": b"pdf"}},
|
||||
positional=2,
|
||||
python=None,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class Form(Enum):
|
||||
SYNC = auto()
|
||||
SYNC_FLAG_FALSE = auto()
|
||||
ASYNC = auto()
|
||||
SYNC_FLAG_TRUE = auto()
|
||||
ASYNC_FLAG_TRUE = auto()
|
||||
|
||||
|
||||
async def call(entry: Entry, form: Form) -> object:
|
||||
match form:
|
||||
case Form.SYNC:
|
||||
return entry.sync(*entry.args(), **entry.kwargs())
|
||||
case Form.SYNC_FLAG_FALSE:
|
||||
return entry.sync(*entry.args(), **entry.kwargs(**{entry.flag: False}))
|
||||
case Form.ASYNC:
|
||||
return await entry.asynchronous(*entry.args(), **entry.kwargs())
|
||||
case Form.SYNC_FLAG_TRUE:
|
||||
return await cast(
|
||||
"Awaitable[object]", entry.sync(*entry.args(), **entry.kwargs(**{entry.flag: True}))
|
||||
) # cast-ok: a sync entry with its flag set returns a coroutine
|
||||
case Form.ASYNC_FLAG_TRUE:
|
||||
return await entry.asynchronous(*entry.args(), **entry.kwargs(**{entry.flag: True}))
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class Native:
|
||||
result: ModelResponse = field(default_factory=ModelResponse)
|
||||
calls: list[tuple[str, NativeCall]] = field(default_factory=list) # mutable-ok: records native calls
|
||||
|
||||
def install(self, entry: Entry) -> None:
|
||||
def sync(native_call: NativeCall) -> ModelResponse:
|
||||
self.calls.append(("sync", native_call))
|
||||
return self.result
|
||||
|
||||
async def asynchronous(native_call: NativeCall) -> ModelResponse:
|
||||
self.calls.append(("async", native_call))
|
||||
return self.result
|
||||
|
||||
entry.sync_binding.override(sync)
|
||||
entry.async_binding.override(asynchronous)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def native(entry: Entry) -> Iterator[Native]:
|
||||
recorder: Final = Native()
|
||||
recorder.install(entry)
|
||||
yield recorder
|
||||
entry.sync_binding.reset()
|
||||
entry.async_binding.reset()
|
||||
|
||||
|
||||
class Outcome(Enum):
|
||||
PYTHON = auto()
|
||||
NATIVE = auto()
|
||||
UNAVAILABLE = auto()
|
||||
NO_PYTHON = auto()
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Scenario:
|
||||
rollout: Rollout | None
|
||||
binding: bool
|
||||
outcome: Outcome
|
||||
|
||||
|
||||
SCENARIOS: Final = {
|
||||
"no-rule": Scenario(None, binding=True, outcome=Outcome.PYTHON),
|
||||
"python-only": Scenario(Rollout.PYTHON_ONLY, binding=True, outcome=Outcome.PYTHON),
|
||||
"required": Scenario(Rollout.RUST_REQUIRED, binding=True, outcome=Outcome.NATIVE),
|
||||
"required-without-extension": Scenario(Rollout.RUST_REQUIRED, binding=False, outcome=Outcome.UNAVAILABLE),
|
||||
}
|
||||
|
||||
|
||||
def expected(entry: Entry, scenario: Scenario) -> Outcome:
|
||||
if entry.python is None and scenario.rollout is not Rollout.RUST_REQUIRED:
|
||||
return Outcome.NO_PYTHON
|
||||
return scenario.outcome
|
||||
|
||||
|
||||
@contextmanager
|
||||
def scenario_rules(entry: Entry, scenario: Scenario) -> Iterator[None]:
|
||||
if not scenario.binding:
|
||||
entry.sync_binding.override(None)
|
||||
entry.async_binding.override(None)
|
||||
with routed(*(() if scenario.rollout is None else (RouteRule(entry.route, scenario.rollout),))):
|
||||
yield
|
||||
|
||||
|
||||
def entry_id(entry: Entry) -> str:
|
||||
return entry.route.value
|
||||
|
||||
|
||||
@pytest.mark.parametrize("form", tuple(Form), ids=lambda form: form.name.lower())
|
||||
@pytest.mark.parametrize("scenario", tuple(SCENARIOS.values()), ids=tuple(SCENARIOS))
|
||||
@pytest.mark.parametrize("entry", ENTRIES, ids=entry_id)
|
||||
async def test_every_public_form_takes_the_path_the_rules_pick(
|
||||
entry: Entry, scenario: Scenario, form: Form, native: Native
|
||||
) -> None:
|
||||
outcome: Final = expected(entry, scenario)
|
||||
with scenario_rules(entry, scenario):
|
||||
if outcome is Outcome.UNAVAILABLE:
|
||||
with pytest.raises(RuntimeError, match=f"Rust {entry.route.value} bridge is unavailable"):
|
||||
await call(entry, form)
|
||||
return
|
||||
if outcome is Outcome.NO_PYTHON:
|
||||
with pytest.raises(NoPythonImplementationError, match=f"{entry.route.value} has no Python implementation"):
|
||||
await call(entry, form)
|
||||
assert native.calls == []
|
||||
return
|
||||
result: Final = await call(entry, form)
|
||||
|
||||
if outcome is Outcome.PYTHON:
|
||||
assert entry.python is not None
|
||||
assert entry.python(result)
|
||||
assert native.calls == []
|
||||
return
|
||||
assert result is native.result
|
||||
[(binding, native_call)] = native.calls
|
||||
assert binding == ("sync" if form in (Form.SYNC, Form.SYNC_FLAG_FALSE) else "async")
|
||||
assert {name: native_call.resolved[name] for name in entry.arguments} == entry.arguments
|
||||
assert entry.flag not in native_call.resolved
|
||||
|
||||
|
||||
@pytest.mark.parametrize("entry", ENTRIES, ids=entry_id)
|
||||
async def test_async_forms_hand_native_the_same_call(entry: Entry, native: Native) -> None:
|
||||
with scenario_rules(entry, SCENARIOS["required"]):
|
||||
for form in (Form.ASYNC, Form.SYNC_FLAG_TRUE, Form.ASYNC_FLAG_TRUE):
|
||||
_ = await call(entry, form)
|
||||
|
||||
assert [binding for binding, _ in native.calls] == ["async", "async", "async"]
|
||||
resolved: Final = [dict(native_call.resolved) for _, native_call in native.calls]
|
||||
assert resolved[1:] == [resolved[0], resolved[0]]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("entry", ENTRIES, ids=entry_id)
|
||||
async def test_sync_entry_with_flag_true_does_nothing_until_awaited(entry: Entry, native: Native) -> None:
|
||||
with scenario_rules(entry, SCENARIOS["required"]):
|
||||
pending: Final = entry.sync(*entry.args(), **entry.kwargs(**{entry.flag: True}))
|
||||
|
||||
assert inspect.iscoroutine(pending)
|
||||
assert native.calls == []
|
||||
assert await pending is native.result
|
||||
assert len(native.calls) == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize("entry", ENTRIES, ids=entry_id)
|
||||
async def test_sync_entry_with_flag_true_raises_setup_errors_at_await(entry: Entry, native: Native) -> None:
|
||||
with scenario_rules(entry, SCENARIOS["required-without-extension"]):
|
||||
pending: Final = entry.sync(*entry.args(), **entry.kwargs(**{entry.flag: True}))
|
||||
|
||||
assert inspect.iscoroutine(pending)
|
||||
with pytest.raises(RuntimeError, match="is unavailable"):
|
||||
await pending
|
||||
|
|
@ -4,6 +4,6 @@ Test what each side of the bridge does, not the rollout policy that picks a side
|
|||
|
||||
Call each path directly with an explicit decision instead. The Python path is the implementation the dispatcher falls back to, e.g. `litellm.responses.main.responses`. The Rust path is the native binding, e.g. `NATIVE_OCR.load()` from `litellm/rust_bridge/ocr/entrypoints.py`, called with the `NativeCall` envelope that dispatch would hand it for every public inference route. When the native side reads a policy-derived setting such as `settings.secret_manager().native`, pin that field in the test instead of deriving it from the catalog. `ocr/test_secrets.py` shows the pattern
|
||||
|
||||
Rollout policy itself, meaning which rule matches and what `LITELLM_RUST` changes, belongs in `test_catalog.py`, `test_configuration.py` and `test_dispatch.py`, tested against rules the test builds rather than the shipped `catalog.RULES`
|
||||
Rollout policy itself, meaning which rule matches and what `LITELLM_RUST` changes, belongs in `test_catalog.py` and `test_configuration.py`, tested against rules the test builds rather than the shipped `catalog.RULES`. Which path a public call takes under those rules is a dispatch test and lives in `tests/test_litellm_rust/` for now
|
||||
|
||||
Before adding a test here, ask whether it checks something Rust cannot. A `host/routes/<route>.py` module is the Python half of a native route: it projects Python-only state (the cost map, `litellm.*` settings, request kwargs) into the plain values the Rust side consumes, and maps native failures back onto public exceptions. Those projections are what belongs here, because a wrong key or an ignored provider prefix ships the wrong value to Rust and no Rust test sees it. `host/routes/test_messages.py` shows the shape. Behavior that lives in Rust (a request transform given its inputs, header assembly, stream relay) is tested in the crate, and the route end to end is tested against a recording server in `tests/test_litellm_rust/`. A test that only re-checks a Python helper the route host happens to call is a duplicate of that helper's own test and should not be added
|
||||
|
|
|
|||
|
|
@ -7,10 +7,12 @@ import pytest
|
|||
|
||||
from litellm.rust_bridge import catalog, configuration
|
||||
from litellm.rust_bridge.catalog import (
|
||||
Feature,
|
||||
Route,
|
||||
RouteContext,
|
||||
RouteRule,
|
||||
Rules,
|
||||
Target,
|
||||
)
|
||||
from litellm.rust_bridge.configuration import Decision, Rollout
|
||||
|
||||
|
|
@ -23,36 +25,6 @@ def isolated_configuration(monkeypatch: pytest.MonkeyPatch) -> Generator[None]:
|
|||
configuration.reset_rust_configuration()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("route", tuple(Route))
|
||||
@pytest.mark.parametrize(
|
||||
"provider", (None, "bedrock", "mistral", "anthropic", "deepseek", "vertex_ai", "openai", "azure_ai", "unknown")
|
||||
)
|
||||
@pytest.mark.parametrize("process", (None, False, True))
|
||||
@pytest.mark.parametrize("environment", (None, "0", "1"))
|
||||
def test_shipped_decisions(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
route: Route,
|
||||
provider: str | None,
|
||||
process: bool | None,
|
||||
environment: str | None,
|
||||
) -> None:
|
||||
configuration.rust(process)
|
||||
if environment is not None:
|
||||
monkeypatch.setenv("LITELLM_RUST", environment)
|
||||
context: Final = RouteContext(route, provider=provider, model="test-model")
|
||||
|
||||
if route is Route.OCR or (route is Route.TRANSCRIPTION and provider == "bedrock"):
|
||||
assert catalog.rollout(context) is Rollout.RUST_REQUIRED
|
||||
assert catalog.decision(context) is Decision.RUST_REQUIRED
|
||||
elif route is Route.MESSAGES and provider in ("anthropic", "deepseek", "vertex_ai"):
|
||||
assert catalog.rollout(context) is Rollout.RUST_OPT_IN
|
||||
opted_in: Final = environment == "1" or (environment is None and process is True)
|
||||
assert catalog.decision(context) is (Decision.RUST_WITH_FALLBACK if opted_in else Decision.PYTHON)
|
||||
else:
|
||||
assert catalog.rollout(context) is Rollout.PYTHON_ONLY
|
||||
assert catalog.decision(context) is Decision.PYTHON
|
||||
|
||||
|
||||
@pytest.mark.parametrize("route", tuple(Route))
|
||||
def test_missing_rule_stays_on_python_even_when_rust_is_enabled(monkeypatch: pytest.MonkeyPatch, route: Route) -> None:
|
||||
configuration.rust(True)
|
||||
|
|
@ -72,6 +44,11 @@ def test_missing_rule_stays_on_python_even_when_rust_is_enabled(monkeypatch: pyt
|
|||
(RouteContext(Route.RESPONSES, provider="openai", model="other"), Decision.PYTHON),
|
||||
(RouteContext(Route.RESPONSES, provider="anthropic", model="m"), Decision.PYTHON),
|
||||
(RouteContext(Route.MESSAGES, provider="openai", model="m"), Decision.PYTHON),
|
||||
(RouteContext(Route.MESSAGES, target=Target.MESSAGES), Decision.RUST_REQUIRED),
|
||||
(RouteContext(Route.MESSAGES, target=Target.CHAT_COMPLETIONS), Decision.PYTHON),
|
||||
(RouteContext(Route.MESSAGES), Decision.PYTHON),
|
||||
(RouteContext(Route.OCR, features=frozenset({Feature.SAFEGUARDS_CLASSIFIER})), Decision.PYTHON),
|
||||
(RouteContext(Route.OCR), Decision.RUST_REQUIRED),
|
||||
),
|
||||
)
|
||||
def test_first_matching_rule_respects_every_constraint(context: RouteContext, expected: Decision) -> None:
|
||||
|
|
@ -83,24 +60,14 @@ def test_first_matching_rule_respects_every_constraint(context: RouteContext, ex
|
|||
models=frozenset({"m"}),
|
||||
),
|
||||
RouteRule(Route.RESPONSES, Rollout.PYTHON_ONLY),
|
||||
RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED, targets=frozenset({Target.MESSAGES})),
|
||||
RouteRule(Route.OCR, Rollout.PYTHON_ONLY, features=frozenset({Feature.SAFEGUARDS_CLASSIFIER})),
|
||||
RouteRule(Route.OCR, Rollout.RUST_REQUIRED),
|
||||
)
|
||||
|
||||
assert catalog.decision(context, rules) is expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize("process", (None, False, True))
|
||||
@pytest.mark.parametrize("environment", (None, "0", "1"))
|
||||
def test_ocr_has_no_python_path_to_opt_out_to(
|
||||
monkeypatch: pytest.MonkeyPatch, process: bool | None, environment: str | None
|
||||
) -> None:
|
||||
configuration.rust(process)
|
||||
if environment is not None:
|
||||
monkeypatch.setenv("LITELLM_RUST", environment)
|
||||
|
||||
assert catalog.decision(RouteContext(Route.OCR, model="m")) is Decision.RUST_REQUIRED
|
||||
assert catalog.decision(RouteContext(Route.OCR, provider="aws_textract", model="m")) is Decision.RUST_REQUIRED
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("context", "expected"),
|
||||
(
|
||||
|
|
@ -154,6 +121,10 @@ def test_empty_constraints_match_nothing() -> None:
|
|||
rules: Final[Rules] = (
|
||||
RouteRule(Route.OCR, Rollout.RUST_REQUIRED, providers=frozenset()),
|
||||
RouteRule(Route.OCR, Rollout.RUST_REQUIRED, models=frozenset()),
|
||||
RouteRule(Route.OCR, Rollout.RUST_REQUIRED, features=frozenset()),
|
||||
)
|
||||
|
||||
assert catalog.decision(RouteContext(Route.OCR), rules) is Decision.PYTHON
|
||||
assert catalog.decision(
|
||||
RouteContext(Route.OCR, features=frozenset({Feature.SAFEGUARDS_CLASSIFIER})), rules
|
||||
) is Decision.PYTHON
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@ from typing import Final
|
|||
|
||||
import pytest
|
||||
|
||||
from litellm.rust_bridge.public_call import native_call, signature
|
||||
from litellm.rust_bridge.public_call import AsyncFlag, native_call, signature
|
||||
|
||||
|
||||
def _messages(
|
||||
|
|
@ -18,6 +18,7 @@ def _messages(
|
|||
|
||||
|
||||
_SIGNATURE: Final = signature(_messages)
|
||||
_FLAG: Final = AsyncFlag("aflag", sync=_SIGNATURE)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("supplied", ({}, {"api_key": None}, {"api_key": "explicit"}))
|
||||
|
|
@ -62,3 +63,16 @@ def test_resolved_lays_the_keywords_over_the_base_and_equals_the_full_binding()
|
|||
}
|
||||
assert call.resolved["metadata"] is metadata
|
||||
assert supplied == {"model": "keyword-model", "metadata": metadata}
|
||||
|
||||
|
||||
def test_flag_keywords_name_positionals_and_drop_only_the_flag() -> None:
|
||||
messages: Final[Sequence[object]] = [{"role": "user", "content": "hello"}]
|
||||
|
||||
keywords: Final = _FLAG.keywords((128, messages), {"model": "m", "stream": True, "aflag": True})
|
||||
|
||||
assert keywords == {"max_tokens": 128, "messages": messages, "model": "m", "stream": True}
|
||||
|
||||
|
||||
def test_flag_keywords_reject_a_parameter_given_twice() -> None:
|
||||
with pytest.raises(TypeError):
|
||||
_FLAG.keywords((128,), {"max_tokens": 64})
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue