From a84f68b6e36072539794e4abb387c65ef04e71af Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 16 Sep 2026 15:02:12 -0700 Subject: [PATCH] refactor(rust_bridge): give chat completions, messages and responses the ocr dispatch shape Each route now has litellm/rust_bridge//{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 --- litellm/__init__.py | 3 + litellm/chat_completions/__init__.py | 3 + litellm/chat_completions/dispatch.py | 114 +++++ litellm/llms/anthropic/chat/handler.py | 86 +--- litellm/llms/custom_httpx/llm_http_handler.py | 99 ---- litellm/messages/__init__.py | 3 + litellm/messages/dispatch.py | 113 +++++ litellm/responses/dispatch.py | 106 +++++ .../rust_bridge/chat_completions/callbacks.py | 19 + .../chat_completions/entrypoints.py | 54 +++ .../rust_bridge/chat_completions/native.py | 445 ------------------ litellm/rust_bridge/failures.py | 37 ++ litellm/rust_bridge/messages/callbacks.py | 23 + litellm/rust_bridge/messages/entrypoints.py | 54 +++ litellm/rust_bridge/messages/native.py | 136 ------ litellm/rust_bridge/ocr/callbacks.py | 31 +- litellm/rust_bridge/public_call.py | 42 ++ litellm/rust_bridge/responses/callbacks.py | 19 + litellm/rust_bridge/responses/entrypoints.py | 54 +++ tests/e2e/e2e_config.py | 2 - .../test_messages_azure_foundry_e2e.py | 10 +- .../test_rust_bridge_messages.py | 237 ---------- .../test_litellm/chat_completions/__init__.py | 0 .../chat_completions/test_dispatch.py | 186 ++++++++ .../chat/test_anthropic_chat_handler.py | 56 +-- .../chat/test_bedrock_converse_handler.py | 41 +- tests/test_litellm/messages/__init__.py | 0 tests/test_litellm/messages/test_dispatch.py | 198 ++++++++ tests/test_litellm/responses/test_dispatch.py | 195 ++++++++ .../chat_completions/test_callbacks.py | 49 ++ .../chat_completions/test_native.py | 300 ------------ .../rust_bridge/messages/__init__.py | 0 .../rust_bridge/messages/test_callbacks.py | 42 ++ .../rust_bridge/responses/__init__.py | 0 .../rust_bridge/responses/test_callbacks.py | 57 +++ .../test_litellm/rust_bridge/test_failures.py | 54 +++ 36 files changed, 1447 insertions(+), 1421 deletions(-) create mode 100644 litellm/chat_completions/__init__.py create mode 100644 litellm/chat_completions/dispatch.py create mode 100644 litellm/messages/__init__.py create mode 100644 litellm/messages/dispatch.py create mode 100644 litellm/responses/dispatch.py create mode 100644 litellm/rust_bridge/chat_completions/callbacks.py create mode 100644 litellm/rust_bridge/chat_completions/entrypoints.py delete mode 100644 litellm/rust_bridge/chat_completions/native.py create mode 100644 litellm/rust_bridge/failures.py create mode 100644 litellm/rust_bridge/messages/callbacks.py create mode 100644 litellm/rust_bridge/messages/entrypoints.py delete mode 100644 litellm/rust_bridge/messages/native.py create mode 100644 litellm/rust_bridge/public_call.py create mode 100644 litellm/rust_bridge/responses/callbacks.py create mode 100644 litellm/rust_bridge/responses/entrypoints.py delete mode 100644 tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py create mode 100644 tests/test_litellm/chat_completions/__init__.py create mode 100644 tests/test_litellm/chat_completions/test_dispatch.py create mode 100644 tests/test_litellm/messages/__init__.py create mode 100644 tests/test_litellm/messages/test_dispatch.py create mode 100644 tests/test_litellm/responses/test_dispatch.py create mode 100644 tests/test_litellm/rust_bridge/chat_completions/test_callbacks.py delete mode 100644 tests/test_litellm/rust_bridge/chat_completions/test_native.py create mode 100644 tests/test_litellm/rust_bridge/messages/__init__.py create mode 100644 tests/test_litellm/rust_bridge/messages/test_callbacks.py create mode 100644 tests/test_litellm/rust_bridge/responses/__init__.py create mode 100644 tests/test_litellm/rust_bridge/responses/test_callbacks.py create mode 100644 tests/test_litellm/rust_bridge/test_failures.py diff --git a/litellm/__init__.py b/litellm/__init__.py index c6f03172d8e..c80720c3677 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -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 * diff --git a/litellm/chat_completions/__init__.py b/litellm/chat_completions/__init__.py new file mode 100644 index 00000000000..b5f139da0c8 --- /dev/null +++ b/litellm/chat_completions/__init__.py @@ -0,0 +1,3 @@ +from .dispatch import acompletion, completion + +__all__ = ("acompletion", "completion") diff --git a/litellm/chat_completions/dispatch.py b/litellm/chat_completions/dispatch.py new file mode 100644 index 00000000000..83e5d956988 --- /dev/null +++ b/litellm/chat_completions/dispatch.py @@ -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 diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py index dff3a0be3fc..73c101ebdef 100644 --- a/litellm/llms/anthropic/chat/handler.py +++ b/litellm/llms/anthropic/chat/handler.py @@ -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: diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 98e74ddce81..27012585a10 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -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, diff --git a/litellm/messages/__init__.py b/litellm/messages/__init__.py new file mode 100644 index 00000000000..7c492ba4c3b --- /dev/null +++ b/litellm/messages/__init__.py @@ -0,0 +1,3 @@ +from .dispatch import anthropic_messages, anthropic_messages_handler + +__all__ = ("anthropic_messages", "anthropic_messages_handler") diff --git a/litellm/messages/dispatch.py b/litellm/messages/dispatch.py new file mode 100644 index 00000000000..af7123046a4 --- /dev/null +++ b/litellm/messages/dispatch.py @@ -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 diff --git a/litellm/responses/dispatch.py b/litellm/responses/dispatch.py new file mode 100644 index 00000000000..a85a7feb542 --- /dev/null +++ b/litellm/responses/dispatch.py @@ -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 diff --git a/litellm/rust_bridge/chat_completions/callbacks.py b/litellm/rust_bridge/chat_completions/callbacks.py new file mode 100644 index 00000000000..9a00ce340ba --- /dev/null +++ b/litellm/rust_bridge/chat_completions/callbacks.py @@ -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)) diff --git a/litellm/rust_bridge/chat_completions/entrypoints.py b/litellm/rust_bridge/chat_completions/entrypoints.py new file mode 100644 index 00000000000..6e41600c42e --- /dev/null +++ b/litellm/rust_bridge/chat_completions/entrypoints.py @@ -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) diff --git a/litellm/rust_bridge/chat_completions/native.py b/litellm/rust_bridge/chat_completions/native.py deleted file mode 100644 index 1e03806f38c..00000000000 --- a/litellm/rust_bridge/chat_completions/native.py +++ /dev/null @@ -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() diff --git a/litellm/rust_bridge/failures.py b/litellm/rust_bridge/failures.py new file mode 100644 index 00000000000..b714341fe43 --- /dev/null +++ b/litellm/rust_bridge/failures.py @@ -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 diff --git a/litellm/rust_bridge/messages/callbacks.py b/litellm/rust_bridge/messages/callbacks.py new file mode 100644 index 00000000000..1aff6c7f75d --- /dev/null +++ b/litellm/rust_bridge/messages/callbacks.py @@ -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)) diff --git a/litellm/rust_bridge/messages/entrypoints.py b/litellm/rust_bridge/messages/entrypoints.py new file mode 100644 index 00000000000..46565bfd46a --- /dev/null +++ b/litellm/rust_bridge/messages/entrypoints.py @@ -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) diff --git a/litellm/rust_bridge/messages/native.py b/litellm/rust_bridge/messages/native.py deleted file mode 100644 index 40d0ddf622b..00000000000 --- a/litellm/rust_bridge/messages/native.py +++ /dev/null @@ -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), - ) diff --git a/litellm/rust_bridge/ocr/callbacks.py b/litellm/rust_bridge/ocr/callbacks.py index 6c2c0573779..c30943e6fd7 100644 --- a/litellm/rust_bridge/ocr/callbacks.py +++ b/litellm/rust_bridge/ocr/callbacks.py @@ -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)) diff --git a/litellm/rust_bridge/public_call.py b/litellm/rust_bridge/public_call.py new file mode 100644 index 00000000000..2a41926a802 --- /dev/null +++ b/litellm/rust_bridge/public_call.py @@ -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 diff --git a/litellm/rust_bridge/responses/callbacks.py b/litellm/rust_bridge/responses/callbacks.py new file mode 100644 index 00000000000..180b89c4412 --- /dev/null +++ b/litellm/rust_bridge/responses/callbacks.py @@ -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)) diff --git a/litellm/rust_bridge/responses/entrypoints.py b/litellm/rust_bridge/responses/entrypoints.py new file mode 100644 index 00000000000..9bba7406b6d --- /dev/null +++ b/litellm/rust_bridge/responses/entrypoints.py @@ -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) diff --git a/tests/e2e/e2e_config.py b/tests/e2e/e2e_config.py index 896cb3e7efe..0a549c44b25 100644 --- a/tests/e2e/e2e_config.py +++ b/tests/e2e/e2e_config.py @@ -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 diff --git a/tests/e2e/llm_translation/test_messages_azure_foundry_e2e.py b/tests/e2e/llm_translation/test_messages_azure_foundry_e2e.py index d8d44820e80..07be68a964b 100644 --- a/tests/e2e/llm_translation/test_messages_azure_foundry_e2e.py +++ b/tests/e2e/llm_translation/test_messages_azure_foundry_e2e.py @@ -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: diff --git a/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py b/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py deleted file mode 100644 index 9f7f1bc86c7..00000000000 --- a/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py +++ /dev/null @@ -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 diff --git a/tests/test_litellm/chat_completions/__init__.py b/tests/test_litellm/chat_completions/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/chat_completions/test_dispatch.py b/tests/test_litellm/chat_completions/test_dispatch.py new file mode 100644 index 00000000000..5892b208302 --- /dev/null +++ b/tests/test_litellm/chat_completions/test_dispatch.py @@ -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 diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py index e45d655ff7f..5667d5ca56c 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py @@ -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 diff --git a/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py b/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py index 79f41a22fe3..67ffe7570a1 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py +++ b/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py @@ -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() diff --git a/tests/test_litellm/messages/__init__.py b/tests/test_litellm/messages/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/messages/test_dispatch.py b/tests/test_litellm/messages/test_dispatch.py new file mode 100644 index 00000000000..840e9dec667 --- /dev/null +++ b/tests/test_litellm/messages/test_dispatch.py @@ -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 diff --git a/tests/test_litellm/responses/test_dispatch.py b/tests/test_litellm/responses/test_dispatch.py new file mode 100644 index 00000000000..3daf2b475fc --- /dev/null +++ b/tests/test_litellm/responses/test_dispatch.py @@ -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 diff --git a/tests/test_litellm/rust_bridge/chat_completions/test_callbacks.py b/tests/test_litellm/rust_bridge/chat_completions/test_callbacks.py new file mode 100644 index 00000000000..94ac358c6d1 --- /dev/null +++ b/tests/test_litellm/rust_bridge/chat_completions/test_callbacks.py @@ -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 diff --git a/tests/test_litellm/rust_bridge/chat_completions/test_native.py b/tests/test_litellm/rust_bridge/chat_completions/test_native.py deleted file mode 100644 index 14f8113924d..00000000000 --- a/tests/test_litellm/rust_bridge/chat_completions/test_native.py +++ /dev/null @@ -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" diff --git a/tests/test_litellm/rust_bridge/messages/__init__.py b/tests/test_litellm/rust_bridge/messages/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/rust_bridge/messages/test_callbacks.py b/tests/test_litellm/rust_bridge/messages/test_callbacks.py new file mode 100644 index 00000000000..8ba0497ffbe --- /dev/null +++ b/tests/test_litellm/rust_bridge/messages/test_callbacks.py @@ -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 diff --git a/tests/test_litellm/rust_bridge/responses/__init__.py b/tests/test_litellm/rust_bridge/responses/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/rust_bridge/responses/test_callbacks.py b/tests/test_litellm/rust_bridge/responses/test_callbacks.py new file mode 100644 index 00000000000..6ecc5bcf0b9 --- /dev/null +++ b/tests/test_litellm/rust_bridge/responses/test_callbacks.py @@ -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 diff --git a/tests/test_litellm/rust_bridge/test_failures.py b/tests/test_litellm/rust_bridge/test_failures.py new file mode 100644 index 00000000000..80057b816d3 --- /dev/null +++ b/tests/test_litellm/rust_bridge/test_failures.py @@ -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"