mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-04 07:27:51 +00:00
fix(rust-bridge): share provider resolution across route dispatch and execution
Some checks failed
LiteLLM Rust / rust-lint (push) Has been cancelled
LiteLLM Rust / rust-test (push) Has been cancelled
LiteLLM Rust / rust-wheel (push) Has been cancelled
Terraform Provider / gofmt, vet, build, test (push) Has been cancelled
Terraform Provider / Provider endpoints vs proxy OpenAPI schema (push) Has been cancelled
Some checks failed
LiteLLM Rust / rust-lint (push) Has been cancelled
LiteLLM Rust / rust-test (push) Has been cancelled
LiteLLM Rust / rust-wheel (push) Has been cancelled
Terraform Provider / gofmt, vet, build, test (push) Has been cancelled
Terraform Provider / Provider endpoints vs proxy OpenAPI schema (push) Has been cancelled
This commit is contained in:
parent
560372456e
commit
b336411232
7 changed files with 155 additions and 65 deletions
|
|
@ -3,10 +3,8 @@ 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.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 import detectors as route_detectors
|
||||
|
|
@ -21,7 +19,6 @@ from litellm.rust_bridge.public_call import (
|
|||
native_call,
|
||||
native_call_hook,
|
||||
optional_sequence,
|
||||
optional_str,
|
||||
signature,
|
||||
)
|
||||
from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse
|
||||
|
|
@ -75,16 +72,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")
|
||||
|
|
@ -119,33 +106,47 @@ _ASYNC_FEATURES: Final[detectors.Detectors] = MappingProxyType(
|
|||
**_FEATURES,
|
||||
Feature.PRE_REQUEST_HOOK: route_detectors.pre_request_hook,
|
||||
Feature.WEBSEARCH_INTERCEPTION: route_detectors.websearch_interception,
|
||||
Feature.MESSAGES_INTERCEPTOR: route_detectors.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 route_detectors.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 import detectors as route_detectors
|
||||
|
|
@ -66,14 +63,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,
|
||||
|
|
@ -106,9 +95,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,18 +60,7 @@ def websearch_interception(_request: NativeCall) -> bool:
|
|||
return any(isinstance(callback, WebSearchInterceptionLogger) for callback in litellm.callbacks)
|
||||
|
||||
|
||||
def _resolved_provider(request: NativeCall) -> str | None:
|
||||
from litellm.exceptions import BadRequestError
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
|
||||
|
||||
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
|
||||
|
||||
|
||||
def intercepted(request: NativeCall) -> bool:
|
||||
def intercepted(request: NativeCall, provider: str | None) -> bool:
|
||||
"""The async Messages handler delegates matching tools to a registered interceptor."""
|
||||
from litellm.llms.anthropic.pass_through.messages.interceptors import get_messages_interceptors
|
||||
|
||||
|
|
@ -79,7 +68,6 @@ def intercepted(request: NativeCall) -> bool:
|
|||
if not tools:
|
||||
return False
|
||||
mappings: Final = tuple(mapping for mapping in map(optional_mapping, tools) if mapping is not None)
|
||||
provider: Final = _resolved_provider(request)
|
||||
return any(interceptor.can_handle(mappings, provider) for interceptor in get_messages_interceptors())
|
||||
|
||||
|
||||
|
|
|
|||
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:
|
||||
|
|
|
|||
|
|
@ -45,22 +45,19 @@ def call(**kwargs: object) -> NativeCall:
|
|||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("model", "provider", "tools", "expected"),
|
||||
("provider", "tools", "expected"),
|
||||
(
|
||||
("deepseek/test-model", None, [{"type": ANTHROPIC_ADVISOR_TOOL_TYPE}], True),
|
||||
("anthropic/test-model", None, [{"type": ANTHROPIC_ADVISOR_TOOL_TYPE}], False),
|
||||
("test-model", "anthropic", [{"type": ANTHROPIC_ADVISOR_TOOL_TYPE}], False),
|
||||
("unresolved-model", None, [{"type": ANTHROPIC_ADVISOR_TOOL_TYPE}], True),
|
||||
("deepseek/test-model", None, [{"type": "function", "name": "f"}], False),
|
||||
("deepseek/test-model", None, ["not-a-tool"], False),
|
||||
("deepseek/test-model", None, [], False),
|
||||
("deepseek/test-model", None, None, False),
|
||||
("deepseek", [{"type": ANTHROPIC_ADVISOR_TOOL_TYPE}], True),
|
||||
("anthropic", [{"type": ANTHROPIC_ADVISOR_TOOL_TYPE}], False),
|
||||
(None, [{"type": ANTHROPIC_ADVISOR_TOOL_TYPE}], True),
|
||||
("deepseek", [{"type": "function", "name": "f"}], False),
|
||||
("deepseek", ["not-a-tool"], False),
|
||||
("deepseek", [], False),
|
||||
("deepseek", None, False),
|
||||
),
|
||||
)
|
||||
def test_messages_interception_matches_tools_and_resolved_provider(
|
||||
model: str, provider: str | None, tools: object, expected: bool
|
||||
) -> None:
|
||||
assert messages_detectors.intercepted(call(model=model, custom_llm_provider=provider, tools=tools)) is expected
|
||||
def test_messages_interception_matches_tools_and_provider(provider: str | None, tools: object, expected: bool) -> None:
|
||||
assert messages_detectors.intercepted(call(model="test-model", tools=tools), provider) is expected
|
||||
|
||||
|
||||
def _wrap_nested_level(value: object, depth: int) -> object:
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue