scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-19 00:01:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 88

refactor(rust_bridge): give chat completions, messages and responses the ocr dispatch shape

Browse source 
Each route now has litellm/rust_bridge/<route>/{entrypoints,callbacks}.py and a
public dispatch module (litellm/chat_completions/dispatch.py,
litellm/responses/dispatch.py, litellm/messages/dispatch.py) that binds the
public call to the legacy Python signature, builds a frozen request, and asks
the runtime to pick Rust or Python from the catalog. The legacy implementations
stay in litellm/main.py, litellm/responses/main.py and the anthropic messages
handler, and litellm/__init__.py re-exports the dispatch names over them the
same way it already does for ocr

The per-handler shims in rust_bridge/chat_completions/native.py and
rust_bridge/messages/native.py are removed along with their call sites in the
anthropic and bedrock chat handlers and the http handler. The exception
mapping that every callbacks module repeated moves to rust_bridge/failures.py
and the signature binding helpers to rust_bridge/public_call.py
This commit is contained in:
Yujong Lee 2026-09-16 15:02:12 -07:00
parent 62c862796a
commit a84f68b6e3
36 changed files with 1447 additions and 1421 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
litellm/__init__.py
View file

@ -1406,7 +1406,9 @@ from .videos.main import *
from .batch_completion.main import *
from .rerank_api.main import *
from .llms.anthropic.experimental_pass_through.messages.handler import *
from .messages.dispatch import *
from .responses.main import *
from .responses.dispatch import *
# Interactions API is available as litellm.interactions module
# Usage: litellm.interactions.create(), litellm.interactions.get(), etc.
@ -1435,6 +1437,7 @@ from .skills.main import (
)
from .containers.main import *
from .ocr.dispatch import *
from .chat_completions.dispatch import *
from .rust_bridge import rust
from .rag.main import *
from .sandbox.main import *

3
litellm/chat_completions/__init__.py Normal file
View file

@ -0,0 +1,3 @@
from .dispatch import acompletion, completion
__all__ = ("acompletion", "completion")

114
litellm/chat_completions/dispatch.py Normal file
View file

@ -0,0 +1,114 @@
import inspect
from collections.abc import Awaitable, Callable, Coroutine, 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 import main
from litellm.rust_bridge.catalog import Context, Delivery, Route
from litellm.rust_bridge.chat_completions.entrypoints import (
NATIVE_ACOMPLETION,
NATIVE_COMPLETION,
LiteLLMChatCompletionsRequest,
NativeAcompletion,
)
from litellm.rust_bridge.public_call import (
bind,
optional_bool,
optional_mapping,
optional_sequence,
optional_str,
signature,
)
from litellm.rust_bridge.runtime import arun, run
from litellm.types.utils import ModelResponse
from litellm.utils import CustomStreamWrapper
__all__ = ("acompletion", "completion")
ChatResult: TypeAlias = ModelResponse | CustomStreamWrapper
PythonCompletion: TypeAlias = Callable[..., ChatResult | Coroutine[object, object, ChatResult]]
PythonAcompletion: TypeAlias = Callable[..., Awaitable[ChatResult]]
def _python_completion() -> PythonCompletion:
return cast( # cast-ok: forward the original call shape through the Python @client decorator
PythonCompletion, main.completion
)
def _python_acompletion() -> PythonAcompletion:
return cast( # cast-ok: forward the original call shape through the Python @client decorator
PythonAcompletion, main.acompletion
)
_COMPLETION: Final = signature(_python_completion())
_ACOMPLETION: Final = signature(_python_acompletion())
def _public_request(
legacy: inspect.Signature, args: tuple[object, ...], kwargs: Mapping[str, object]
) -> LiteLLMChatCompletionsRequest | None:
fields: Final = bind(legacy, args, kwargs)
if fields is None:
return None
model: Final = fields.get("model")
messages: Final = optional_sequence(fields.get("messages"))
extra: Final = optional_mapping(fields.get("kwargs")) or MappingProxyType({})
if not isinstance(model, str) or messages is None:
return None
return LiteLLMChatCompletionsRequest(
model=model,
messages=messages,
stream=optional_bool(fields.get("stream")),
api_key=optional_str(fields.get("api_key")),
api_base=optional_str(extra.get("api_base")) or optional_str(fields.get("base_url")),
custom_llm_provider=optional_str(extra.get("custom_llm_provider")),
extra_headers=optional_mapping(fields.get("extra_headers")),
kwargs=extra,
)
def completion(
*args: object,
**kwargs: object, # kwargs-ok: preserve the public chat completions call shape
) -> ChatResult | Coroutine[object, object, ChatResult]:
python: Final = _python_completion()
request: Final = _public_request(_COMPLETION, args, kwargs)
if request is None or request.kwargs.get("acompletion") is True:
return python(*args, **kwargs)
return run(
_context(request),
binding=NATIVE_COMPLETION,
native=lambda hook: hook(request, args, kwargs),
python=lambda: python(*args, **kwargs),
)
async def acompletion(*args: object, **kwargs: object) -> ChatResult: # kwargs-ok: preserve the public call shape
python: Final = _python_acompletion()
request: Final = _public_request(_ACOMPLETION, args, kwargs)
if request is None:
return await python(*args, **kwargs)
async def native(hook: NativeAcompletion) -> ChatResult:
return await hook(request, args, kwargs)
return await arun(
_context(request), binding=NATIVE_ACOMPLETION, native=native, python=lambda: python(*args, **kwargs)
)
def _context(request: LiteLLMChatCompletionsRequest) -> Context:
return Context(
Route.CHAT_COMPLETIONS,
provider=request.custom_llm_provider,
model=request.model,
delivery=Delivery.STREAMING if request.stream else Delivery.COMPLETED,
)
completion.__doc__ = _python_completion().__doc__
completion.__wrapped__ = _python_completion() # pyright: ignore[reportFunctionMemberAccess] # inspect.signature follows __wrapped__ to the legacy signature
acompletion.__doc__ = _python_acompletion().__doc__
acompletion.__wrapped__ = _python_acompletion() # pyright: ignore[reportFunctionMemberAccess] # inspect.signature follows __wrapped__ to the legacy signature

86
litellm/llms/anthropic/chat/handler.py
View file

@ -25,8 +25,6 @@ from litellm.llms.custom_httpx.http_handler import (
_get_httpx_client,
get_async_httpx_client,
)
from litellm.rust_bridge.chat_completions import native as rust_chat_completions_bridge
from litellm.rust_bridge.chat_completions.native import rust_chat_completions_accepts
from litellm.types.llms.anthropic import (
ContentBlockDelta,
ContentBlockStart,
@ -375,24 +373,22 @@ class AnthropicChatCompletion(BaseLLM):
"""Filter beta headers and emit pre_call, returning `(headers, data)`.
The pair stays mutable because the streaming path rewrites it in
place (`data["stream"] = True`) before sending. A Rust attempt that
declined already emitted pre_call for this request, so skip it there.
place (`data["stream"] = True`) before sending.
"""
request_headers, data = update_request_with_filtered_beta(
headers=headers,
request_data=request_data,
provider=custom_llm_provider,
)
if not serves_via_rust:
logging_obj.pre_call(
input=messages,
api_key=api_key,
additional_args={
"complete_input_dict": data,
"api_base": api_base,
"headers": request_headers,
},
)
logging_obj.pre_call(
input=messages,
api_key=api_key,
additional_args={
"complete_input_dict": data,
"api_base": api_base,
"headers": request_headers,
},
)
print_verbose(f"_is_function_call: {_is_function_call}")
return request_headers, data
@ -456,68 +452,6 @@ class AnthropicChatCompletion(BaseLLM):
timeout=timeout,
)
# The Rust core owns the whole call for the subset it accepts, so ask
# before transforming: whichever path runs emits pre_call exactly once.
# `get_config` merges the class-level defaults (Anthropic's required
# `max_tokens` among them) that `transform_request` would have applied.
rust_optional_params: Final = { # mutable-ok: json.dumps in the bridge rejects a mappingproxy
**AnthropicConfig.get_config(model=model),
**optional_params,
}
serves_via_rust: Final = rust_chat_completions_accepts(
model=model,
messages=messages,
optional_params=rust_optional_params,
custom_llm_provider=custom_llm_provider,
litellm_params=litellm_params,
stream=stream,
)
if serves_via_rust:
rust_logging_args: Final = { # mutable-ok: logging callbacks read additional_args as a plain dict
"complete_input_dict": { # mutable-ok: same, and it is serialized alongside its parent
"model": model,
"messages": messages,
**rust_optional_params,
},
"api_base": api_base,
"headers": headers,
}
logging_obj.pre_call(input=messages, api_key=api_key, additional_args=rust_logging_args)
log_rust_post_call: Final = rust_chat_completions_bridge.response_logger(
logging_obj=logging_obj,
messages=messages,
api_key=api_key,
additional_args=rust_logging_args,
)
if acompletion is True:
return rust_chat_completions_bridge.achat_completions_or_fallback(
model=model,
messages=messages,
optional_params=rust_optional_params,
model_response=model_response,
api_key=api_key,
api_base=api_base,
custom_llm_provider=custom_llm_provider,
extra_headers=headers,
timeout=timeout,
on_response=log_rust_post_call,
python_fallback=acompletion_dispatch,
)
rust_response: Final = rust_chat_completions_bridge.chat_completions(
model=model,
messages=messages,
optional_params=rust_optional_params,
model_response=model_response,
api_key=api_key,
api_base=api_base,
custom_llm_provider=custom_llm_provider,
extra_headers=headers,
timeout=timeout,
on_response=log_rust_post_call,
)
if rust_response is not None:
return rust_response
if acompletion is True:
return acompletion_dispatch()
else:

99
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -185,9 +185,6 @@ if TYPE_CHECKING:
from litellm.llms.anthropic.experimental_pass_through.messages.fake_stream_iterator import (
FakeAnthropicMessagesStreamIterator,
)
from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import (
AnthropicMessagesStreamingResponse,
)
from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig
from litellm.types.llms.openai_evals import (
CancelEvalResponse,
@ -2285,36 +2282,6 @@ class BaseLLMHTTPHandler:
},
)
rust_messages_response: Final = await self._maybe_rust_anthropic_messages(
custom_llm_provider=custom_llm_provider,
litellm_params=litellm_params,
has_agentic_hook=self._has_agentic_completion_hook(logging_obj),
model=model,
api_key=api_key,
api_base=api_base,
headers=headers,
request_body=request_body,
timeout=self._resolve_anthropic_messages_timeout(
litellm_params=litellm_params,
stream=stream or False,
custom_llm_provider=custom_llm_provider,
),
)
if rust_messages_response is not None:
if stream:
return self._rust_anthropic_messages_fake_stream(rust_messages_response)
return await self._finalize_anthropic_messages_response(
initial_response=rust_messages_response,
model=model,
messages=messages,
anthropic_messages_provider_config=anthropic_messages_provider_config,
anthropic_messages_optional_request_params=anthropic_messages_optional_request_params,
logging_obj=logging_obj,
custom_llm_provider=custom_llm_provider,
api_key=api_key,
kwargs=kwargs,
)
response: Final = await self._async_post_anthropic_messages_with_http_error_retry(
async_httpx_client=async_httpx_client,
request_url=request_url,
@ -2443,72 +2410,6 @@ class BaseLLMHTTPHandler:
"anthropic_messages",
)
@staticmethod
async def _maybe_rust_anthropic_messages(
*,
custom_llm_provider: str,
litellm_params: GenericLiteLLMParams,
has_agentic_hook: bool,
model: str,
api_key: str | None,
api_base: str | None,
headers: dict,
request_body: dict,
timeout: float | httpx.Timeout | None,
) -> AnthropicMessagesResponse | None:
from litellm.rust_bridge.catalog import Context, Route, decision
from litellm.rust_bridge.configuration import Decision
if decision(Context(Route.MESSAGES, provider=custom_llm_provider, model=model)) is Decision.PYTHON:
return None
if has_agentic_hook:
return None
from litellm.rust_bridge.messages import native as rust_messages_bridge
upstream_body: Final = {key: value for key, value in request_body.items() if key != "stream"}
try:
rust_response: Final = await rust_messages_bridge.amessages(
model=model,
body=upstream_body,
api_key=api_key,
api_base=api_base,
custom_llm_provider=custom_llm_provider,
extra_headers=headers,
timeout=timeout,
)
except Exception as rust_error: # noqa: BLE001 # rollout-safety fallback: any Rust bridge failure must fall back to the Python path
verbose_logger.debug(
"Rust Anthropic messages bridge raised %s; falling back to Python path",
type(rust_error).__name__,
)
return None
if rust_response is None:
return None
response_obj: Final = cast(AnthropicMessagesResponse, dict(rust_response))
response_obj["_hidden_params"] = {"additional_headers": {"x-litellm-rust": "true"}}
return response_obj
@staticmethod
def _rust_anthropic_messages_fake_stream(
rust_response: AnthropicMessagesResponse,
) -> "AnthropicMessagesStreamingResponse":
from litellm.llms.anthropic.experimental_pass_through.messages.fake_stream_iterator import (
FakeAnthropicMessagesStreamIterator,
)
from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import (
AnthropicMessagesStreamHiddenParams,
AnthropicMessagesStreamingResponse,
)
completion_stream = cast(AsyncIterator[bytes], FakeAnthropicMessagesStreamIterator(response=rust_response))
hidden_params: Final = AnthropicMessagesStreamHiddenParams(additional_headers={"x-litellm-rust": "true"})
return AnthropicMessagesStreamingResponse(
completion_stream=completion_stream,
hidden_params=hidden_params,
)
def anthropic_messages_handler(
self,
model: str,

3
litellm/messages/__init__.py Normal file
View file

@ -0,0 +1,3 @@
from .dispatch import anthropic_messages, anthropic_messages_handler
__all__ = ("anthropic_messages", "anthropic_messages_handler")

113
litellm/messages/dispatch.py Normal file
View file

@ -0,0 +1,113 @@
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.llms.anthropic.experimental_pass_through.messages import handler as main
from litellm.rust_bridge.catalog import Context, Delivery, Route
from litellm.rust_bridge.messages.entrypoints import (
NATIVE_AMESSAGES,
NATIVE_MESSAGES,
LiteLLMMessagesRequest,
NativeAmessages,
)
from litellm.rust_bridge.public_call import (
bind,
optional_bool,
optional_mapping,
optional_sequence,
optional_str,
signature,
)
from litellm.rust_bridge.runtime import arun, run
from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse
__all__ = ("anthropic_messages", "anthropic_messages_handler")
MessagesResult: TypeAlias = AnthropicMessagesResponse | Iterator[bytes] | AsyncIterator[object]
PythonMessages: TypeAlias = Callable[..., MessagesResult | Coroutine[object, object, MessagesResult]]
PythonAmessages: TypeAlias = Callable[..., Awaitable[MessagesResult]]
def _python_messages() -> PythonMessages:
return cast( # cast-ok: forward the original call shape through the legacy handler
PythonMessages, main.anthropic_messages_handler
)
def _python_amessages() -> PythonAmessages:
return cast( # cast-ok: forward the original call shape through the Python @client decorator
PythonAmessages, main.anthropic_messages
)
_MESSAGES: Final = signature(_python_messages())
_AMESSAGES: Final = signature(_python_amessages())
def _public_request(
legacy: inspect.Signature, args: tuple[object, ...], kwargs: Mapping[str, object]
) -> LiteLLMMessagesRequest | None:
fields: Final = bind(legacy, args, kwargs)
if fields is None:
return None
model: Final = fields.get("model")
messages: Final = optional_sequence(fields.get("messages"))
max_tokens: Final = fields.get("max_tokens")
if not isinstance(model, str) or messages is None or not isinstance(max_tokens, int):
return None
return LiteLLMMessagesRequest(
model=model,
messages=messages,
max_tokens=max_tokens,
stream=optional_bool(fields.get("stream")),
api_key=optional_str(fields.get("api_key")),
api_base=optional_str(fields.get("api_base")),
custom_llm_provider=optional_str(fields.get("custom_llm_provider")),
kwargs=optional_mapping(fields.get("kwargs")) or MappingProxyType({}),
)
def anthropic_messages_handler(
*args: object,
**kwargs: object, # kwargs-ok: preserve the public Anthropic Messages call shape
) -> MessagesResult | Coroutine[object, object, MessagesResult]:
python: Final = _python_messages()
request: Final = _public_request(_MESSAGES, args, kwargs)
if request is None or request.kwargs.get("is_async") is True:
return python(*args, **kwargs)
return run(
_context(request),
binding=NATIVE_MESSAGES,
native=lambda hook: hook(request, args, kwargs),
python=lambda: python(*args, **kwargs),
)
async def anthropic_messages(*args: object, **kwargs: object) -> MessagesResult: # kwargs-ok: public call shape
python: Final = _python_amessages()
request: Final = _public_request(_AMESSAGES, args, kwargs)
if request is None:
return await python(*args, **kwargs)
async def native(hook: NativeAmessages) -> MessagesResult:
return await hook(request, args, kwargs)
return await arun(
_context(request), binding=NATIVE_AMESSAGES, native=native, python=lambda: python(*args, **kwargs)
)
def _context(request: LiteLLMMessagesRequest) -> Context:
return Context(
Route.MESSAGES,
provider=request.custom_llm_provider,
model=request.model,
delivery=Delivery.STREAMING if request.stream else Delivery.COMPLETED,
)
anthropic_messages_handler.__doc__ = _python_messages().__doc__
anthropic_messages_handler.__wrapped__ = _python_messages() # pyright: ignore[reportFunctionMemberAccess] # inspect.signature follows __wrapped__ to the legacy signature
anthropic_messages.__doc__ = _python_amessages().__doc__
anthropic_messages.__wrapped__ = _python_amessages() # pyright: ignore[reportFunctionMemberAccess] # inspect.signature follows __wrapped__ to the legacy signature

106
litellm/responses/dispatch.py Normal file
View file

@ -0,0 +1,106 @@
import inspect
from collections.abc import Awaitable, Callable, Coroutine, 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.responses import main
from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator
from litellm.rust_bridge.catalog import Context, Delivery, Route
from litellm.rust_bridge.public_call import bind, optional_bool, optional_mapping, optional_str, signature
from litellm.rust_bridge.responses.entrypoints import (
NATIVE_ARESPONSES,
NATIVE_RESPONSES,
LiteLLMResponsesRequest,
NativeAresponses,
)
from litellm.rust_bridge.runtime import arun, run
from litellm.types.llms.openai import ResponsesAPIResponse
__all__ = ("aresponses", "responses")
ResponsesResult: TypeAlias = ResponsesAPIResponse | BaseResponsesAPIStreamingIterator
PythonResponses: TypeAlias = Callable[..., ResponsesResult | Coroutine[object, object, ResponsesResult]]
PythonAresponses: TypeAlias = Callable[..., Awaitable[ResponsesResult]]
def _python_responses() -> PythonResponses:
return cast( # cast-ok: forward the original call shape through the Python @client decorator
PythonResponses, main.responses
)
def _python_aresponses() -> PythonAresponses:
return cast( # cast-ok: forward the original call shape through the Python @client decorator
PythonAresponses, main.aresponses
)
_RESPONSES: Final = signature(_python_responses())
_ARESPONSES: Final = signature(_python_aresponses())
def _public_request(
legacy: inspect.Signature, args: tuple[object, ...], kwargs: Mapping[str, object]
) -> LiteLLMResponsesRequest | None:
fields: Final = bind(legacy, args, kwargs)
if fields is None:
return None
model: Final = fields.get("model")
extra: Final = optional_mapping(fields.get("kwargs")) or MappingProxyType({})
if not isinstance(model, str):
return None
return LiteLLMResponsesRequest(
model=model,
input=fields.get("input"),
stream=optional_bool(fields.get("stream")),
api_key=optional_str(extra.get("api_key")),
api_base=optional_str(extra.get("api_base")) or optional_str(extra.get("base_url")),
custom_llm_provider=optional_str(fields.get("custom_llm_provider")),
extra_headers=optional_mapping(fields.get("extra_headers")),
kwargs=extra,
)
def responses(
*args: object,
**kwargs: object, # kwargs-ok: preserve the public Responses call shape
) -> ResponsesResult | Coroutine[object, object, ResponsesResult]:
python: Final = _python_responses()
request: Final = _public_request(_RESPONSES, args, kwargs)
if request is None or request.kwargs.get("aresponses") is True:
return python(*args, **kwargs)
return run(
_context(request),
binding=NATIVE_RESPONSES,
native=lambda hook: hook(request, args, kwargs),
python=lambda: python(*args, **kwargs),
)
async def aresponses(*args: object, **kwargs: object) -> ResponsesResult: # kwargs-ok: preserve the public call shape
python: Final = _python_aresponses()
request: Final = _public_request(_ARESPONSES, args, kwargs)
if request is None:
return await python(*args, **kwargs)
async def native(hook: NativeAresponses) -> ResponsesResult:
return await hook(request, args, kwargs)
return await arun(
_context(request), binding=NATIVE_ARESPONSES, native=native, python=lambda: python(*args, **kwargs)
)
def _context(request: LiteLLMResponsesRequest) -> Context:
return Context(
Route.RESPONSES,
provider=request.custom_llm_provider,
model=request.model,
delivery=Delivery.STREAMING if request.stream else Delivery.COMPLETED,
)
responses.__doc__ = _python_responses().__doc__
responses.__wrapped__ = _python_responses() # pyright: ignore[reportFunctionMemberAccess] # inspect.signature follows __wrapped__ to the legacy signature
aresponses.__doc__ = _python_aresponses().__doc__
aresponses.__wrapped__ = _python_aresponses() # pyright: ignore[reportFunctionMemberAccess] # inspect.signature follows __wrapped__ to the legacy signature

19
litellm/rust_bridge/chat_completions/callbacks.py Normal file
View file

@ -0,0 +1,19 @@
from __future__ import annotations
from collections.abc import Mapping
from litellm.rust_bridge import failures
from litellm.rust_bridge.chat_completions.entrypoints import LiteLLMChatCompletionsRequest
from litellm.types.utils import ModelResponse
def response(value: Mapping[str, object]) -> ModelResponse:
return ModelResponse(**value)
def arguments(request: LiteLLMChatCompletionsRequest) -> Mapping[str, object]:
return request.kwargs
def map_failure(error: Exception, request: LiteLLMChatCompletionsRequest, request_provider: str) -> Exception:
return failures.map_failure(error, request.model, request_provider, arguments(request))

54
litellm/rust_bridge/chat_completions/entrypoints.py Normal file
View file

@ -0,0 +1,54 @@
from __future__ import annotations
from collections.abc import Awaitable, Mapping, Sequence
from dataclasses import dataclass
from typing import Final, Protocol, cast # noqa: TID251 # validates dynamically loaded native callables
from litellm.rust_bridge.bindings import NativeBinding
from litellm.types.utils import ModelResponse
@dataclass(frozen=True, slots=True)
class LiteLLMChatCompletionsRequest:
model: str
messages: Sequence[object]
stream: bool | None
api_key: str | None
api_base: str | None
custom_llm_provider: str | None
extra_headers: Mapping[str, object] | None
kwargs: Mapping[str, object]
class NativeCompletion(Protocol):
def __call__(
self,
request: LiteLLMChatCompletionsRequest,
args: tuple[object, ...],
kwargs: Mapping[str, object],
) -> ModelResponse: ...
class NativeAcompletion(Protocol):
def __call__(
self,
request: LiteLLMChatCompletionsRequest,
args: tuple[object, ...],
kwargs: Mapping[str, object],
) -> Awaitable[ModelResponse]: ...
def _completion_binding(value: object) -> NativeCompletion | None:
if not callable(value):
return None
return cast("NativeCompletion", value) # cast-ok: callable validated at the native binding boundary
def _acompletion_binding(value: object) -> NativeAcompletion | None:
if not callable(value):
return None
return cast("NativeAcompletion", value) # cast-ok: callable validated at the native binding boundary
NATIVE_COMPLETION: Final = NativeBinding("completion", validate=_completion_binding)
NATIVE_ACOMPLETION: Final = NativeBinding("acompletion", validate=_acompletion_binding)

445
litellm/rust_bridge/chat_completions/native.py
View file

@ -1,445 +0,0 @@
"""Thin Python wrapper for the native Rust chat completions bridge.
The Rust core owns the conversation translation, the provider call, and the
response normalization for the subset of `/chat/completions` requests it
accepts. This module only marshals inputs and hands the normalized result to
LiteLLM's existing `ModelResponse` builder.
``None`` means the provider was never called, so the caller is free to serve the
request on the Python path. A failure after the call was issued raises instead:
retrying it there would bill the customer for the same work twice.
"""
from __future__ import annotations
import json
from collections.abc import Awaitable, Callable, Mapping, Sequence
from dataclasses import dataclass
from typing import TYPE_CHECKING, Final, Protocol
import httpx
from pydantic import TypeAdapter, ValidationError
from litellm._logging import verbose_logger
from litellm.exceptions import APIError
from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import (
convert_to_model_response_object,
)
from litellm.llms.bedrock.request_metadata import bedrock_request_metadata_is_owned
from litellm.rust_bridge.catalog import Context, Delivery, Route, decision
from litellm.rust_bridge.configuration import Decision
from litellm.rust_bridge.loader import get_native_bridge
from litellm.rust_bridge.timeouts import timeout_to_seconds
from litellm.types.utils import ModelResponse
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
# `litellm_params` values are `object`, so validate the one this module reads
# rather than narrowing an unparameterized `Mapping` and typing the result Any.
_LITELLM_METADATA_ADAPTER: Final = TypeAdapter(Mapping[str, object])
RUST_RESPONSE_HEADER: Final = "x-litellm-rust"
class RustChatCompletions(Protocol):
def __call__(
self,
model: str,
messages: Sequence[object],
optional_params: Mapping[str, object] | None,
api_key: str | None,
api_base: str | None,
custom_llm_provider: str | None,
extra_headers: Mapping[str, object] | None,
timeout_seconds: float | None,
) -> Mapping[str, object]:
raise NotImplementedError
class RustAchatCompletions(Protocol):
def __call__(
self,
model: str,
messages: Sequence[object],
optional_params: Mapping[str, object] | None,
api_key: str | None,
api_base: str | None,
custom_llm_provider: str | None,
extra_headers: Mapping[str, object] | None,
timeout_seconds: float | None,
) -> Awaitable[Mapping[str, object]]:
raise NotImplementedError
class RustChatCompletionsDecline(Protocol):
def __call__(
self,
model: str,
messages: Sequence[object],
optional_params: Mapping[str, object] | None,
custom_llm_provider: str | None,
) -> str | None:
raise NotImplementedError
class ResponseObserver(Protocol):
"""Invoked with the payload the core returned, on success only.
Lets the caller emit its own `post_call` on whichever path served the
request. Both entry points call it, so the synchronous and asynchronous
paths cannot drift apart the way the pre_call suppression once did.
"""
def __call__(self, rust_response: Mapping[str, object], /) -> None:
raise NotImplementedError
def response_logger(
*,
logging_obj: LiteLLMLoggingObj,
messages: Sequence[object],
api_key: str,
additional_args: Mapping[str, object],
) -> ResponseObserver:
"""A `ResponseObserver` that emits the caller's `post_call` for a Rust-served
request.
The core owns the provider call, so the Python transform that normally
raises this event never runs; without it every `post_call` callback goes
silent on a Rust-served request and `original_response` stays unset. The
payload is the core's normalized response rather than the provider's wire
body, which is the closest thing that crosses the bridge.
"""
def log(rust_response: Mapping[str, object], /) -> None:
logging_obj.post_call(
input=messages,
api_key=api_key,
original_response=json.dumps(rust_response),
additional_args=additional_args,
)
return log
class _Unset:
pass
_UNSET: Final[_Unset] = _Unset()
@dataclass(slots=True)
class _RustChatCompletionsState:
chat_completions: RustChatCompletions | None = None
achat_completions: RustAchatCompletions | None = None
decline: RustChatCompletionsDecline | None = None
_STATE: Final[_RustChatCompletionsState] = _RustChatCompletionsState()
def set_rust_chat_completions(
*,
chat_completions: RustChatCompletions | None | _Unset = _UNSET,
achat_completions: RustAchatCompletions | None | _Unset = _UNSET,
decline: RustChatCompletionsDecline | None | _Unset = _UNSET,
) -> None:
"""Inject the native callables, so tests can supply a double instead of
patching module attributes."""
if not isinstance(chat_completions, _Unset):
_STATE.chat_completions = chat_completions
if not isinstance(achat_completions, _Unset):
_STATE.achat_completions = achat_completions
if not isinstance(decline, _Unset):
_STATE.decline = decline
def load_rust_chat_completions() -> RustChatCompletions | None:
if _STATE.chat_completions is not None:
return _STATE.chat_completions
native_bridge: Final = get_native_bridge()
if native_bridge is None:
return None
loaded: RustChatCompletions | None = getattr(native_bridge, "chat_completions", None)
return loaded
def load_rust_achat_completions() -> RustAchatCompletions | None:
if _STATE.achat_completions is not None:
return _STATE.achat_completions
native_bridge: Final = get_native_bridge()
if native_bridge is None:
return None
loaded: RustAchatCompletions | None = getattr(native_bridge, "achat_completions", None)
return loaded
def _load_rust_decline() -> RustChatCompletionsDecline | None:
if _STATE.decline is not None:
return _STATE.decline
native_bridge: Final = get_native_bridge()
if native_bridge is None:
return None
loaded: RustChatCompletionsDecline | None = getattr(native_bridge, "chat_completions_decline", None)
return loaded
def _anthropic_user_id_reaches_the_body(litellm_params: Mapping[str, object] | None) -> bool:
metadata: Final = litellm_params.get("metadata") if litellm_params is not None else None
try:
entries: Final = _LITELLM_METADATA_ADAPTER.validate_python(metadata)
except ValidationError:
return False
return entries.get("user_id") is not None
def _litellm_metadata_reaches_the_provider(
custom_llm_provider: str | None, litellm_params: Mapping[str, object] | None
) -> bool:
"""Whether the Python transform would promote proxy-owned attribution into the
provider request, below this gate and inside the function the Rust route replaces.
`AnthropicConfig.transform_request` promotes a valid `metadata["user_id"]`
into the Messages body, so the core never sees the key and would send the
request to Anthropic with the abuse-detection attribution missing.
`AmazonConverseConfig` resolves proxy-owned `requestMetadata` onto the
Converse body whenever the operator armed `bedrock_request_metadata_fields`.
Owning that field also means evicting a caller-supplied one, which the core
cannot do either, so ownership alone is the condition rather than whether
anything resolved.
Deliberately a superset of Python's condition in both cases: declining a
request Python would not have attributed anyway costs only the Rust path,
while missing one loses the attribution silently.
"""
match custom_llm_provider:
case "anthropic":
return _anthropic_user_id_reaches_the_body(litellm_params)
case "bedrock":
return bedrock_request_metadata_is_owned()
case _:
return False
def rust_chat_completions_accepts(
*,
model: str,
messages: Sequence[object],
optional_params: Mapping[str, object],
custom_llm_provider: str | None,
litellm_params: Mapping[str, object] | None,
stream: object,
) -> bool:
"""Whether the Rust path will serve this request.
Asked before the caller commits to either path, so pre-call logging is
emitted exactly once, on whichever path actually runs. The core's own
capability gate answers the second half; it resolves no credentials and
performs no I/O.
"""
context: Final = Context(
Route.CHAT_COMPLETIONS,
provider=custom_llm_provider,
model=model,
delivery=Delivery.STREAMING if stream else Delivery.COMPLETED,
)
if decision(context) is Decision.PYTHON:
return False
if _litellm_metadata_reaches_the_provider(custom_llm_provider, litellm_params):
verbose_logger.debug("Rust chat completions declined (litellm metadata user_id); using the Python path")
return False
decline: Final = _load_rust_decline()
if decline is None:
return False
try:
reason: Final = decline(
model=model,
messages=messages,
optional_params=optional_params,
custom_llm_provider=custom_llm_provider,
)
except Exception as rust_error: # noqa: BLE001 # rollout-safety fallback: any Rust bridge failure must fall back to the Python path
verbose_logger.debug(
"Rust chat completions gate raised %s; staying on the Python path",
type(rust_error).__name__,
)
return False
if reason is not None:
verbose_logger.debug("Rust chat completions declined (%s); using the Python path", reason)
return False
return True
def _rust_bridge_exceptions() -> tuple[type[BaseException], type[BaseException]] | None:
"""`(declined, upstream_failed)` from the native module, or None when absent."""
native_bridge: Final = get_native_bridge()
if native_bridge is None:
return None
declined: Final = getattr(native_bridge, "RustBridgeDeclined", None)
upstream: Final = getattr(native_bridge, "RustUpstreamError", None)
if declined is None or upstream is None:
return None
return declined, upstream
def _reraise_or_decline(
rust_error: BaseException,
*,
model: str,
custom_llm_provider: str | None,
) -> None:
"""Re-raise a failure the provider already saw, or return so the caller declines.
A request that never reached the provider is safe to serve on the Python
path. One that did is not: the provider has already done the work, so a
second attempt bills for it twice. Those surface as an `APIError` carrying
the upstream status, which LiteLLM's exception mapping already understands.
"""
exceptions: Final = _rust_bridge_exceptions()
if exceptions is None:
verbose_logger.debug(
"Rust chat completions bridge raised %s; falling back to Python path",
type(rust_error).__name__,
)
return
declined, upstream_failed = exceptions
if isinstance(rust_error, upstream_failed):
args: Final = rust_error.args
status: Final = args[0] if args else 0
message: Final = args[1] if len(args) > 1 else ""
raise APIError(
status_code=int(status) or 500,
message=f"litellm rust chat completions: {message}",
llm_provider=custom_llm_provider or "",
model=model,
)
if not isinstance(rust_error, declined):
raise rust_error
verbose_logger.debug(
"Rust chat completions declined before calling the provider (%s); using the Python path",
rust_error,
)
def _build_model_response(
rust_response: Mapping[str, object],
model_response: ModelResponse,
) -> ModelResponse:
built: Final = convert_to_model_response_object(
response_object=dict(rust_response), # mutable-ok: the converter takes a real dict and rewrites it
model_response_object=model_response,
hidden_params={"additional_headers": {RUST_RESPONSE_HEADER: "true"}}, # mutable-ok: rewritten by the converter
)
if not isinstance(built, ModelResponse):
raise TypeError(f"expected a ModelResponse from the rust path, got {type(built).__name__}")
return built
def chat_completions(
*,
model: str,
messages: Sequence[object],
optional_params: Mapping[str, object],
model_response: ModelResponse,
api_key: str | None,
api_base: str | None,
custom_llm_provider: str | None,
extra_headers: Mapping[str, object] | None,
timeout: float | httpx.Timeout | None,
on_response: ResponseObserver,
) -> ModelResponse | None:
rust_chat_completions: Final = load_rust_chat_completions()
if rust_chat_completions is None:
return None
try:
rust_response: Final = rust_chat_completions(
model=model,
messages=messages,
optional_params=optional_params,
api_key=api_key,
api_base=api_base,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
timeout_seconds=timeout_to_seconds(timeout),
)
except Exception as rust_error: # noqa: BLE001 # rollout safety: the helper re-raises anything the provider already saw
_reraise_or_decline(rust_error, model=model, custom_llm_provider=custom_llm_provider)
return None
on_response(rust_response)
return _build_model_response(rust_response, model_response)
async def achat_completions(
*,
model: str,
messages: Sequence[object],
optional_params: Mapping[str, object],
model_response: ModelResponse,
api_key: str | None,
api_base: str | None,
custom_llm_provider: str | None,
extra_headers: Mapping[str, object] | None,
timeout: float | httpx.Timeout | None,
on_response: ResponseObserver,
) -> ModelResponse | None:
rust_achat_completions: Final = load_rust_achat_completions()
if rust_achat_completions is None:
return None
try:
rust_response: Final = await rust_achat_completions(
model=model,
messages=messages,
optional_params=optional_params,
api_key=api_key,
api_base=api_base,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
timeout_seconds=timeout_to_seconds(timeout),
)
except Exception as rust_error: # noqa: BLE001 # rollout safety: the helper re-raises anything the provider already saw
_reraise_or_decline(rust_error, model=model, custom_llm_provider=custom_llm_provider)
return None
on_response(rust_response)
return _build_model_response(rust_response, model_response)
async def achat_completions_or_fallback(
*,
model: str,
messages: Sequence[object],
optional_params: Mapping[str, object],
model_response: ModelResponse,
api_key: str | None,
api_base: str | None,
custom_llm_provider: str | None,
extra_headers: Mapping[str, object] | None,
timeout: float | httpx.Timeout | None,
on_response: ResponseObserver,
python_fallback: Callable[[], Awaitable[object]],
) -> object:
"""Await the Rust path, falling back to the caller's own Python path when
the bridge is unavailable or the call fails.
The caller supplies the fallback, so the bridge stays free of provider
dispatch. This exists because a caller that dispatches asynchronously has
already returned a coroutine by the time a Rust failure surfaces, and so
cannot fall back on its own.
"""
response: Final = await achat_completions(
model=model,
messages=messages,
optional_params=optional_params,
model_response=model_response,
api_key=api_key,
api_base=api_base,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
timeout=timeout,
on_response=on_response,
)
if response is not None:
return response
return await python_fallback()

37
litellm/rust_bridge/failures.py Normal file
View file

@ -0,0 +1,37 @@
"""Map a native failure onto LiteLLM's public exception contract."""
from __future__ import annotations
from collections.abc import Mapping
from typing import Final, Protocol, cast # noqa: TID251 # adapts the public exception mapper
import litellm
class ExceptionMapper(Protocol):
def __call__(
self,
*,
model: str,
custom_llm_provider: str | None,
original_exception: Exception,
completion_kwargs: dict[str, object], # mutable-ok: the legacy public exception mapper mutates its kwargs
extra_kwargs: dict[str, object], # mutable-ok: the legacy public exception mapper mutates its kwargs
) -> Exception: ...
def map_failure(error: Exception, model: str, request_provider: str, kwargs: Mapping[str, object]) -> Exception:
mapper: Final = cast( # cast-ok: bounded adapter for the legacy public exception mapper
ExceptionMapper, litellm.exception_type
)
try:
return mapper(
model=model.removeprefix(f"{request_provider}/"),
custom_llm_provider=request_provider,
original_exception=error,
completion_kwargs=dict(kwargs), # mutable-ok: exception mapper requires owned kwargs
extra_kwargs=dict(kwargs), # mutable-ok: exception mapper requires owned kwargs
)
except Exception as public_error:
public_error.__context__ = error
return public_error

23
litellm/rust_bridge/messages/callbacks.py Normal file
View file

@ -0,0 +1,23 @@
from __future__ import annotations
from collections.abc import Mapping
from typing import cast # noqa: TID251 # narrows the normalized native payload to the public TypedDict
from litellm.rust_bridge import failures
from litellm.rust_bridge.messages.entrypoints import LiteLLMMessagesRequest
from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse
def response(value: Mapping[str, object]) -> AnthropicMessagesResponse:
return cast( # cast-ok: AnthropicMessagesResponse is a TypedDict over the normalized native payload
AnthropicMessagesResponse,
dict(value), # mutable-ok: the public Messages response is a TypedDict the caller may annotate in place
)
def arguments(request: LiteLLMMessagesRequest) -> Mapping[str, object]:
return request.kwargs
def map_failure(error: Exception, request: LiteLLMMessagesRequest, request_provider: str) -> Exception:
return failures.map_failure(error, request.model, request_provider, arguments(request))

54
litellm/rust_bridge/messages/entrypoints.py Normal file
View file

@ -0,0 +1,54 @@
from __future__ import annotations
from collections.abc import Awaitable, Mapping, Sequence
from dataclasses import dataclass
from typing import Final, Protocol, cast # noqa: TID251 # validates dynamically loaded native callables
from litellm.rust_bridge.bindings import NativeBinding
from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse
@dataclass(frozen=True, slots=True)
class LiteLLMMessagesRequest:
model: str
messages: Sequence[object]
max_tokens: int
stream: bool | None
api_key: str | None
api_base: str | None
custom_llm_provider: str | None
kwargs: Mapping[str, object]
class NativeMessages(Protocol):
def __call__(
self,
request: LiteLLMMessagesRequest,
args: tuple[object, ...],
kwargs: Mapping[str, object],
) -> AnthropicMessagesResponse: ...
class NativeAmessages(Protocol):
def __call__(
self,
request: LiteLLMMessagesRequest,
args: tuple[object, ...],
kwargs: Mapping[str, object],
) -> Awaitable[AnthropicMessagesResponse]: ...
def _messages_binding(value: object) -> NativeMessages | None:
if not callable(value):
return None
return cast("NativeMessages", value) # cast-ok: callable validated at the native binding boundary
def _amessages_binding(value: object) -> NativeAmessages | None:
if not callable(value):
return None
return cast("NativeAmessages", value) # cast-ok: callable validated at the native binding boundary
NATIVE_MESSAGES: Final = NativeBinding("anthropic_messages_handler", validate=_messages_binding)
NATIVE_AMESSAGES: Final = NativeBinding("anthropic_messages", validate=_amessages_binding)

136
litellm/rust_bridge/messages/native.py
View file

@ -1,136 +0,0 @@
"""Thin Python wrapper for the native Rust Anthropic Messages bridge."""
from __future__ import annotations
from collections.abc import Awaitable
from dataclasses import dataclass
from typing import Final, Protocol, cast
import httpx
from litellm.rust_bridge.timeouts import timeout_to_seconds
class RustMessages(Protocol):
def __call__(
self,
model: str,
body: dict[str, object],
api_key: str | None,
api_base: str | None,
custom_llm_provider: str | None,
extra_headers: dict[str, object] | None,
timeout_seconds: float | None,
) -> dict[str, object]:
raise NotImplementedError
class RustAmessages(Protocol):
def __call__(
self,
model: str,
body: dict[str, object],
api_key: str | None,
api_base: str | None,
custom_llm_provider: str | None,
extra_headers: dict[str, object] | None,
timeout_seconds: float | None,
) -> Awaitable[dict[str, object]]:
raise NotImplementedError
class _Unset:
pass
_UNSET: Final[_Unset] = _Unset()
@dataclass(slots=True)
class _RustMessagesState:
messages: RustMessages | None = None
amessages: RustAmessages | None = None
_STATE: Final[_RustMessagesState] = _RustMessagesState()
def set_rust_messages(
*,
messages: RustMessages | None | _Unset = _UNSET,
amessages: RustAmessages | None | _Unset = _UNSET,
) -> None:
if not isinstance(messages, _Unset):
_STATE.messages = messages
if not isinstance(amessages, _Unset):
_STATE.amessages = amessages
def load_rust_messages() -> RustMessages | None:
if _STATE.messages is not None:
return _STATE.messages
from litellm.rust_bridge import get_native_bridge
native_bridge: Final = get_native_bridge()
if native_bridge is None:
return None
return cast(RustMessages, getattr(native_bridge, "messages", None))
def load_rust_amessages() -> RustAmessages | None:
if _STATE.amessages is not None:
return _STATE.amessages
from litellm.rust_bridge import get_native_bridge
native_bridge: Final = get_native_bridge()
if native_bridge is None:
return None
return cast(RustAmessages, getattr(native_bridge, "amessages", None))
def messages(
*,
model: str,
body: dict[str, object],
api_key: str | None,
api_base: str | None,
custom_llm_provider: str | None,
extra_headers: dict[str, object] | None,
timeout: float | httpx.Timeout | None,
) -> dict[str, object] | None:
rust_messages: Final = load_rust_messages()
if rust_messages is None:
return None
return rust_messages(
model=model,
body=body,
api_key=api_key,
api_base=api_base,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
timeout_seconds=timeout_to_seconds(timeout),
)
async def amessages(
*,
model: str,
body: dict[str, object],
api_key: str | None,
api_base: str | None,
custom_llm_provider: str | None,
extra_headers: dict[str, object] | None,
timeout: float | httpx.Timeout | None,
) -> dict[str, object] | None:
rust_amessages: Final = load_rust_amessages()
if rust_amessages is None:
return None
return await rust_amessages(
model=model,
body=body,
api_key=api_key,
api_base=api_base,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
timeout_seconds=timeout_to_seconds(timeout),
)

31
litellm/rust_bridge/ocr/callbacks.py
View file

@ -2,29 +2,17 @@ from __future__ import annotations
from collections.abc import Mapping
from types import MappingProxyType
from typing import Final, Protocol, cast # noqa: TID251 # adapts the public exception mapper
from typing import Final
from pydantic import TypeAdapter
import litellm
from litellm.llms.base_llm.ocr.transformation import PROVIDER_NATIVE_RESPONSE_KEY, OCRResponse
from litellm.rust_bridge import failures
from litellm.rust_bridge.ocr.entrypoints import LiteLLMOcrRequest
_RESPONSE_ADAPTER: Final = TypeAdapter(dict[str, object])
class ExceptionMapper(Protocol):
def __call__(
self,
*,
model: str,
custom_llm_provider: str | None,
original_exception: Exception,
completion_kwargs: dict[str, object],
extra_kwargs: dict[str, object],
) -> Exception: ...
def response(value: Mapping[str, object]) -> OCRResponse:
provider_native_response: Final = value.get(PROVIDER_NATIVE_RESPONSE_KEY)
normalized: Final = OCRResponse.model_validate(
@ -40,17 +28,4 @@ def arguments(request: LiteLLMOcrRequest) -> Mapping[str, object]:
def map_failure(error: Exception, request: LiteLLMOcrRequest, request_provider: str) -> Exception:
mapper: Final = cast( # cast-ok: bounded adapter for the legacy public exception mapper
ExceptionMapper, litellm.exception_type
)
try:
return mapper(
model=request.model.removeprefix(f"{request_provider}/"),
custom_llm_provider=request_provider,
original_exception=error,
completion_kwargs=dict(arguments(request)), # mutable-ok: exception mapper requires owned kwargs
extra_kwargs=dict(request.kwargs), # mutable-ok: exception mapper requires owned kwargs
)
except Exception as public_error:
public_error.__context__ = error
return public_error
return failures.map_failure(error, request.model, request_provider, arguments(request))

42
litellm/rust_bridge/public_call.py Normal file
View file

@ -0,0 +1,42 @@
"""Bind a public LiteLLM call to its legacy Python signature without running it."""
from __future__ import annotations
import inspect
from collections.abc import Callable, Mapping, Sequence
from typing import Final, cast # noqa: TID251 # narrows caller-owned containers without copying them
def signature(legacy: Callable[..., object]) -> inspect.Signature:
return inspect.signature(legacy)
def bind(
legacy: inspect.Signature, args: tuple[object, ...], kwargs: Mapping[str, object]
) -> Mapping[str, object] | None:
try:
bound: Final = legacy.bind(*args, **kwargs)
except TypeError:
return None
bound.apply_defaults()
return bound.arguments
def optional_str(value: object) -> str | None:
return value if isinstance(value, str) else None
def optional_bool(value: object) -> bool | None:
return value if isinstance(value, bool) else None
def optional_mapping(value: object) -> Mapping[str, object] | None:
if not isinstance(value, Mapping):
return None
return cast("Mapping[str, object]", value) # cast-ok: the same caller-owned object is handed on unchanged
def optional_sequence(value: object) -> Sequence[object] | None:
if isinstance(value, str | bytes) or not isinstance(value, Sequence):
return None
return cast("Sequence[object]", value) # cast-ok: the same caller-owned object is handed on unchanged

19
litellm/rust_bridge/responses/callbacks.py Normal file
View file

@ -0,0 +1,19 @@
from __future__ import annotations
from collections.abc import Mapping
from litellm.rust_bridge import failures
from litellm.rust_bridge.responses.entrypoints import LiteLLMResponsesRequest
from litellm.types.llms.openai import ResponsesAPIResponse
def response(value: Mapping[str, object]) -> ResponsesAPIResponse:
return ResponsesAPIResponse.model_validate(value)
def arguments(request: LiteLLMResponsesRequest) -> Mapping[str, object]:
return request.kwargs
def map_failure(error: Exception, request: LiteLLMResponsesRequest, request_provider: str) -> Exception:
return failures.map_failure(error, request.model, request_provider, arguments(request))

54
litellm/rust_bridge/responses/entrypoints.py Normal file
View file

@ -0,0 +1,54 @@
from __future__ import annotations
from collections.abc import Awaitable, Mapping
from dataclasses import dataclass
from typing import Final, Protocol, cast # noqa: TID251 # validates dynamically loaded native callables
from litellm.rust_bridge.bindings import NativeBinding
from litellm.types.llms.openai import ResponsesAPIResponse
@dataclass(frozen=True, slots=True)
class LiteLLMResponsesRequest:
model: str
input: object
stream: bool | None
api_key: str | None
api_base: str | None
custom_llm_provider: str | None
extra_headers: Mapping[str, object] | None
kwargs: Mapping[str, object]
class NativeResponses(Protocol):
def __call__(
self,
request: LiteLLMResponsesRequest,
args: tuple[object, ...],
kwargs: Mapping[str, object],
) -> ResponsesAPIResponse: ...
class NativeAresponses(Protocol):
def __call__(
self,
request: LiteLLMResponsesRequest,
args: tuple[object, ...],
kwargs: Mapping[str, object],
) -> Awaitable[ResponsesAPIResponse]: ...
def _responses_binding(value: object) -> NativeResponses | None:
if not callable(value):
return None
return cast("NativeResponses", value) # cast-ok: callable validated at the native binding boundary
def _aresponses_binding(value: object) -> NativeAresponses | None:
if not callable(value):
return None
return cast("NativeAresponses", value) # cast-ok: callable validated at the native binding boundary
NATIVE_RESPONSES: Final = NativeBinding("responses", validate=_responses_binding)
NATIVE_ARESPONSES: Final = NativeBinding("aresponses", validate=_aresponses_binding)

2
tests/e2e/e2e_config.py
View file

@ -101,8 +101,6 @@ SLOW_PROVIDER_TIMEOUT_SECONDS = float(os.environ.get("E2E_SLOW_PROVIDER_TIMEOUT"
# fresh connection and the next call re-rolls. See ProxyClient._await_model_servable.
PROPAGATION_TIMEOUT = float(os.environ.get("E2E_PROPAGATION_TIMEOUT", "15"))
EXPECT_RUST = os.environ.get("E2E_EXPECT_RUST", "").strip().lower() in ("1", "true", "yes")
# Record/replay fixture selection (see fixture_mode.py and provider_edge.py).
# The raw mode value is parsed and validated there; "live" (the default, also
# for empty values) means the harness behaves exactly as before this knob

10
tests/e2e/llm_translation/test_messages_azure_foundry_e2e.py
View file

@ -11,8 +11,7 @@ sent in the request.
from __future__ import annotations
import pytest
from e2e_config import EXPECT_RUST, unique_marker
from e2e_config import unique_marker
from e2e_http import StreamingResponse, require_successful_call, unwrap
from endpoints_client import EndpointsClient
from lifecycle import ResourceManager
@ -50,13 +49,6 @@ def _assert_streamed_ok(result: StreamingResponse) -> None:
assert any("message_stop" in event for event in result.stream_events), (
"stream never reached message_stop"
)
if EXPECT_RUST:
assert result.headers.get("x-litellm-rust") == "true", (
"E2E_EXPECT_RUST is set, so this gateway must serve /v1/messages through the "
"Rust path, but the response carried no x-litellm-rust marker. The request "
"still succeeded, which is exactly the failure mode: a gateway whose native "
f"extension is unavailable falls back to Python silently. headers={result.headers}"
)
class TestAzureFoundryMessages:

237
tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py
View file

@ -1,237 +0,0 @@
"""Tests for the optional Rust-backed Anthropic Messages path."""
import importlib
from typing import cast
import httpx
import pytest
import litellm
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
from litellm.rust_bridge import configuration
from litellm.types.llms.anthropic_messages.anthropic_response import (
AnthropicMessagesResponse,
)
from litellm.types.router import GenericLiteLLMParams
rust_messages = importlib.import_module("litellm.rust_bridge.messages.native")
rust_bridge_loader = importlib.import_module("litellm.rust_bridge.loader")
FAKE_MESSAGES_RESPONSE: dict[str, object] = {
"id": "msg_123",
"type": "message",
"role": "assistant",
"model": "claude-sonnet-4-5-20250929",
"content": [{"type": "text", "text": "hello world"}],
"stop_reason": "end_turn",
"usage": {"input_tokens": 5, "output_tokens": 3},
}
REQUEST_BODY: dict[str, object] = {
"model": "claude-sonnet-4-5",
"max_tokens": 64,
"messages": [{"role": "user", "content": "hi"}],
}
class RecordingMessages:
def __init__(self) -> None:
self.calls: list[dict[str, object]] = []
def __call__(
self,
model: str,
body: dict[str, object],
api_key: str | None,
api_base: str | None,
custom_llm_provider: str | None,
extra_headers: dict[str, object] | None,
timeout_seconds: float | None,
) -> dict[str, object]:
self.calls.append(
{
"model": model,
"body": body,
"api_key": api_key,
"api_base": api_base,
"custom_llm_provider": custom_llm_provider,
"extra_headers": extra_headers,
"timeout_seconds": timeout_seconds,
}
)
return dict(FAKE_MESSAGES_RESPONSE)
class RecordingAsyncMessages:
def __init__(self) -> None:
self.calls: list[dict[str, object]] = []
async def __call__(
self,
model: str,
body: dict[str, object],
api_key: str | None,
api_base: str | None,
custom_llm_provider: str | None,
extra_headers: dict[str, object] | None,
timeout_seconds: float | None,
) -> dict[str, object]:
self.calls.append(
{
"model": model,
"body": body,
"api_key": api_key,
"api_base": api_base,
"custom_llm_provider": custom_llm_provider,
"extra_headers": extra_headers,
"timeout_seconds": timeout_seconds,
}
)
return dict(FAKE_MESSAGES_RESPONSE)
class ExplodingAsyncMessages:
def __init__(self) -> None:
self.calls = 0
async def __call__(self, **kwargs: object) -> dict[str, object]:
self.calls += 1
raise AssertionError("bridge must not be called")
@pytest.fixture(autouse=True)
def _reset_rust_flag():
rust_messages.set_rust_messages(messages=None, amessages=None)
configuration.reset_rust_configuration()
rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL
yield
rust_messages.set_rust_messages(messages=None, amessages=None)
configuration.reset_rust_configuration()
rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL
def test_load_rust_messages_returns_injected_impl():
bridge = RecordingMessages()
litellm.rust(True)
rust_messages.set_rust_messages(messages=bridge)
assert rust_messages.load_rust_messages() is bridge
def test_load_rust_amessages_returns_injected_impl():
bridge = RecordingAsyncMessages()
litellm.rust(True)
rust_messages.set_rust_messages(amessages=bridge)
assert rust_messages.load_rust_amessages() is bridge
def test_messages_wrapper_returns_none_when_bridge_absent(monkeypatch):
monkeypatch.setattr(
importlib.import_module("litellm.rust_bridge"),
"get_native_bridge",
lambda: None,
)
litellm.rust(True)
assert rust_messages.load_rust_messages() is None
result = rust_messages.messages(
model="claude",
body=REQUEST_BODY,
api_key="k",
api_base="b",
custom_llm_provider="azure_ai",
extra_headers={},
timeout=30.0,
)
assert result is None
def test_messages_wrapper_forwards_args_and_converts_timeout():
bridge = RecordingMessages()
litellm.rust(True)
rust_messages.set_rust_messages(messages=bridge)
response = rust_messages.messages(
model="claude-sonnet-4-5",
body=REQUEST_BODY,
api_key="sk-azure",
api_base="https://resource.services.ai.azure.com/anthropic",
custom_llm_provider="azure_ai",
extra_headers={"anthropic-beta": "token-efficient-tools-2025-02-19"},
timeout=httpx.Timeout(600.0, read=42.0),
)
assert response == FAKE_MESSAGES_RESPONSE
assert bridge.calls[0] == {
"model": "claude-sonnet-4-5",
"body": REQUEST_BODY,
"api_key": "sk-azure",
"api_base": "https://resource.services.ai.azure.com/anthropic",
"custom_llm_provider": "azure_ai",
"extra_headers": {"anthropic-beta": "token-efficient-tools-2025-02-19"},
"timeout_seconds": 42.0,
}
@pytest.mark.asyncio
async def test_amessages_wrapper_forwards_args():
bridge = RecordingAsyncMessages()
litellm.rust(True)
rust_messages.set_rust_messages(amessages=bridge)
response = await rust_messages.amessages(
model="claude-sonnet-4-5",
body=REQUEST_BODY,
api_key="sk-azure",
api_base="https://resource.services.ai.azure.com/anthropic",
custom_llm_provider="azure_ai",
extra_headers=None,
timeout=12.5,
)
assert response == FAKE_MESSAGES_RESPONSE
assert bridge.calls[0]["model"] == "claude-sonnet-4-5"
assert bridge.calls[0]["timeout_seconds"] == 12.5
def _gate(**overrides):
kwargs = {
"custom_llm_provider": "azure_ai",
"litellm_params": GenericLiteLLMParams(api_key="sk-azure"),
"has_agentic_hook": False,
"model": "claude-sonnet-4-5",
"api_key": "sk-azure",
"api_base": "https://resource.services.ai.azure.com/anthropic",
"headers": {"x-api-key": "sk-azure", "anthropic-version": "2023-06-01"},
"request_body": dict(REQUEST_BODY),
"timeout": 30.0,
}
kwargs.update(overrides)
return BaseLLMHTTPHandler._maybe_rust_anthropic_messages(**kwargs)
@pytest.mark.asyncio
@pytest.mark.parametrize("custom_llm_provider", ("azure_ai", "anthropic", "openai"))
async def test_gate_stays_on_python_with_the_switch_on(custom_llm_provider):
bridge = ExplodingAsyncMessages()
litellm.rust(True)
rust_messages.set_rust_messages(amessages=bridge)
response = await _gate(custom_llm_provider=custom_llm_provider)
assert response is None
assert bridge.calls == 0
@pytest.mark.asyncio
async def test_fake_stream_wraps_rust_response_as_anthropic_sse():
response = cast(AnthropicMessagesResponse, dict(FAKE_MESSAGES_RESPONSE))
stream = BaseLLMHTTPHandler._rust_anthropic_messages_fake_stream(response)
assert stream._hidden_params["additional_headers"] == {"x-litellm-rust": "true"}
chunks = [chunk async for chunk in stream]
joined = b"".join(chunks)
assert b"event: message_start" in joined
assert b"event: content_block_delta" in joined
assert b"hello world" in joined
assert b"event: message_stop" in joined

0
tests/test_litellm/chat_completions/__init__.py Normal file
View file

186
tests/test_litellm/chat_completions/test_dispatch.py Normal file
View file

@ -0,0 +1,186 @@
import inspect
from collections.abc import Generator, Mapping
from typing import Final
from unittest.mock import AsyncMock, Mock
import pytest
import litellm
from litellm import main as python_chat
from litellm.rust_bridge import configuration, runtime
from litellm.rust_bridge.catalog import Route, Rule, decision
from litellm.rust_bridge.chat_completions.entrypoints import (
NATIVE_ACOMPLETION,
NATIVE_COMPLETION,
LiteLLMChatCompletionsRequest,
)
from litellm.rust_bridge.configuration import Rollout
from litellm.types.utils import ModelResponse
MESSAGES: Final = [{"role": "user", "content": "hi"}]
RUST_RULES: Final = (Rule(Route.CHAT_COMPLETIONS, Rollout.RUST_OPT_OUT),)
@pytest.fixture(autouse=True)
def isolated_configuration(monkeypatch: pytest.MonkeyPatch) -> Generator[None]:
monkeypatch.delenv("LITELLM_RUST", raising=False)
configuration.reset_rust_configuration()
yield
NATIVE_COMPLETION.reset()
NATIVE_ACOMPLETION.reset()
configuration.reset_rust_configuration()
@pytest.fixture
def rust_route(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(runtime, "decision", lambda context, rules=RUST_RULES: decision(context, RUST_RULES))
def test_public_signature_is_the_legacy_signature() -> None:
assert inspect.signature(litellm.completion) == inspect.signature(python_chat.completion)
assert inspect.signature(litellm.acompletion) == inspect.signature(python_chat.acompletion)
@pytest.mark.asyncio
@pytest.mark.parametrize("asynchronous", [False, True])
async def test_python_only_route_never_loads_native(monkeypatch: pytest.MonkeyPatch, asynchronous: bool) -> None:
response: Final = ModelResponse()
fallback: Final = AsyncMock(return_value=response) if asynchronous else Mock(return_value=response)
monkeypatch.setattr(python_chat, "acompletion" if asynchronous else "completion", fallback)
monkeypatch.setattr(
NATIVE_ACOMPLETION if asynchronous else NATIVE_COMPLETION,
"load",
Mock(side_effect=AssertionError("native must not be loaded")),
)
litellm.rust(True)
result: Final = (
await litellm.acompletion("gpt-4o", MESSAGES, temperature=0.1)
if asynchronous
else litellm.completion("gpt-4o", MESSAGES, temperature=0.1)
)
assert result is response
fallback.assert_called_once_with("gpt-4o", MESSAGES, temperature=0.1)
@pytest.mark.asyncio
@pytest.mark.parametrize("asynchronous", [False, True])
async def test_unavailable_native_uses_python(
monkeypatch: pytest.MonkeyPatch, rust_route: None, asynchronous: bool
) -> None:
response: Final = ModelResponse()
fallback: Final = AsyncMock(return_value=response) if asynchronous else Mock(return_value=response)
monkeypatch.setattr(python_chat, "acompletion" if asynchronous else "completion", fallback)
(NATIVE_ACOMPLETION if asynchronous else NATIVE_COMPLETION).override(None)
result: Final = (
await litellm.acompletion("gpt-4o", MESSAGES, temperature=0.1)
if asynchronous
else litellm.completion("gpt-4o", MESSAGES, temperature=0.1)
)
assert result is response
fallback.assert_called_once_with("gpt-4o", MESSAGES, temperature=0.1)
def test_native_receives_the_bound_request_and_original_call_shape(rust_route: None) -> None:
captured: Final[list[tuple[LiteLLMChatCompletionsRequest, tuple[object, ...], Mapping[str, object]]]] = []
def native(
request: LiteLLMChatCompletionsRequest,
args: tuple[object, ...],
kwargs: Mapping[str, object],
) -> ModelResponse:
captured.append((request, args, kwargs))
return ModelResponse(model=request.model)
NATIVE_COMPLETION.override(native)
response: Final = litellm.completion(
"anthropic/claude-sonnet-4-5",
MESSAGES,
stream=True,
api_key="sk-test",
base_url="https://example.invalid",
extra_headers={"x-test": "1"},
custom_llm_provider="anthropic",
metadata={"user_id": "u"},
)
request, call_args, hook_kwargs = captured[0]
assert isinstance(response, ModelResponse)
assert response.model == "anthropic/claude-sonnet-4-5"
assert request.model == "anthropic/claude-sonnet-4-5"
assert request.messages is MESSAGES
assert request.stream is True
assert request.api_key == "sk-test"
assert request.api_base == "https://example.invalid"
assert request.custom_llm_provider == "anthropic"
assert request.extra_headers == {"x-test": "1"}
assert request.kwargs == {"custom_llm_provider": "anthropic", "metadata": {"user_id": "u"}}
assert call_args == ("anthropic/claude-sonnet-4-5", MESSAGES)
assert hook_kwargs["metadata"] == {"user_id": "u"}
assert "temperature" not in hook_kwargs
def test_internal_async_dispatch_marker_stays_on_python(monkeypatch: pytest.MonkeyPatch, rust_route: None) -> None:
native: Final = Mock(side_effect=AssertionError("acompletion's inner completion() call must stay on Python"))
NATIVE_COMPLETION.override(native)
response: Final = ModelResponse()
fallback: Final = Mock(return_value=response)
monkeypatch.setattr(python_chat, "completion", fallback)
assert litellm.completion("gpt-4o", MESSAGES, acompletion=True) is response
native.assert_not_called()
@pytest.mark.parametrize("enabled", [False, True], ids=["flag-disabled", "flag-enabled"])
def test_public_binding_errors_do_not_depend_on_native_selection(rust_route: None, enabled: bool) -> None:
native: Final = Mock(side_effect=AssertionError("binding errors precede admission"))
litellm.rust(enabled)
NATIVE_COMPLETION.override(native)
with pytest.raises(TypeError, match=r"completion\(\) got multiple values for argument 'model'"):
litellm.completion("gpt-4o", MESSAGES, model="duplicate")
with pytest.raises(TypeError, match=r"completion\(\) missing 1 required positional argument: 'model'"):
litellm.completion()
native.assert_not_called()
class Declined(Exception):
pass
class Upstream(Exception):
pass
@pytest.mark.asyncio
@pytest.mark.parametrize("asynchronous", [False, True])
@pytest.mark.parametrize("declined", [False, True])
async def test_only_native_declines_replay_on_python(
monkeypatch: pytest.MonkeyPatch, rust_route: None, asynchronous: bool, declined: bool
) -> None:
failure: Final = Declined("unsupported") if declined else RuntimeError("provider already called")
native: Final = AsyncMock(side_effect=failure) if asynchronous else Mock(side_effect=failure)
(NATIVE_ACOMPLETION if asynchronous else NATIVE_COMPLETION).override(native)
monkeypatch.setattr(runtime, "native_exception_types", lambda: (Declined, Upstream))
response: Final = ModelResponse()
fallback: Final = AsyncMock(return_value=response) if asynchronous else Mock(return_value=response)
monkeypatch.setattr(python_chat, "acompletion" if asynchronous else "completion", fallback)
async def call() -> object:
if asynchronous:
return await litellm.acompletion("gpt-4o", MESSAGES)
return litellm.completion("gpt-4o", MESSAGES)
if declined:
assert await call() is response
fallback.assert_called_once_with("gpt-4o", MESSAGES)
else:
with pytest.raises(RuntimeError) as caught:
await call()
assert caught.value is failure
fallback.assert_not_called()
assert native.call_count == 1

56
tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py
View file

@ -8,9 +8,9 @@ import httpx
import pytest
import litellm
from litellm._uuid import uuid
from litellm.constants import RESPONSE_FORMAT_TOOL_NAME
from litellm.llms.anthropic.chat.handler import ModelResponseIterator, make_call
from litellm._uuid import uuid
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.types.llms.openai import (
ChatCompletionToolCallChunk,
@ -2333,22 +2333,7 @@ def test_non_bash_tool_result_skipped():
), f"Expected 0 code_interpreter_results for text_editor result, got {len(code_results)}"
class TestRustChatCompletionsHook:
"""The catalog keeps Anthropic chat completions on the Python path, so the
injected native callables are never consulted even with the switch on."""
@pytest.fixture(autouse=True)
def _reset_bridge(self, monkeypatch):
from litellm.rust_bridge.chat_completions import native as bridge
from litellm.rust_bridge import configuration
monkeypatch.setenv("LITELLM_RUST", "1")
configuration.reset_rust_configuration()
bridge.set_rust_chat_completions(chat_completions=None, achat_completions=None, decline=None)
yield
bridge.set_rust_chat_completions(chat_completions=None, achat_completions=None, decline=None)
configuration.reset_rust_configuration()
class TestAnthropicChatCompletionPreCallLogging:
@staticmethod
def _completion_kwargs(**overrides):
from litellm.types.utils import ModelResponse
@ -2374,45 +2359,10 @@ class TestRustChatCompletionsHook:
kwargs.update(overrides)
return kwargs
@staticmethod
def _inject():
from litellm.rust_bridge.chat_completions import native as bridge
seen = {"gate": [], "call": []}
def gate(**kwargs):
seen["gate"].append(kwargs)
def native(**kwargs):
seen["call"].append(kwargs)
raise AssertionError("the native call must not run for a python-only route")
bridge.set_rust_chat_completions(decline=gate, chat_completions=native)
return seen
def test_the_python_only_route_never_consults_the_core(self):
from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
seen = self._inject()
with patch.object(
AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []}
) as transform:
try:
AnthropicChatCompletion().completion(**self._completion_kwargs())
except Exception:
# The Python path goes on to make an HTTP call; reaching it is
# the assertion, so the network failure below is expected.
pass
assert seen["gate"] == []
assert seen["call"] == []
assert transform.called
def test_pre_call_logging_fires_once_on_the_python_path(self):
from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
self._inject()
calls = {"pre_call": []}
logging_obj = MagicMock()
logging_obj.pre_call.side_effect = lambda **kwargs: calls["pre_call"].append(kwargs)
@ -2422,6 +2372,8 @@ class TestRustChatCompletionsHook:
try:
AnthropicChatCompletion().completion(**self._completion_kwargs(logging_obj=logging_obj))
except Exception:
# The Python path goes on to make an HTTP call; reaching it is
# the assertion, so the network failure below is expected.
pass
assert len(calls["pre_call"]) == 1

41
tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py
View file

@ -1,8 +1,6 @@
"""Tests for `BedrockConverseLLM.completion`.
The catalog keeps Bedrock chat completions on the Python path, so the injected
native callables are never consulted. AWS credential resolution is stubbed so
nothing reaches STS.
AWS credential resolution is stubbed so nothing reaches STS.
"""
from __future__ import annotations
@ -21,7 +19,6 @@ from botocore.exceptions import ClientError
from litellm.llms.bedrock.chat.converse_handler import BedrockConverseLLM
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.rust_bridge import configuration
from litellm.rust_bridge.chat_completions import native as bridge
from litellm.types.utils import ModelResponse
from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe
@ -33,33 +30,13 @@ RESOLVED_CREDENTIALS = Credentials(
@pytest.fixture(autouse=True)
def reset_bridge(monkeypatch):
def reset_rust_configuration(monkeypatch):
monkeypatch.setenv("LITELLM_RUST", "1")
configuration.reset_rust_configuration()
bridge.set_rust_chat_completions(
chat_completions=None, achat_completions=None, decline=None
)
yield
bridge.set_rust_chat_completions(
chat_completions=None, achat_completions=None, decline=None
)
configuration.reset_rust_configuration()
def _inject():
seen: dict[str, list[dict]] = {"gate": [], "call": []}
def gate(**kwargs):
seen["gate"].append(kwargs)
def native(**kwargs):
seen["call"].append(kwargs)
raise AssertionError("the native call must not run for a python-only route")
bridge.set_rust_chat_completions(decline=gate, chat_completions=native)
return seen
def _completion_kwargs(**overrides):
kwargs = {
"model": "bedrock/us-east-1/anthropic.claude-sonnet-4-5-v1:0",
@ -199,17 +176,7 @@ def _sync_client_returning_converse_response():
return client
def test_the_python_only_route_never_consults_the_core():
seen = _inject()
response = _run(client=_sync_client_returning_converse_response())
assert response.choices[0].message.content == "hi"
assert seen["gate"] == []
assert seen["call"] == []
def test_the_sync_python_path_logs_pre_call_once():
_inject()
logging_obj = MagicMock()
response = _run(
logging_obj=logging_obj,
@ -222,8 +189,8 @@ def test_the_sync_python_path_logs_pre_call_once():
def test_bearer_token_auth_serves_when_boto3_resolves_no_sigv4_credentials(monkeypatch):
"""With only `AWS_BEARER_TOKEN_BEDROCK` configured boto3 resolves no
credentials at all. Preparing the Rust handoff must not dereference that
None: the bearer token signs the request on its own."""
credentials at all. The handler must not dereference that None: the bearer
token signs the request on its own."""
monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "bedrock-bearer-token")
client = _sync_client_returning_converse_response()

0
tests/test_litellm/messages/__init__.py Normal file
View file

198
tests/test_litellm/messages/test_dispatch.py Normal file
View file

@ -0,0 +1,198 @@
import inspect
from collections.abc import Generator, Mapping
from typing import Final
from unittest.mock import AsyncMock, Mock
import pytest
import litellm
from litellm.llms.anthropic.experimental_pass_through.messages import handler as python_messages
from litellm.rust_bridge import configuration, runtime
from litellm.rust_bridge.catalog import Route, Rule, decision
from litellm.rust_bridge.configuration import Rollout
from litellm.rust_bridge.messages.entrypoints import (
NATIVE_AMESSAGES,
NATIVE_MESSAGES,
LiteLLMMessagesRequest,
)
from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse
MESSAGES: Final = [{"role": "user", "content": "hi"}]
RUST_RULES: Final = (Rule(Route.MESSAGES, Rollout.RUST_OPT_OUT),)
def _response(model: str = "claude-sonnet-4-5") -> AnthropicMessagesResponse:
return AnthropicMessagesResponse(id="msg_test", type="message", role="assistant", model=model, content=[])
@pytest.fixture(autouse=True)
def isolated_configuration(monkeypatch: pytest.MonkeyPatch) -> Generator[None]:
monkeypatch.delenv("LITELLM_RUST", raising=False)
configuration.reset_rust_configuration()
yield
NATIVE_MESSAGES.reset()
NATIVE_AMESSAGES.reset()
configuration.reset_rust_configuration()
@pytest.fixture
def rust_route(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(runtime, "decision", lambda context, rules=RUST_RULES: decision(context, RUST_RULES))
def test_public_signature_is_the_legacy_signature() -> None:
assert inspect.signature(litellm.anthropic_messages_handler) == inspect.signature(
python_messages.anthropic_messages_handler
)
assert inspect.signature(litellm.anthropic_messages) == inspect.signature(python_messages.anthropic_messages)
@pytest.mark.asyncio
@pytest.mark.parametrize("asynchronous", [False, True])
async def test_python_only_route_never_loads_native(monkeypatch: pytest.MonkeyPatch, asynchronous: bool) -> None:
response: Final = _response()
fallback: Final = AsyncMock(return_value=response) if asynchronous else Mock(return_value=response)
monkeypatch.setattr(
python_messages, "anthropic_messages" if asynchronous else "anthropic_messages_handler", fallback
)
monkeypatch.setattr(
NATIVE_AMESSAGES if asynchronous else NATIVE_MESSAGES,
"load",
Mock(side_effect=AssertionError("native must not be loaded")),
)
litellm.rust(True)
result: Final = (
await litellm.anthropic_messages(16, MESSAGES, "claude-sonnet-4-5", temperature=0.1)
if asynchronous
else litellm.anthropic_messages_handler(16, MESSAGES, "claude-sonnet-4-5", temperature=0.1)
)
assert result is response
fallback.assert_called_once_with(16, MESSAGES, "claude-sonnet-4-5", temperature=0.1)
@pytest.mark.asyncio
@pytest.mark.parametrize("asynchronous", [False, True])
async def test_unavailable_native_uses_python(
monkeypatch: pytest.MonkeyPatch, rust_route: None, asynchronous: bool
) -> None:
response: Final = _response()
fallback: Final = AsyncMock(return_value=response) if asynchronous else Mock(return_value=response)
monkeypatch.setattr(
python_messages, "anthropic_messages" if asynchronous else "anthropic_messages_handler", fallback
)
(NATIVE_AMESSAGES if asynchronous else NATIVE_MESSAGES).override(None)
result: Final = (
await litellm.anthropic_messages(16, MESSAGES, "claude-sonnet-4-5", temperature=0.1)
if asynchronous
else litellm.anthropic_messages_handler(16, MESSAGES, "claude-sonnet-4-5", temperature=0.1)
)
assert result is response
fallback.assert_called_once_with(16, MESSAGES, "claude-sonnet-4-5", temperature=0.1)
def test_native_receives_the_bound_request_and_original_call_shape(rust_route: None) -> None:
captured: Final[list[tuple[LiteLLMMessagesRequest, tuple[object, ...], Mapping[str, object]]]] = []
def native(
request: LiteLLMMessagesRequest,
args: tuple[object, ...],
kwargs: Mapping[str, object],
) -> AnthropicMessagesResponse:
captured.append((request, args, kwargs))
return _response(request.model)
NATIVE_MESSAGES.override(native)
response: Final = litellm.anthropic_messages_handler(
16,
MESSAGES,
"anthropic/claude-sonnet-4-5",
stream=True,
api_key="sk-test",
api_base="https://example.invalid",
custom_llm_provider="anthropic",
litellm_metadata={"user_id": "u"},
)
request, call_args, hook_kwargs = captured[0]
assert isinstance(response, dict)
assert response["model"] == "anthropic/claude-sonnet-4-5"
assert request.model == "anthropic/claude-sonnet-4-5"
assert request.messages is MESSAGES
assert request.max_tokens == 16
assert request.stream is True
assert request.api_key == "sk-test"
assert request.api_base == "https://example.invalid"
assert request.custom_llm_provider == "anthropic"
assert request.kwargs == {"litellm_metadata": {"user_id": "u"}}
assert call_args == (16, MESSAGES, "anthropic/claude-sonnet-4-5")
assert hook_kwargs["litellm_metadata"] == {"user_id": "u"}
assert "temperature" not in hook_kwargs
def test_internal_async_dispatch_marker_stays_on_python(monkeypatch: pytest.MonkeyPatch, rust_route: None) -> None:
native: Final = Mock(side_effect=AssertionError("the async handler's inner sync call must stay on Python"))
NATIVE_MESSAGES.override(native)
response: Final = _response()
fallback: Final = Mock(return_value=response)
monkeypatch.setattr(python_messages, "anthropic_messages_handler", fallback)
assert litellm.anthropic_messages_handler(16, MESSAGES, "claude-sonnet-4-5", is_async=True) is response
native.assert_not_called()
@pytest.mark.parametrize("enabled", [False, True], ids=["flag-disabled", "flag-enabled"])
def test_public_binding_errors_do_not_depend_on_native_selection(rust_route: None, enabled: bool) -> None:
native: Final = Mock(side_effect=AssertionError("binding errors precede admission"))
litellm.rust(enabled)
NATIVE_MESSAGES.override(native)
with pytest.raises(TypeError, match=r"anthropic_messages_handler\(\) got multiple values for argument 'model'"):
litellm.anthropic_messages_handler(16, MESSAGES, "claude-sonnet-4-5", model="duplicate")
with pytest.raises(TypeError, match=r"anthropic_messages_handler\(\) missing 3 required positional arguments"):
litellm.anthropic_messages_handler()
native.assert_not_called()
class Declined(Exception):
pass
class Upstream(Exception):
pass
@pytest.mark.asyncio
@pytest.mark.parametrize("asynchronous", [False, True])
@pytest.mark.parametrize("declined", [False, True])
async def test_only_native_declines_replay_on_python(
monkeypatch: pytest.MonkeyPatch, rust_route: None, asynchronous: bool, declined: bool
) -> None:
failure: Final = Declined("unsupported") if declined else RuntimeError("provider already called")
native: Final = AsyncMock(side_effect=failure) if asynchronous else Mock(side_effect=failure)
(NATIVE_AMESSAGES if asynchronous else NATIVE_MESSAGES).override(native)
monkeypatch.setattr(runtime, "native_exception_types", lambda: (Declined, Upstream))
response: Final = _response()
fallback: Final = AsyncMock(return_value=response) if asynchronous else Mock(return_value=response)
monkeypatch.setattr(
python_messages, "anthropic_messages" if asynchronous else "anthropic_messages_handler", fallback
)
async def call() -> object:
if asynchronous:
return await litellm.anthropic_messages(16, MESSAGES, "claude-sonnet-4-5")
return litellm.anthropic_messages_handler(16, MESSAGES, "claude-sonnet-4-5")
if declined:
assert await call() is response
fallback.assert_called_once_with(16, MESSAGES, "claude-sonnet-4-5")
else:
with pytest.raises(RuntimeError) as caught:
await call()
assert caught.value is failure
fallback.assert_not_called()
assert native.call_count == 1

195
tests/test_litellm/responses/test_dispatch.py Normal file
View file

@ -0,0 +1,195 @@
import inspect
from collections.abc import Generator, Mapping
from typing import Final
from unittest.mock import AsyncMock, Mock
import pytest
import litellm
from litellm.responses import main as python_responses
from litellm.rust_bridge import configuration, runtime
from litellm.rust_bridge.catalog import Route, Rule, decision
from litellm.rust_bridge.configuration import Rollout
from litellm.rust_bridge.responses.entrypoints import (
NATIVE_ARESPONSES,
NATIVE_RESPONSES,
LiteLLMResponsesRequest,
)
from litellm.types.llms.openai import ResponsesAPIResponse
RUST_RULES: Final = (Rule(Route.RESPONSES, Rollout.RUST_OPT_OUT),)
def _response(model: str = "gpt-4o") -> ResponsesAPIResponse:
return ResponsesAPIResponse(
id="resp_test", object="response", created_at=0, model=model, output=[], status="completed"
)
@pytest.fixture(autouse=True)
def isolated_configuration(monkeypatch: pytest.MonkeyPatch) -> Generator[None]:
monkeypatch.delenv("LITELLM_RUST", raising=False)
configuration.reset_rust_configuration()
yield
NATIVE_RESPONSES.reset()
NATIVE_ARESPONSES.reset()
configuration.reset_rust_configuration()
@pytest.fixture
def rust_route(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(runtime, "decision", lambda context, rules=RUST_RULES: decision(context, RUST_RULES))
def test_public_signature_is_the_legacy_signature() -> None:
assert inspect.signature(litellm.responses) == inspect.signature(python_responses.responses)
assert inspect.signature(litellm.aresponses) == inspect.signature(python_responses.aresponses)
@pytest.mark.asyncio
@pytest.mark.parametrize("asynchronous", [False, True])
async def test_python_only_route_never_loads_native(monkeypatch: pytest.MonkeyPatch, asynchronous: bool) -> None:
response: Final = _response()
fallback: Final = AsyncMock(return_value=response) if asynchronous else Mock(return_value=response)
monkeypatch.setattr(python_responses, "aresponses" if asynchronous else "responses", fallback)
monkeypatch.setattr(
NATIVE_ARESPONSES if asynchronous else NATIVE_RESPONSES,
"load",
Mock(side_effect=AssertionError("native must not be loaded")),
)
litellm.rust(True)
result: Final = (
await litellm.aresponses("hi", "gpt-4o", temperature=0.1)
if asynchronous
else litellm.responses("hi", "gpt-4o", temperature=0.1)
)
assert result is response
fallback.assert_called_once_with("hi", "gpt-4o", temperature=0.1)
@pytest.mark.asyncio
@pytest.mark.parametrize("asynchronous", [False, True])
async def test_unavailable_native_uses_python(
monkeypatch: pytest.MonkeyPatch, rust_route: None, asynchronous: bool
) -> None:
response: Final = _response()
fallback: Final = AsyncMock(return_value=response) if asynchronous else Mock(return_value=response)
monkeypatch.setattr(python_responses, "aresponses" if asynchronous else "responses", fallback)
(NATIVE_ARESPONSES if asynchronous else NATIVE_RESPONSES).override(None)
result: Final = (
await litellm.aresponses("hi", "gpt-4o", temperature=0.1)
if asynchronous
else litellm.responses("hi", "gpt-4o", temperature=0.1)
)
assert result is response
fallback.assert_called_once_with("hi", "gpt-4o", temperature=0.1)
def test_native_receives_the_bound_request_and_original_call_shape(rust_route: None) -> None:
captured: Final[list[tuple[LiteLLMResponsesRequest, tuple[object, ...], Mapping[str, object]]]] = []
def native(
request: LiteLLMResponsesRequest,
args: tuple[object, ...],
kwargs: Mapping[str, object],
) -> ResponsesAPIResponse:
captured.append((request, args, kwargs))
return _response(request.model)
NATIVE_RESPONSES.override(native)
response: Final = litellm.responses(
"hi",
"anthropic/claude-sonnet-4-5",
stream=True,
api_key="sk-test",
api_base="https://example.invalid",
extra_headers={"x-test": "1"},
custom_llm_provider="anthropic",
litellm_metadata={"user_id": "u"},
)
request, call_args, hook_kwargs = captured[0]
assert isinstance(response, ResponsesAPIResponse)
assert response.model == "anthropic/claude-sonnet-4-5"
assert request.model == "anthropic/claude-sonnet-4-5"
assert request.input == "hi"
assert request.stream is True
assert request.api_key == "sk-test"
assert request.api_base == "https://example.invalid"
assert request.custom_llm_provider == "anthropic"
assert request.extra_headers == {"x-test": "1"}
assert request.kwargs == {
"api_key": "sk-test",
"api_base": "https://example.invalid",
"litellm_metadata": {"user_id": "u"},
}
assert call_args == ("hi", "anthropic/claude-sonnet-4-5")
assert hook_kwargs["litellm_metadata"] == {"user_id": "u"}
assert "temperature" not in hook_kwargs
def test_internal_async_dispatch_marker_stays_on_python(monkeypatch: pytest.MonkeyPatch, rust_route: None) -> None:
native: Final = Mock(side_effect=AssertionError("aresponses's inner responses() call must stay on Python"))
NATIVE_RESPONSES.override(native)
response: Final = _response()
fallback: Final = Mock(return_value=response)
monkeypatch.setattr(python_responses, "responses", fallback)
assert litellm.responses("hi", "gpt-4o", aresponses=True) is response
native.assert_not_called()
@pytest.mark.parametrize("enabled", [False, True], ids=["flag-disabled", "flag-enabled"])
def test_public_binding_errors_do_not_depend_on_native_selection(rust_route: None, enabled: bool) -> None:
native: Final = Mock(side_effect=AssertionError("binding errors precede admission"))
litellm.rust(enabled)
NATIVE_RESPONSES.override(native)
with pytest.raises(TypeError, match=r"responses\(\) got multiple values for argument 'model'"):
litellm.responses("hi", "gpt-4o", model="duplicate")
with pytest.raises(TypeError, match=r"responses\(\) missing 2 required positional arguments: 'input' and 'model'"):
litellm.responses()
native.assert_not_called()
class Declined(Exception):
pass
class Upstream(Exception):
pass
@pytest.mark.asyncio
@pytest.mark.parametrize("asynchronous", [False, True])
@pytest.mark.parametrize("declined", [False, True])
async def test_only_native_declines_replay_on_python(
monkeypatch: pytest.MonkeyPatch, rust_route: None, asynchronous: bool, declined: bool
) -> None:
failure: Final = Declined("unsupported") if declined else RuntimeError("provider already called")
native: Final = AsyncMock(side_effect=failure) if asynchronous else Mock(side_effect=failure)
(NATIVE_ARESPONSES if asynchronous else NATIVE_RESPONSES).override(native)
monkeypatch.setattr(runtime, "native_exception_types", lambda: (Declined, Upstream))
response: Final = _response()
fallback: Final = AsyncMock(return_value=response) if asynchronous else Mock(return_value=response)
monkeypatch.setattr(python_responses, "aresponses" if asynchronous else "responses", fallback)
async def call() -> object:
if asynchronous:
return await litellm.aresponses("hi", "gpt-4o")
return litellm.responses("hi", "gpt-4o")
if declined:
assert await call() is response
fallback.assert_called_once_with("hi", "gpt-4o")
else:
with pytest.raises(RuntimeError) as caught:
await call()
assert caught.value is failure
fallback.assert_not_called()
assert native.call_count == 1

49
tests/test_litellm/rust_bridge/chat_completions/test_callbacks.py Normal file
View file

@ -0,0 +1,49 @@
from types import MappingProxyType
from typing import Final
from litellm.rust_bridge.chat_completions.callbacks import arguments, response
from litellm.rust_bridge.chat_completions.entrypoints import LiteLLMChatCompletionsRequest
from litellm.types.utils import ModelResponse
def test_response_builds_the_public_model_response() -> None:
built: Final = response(
MappingProxyType(
{
"id": "chatcmpl-native",
"object": "chat.completion",
"created": 1,
"model": "claude-sonnet-4-5",
"choices": [
{
"index": 0,
"finish_reason": "stop",
"message": {"role": "assistant", "content": "native"},
}
],
"usage": {"prompt_tokens": 2, "completion_tokens": 3, "total_tokens": 5},
}
)
)
assert isinstance(built, ModelResponse)
assert built.id == "chatcmpl-native"
assert built.choices[0].message.content == "native"
assert built.usage is not None
assert built.usage.total_tokens == 5
def test_arguments_are_the_public_kwargs_view() -> None:
kwargs: Final = MappingProxyType({"metadata": {"user_id": "u"}})
request: Final = LiteLLMChatCompletionsRequest(
model="anthropic/claude-sonnet-4-5",
messages=[{"role": "user", "content": "hi"}],
stream=None,
api_key=None,
api_base=None,
custom_llm_provider="anthropic",
extra_headers=None,
kwargs=kwargs,
)
assert arguments(request) is kwargs

300
tests/test_litellm/rust_bridge/chat_completions/test_native.py
View file

@ -1,300 +0,0 @@
"""Tests for the Rust chat completions bridge.
The native callables are dependency-injected through
``set_rust_chat_completions`` rather than patched, so these run without the
compiled extension present.
"""
from __future__ import annotations
import pytest
from litellm.rust_bridge import configuration
from litellm.rust_bridge.chat_completions import native as bridge
from litellm.types.utils import ModelResponse
RUST_RESPONSE = {
"created": 1_700_000_000,
"model": "claude-sonnet-4-5-20260101",
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": "hello from rust"},
"finish_reason": "stop",
}
],
"usage": {
"prompt_tokens": 11,
"completion_tokens": 4,
"total_tokens": 15,
"prompt_tokens_details": {
"cached_tokens": 0,
"cache_creation_tokens": 0,
"text_tokens": 11,
},
},
}
MESSAGES = [{"role": "user", "content": "hi"}]
class _FakeDeclined(Exception):
"""Stands in for the native `RustBridgeDeclined`."""
class _FakeUpstream(Exception):
"""Stands in for the native `RustUpstreamError`; args are (status, message)."""
class _FakeNative:
RustBridgeDeclined = _FakeDeclined
RustUpstreamError = _FakeUpstream
def _fake_native_bridge(monkeypatch):
"""Expose the bridge's exception classes without the compiled extension."""
monkeypatch.setattr(bridge, "get_native_bridge", lambda: _FakeNative())
def _hide_native_bridge(monkeypatch):
"""Simulate a wheel built without the compiled extension.
There is no injection seam for "the .so is absent", so the loader itself is
replaced; every other case here uses `set_rust_chat_completions`.
"""
monkeypatch.setattr(bridge, "get_native_bridge", lambda: None)
@pytest.fixture(autouse=True)
def reset_bridge(monkeypatch):
"""Every test starts with no injected callables, and leaves none behind."""
bridge.set_rust_chat_completions(chat_completions=None, achat_completions=None, decline=None)
configuration.reset_rust_configuration()
monkeypatch.setenv("LITELLM_RUST", "1")
yield
bridge.set_rust_chat_completions(chat_completions=None, achat_completions=None, decline=None)
configuration.reset_rust_configuration()
class _RecordingDecline:
"""A stand-in for the native gate that records what it was asked."""
def __init__(self, reason: str | None = None):
self.reason = reason
self.calls: list[dict] = []
def __call__(self, **kwargs):
self.calls.append(kwargs)
return self.reason
class _RecordingCall:
def __init__(self, result=None, error: Exception | None = None):
self.result = result if result is not None else dict(RUST_RESPONSE)
self.error = error
self.calls: list[dict] = []
def __call__(self, **kwargs):
self.calls.append(kwargs)
if self.error is not None:
raise self.error
return self.result
class _RecordingAsyncCall(_RecordingCall):
async def __call__(self, **kwargs):
return _RecordingCall.__call__(self, **kwargs)
def _accepts(**overrides) -> bool:
kwargs = {
"model": "claude-sonnet-4-5",
"messages": MESSAGES,
"optional_params": {"max_tokens": 16},
"custom_llm_provider": "anthropic",
"litellm_params": {},
"stream": None,
}
kwargs.update(overrides)
return bridge.rust_chat_completions_accepts(**kwargs)
class TestGate:
@pytest.mark.parametrize("custom_llm_provider", ("anthropic", "bedrock", "openai", None))
def test_the_python_only_route_never_consults_the_core(self, custom_llm_provider):
gate = _RecordingDecline()
bridge.set_rust_chat_completions(decline=gate)
configuration.rust(True)
assert _accepts(custom_llm_provider=custom_llm_provider) is False
assert _accepts(custom_llm_provider=custom_llm_provider, stream=True) is False
assert gate.calls == []
def _call_kwargs(model_response: ModelResponse) -> dict:
return {
"model": "claude-sonnet-4-5",
"messages": MESSAGES,
"optional_params": {"max_tokens": 16},
"model_response": model_response,
"api_key": "sk-test",
"api_base": None,
"custom_llm_provider": "anthropic",
"extra_headers": {},
"timeout": 30.0,
"on_response": lambda _rust_response: None,
}
class TestSyncCall:
def test_builds_a_model_response_and_stamps_the_rust_header(self):
native = _RecordingCall()
bridge.set_rust_chat_completions(chat_completions=native)
model_response = ModelResponse()
original_id = model_response.id
result = bridge.chat_completions(**_call_kwargs(model_response))
assert result is not None
assert result.choices[0].message.content == "hello from rust"
assert result.choices[0].finish_reason == "stop"
assert result.model == "claude-sonnet-4-5-20260101"
assert result.usage.prompt_tokens == 11
assert result.usage.completion_tokens == 4
assert result.usage.total_tokens == 15
assert result._hidden_params["additional_headers"] == {"x-litellm-rust": "true"}
assert result.id == original_id, "the rust path must keep the chatcmpl id litellm already minted"
def test_passes_the_timeout_through_as_seconds(self):
native = _RecordingCall()
bridge.set_rust_chat_completions(chat_completions=native)
bridge.chat_completions(**_call_kwargs(ModelResponse()))
assert native.calls[0]["timeout_seconds"] == 30.0
def test_falls_back_when_the_bridge_is_unavailable(self, monkeypatch):
_hide_native_bridge(monkeypatch)
assert bridge.chat_completions(**_call_kwargs(ModelResponse())) is None
def test_falls_back_when_the_core_declines_before_calling_the_provider(self, monkeypatch):
_fake_native_bridge(monkeypatch)
bridge.set_rust_chat_completions(chat_completions=_RecordingCall(error=_FakeDeclined("streaming")))
assert bridge.chat_completions(**_call_kwargs(ModelResponse())) is None
class TestAsyncCall:
@pytest.mark.asyncio
async def test_builds_a_model_response(self):
bridge.set_rust_chat_completions(achat_completions=_RecordingAsyncCall())
result = await bridge.achat_completions(**_call_kwargs(ModelResponse()))
assert result is not None
assert result.choices[0].message.content == "hello from rust"
assert result._hidden_params["additional_headers"] == {"x-litellm-rust": "true"}
@pytest.mark.asyncio
async def test_falls_back_when_the_bridge_is_unavailable(self, monkeypatch):
_hide_native_bridge(monkeypatch)
assert await bridge.achat_completions(**_call_kwargs(ModelResponse())) is None
@pytest.mark.asyncio
async def test_falls_back_when_the_core_declines_before_calling_the_provider(self, monkeypatch):
_fake_native_bridge(monkeypatch)
bridge.set_rust_chat_completions(achat_completions=_RecordingAsyncCall(error=_FakeDeclined("streaming")))
assert await bridge.achat_completions(**_call_kwargs(ModelResponse())) is None
class TestAsyncFallbackWrapper:
@pytest.mark.asyncio
async def test_returns_the_rust_response_without_running_the_fallback(self):
bridge.set_rust_chat_completions(achat_completions=_RecordingAsyncCall())
ran = []
async def fallback():
ran.append(True)
return "python"
result = await bridge.achat_completions_or_fallback(**_call_kwargs(ModelResponse()), python_fallback=fallback)
assert result.choices[0].message.content == "hello from rust"
assert ran == []
@pytest.mark.asyncio
async def test_runs_the_fallback_when_the_core_declines(self, monkeypatch):
_fake_native_bridge(monkeypatch)
bridge.set_rust_chat_completions(achat_completions=_RecordingAsyncCall(error=_FakeDeclined("streaming")))
async def fallback():
return "python"
result = await bridge.achat_completions_or_fallback(**_call_kwargs(ModelResponse()), python_fallback=fallback)
assert result == "python"
@pytest.mark.asyncio
async def test_runs_the_fallback_when_the_bridge_is_unavailable(self, monkeypatch):
_hide_native_bridge(monkeypatch)
async def fallback():
return "python"
result = await bridge.achat_completions_or_fallback(**_call_kwargs(ModelResponse()), python_fallback=fallback)
assert result == "python"
class TestFailureClassification:
"""A failure the provider already saw must not be retried on the Python
path: it would bill the customer for the same work twice."""
@pytest.fixture(autouse=True)
def _native_exceptions(self, monkeypatch):
_fake_native_bridge(monkeypatch)
def test_a_decline_falls_back_because_nothing_was_sent(self):
bridge.set_rust_chat_completions(chat_completions=_RecordingCall(error=_FakeDeclined("streaming")))
assert bridge.chat_completions(**_call_kwargs(ModelResponse())) is None
def test_an_upstream_failure_is_surfaced_with_its_status(self):
from litellm.exceptions import APIError
bridge.set_rust_chat_completions(chat_completions=_RecordingCall(error=_FakeUpstream(429, "429: rate limited")))
with pytest.raises(APIError) as raised:
bridge.chat_completions(**_call_kwargs(ModelResponse()))
assert raised.value.status_code == 429
assert "rate limited" in str(raised.value)
def test_a_transport_failure_with_no_response_surfaces_as_a_500(self):
from litellm.exceptions import APIError
bridge.set_rust_chat_completions(chat_completions=_RecordingCall(error=_FakeUpstream(0, "connection reset")))
with pytest.raises(APIError) as raised:
bridge.chat_completions(**_call_kwargs(ModelResponse()))
assert raised.value.status_code == 500
def test_an_unrecognized_error_is_not_swallowed(self):
bridge.set_rust_chat_completions(chat_completions=_RecordingCall(error=RuntimeError("something else")))
with pytest.raises(RuntimeError):
bridge.chat_completions(**_call_kwargs(ModelResponse()))
@pytest.mark.asyncio
async def test_the_async_wrapper_does_not_fall_back_on_an_upstream_failure(self):
from litellm.exceptions import APIError
bridge.set_rust_chat_completions(achat_completions=_RecordingAsyncCall(error=_FakeUpstream(500, "500: boom")))
ran = []
async def fallback():
ran.append(True)
return "python"
with pytest.raises(APIError):
await bridge.achat_completions_or_fallback(**_call_kwargs(ModelResponse()), python_fallback=fallback)
assert ran == [], "a request the provider already served must not be re-issued"
@pytest.mark.asyncio
async def test_the_async_wrapper_falls_back_on_a_decline(self):
bridge.set_rust_chat_completions(
achat_completions=_RecordingAsyncCall(error=_FakeDeclined("blank message text"))
)
async def fallback():
return "python"
result = await bridge.achat_completions_or_fallback(**_call_kwargs(ModelResponse()), python_fallback=fallback)
assert result == "python"

0
tests/test_litellm/rust_bridge/messages/__init__.py Normal file
View file

42
tests/test_litellm/rust_bridge/messages/test_callbacks.py Normal file
View file

@ -0,0 +1,42 @@
from types import MappingProxyType
from typing import Final
from litellm.rust_bridge.messages.callbacks import arguments, response
from litellm.rust_bridge.messages.entrypoints import LiteLLMMessagesRequest
def test_response_is_a_detached_public_messages_dict() -> None:
native: Final = MappingProxyType(
{
"id": "msg_native",
"type": "message",
"role": "assistant",
"model": "claude-sonnet-4-5",
"content": [{"type": "text", "text": "native"}],
"stop_reason": "end_turn",
"usage": {"input_tokens": 2, "output_tokens": 3},
}
)
built: Final = response(native)
assert built == dict(native)
assert isinstance(built, dict)
built["_hidden_params"] = {"annotated": True}
assert "_hidden_params" not in native
def test_arguments_are_the_public_kwargs_view() -> None:
kwargs: Final = MappingProxyType({"litellm_metadata": {"user_id": "u"}})
request: Final = LiteLLMMessagesRequest(
model="claude-sonnet-4-5",
messages=[{"role": "user", "content": "hi"}],
max_tokens=16,
stream=None,
api_key=None,
api_base=None,
custom_llm_provider="anthropic",
kwargs=kwargs,
)
assert arguments(request) is kwargs

0
tests/test_litellm/rust_bridge/responses/__init__.py Normal file
View file

57
tests/test_litellm/rust_bridge/responses/test_callbacks.py Normal file
View file

@ -0,0 +1,57 @@
from types import MappingProxyType
from typing import Final
import pytest
from pydantic import ValidationError
from litellm.rust_bridge.responses.callbacks import arguments, response
from litellm.rust_bridge.responses.entrypoints import LiteLLMResponsesRequest
from litellm.types.llms.openai import ResponsesAPIResponse
def test_response_validates_into_the_public_responses_model() -> None:
built: Final = response(
MappingProxyType(
{
"id": "resp_native",
"object": "response",
"created_at": 1,
"model": "gpt-4o",
"status": "completed",
"output": [
{
"type": "message",
"id": "msg_native",
"role": "assistant",
"status": "completed",
"content": [{"type": "output_text", "text": "native", "annotations": []}],
}
],
}
)
)
assert isinstance(built, ResponsesAPIResponse)
assert built.id == "resp_native"
assert built.output[0].content[0].text == "native"
def test_response_rejects_a_payload_missing_required_fields() -> None:
with pytest.raises(ValidationError):
response(MappingProxyType({"object": "response"}))
def test_arguments_are_the_public_kwargs_view() -> None:
kwargs: Final = MappingProxyType({"litellm_metadata": {"user_id": "u"}})
request: Final = LiteLLMResponsesRequest(
model="gpt-4o",
input="hi",
stream=None,
api_key=None,
api_base=None,
custom_llm_provider="openai",
extra_headers=None,
kwargs=kwargs,
)
assert arguments(request) is kwargs

54
tests/test_litellm/rust_bridge/test_failures.py Normal file
View file

@ -0,0 +1,54 @@
from types import MappingProxyType
from typing import Final
import pytest
import litellm
from litellm.rust_bridge import failures
class UpstreamRateLimited(Exception):
status_code = 429
message = "rate limited"
def test_upstream_status_maps_onto_the_public_exception_contract() -> None:
upstream: Final = UpstreamRateLimited("rate limited")
mapped: Final = failures.map_failure(upstream, "anthropic/claude-sonnet-4-5", "anthropic", MappingProxyType({}))
assert isinstance(mapped, litellm.RateLimitError)
assert mapped.llm_provider == "anthropic"
assert mapped.model == "claude-sonnet-4-5"
def test_mapper_failure_keeps_the_native_error_as_context(monkeypatch: pytest.MonkeyPatch) -> None:
def explode(**_kwargs: object) -> Exception:
raise ValueError("mapper broke")
monkeypatch.setattr(litellm, "exception_type", explode)
native_error: Final = RuntimeError("native")
mapped: Final = failures.map_failure(native_error, "mistral/mistral-ocr-latest", "mistral", MappingProxyType({}))
assert isinstance(mapped, ValueError)
assert mapped.__context__ is native_error
def test_kwargs_are_handed_to_the_mapper_as_owned_copies(monkeypatch: pytest.MonkeyPatch) -> None:
seen: Final[list[dict[str, object]]] = []
def record(**kwargs: object) -> Exception:
seen.append(dict(kwargs))
return RuntimeError("mapped")
monkeypatch.setattr(litellm, "exception_type", record)
request_kwargs: Final = MappingProxyType({"metadata": {"user_id": "u"}})
failures.map_failure(RuntimeError("native"), "gpt-4o", "openai", request_kwargs)
assert seen[0]["completion_kwargs"] == {"metadata": {"user_id": "u"}}
assert seen[0]["extra_kwargs"] == {"metadata": {"user_id": "u"}}
assert seen[0]["completion_kwargs"] is not request_kwargs
assert seen[0]["model"] == "gpt-4o"
assert seen[0]["custom_llm_provider"] == "openai"