mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-24 05:45:11 +00:00
fix(rust-bridge): share provider resolution across route dispatch and execution
This commit is contained in:
parent
7657618605
commit
c04eba9acb
5 changed files with 145 additions and 41 deletions
|
|
@ -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),
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
)
|
||||
|
|
|
|||
60
litellm/rust_bridge/provider.py
Normal file
60
litellm/rust_bridge/provider.py
Normal file
|
|
@ -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)
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue