diff --git a/litellm/messages/dispatch.py b/litellm/messages/dispatch.py index 77d6257e3e7..f0fdfaecddc 100644 --- a/litellm/messages/dispatch.py +++ b/litellm/messages/dispatch.py @@ -3,12 +3,10 @@ from collections.abc import AsyncIterator, Awaitable, Callable, Coroutine, Itera from types import MappingProxyType from typing import Final, TypeAlias, cast # noqa: TID251 # native binding selects a sync result or an async awaitable -from litellm.exceptions import BadRequestError -from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.llms.anthropic.pass_through.messages import handler as main from litellm.llms.anthropic.pass_through.messages.interceptors import get_messages_interceptors from litellm.llms.anthropic.pass_through.safeguards import native_route_classifies -from litellm.rust_bridge import detectors +from litellm.rust_bridge import detectors, provider from litellm.rust_bridge.catalog import Feature, Route, RouteContext, Target from litellm.rust_bridge.dispatch import PublicDispatch from litellm.rust_bridge.messages.entrypoints import ( @@ -23,7 +21,6 @@ from litellm.rust_bridge.public_call import ( native_call_hook, optional_mapping, optional_sequence, - optional_str, signature, ) from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse @@ -77,16 +74,6 @@ def _public_request( return native_call(legacy, args, _native_kwargs(fields, kwargs)) -def _resolved(request: NativeCall) -> tuple[str, str | None]: - model: Final = str(request.resolved["model"]) - provider: Final = optional_str(request.resolved.get("custom_llm_provider")) - try: - resolved_model, resolved_provider, _, _ = get_llm_provider(model, provider) - except BadRequestError: - return model, provider - return resolved_model, resolved_provider - - def _target(request: NativeCall, model: str, provider: str | None) -> Target: requested_model: Final = str(request.resolved["model"]) model_info: Final = request.resolved.get("model_info") @@ -102,12 +89,11 @@ def _target(request: NativeCall, model: str, provider: str | None) -> Target: return Target.CHAT_COMPLETIONS -def _intercepted(request: NativeCall) -> bool: +def _intercepted(request: NativeCall, provider: str | None) -> bool: tools: Final = optional_sequence(request.resolved.get("tools")) if not tools: return False mappings: Final = tuple(mapping for mapping in map(optional_mapping, tools) if mapping is not None) - provider: Final = _resolved(request)[1] return any(interceptor.can_handle(mappings, provider) for interceptor in get_messages_interceptors()) @@ -131,33 +117,47 @@ _ASYNC_FEATURES: Final[detectors.Detectors] = MappingProxyType( **_FEATURES, Feature.PRE_REQUEST_HOOK: detectors.pre_request_hook, Feature.WEBSEARCH_INTERCEPTION: detectors.websearch_interception, - Feature.MESSAGES_INTERCEPTOR: _intercepted, } ) -def _context(request: NativeCall, features: detectors.Detectors) -> RouteContext: +def _features( + request: NativeCall, provider: str | None, features: detectors.Detectors, intercepting: bool +) -> frozenset[Feature]: + detected: Final = detectors.detected(features, request) + if not intercepting or not _intercepted(request, provider): + return detected + return detected | {Feature.MESSAGES_INTERCEPTOR} + + +def _context(request: NativeCall, features: detectors.Detectors, intercepting: bool) -> RouteContext: requested_model: Final = str(request.resolved["model"]) - model, provider = _resolved(request) - return RouteContext( - Route.MESSAGES, - provider=provider, - model=requested_model, - target=_target(request, model, provider), - features=detectors.detected(features, request), - ) + resolution: Final = provider.resolve(request) + match resolution: + case provider.Resolved(model=resolved_model, provider=resolved_provider): + return RouteContext( + Route.MESSAGES, + provider=resolved_provider, + model=requested_model, + target=_target(request, resolved_model, resolved_provider), + features=_features(request, resolved_provider, features, intercepting), + ) + case provider.Unresolved(): + return RouteContext( + Route.MESSAGES, model=requested_model, features=_features(request, None, features, intercepting) + ) _DISPATCH: Final = PublicDispatch( route=Route.MESSAGES, request=lambda args, kwargs: _public_request(_MESSAGES, args, kwargs), - context=lambda request: _context(request, _FEATURES), + context=lambda request: _context(request, _FEATURES, intercepting=False), ) _ADISPATCH: Final = PublicDispatch( route=Route.MESSAGES, request=lambda args, kwargs: _public_request(_AMESSAGES, args, kwargs), - context=lambda request: _context(request, _ASYNC_FEATURES), + context=lambda request: _context(request, _ASYNC_FEATURES, intercepting=True), ) diff --git a/litellm/responses/dispatch.py b/litellm/responses/dispatch.py index 0cac90f4361..f857eb5010f 100644 --- a/litellm/responses/dispatch.py +++ b/litellm/responses/dispatch.py @@ -3,11 +3,9 @@ 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.exceptions import BadRequestError -from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.responses import main from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator -from litellm.rust_bridge import detectors +from litellm.rust_bridge import detectors, provider from litellm.rust_bridge.catalog import Feature, Route, RouteContext from litellm.rust_bridge.dispatch import PublicDispatch from litellm.rust_bridge.public_call import ( @@ -16,7 +14,6 @@ from litellm.rust_bridge.public_call import ( bind, native_call, native_call_hook, - optional_str, signature, ) from litellm.rust_bridge.responses.entrypoints import ( @@ -65,14 +62,6 @@ def _public_request( return native_call(legacy, args, kwargs) -def _provider(request: NativeCall) -> str | None: - provider: Final = optional_str(request.resolved.get("custom_llm_provider")) - try: - return get_llm_provider(str(request.resolved["model"]), provider)[1] - except BadRequestError: - return provider - - _FEATURES: Final[detectors.Detectors] = MappingProxyType( { **detectors.COMMON, @@ -105,9 +94,10 @@ _SYNC_FEATURES: Final[detectors.Detectors] = MappingProxyType( def _context(request: NativeCall, features: detectors.Detectors) -> RouteContext: + resolution: Final = provider.resolve(request) return RouteContext( Route.RESPONSES, - provider=_provider(request), + provider=resolution.provider if isinstance(resolution, provider.Resolved) else None, model=str(request.resolved["model"]), features=detectors.detected(features, request), ) diff --git a/litellm/rust_bridge/provider.py b/litellm/rust_bridge/provider.py new file mode 100644 index 00000000000..61964dc3812 --- /dev/null +++ b/litellm/rust_bridge/provider.py @@ -0,0 +1,60 @@ +from __future__ import annotations + +from dataclasses import dataclass +from types import MappingProxyType +from typing import Final, TypeAlias + +from litellm.exceptions import BadRequestError +from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider +from litellm.rust_bridge.public_call import NativeCall, optional_str +from litellm.types.router import GenericLiteLLMParams + +__all__ = ("Resolution", "Resolved", "Unresolved", "resolve") + +_HANDLER_FIELDS: Final = frozenset({"api_base", "api_key", "custom_llm_provider"}) + + +@dataclass(frozen=True, slots=True) +class Resolved: + model: str + provider: str + + +@dataclass(frozen=True, slots=True) +class Unresolved: + pass + + +Resolution: TypeAlias = Resolved | Unresolved + + +def resolve(request: NativeCall) -> Resolution: + resolved: Final = request.resolved + model: Final = str(resolved["model"]) + custom_llm_provider: Final = optional_str(resolved.get("custom_llm_provider")) + api_base: Final = optional_str(resolved.get("api_base")) + api_key: Final = optional_str(resolved.get("api_key")) + extra: Final = MappingProxyType( + { + name: value + for name, value in request.kwargs.items() + if name not in request.base and name not in _HANDLER_FIELDS + } + ) + litellm_params: Final = GenericLiteLLMParams( + **extra, + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + ) + try: + resolved_model, resolved_provider, _, _ = get_llm_provider( + model=model, + custom_llm_provider=custom_llm_provider, + api_base=api_base, + api_key=api_key, + litellm_params=litellm_params, + ) + except BadRequestError: + return Unresolved() + return Resolved(resolved_model, resolved_provider) diff --git a/tests/test_litellm_rust/messages/test_dispatch.py b/tests/test_litellm_rust/messages/test_dispatch.py index 50f71569d73..309a0e3776c 100644 --- a/tests/test_litellm_rust/messages/test_dispatch.py +++ b/tests/test_litellm_rust/messages/test_dispatch.py @@ -601,3 +601,55 @@ async def test_interceptor_feature_follows_the_resolved_provider() -> None: rules=rules, ) assert result is native_result + + +@pytest.mark.parametrize( + "kwargs", + ( + {"api_base": "https://api.groq.com/openai/v1"}, + {"api_base": "https://api.deepseek.com/v1"}, + {"use_litellm_proxy": True}, + ), +) +def test_provider_rule_uses_legacy_handler_provider_inputs(kwargs: Mapping[str, object]) -> None: + python_result: Final = response() + + def python(*call_args: object, **call_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: fallback + return python_result + + def native(request: NativeCall) -> AnthropicMessagesResponse: + pytest.fail("An anthropic provider rule must not route providers resolved from other handler inputs") + + result: Final = _DISPATCH.run( + (16, MESSAGES, "claude-sonnet-4-5"), + kwargs, + python=python, + binding=messages_binding(native), + native=lambda hook, request, call_args, call_kwargs: hook(request), + rules=(RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED, providers=frozenset({"anthropic"})),), + ) + assert result is python_result + + +@pytest.mark.parametrize( + "kwargs", + ( + {}, + {"custom_llm_provider": "anthropic"}, + ), +) +def test_provider_rule_keeps_anthropic_requests_native(kwargs: Mapping[str, object]) -> None: + native_result: Final = response() + + def python(*call_args: object, **call_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: fallback + return response() + + result: Final = _DISPATCH.run( + (16, MESSAGES, "claude-sonnet-4-5"), + kwargs, + python=python, + binding=messages_binding(lambda request: native_result), + native=lambda hook, request, call_args, call_kwargs: hook(request), + rules=(RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED, providers=frozenset({"anthropic"})),), + ) + assert result is native_result diff --git a/tests/test_litellm_rust/responses/test_dispatch.py b/tests/test_litellm_rust/responses/test_dispatch.py index 4498dbdf39d..226e64b2742 100644 --- a/tests/test_litellm_rust/responses/test_dispatch.py +++ b/tests/test_litellm_rust/responses/test_dispatch.py @@ -388,6 +388,8 @@ async def test_deployment_hook_keeps_only_sync_calls_on_python(monkeypatch: pyte ({"model": "anthropic/claude-sonnet-4-5"}, False), ({"model": "claude-sonnet-4-5"}, False), ({"model": "gpt-5", "custom_llm_provider": "azure"}, False), + ({"model": "gpt-5", "api_base": "https://api.groq.com/openai/v1"}, False), + ({"model": "gpt-5", "use_litellm_proxy": True}, False), ), ) def test_provider_rule_sees_the_resolved_provider(kwargs: Mapping[str, object], native: bool) -> None: