mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-20 00:11:50 +00:00
refactor(rust_bridge): keep every route but OCR and Bedrock transcription on Python
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Yujong Lee
parent
f9d423827f
commit
803baead7a
7 changed files with 56 additions and 956 deletions
|
|
@ -63,27 +63,12 @@ class Rule:
|
|||
|
||||
Rules: TypeAlias = tuple[Rule, ...]
|
||||
|
||||
_COMPLETED: Final = frozenset({Delivery.COMPLETED})
|
||||
|
||||
RULES: Final[Rules] = (
|
||||
Rule(Route.OCR, Rollout.RUST_OPT_OUT),
|
||||
Rule(Route.TRANSCRIPTION, Rollout.RUST_REQUIRED, providers=frozenset({"bedrock"})),
|
||||
Rule(Route.TRANSCRIPTION, Rollout.PYTHON_ONLY),
|
||||
Rule(
|
||||
Route.CHAT_COMPLETIONS,
|
||||
Rollout.RUST_OPT_IN,
|
||||
providers=frozenset({"anthropic", "bedrock"}),
|
||||
deliveries=_COMPLETED,
|
||||
),
|
||||
Rule(Route.CHAT_COMPLETIONS, Rollout.PYTHON_ONLY),
|
||||
Rule(Route.MESSAGES, Rollout.RUST_OPT_IN, providers=frozenset({"anthropic", "azure_ai"})),
|
||||
Rule(Route.MESSAGES, Rollout.PYTHON_ONLY),
|
||||
Rule(
|
||||
Route.RESPONSES,
|
||||
Rollout.RUST_OPT_IN,
|
||||
providers=frozenset({"openai"}),
|
||||
deliveries=frozenset({Delivery.WEBSOCKET}),
|
||||
),
|
||||
Rule(Route.RESPONSES, Rollout.PYTHON_ONLY),
|
||||
Rule(Route.EMBEDDING, Rollout.PYTHON_ONLY),
|
||||
Rule(Route.RERANK, Rollout.PYTHON_ONLY),
|
||||
|
|
|
|||
|
|
@ -99,15 +99,6 @@ class ExplodingAsyncMessages:
|
|||
raise AssertionError("bridge must not be called")
|
||||
|
||||
|
||||
class RaisingAsyncMessages:
|
||||
def __init__(self) -> None:
|
||||
self.calls = 0
|
||||
|
||||
async def __call__(self, **kwargs: object) -> dict[str, object]:
|
||||
self.calls += 1
|
||||
raise RuntimeError("upstream request failed with status 400: bad request")
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _reset_rust_flag():
|
||||
rust_messages.set_rust_messages(messages=None, amessages=None)
|
||||
|
|
@ -218,152 +209,18 @@ def _gate(**overrides):
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gate_invokes_rust_and_marks_response_header():
|
||||
bridge = RecordingAsyncMessages()
|
||||
litellm.rust(True)
|
||||
rust_messages.set_rust_messages(amessages=bridge)
|
||||
|
||||
response = await _gate()
|
||||
|
||||
assert response is not None
|
||||
assert response["id"] == "msg_123"
|
||||
assert response["_hidden_params"]["additional_headers"] == {"x-litellm-rust": "true"}
|
||||
call = bridge.calls[0]
|
||||
assert call["model"] == "claude-sonnet-4-5"
|
||||
assert call["body"] == REQUEST_BODY
|
||||
assert call["api_key"] == "sk-azure"
|
||||
assert call["api_base"] == "https://resource.services.ai.azure.com/anthropic"
|
||||
assert call["extra_headers"] == {"x-api-key": "sk-azure", "anthropic-version": "2023-06-01"}
|
||||
assert call["timeout_seconds"] == 30.0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gate_falls_back_to_python_when_bridge_raises():
|
||||
bridge = RaisingAsyncMessages()
|
||||
litellm.rust(True)
|
||||
rust_messages.set_rust_messages(amessages=bridge)
|
||||
|
||||
response = await _gate()
|
||||
|
||||
assert response is None
|
||||
assert bridge.calls == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gate_skips_rust_when_flag_absent():
|
||||
bridge = ExplodingAsyncMessages()
|
||||
rust_messages.set_rust_messages(amessages=bridge)
|
||||
|
||||
response = await _gate(litellm_params=GenericLiteLLMParams(api_key="sk-azure"))
|
||||
|
||||
assert response is None
|
||||
assert bridge.calls == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gate_uses_process_enable_without_request_override():
|
||||
bridge = RecordingAsyncMessages()
|
||||
rust_messages.set_rust_messages(amessages=bridge)
|
||||
litellm.rust(True)
|
||||
|
||||
response = await _gate(litellm_params=GenericLiteLLMParams(api_key="sk-azure"))
|
||||
|
||||
assert response is not None
|
||||
assert bridge.calls[0]["custom_llm_provider"] == "azure_ai"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gate_invokes_rust_for_native_anthropic_provider():
|
||||
bridge = RecordingAsyncMessages()
|
||||
litellm.rust(True)
|
||||
rust_messages.set_rust_messages(amessages=bridge)
|
||||
|
||||
response = await _gate(
|
||||
custom_llm_provider="anthropic",
|
||||
litellm_params=GenericLiteLLMParams(api_key="sk-ant"),
|
||||
api_key="sk-ant",
|
||||
api_base="https://api.anthropic.com",
|
||||
headers={"x-api-key": "sk-ant", "anthropic-version": "2023-06-01"},
|
||||
)
|
||||
|
||||
assert response is not None
|
||||
assert response["_hidden_params"]["additional_headers"] == {"x-litellm-rust": "true"}
|
||||
assert bridge.calls[0]["custom_llm_provider"] == "anthropic"
|
||||
assert bridge.calls[0]["api_key"] == "sk-ant"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gate_invokes_rust_when_env_var_set(monkeypatch):
|
||||
bridge = RecordingAsyncMessages()
|
||||
rust_messages.set_rust_messages(amessages=bridge)
|
||||
monkeypatch.setenv("LITELLM_RUST", "1")
|
||||
|
||||
response = await _gate(
|
||||
custom_llm_provider="anthropic",
|
||||
litellm_params=GenericLiteLLMParams(api_key="sk-ant"),
|
||||
)
|
||||
|
||||
assert response is not None
|
||||
assert bridge.calls[0]["custom_llm_provider"] == "anthropic"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gate_env_var_falsey_does_not_enable(monkeypatch):
|
||||
bridge = ExplodingAsyncMessages()
|
||||
rust_messages.set_rust_messages(amessages=bridge)
|
||||
monkeypatch.setenv("LITELLM_RUST", "0")
|
||||
|
||||
response = await _gate(
|
||||
custom_llm_provider="anthropic",
|
||||
litellm_params=GenericLiteLLMParams(api_key="sk-ant"),
|
||||
)
|
||||
|
||||
assert response is None
|
||||
assert bridge.calls == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gate_skips_rust_for_unsupported_provider():
|
||||
@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="openai")
|
||||
response = await _gate(custom_llm_provider=custom_llm_provider)
|
||||
|
||||
assert response is None
|
||||
assert bridge.calls == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gate_skips_rust_for_agentic_hook():
|
||||
bridge = ExplodingAsyncMessages()
|
||||
litellm.rust(True)
|
||||
rust_messages.set_rust_messages(amessages=bridge)
|
||||
|
||||
response = await _gate(has_agentic_hook=True)
|
||||
|
||||
assert response is None
|
||||
assert bridge.calls == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gate_streams_through_rust_when_eligible_and_strips_stream_flag():
|
||||
bridge = RecordingAsyncMessages()
|
||||
litellm.rust(True)
|
||||
rust_messages.set_rust_messages(amessages=bridge)
|
||||
|
||||
streaming_body = {**REQUEST_BODY, "stream": True}
|
||||
response = await _gate(
|
||||
has_agentic_hook=False,
|
||||
request_body=streaming_body,
|
||||
)
|
||||
|
||||
assert response is not None
|
||||
assert response["_hidden_params"]["additional_headers"] == {"x-litellm-rust": "true"}
|
||||
assert "stream" not in bridge.calls[0]["body"]
|
||||
assert bridge.calls[0]["body"] == REQUEST_BODY
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fake_stream_wraps_rust_response_as_anthropic_sse():
|
||||
response = cast(AnthropicMessagesResponse, dict(FAKE_MESSAGES_RESPONSE))
|
||||
|
|
@ -378,17 +235,3 @@ async def test_fake_stream_wraps_rust_response_as_anthropic_sse():
|
|||
assert b"event: content_block_delta" in joined
|
||||
assert b"hello world" in joined
|
||||
assert b"event: message_stop" in joined
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gate_falls_back_when_bridge_unavailable(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
importlib.import_module("litellm.rust_bridge"),
|
||||
"get_native_bridge",
|
||||
lambda: None,
|
||||
)
|
||||
litellm.rust(True)
|
||||
|
||||
response = await _gate()
|
||||
|
||||
assert response is None
|
||||
|
|
|
|||
|
|
@ -2334,46 +2334,20 @@ def test_non_bash_tool_result_skipped():
|
|||
|
||||
|
||||
class TestRustChatCompletionsHook:
|
||||
"""The `rust: true` opt-in on `/chat/completions` for the Anthropic provider.
|
||||
|
||||
The native callables are dependency-injected, so these run without the
|
||||
compiled extension.
|
||||
"""
|
||||
|
||||
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,
|
||||
},
|
||||
},
|
||||
}
|
||||
"""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 import chat_completions as bridge
|
||||
from litellm.rust_bridge import configuration
|
||||
|
||||
monkeypatch.setenv("LITELLM_RUST", "1")
|
||||
bridge.set_rust_chat_completions(
|
||||
chat_completions=None, achat_completions=None, decline=None
|
||||
)
|
||||
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
|
||||
)
|
||||
bridge.set_rust_chat_completions(chat_completions=None, achat_completions=None, decline=None)
|
||||
configuration.reset_rust_configuration()
|
||||
|
||||
@staticmethod
|
||||
def _completion_kwargs(**overrides):
|
||||
|
|
@ -2401,96 +2375,31 @@ class TestRustChatCompletionsHook:
|
|||
return kwargs
|
||||
|
||||
@staticmethod
|
||||
def _recording_logging_obj():
|
||||
"""A logging object that keeps each hook's payload in a real list, so a
|
||||
test can assert which path logged and what it carried."""
|
||||
calls = {"pre_call": [], "post_call": []}
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.pre_call.side_effect = lambda **kwargs: calls["pre_call"].append(kwargs)
|
||||
logging_obj.post_call.side_effect = lambda **kwargs: calls["post_call"].append(kwargs)
|
||||
return logging_obj, calls
|
||||
|
||||
def _inject(self, *, decline_reason=None, sync_result=None, sync_error=None):
|
||||
def _inject():
|
||||
from litellm.rust_bridge import chat_completions as bridge
|
||||
|
||||
seen = {"gate": [], "call": []}
|
||||
|
||||
def gate(**kwargs):
|
||||
seen["gate"].append(kwargs)
|
||||
return decline_reason
|
||||
|
||||
def native(**kwargs):
|
||||
seen["call"].append(kwargs)
|
||||
if sync_error is not None:
|
||||
raise sync_error
|
||||
return dict(sync_result if sync_result is not None else self.RUST_RESPONSE)
|
||||
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_rust_true_serves_the_call_and_stamps_the_header(self):
|
||||
from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion
|
||||
|
||||
seen = self._inject()
|
||||
response = AnthropicChatCompletion().completion(**self._completion_kwargs())
|
||||
|
||||
assert response.choices[0].message.content == "hello from rust"
|
||||
assert response._hidden_params["additional_headers"] == {"x-litellm-rust": "true"}
|
||||
assert len(seen["call"]) == 1
|
||||
|
||||
def test_the_core_receives_the_untranslated_openai_messages(self):
|
||||
"""Rust owns the translation, so the handler must not pre-translate."""
|
||||
from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion
|
||||
|
||||
seen = self._inject()
|
||||
AnthropicChatCompletion().completion(
|
||||
**self._completion_kwargs(
|
||||
messages=[
|
||||
{"role": "system", "content": "be terse"},
|
||||
{"role": "user", "content": "hi"},
|
||||
]
|
||||
)
|
||||
)
|
||||
assert seen["call"][0]["messages"] == [
|
||||
{"role": "system", "content": "be terse"},
|
||||
{"role": "user", "content": "hi"},
|
||||
]
|
||||
|
||||
def test_the_anthropic_max_tokens_default_is_merged_in_before_the_gate(self):
|
||||
"""`transform_request` applies `AnthropicConfig.get_config`; the Rust
|
||||
path skips it, so the handler has to merge it or Anthropic 400s on a
|
||||
request that omits `max_tokens`."""
|
||||
from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion
|
||||
|
||||
seen = self._inject()
|
||||
AnthropicChatCompletion().completion(**self._completion_kwargs(optional_params={}))
|
||||
assert "max_tokens" in seen["gate"][0]["optional_params"]
|
||||
assert seen["call"][0]["optional_params"]["max_tokens"] > 0
|
||||
|
||||
def test_a_caller_supplied_max_tokens_outranks_the_default(self):
|
||||
from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion
|
||||
|
||||
seen = self._inject()
|
||||
AnthropicChatCompletion().completion(
|
||||
**self._completion_kwargs(optional_params={"max_tokens": 7})
|
||||
)
|
||||
assert seen["call"][0]["optional_params"]["max_tokens"] == 7
|
||||
|
||||
def test_without_the_opt_in_the_core_is_never_consulted(self, monkeypatch):
|
||||
monkeypatch.setenv("LITELLM_RUST", "0")
|
||||
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, patch.object(
|
||||
AnthropicChatCompletion, "acompletion_function"
|
||||
):
|
||||
) as transform:
|
||||
try:
|
||||
AnthropicChatCompletion().completion(
|
||||
**self._completion_kwargs(litellm_params={})
|
||||
)
|
||||
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.
|
||||
|
|
@ -2499,218 +2408,19 @@ class TestRustChatCompletionsHook:
|
|||
assert seen["call"] == []
|
||||
assert transform.called
|
||||
|
||||
def test_a_declined_request_never_reaches_the_native_call(self):
|
||||
from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion
|
||||
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
|
||||
|
||||
seen = self._inject(decline_reason="unrecognized request parameter")
|
||||
with patch.object(
|
||||
AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []}
|
||||
):
|
||||
try:
|
||||
AnthropicChatCompletion().completion(**self._completion_kwargs())
|
||||
except Exception:
|
||||
pass
|
||||
assert len(seen["gate"]) == 1
|
||||
assert seen["call"] == []
|
||||
|
||||
def test_streaming_stays_on_the_python_path(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": []}
|
||||
):
|
||||
try:
|
||||
AnthropicChatCompletion().completion(
|
||||
**self._completion_kwargs(optional_params={"max_tokens": 16, "stream": True})
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
assert seen["gate"] == []
|
||||
|
||||
def test_pre_call_logging_fires_exactly_once_on_the_rust_path(self):
|
||||
from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion
|
||||
|
||||
seen = self._inject()
|
||||
logging_obj = MagicMock()
|
||||
AnthropicChatCompletion().completion(
|
||||
**self._completion_kwargs(logging_obj=logging_obj)
|
||||
)
|
||||
assert logging_obj.pre_call.call_count == 1
|
||||
assert len(seen["call"]) == 1
|
||||
|
||||
def test_post_call_logging_fires_on_the_rust_path(self):
|
||||
"""The Rust core owns the provider call, so the Python transform that
|
||||
normally raises `post_call` never runs. Without the bridge hook every
|
||||
post_call callback goes silent and `original_response` stays unset."""
|
||||
import json
|
||||
|
||||
from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion
|
||||
|
||||
self._inject()
|
||||
logging_obj = MagicMock()
|
||||
AnthropicChatCompletion().completion(
|
||||
**self._completion_kwargs(logging_obj=logging_obj)
|
||||
)
|
||||
|
||||
assert logging_obj.post_call.call_count == 1
|
||||
logged = logging_obj.post_call.call_args.kwargs["original_response"]
|
||||
assert json.loads(logged)["choices"][0]["message"]["content"] == "hello from rust"
|
||||
|
||||
def test_post_call_is_not_logged_twice_when_the_sync_rust_call_declines(self, monkeypatch):
|
||||
"""A decline never reached the provider, so the Python path serves the
|
||||
request and owns the only post_call. Firing the hook there too would
|
||||
double every post_call callback for one request."""
|
||||
from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion
|
||||
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
|
||||
from litellm.rust_bridge import chat_completions as bridge
|
||||
|
||||
class _Declined(Exception):
|
||||
pass
|
||||
|
||||
class _FakeNative:
|
||||
RustBridgeDeclined = _Declined
|
||||
RustUpstreamError = type("_Upstream", (Exception,), {})
|
||||
|
||||
def declining_native(**_kwargs):
|
||||
raise _Declined("blank message text")
|
||||
|
||||
monkeypatch.setattr(bridge, "get_native_bridge", lambda: _FakeNative())
|
||||
bridge.set_rust_chat_completions(
|
||||
decline=lambda **_kwargs: None, chat_completions=declining_native
|
||||
)
|
||||
|
||||
logging_obj, calls = self._recording_logging_obj()
|
||||
with patch.object(
|
||||
AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []}
|
||||
):
|
||||
try:
|
||||
AnthropicChatCompletion().completion(
|
||||
**self._completion_kwargs(logging_obj=logging_obj)
|
||||
)
|
||||
except Exception:
|
||||
# The Python path goes on to make an HTTP call; the log count is
|
||||
# the assertion, so a failure past this point is expected.
|
||||
pass
|
||||
|
||||
assert calls["post_call"] == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_the_async_path_falls_back_when_the_core_declines(self, monkeypatch):
|
||||
from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion
|
||||
from litellm.rust_bridge import chat_completions as bridge
|
||||
|
||||
class _Declined(Exception):
|
||||
pass
|
||||
|
||||
class _FakeNative:
|
||||
RustBridgeDeclined = _Declined
|
||||
RustUpstreamError = type("_Upstream", (Exception,), {})
|
||||
|
||||
monkeypatch.setattr(bridge, "get_native_bridge", lambda: _FakeNative())
|
||||
|
||||
async def declining_native(**_kwargs):
|
||||
raise _Declined("blank message text")
|
||||
|
||||
bridge.set_rust_chat_completions(
|
||||
decline=lambda **_kwargs: None, achat_completions=declining_native
|
||||
)
|
||||
|
||||
sentinel = object()
|
||||
|
||||
async def python_path(**_kwargs):
|
||||
return sentinel
|
||||
|
||||
with patch.object(
|
||||
AnthropicChatCompletion, "acompletion_function", side_effect=python_path
|
||||
) as python_call:
|
||||
result = await AnthropicChatCompletion().completion(
|
||||
**self._completion_kwargs(acompletion=True)
|
||||
)
|
||||
|
||||
assert result is sentinel
|
||||
assert python_call.called, "a failing rust call must re-enter the python path"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_the_async_path_serves_the_rust_response_without_the_fallback(self):
|
||||
from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion
|
||||
from litellm.rust_bridge import chat_completions as bridge
|
||||
|
||||
async def native(**_kwargs):
|
||||
return dict(self.RUST_RESPONSE)
|
||||
|
||||
bridge.set_rust_chat_completions(
|
||||
decline=lambda **_kwargs: None, achat_completions=native
|
||||
)
|
||||
|
||||
with patch.object(AnthropicChatCompletion, "acompletion_function") as python_call:
|
||||
result = await AnthropicChatCompletion().completion(
|
||||
**self._completion_kwargs(acompletion=True)
|
||||
)
|
||||
|
||||
assert result.choices[0].message.content == "hello from rust"
|
||||
assert result._hidden_params["additional_headers"] == {"x-litellm-rust": "true"}
|
||||
assert not python_call.called
|
||||
|
||||
|
||||
def test_pre_call_logging_fires_once_when_the_sync_rust_call_declines(self, monkeypatch):
|
||||
"""One request, one pre_call, on the synchronous path too. Without the
|
||||
suppression the Python path logs a second time for the same attempt."""
|
||||
from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion
|
||||
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
|
||||
from litellm.rust_bridge import chat_completions as bridge
|
||||
|
||||
class _Declined(Exception):
|
||||
pass
|
||||
|
||||
class _FakeNative:
|
||||
RustBridgeDeclined = _Declined
|
||||
RustUpstreamError = type("_Upstream", (Exception,), {})
|
||||
|
||||
monkeypatch.setattr(bridge, "get_native_bridge", lambda: _FakeNative())
|
||||
|
||||
def declining_native(**_kwargs):
|
||||
raise _Declined("blank message text")
|
||||
|
||||
bridge.set_rust_chat_completions(
|
||||
decline=lambda **_kwargs: None, chat_completions=declining_native
|
||||
)
|
||||
|
||||
logging_obj, calls = self._recording_logging_obj()
|
||||
with patch.object(
|
||||
AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []}
|
||||
):
|
||||
try:
|
||||
AnthropicChatCompletion().completion(
|
||||
**self._completion_kwargs(logging_obj=logging_obj)
|
||||
)
|
||||
except Exception:
|
||||
# The Python path goes on to make an HTTP call; the log count is
|
||||
# the assertion, so a failure past this point is expected.
|
||||
pass
|
||||
|
||||
assert len(calls["pre_call"]) == 1
|
||||
assert calls["pre_call"][0]["additional_args"]["complete_input_dict"]["model"] == (
|
||||
"claude-sonnet-4-5"
|
||||
)
|
||||
|
||||
def test_pre_call_logging_still_fires_when_rust_is_not_involved(self, monkeypatch):
|
||||
"""The suppression must not swallow the log on the ordinary path."""
|
||||
monkeypatch.setenv("LITELLM_RUST", "0")
|
||||
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()
|
||||
logging_obj, calls = self._recording_logging_obj()
|
||||
calls = {"pre_call": []}
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.pre_call.side_effect = lambda **kwargs: calls["pre_call"].append(kwargs)
|
||||
with patch.object(
|
||||
AnthropicConfig, "transform_request", return_value={"model": "m", "messages": []}
|
||||
):
|
||||
try:
|
||||
AnthropicChatCompletion().completion(
|
||||
**self._completion_kwargs(litellm_params={}, logging_obj=logging_obj)
|
||||
)
|
||||
AnthropicChatCompletion().completion(**self._completion_kwargs(logging_obj=logging_obj))
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
|
|
|||
|
|
@ -1,7 +1,8 @@
|
|||
"""Tests for `BedrockConverseLLM.completion`'s Rust chat completions hook.
|
||||
"""Tests for `BedrockConverseLLM.completion`.
|
||||
|
||||
The native callables are dependency-injected, so these run without the compiled
|
||||
extension, and AWS credential resolution is stubbed so nothing reaches STS.
|
||||
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.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
|
@ -19,31 +20,10 @@ 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 chat_completions as bridge
|
||||
from litellm.rust_bridge import configuration
|
||||
from litellm.types.utils import ModelResponse
|
||||
from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe
|
||||
|
||||
RUST_RESPONSE = {
|
||||
"created": 1_700_000_000,
|
||||
"model": "anthropic.claude-sonnet-4-5-v1:0",
|
||||
"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,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
RESOLVED_CREDENTIALS = Credentials(
|
||||
access_key="AKIARESOLVED",
|
||||
secret_key="resolved-secret",
|
||||
|
|
@ -54,6 +34,7 @@ RESOLVED_CREDENTIALS = Credentials(
|
|||
@pytest.fixture(autouse=True)
|
||||
def reset_bridge(monkeypatch):
|
||||
monkeypatch.setenv("LITELLM_RUST", "1")
|
||||
configuration.reset_rust_configuration()
|
||||
bridge.set_rust_chat_completions(
|
||||
chat_completions=None, achat_completions=None, decline=None
|
||||
)
|
||||
|
|
@ -61,20 +42,18 @@ def reset_bridge(monkeypatch):
|
|||
bridge.set_rust_chat_completions(
|
||||
chat_completions=None, achat_completions=None, decline=None
|
||||
)
|
||||
configuration.reset_rust_configuration()
|
||||
|
||||
|
||||
def _inject(*, decline_reason=None, error: Exception | None = None):
|
||||
def _inject():
|
||||
seen: dict[str, list[dict]] = {"gate": [], "call": []}
|
||||
|
||||
def gate(**kwargs):
|
||||
seen["gate"].append(kwargs)
|
||||
return decline_reason
|
||||
|
||||
def native(**kwargs):
|
||||
seen["call"].append(kwargs)
|
||||
if error is not None:
|
||||
raise error
|
||||
return dict(RUST_RESPONSE)
|
||||
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
|
||||
|
|
@ -106,206 +85,6 @@ def _run(*, credentials: Credentials | None = RESOLVED_CREDENTIALS, **overrides)
|
|||
return BedrockConverseLLM().completion(**_completion_kwargs(**overrides))
|
||||
|
||||
|
||||
def _recording_logging_obj():
|
||||
"""A logging object that keeps each hook's payload in a real list, so a test
|
||||
can assert which path logged and what it carried."""
|
||||
calls = {"pre_call": [], "post_call": []}
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.pre_call.side_effect = lambda **kwargs: calls["pre_call"].append(kwargs)
|
||||
logging_obj.post_call.side_effect = lambda **kwargs: calls["post_call"].append(kwargs)
|
||||
return logging_obj, calls
|
||||
|
||||
|
||||
def test_rust_true_serves_the_call_and_stamps_the_header():
|
||||
seen = _inject()
|
||||
response = _run()
|
||||
|
||||
assert response.choices[0].message.content == "hello from rust"
|
||||
assert response._hidden_params["additional_headers"] == {"x-litellm-rust": "true"}
|
||||
assert len(seen["call"]) == 1
|
||||
|
||||
|
||||
def test_the_core_receives_the_credentials_this_handler_already_resolved():
|
||||
"""Both paths must sign as the same principal, so the resolved credentials
|
||||
are handed down rather than re-derived from ambient AWS state."""
|
||||
seen = _inject()
|
||||
_run()
|
||||
|
||||
params = seen["call"][0]["optional_params"]
|
||||
assert params["aws_access_key_id"] == "AKIARESOLVED"
|
||||
assert params["aws_secret_access_key"] == "resolved-secret"
|
||||
assert params["aws_session_token"] == "resolved-token"
|
||||
assert params["aws_region_name"] == "us-east-1"
|
||||
|
||||
|
||||
def test_the_core_receives_the_converse_url_this_handler_already_built():
|
||||
seen = _inject()
|
||||
_run()
|
||||
|
||||
assert seen["call"][0]["api_base"].endswith(
|
||||
"/model/anthropic.claude-sonnet-4-5-v1%3A0/converse"
|
||||
)
|
||||
assert "bedrock-runtime.us-east-1.amazonaws.com" in seen["call"][0]["api_base"]
|
||||
|
||||
|
||||
def test_the_core_receives_the_untranslated_openai_messages():
|
||||
seen = _inject()
|
||||
_run(
|
||||
messages=[
|
||||
{"role": "system", "content": "be terse"},
|
||||
{"role": "user", "content": "hi"},
|
||||
]
|
||||
)
|
||||
assert seen["call"][0]["messages"] == [
|
||||
{"role": "system", "content": "be terse"},
|
||||
{"role": "user", "content": "hi"},
|
||||
]
|
||||
|
||||
|
||||
def test_without_the_opt_in_the_core_is_never_consulted(monkeypatch):
|
||||
monkeypatch.setenv("LITELLM_RUST", "0")
|
||||
seen = _inject()
|
||||
try:
|
||||
_run(litellm_params={})
|
||||
except Exception:
|
||||
# The Python path goes on to make an HTTP call; not reaching the gate
|
||||
# is the assertion, so a failure past this point is expected.
|
||||
pass
|
||||
assert seen["gate"] == []
|
||||
assert seen["call"] == []
|
||||
|
||||
|
||||
def test_streaming_stays_on_the_python_path():
|
||||
seen = _inject()
|
||||
try:
|
||||
_run(optional_params={"maxTokens": 16, "stream": True})
|
||||
except Exception:
|
||||
pass
|
||||
assert seen["gate"] == []
|
||||
|
||||
|
||||
def test_a_declined_request_never_reaches_the_native_call():
|
||||
seen = _inject(decline_reason="unrecognized request parameter")
|
||||
try:
|
||||
_run()
|
||||
except Exception:
|
||||
pass
|
||||
assert len(seen["gate"]) == 1
|
||||
assert seen["call"] == []
|
||||
|
||||
|
||||
def test_pre_call_logging_fires_exactly_once_on_the_rust_path():
|
||||
_inject()
|
||||
logging_obj = MagicMock()
|
||||
_run(logging_obj=logging_obj)
|
||||
assert logging_obj.pre_call.call_count == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_the_async_path_falls_back_when_the_core_declines(monkeypatch):
|
||||
class _Declined(Exception):
|
||||
pass
|
||||
|
||||
class _FakeNative:
|
||||
RustBridgeDeclined = _Declined
|
||||
RustUpstreamError = type("_Upstream", (Exception,), {})
|
||||
|
||||
monkeypatch.setattr(bridge, "get_native_bridge", lambda: _FakeNative())
|
||||
|
||||
async def declining_native(**_kwargs):
|
||||
raise _Declined("blank message text")
|
||||
|
||||
bridge.set_rust_chat_completions(
|
||||
decline=lambda **_kwargs: None, achat_completions=declining_native
|
||||
)
|
||||
|
||||
sentinel = object()
|
||||
|
||||
async def python_path(**_kwargs):
|
||||
return sentinel
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
BedrockConverseLLM, "get_credentials", return_value=RESOLVED_CREDENTIALS
|
||||
),
|
||||
patch.object(
|
||||
BedrockConverseLLM, "async_completion", side_effect=python_path
|
||||
) as python_call,
|
||||
):
|
||||
result = await BedrockConverseLLM().completion(
|
||||
**_completion_kwargs(acompletion=True)
|
||||
)
|
||||
|
||||
assert result is sentinel
|
||||
assert python_call.called, "a failing rust call must re-enter the python path"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_the_async_path_serves_the_rust_response_without_the_fallback():
|
||||
async def native(**_kwargs):
|
||||
return dict(RUST_RESPONSE)
|
||||
|
||||
bridge.set_rust_chat_completions(
|
||||
decline=lambda **_kwargs: None, achat_completions=native
|
||||
)
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
BedrockConverseLLM, "get_credentials", return_value=RESOLVED_CREDENTIALS
|
||||
),
|
||||
patch.object(BedrockConverseLLM, "async_completion") as python_call,
|
||||
):
|
||||
result = await BedrockConverseLLM().completion(
|
||||
**_completion_kwargs(acompletion=True)
|
||||
)
|
||||
|
||||
assert result.choices[0].message.content == "hello from rust"
|
||||
assert result._hidden_params["additional_headers"] == {"x-litellm-rust": "true"}
|
||||
assert not python_call.called
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pre_call_logging_fires_once_even_when_the_rust_path_declines():
|
||||
"""One request, one pre_call. Without the suppression the Python fallback
|
||||
logs a second one and non-idempotent callbacks run twice."""
|
||||
|
||||
class _Declined(Exception):
|
||||
pass
|
||||
|
||||
class _FakeNative:
|
||||
RustBridgeDeclined = _Declined
|
||||
RustUpstreamError = type("_Upstream", (Exception,), {})
|
||||
|
||||
async def declining_native(**_kwargs):
|
||||
raise _Declined("blank message text")
|
||||
|
||||
logging_obj = MagicMock()
|
||||
served = []
|
||||
|
||||
async def python_path(**kwargs):
|
||||
served.append(kwargs)
|
||||
return ModelResponse()
|
||||
|
||||
with (
|
||||
patch.object(bridge, "get_native_bridge", lambda: _FakeNative()),
|
||||
patch.object(
|
||||
BedrockConverseLLM, "get_credentials", return_value=RESOLVED_CREDENTIALS
|
||||
),
|
||||
patch.object(
|
||||
BedrockConverseLLM, "async_completion", side_effect=python_path
|
||||
),
|
||||
):
|
||||
bridge.set_rust_chat_completions(
|
||||
decline=lambda **_kwargs: None, achat_completions=declining_native
|
||||
)
|
||||
await BedrockConverseLLM().completion(
|
||||
**_completion_kwargs(acompletion=True, logging_obj=logging_obj)
|
||||
)
|
||||
|
||||
assert logging_obj.pre_call.call_count == 1
|
||||
assert served and served[0]["skip_pre_call_logging"] is True
|
||||
|
||||
|
||||
CONVERSE_RESPONSE = {
|
||||
"output": {"message": {"role": "assistant", "content": [{"text": "hi"}]}},
|
||||
"stopReason": "end_turn",
|
||||
|
|
@ -392,48 +171,20 @@ def _sync_client_returning_converse_response():
|
|||
return client
|
||||
|
||||
|
||||
def test_pre_call_logging_fires_once_when_the_sync_rust_path_declines():
|
||||
"""One request, one pre_call, on the synchronous path too.
|
||||
|
||||
The gate accepts and logs, then the native call declines before the
|
||||
provider is reached, so execution continues into the Python path below.
|
||||
That is the same attempt continuing; without the suppression it logs a
|
||||
second pre_call and non-idempotent callbacks run twice for one request.
|
||||
"""
|
||||
|
||||
class _Declined(Exception):
|
||||
pass
|
||||
|
||||
class _FakeNative:
|
||||
RustBridgeDeclined = _Declined
|
||||
RustUpstreamError = type("_Upstream", (Exception,), {})
|
||||
|
||||
def declining_native(**_kwargs):
|
||||
raise _Declined("blank message text")
|
||||
|
||||
logging_obj = MagicMock()
|
||||
|
||||
with patch.object(bridge, "get_native_bridge", lambda: _FakeNative()):
|
||||
bridge.set_rust_chat_completions(
|
||||
decline=lambda **_kwargs: None, chat_completions=declining_native
|
||||
)
|
||||
response = _run(
|
||||
logging_obj=logging_obj,
|
||||
client=_sync_client_returning_converse_response(),
|
||||
)
|
||||
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 logging_obj.pre_call.call_count == 1
|
||||
assert seen["gate"] == []
|
||||
assert seen["call"] == []
|
||||
|
||||
|
||||
def test_the_sync_python_path_still_logs_pre_call_without_the_opt_in(monkeypatch):
|
||||
"""The suppression must not swallow the log on a request the gate declined,
|
||||
so a deployment with no `rust` flag keeps exactly the log it always had."""
|
||||
monkeypatch.setenv("LITELLM_RUST", "0")
|
||||
def test_the_sync_python_path_logs_pre_call_once():
|
||||
_inject()
|
||||
logging_obj = MagicMock()
|
||||
response = _run(
|
||||
logging_obj=logging_obj,
|
||||
litellm_params={},
|
||||
client=_sync_client_returning_converse_response(),
|
||||
)
|
||||
|
||||
|
|
@ -441,83 +192,10 @@ def test_the_sync_python_path_still_logs_pre_call_without_the_opt_in(monkeypatch
|
|||
assert logging_obj.pre_call.call_count == 1
|
||||
|
||||
|
||||
def test_post_call_logging_fires_on_the_sync_rust_path():
|
||||
"""The Rust core owns the provider call, so the Converse transform that
|
||||
normally raises `post_call` never runs. Without the bridge hook every
|
||||
post_call callback goes silent and `original_response` stays unset."""
|
||||
import json
|
||||
|
||||
_inject()
|
||||
logging_obj = MagicMock()
|
||||
_run(logging_obj=logging_obj)
|
||||
|
||||
assert logging_obj.post_call.call_count == 1
|
||||
logged = logging_obj.post_call.call_args.kwargs["original_response"]
|
||||
assert json.loads(logged)["choices"][0]["message"]["content"] == "hello from rust"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_logging_fires_on_the_async_rust_path():
|
||||
"""The asynchronous path runs through the same hook, so the two paths
|
||||
cannot drift apart the way the pre_call suppression once did."""
|
||||
import json
|
||||
|
||||
async def native(**_kwargs):
|
||||
return dict(RUST_RESPONSE)
|
||||
|
||||
bridge.set_rust_chat_completions(
|
||||
decline=lambda **_kwargs: None, achat_completions=native
|
||||
)
|
||||
logging_obj = MagicMock()
|
||||
|
||||
with patch.object(
|
||||
BedrockConverseLLM, "get_credentials", return_value=RESOLVED_CREDENTIALS
|
||||
):
|
||||
await BedrockConverseLLM().completion(
|
||||
**_completion_kwargs(acompletion=True, logging_obj=logging_obj)
|
||||
)
|
||||
|
||||
assert logging_obj.post_call.call_count == 1
|
||||
logged = logging_obj.post_call.call_args.kwargs["original_response"]
|
||||
assert json.loads(logged)["choices"][0]["message"]["content"] == "hello from rust"
|
||||
|
||||
|
||||
def test_post_call_is_not_logged_twice_when_the_sync_rust_call_declines():
|
||||
"""A decline never reached the provider, so the Python path serves the
|
||||
request and owns the only post_call. Firing the hook there too would double
|
||||
every post_call callback for one request."""
|
||||
|
||||
class _Declined(Exception):
|
||||
pass
|
||||
|
||||
class _FakeNative:
|
||||
RustBridgeDeclined = _Declined
|
||||
RustUpstreamError = type("_Upstream", (Exception,), {})
|
||||
|
||||
def declining_native(**_kwargs):
|
||||
raise _Declined("blank message text")
|
||||
|
||||
logging_obj, calls = _recording_logging_obj()
|
||||
|
||||
with patch.object(bridge, "get_native_bridge", lambda: _FakeNative()):
|
||||
bridge.set_rust_chat_completions(
|
||||
decline=lambda **_kwargs: None, chat_completions=declining_native
|
||||
)
|
||||
response = _run(
|
||||
logging_obj=logging_obj,
|
||||
client=_sync_client_returning_converse_response(),
|
||||
)
|
||||
|
||||
assert response.choices[0].message.content == "hi"
|
||||
assert len(calls["post_call"]) == 1
|
||||
assert "hi" in calls["post_call"][0]["original_response"]
|
||||
|
||||
|
||||
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."""
|
||||
monkeypatch.setenv("LITELLM_RUST", "0")
|
||||
monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "bedrock-bearer-token")
|
||||
client = _sync_client_returning_converse_response()
|
||||
|
||||
|
|
@ -528,26 +206,11 @@ def test_bearer_token_auth_serves_when_boto3_resolves_no_sigv4_credentials(monke
|
|||
assert sent_headers["Authorization"] == "Bearer bedrock-bearer-token"
|
||||
|
||||
|
||||
def test_the_rust_opt_in_needs_no_sigv4_principal():
|
||||
"""The core resolves the bearer token itself, so a bearer-only deployment
|
||||
keeps its opt-in and the gate sees no aws_* credential keys to sign with."""
|
||||
seen = _inject()
|
||||
|
||||
response = _run(credentials=None, api_key="bedrock-bearer-token")
|
||||
|
||||
assert response.choices[0].message.content == "hello from rust"
|
||||
params = seen["call"][0]["optional_params"]
|
||||
assert not {"aws_access_key_id", "aws_secret_access_key", "aws_session_token"} & params.keys()
|
||||
assert params["aws_region_name"] == "us-east-1"
|
||||
assert seen["call"][0]["api_key"] == "bedrock-bearer-token"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("configured_through", ["env_var", "api_key"])
|
||||
def test_bearer_token_auth_never_runs_the_sigv4_credential_chain(monkeypatch, configured_through):
|
||||
"""The deployment's AWS profile does not exist, so resolving SigV4 credentials
|
||||
raises; a bearer-token deployment must still serve the request, since the
|
||||
bearer token alone signs it."""
|
||||
monkeypatch.setenv("LITELLM_RUST", "0")
|
||||
if configured_through == "env_var":
|
||||
monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "bedrock-bearer-token")
|
||||
else:
|
||||
|
|
@ -569,7 +232,6 @@ def test_bearer_token_auth_never_runs_the_sigv4_credential_chain(monkeypatch, co
|
|||
|
||||
def test_session_tags_sign_the_request_and_stay_out_of_the_body(monkeypatch):
|
||||
"""The tagged STS session signs the Converse call and the tags never reach the request body (#34069)."""
|
||||
monkeypatch.setenv("LITELLM_RUST", "0")
|
||||
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
|
||||
monkeypatch.delenv("AWS_WEB_IDENTITY_TOKEN_FILE", raising=False)
|
||||
monkeypatch.delenv("AWS_ROLE_ARN", raising=False)
|
||||
|
|
|
|||
|
|
@ -2912,19 +2912,13 @@ async def test_generic_http_handler_async_streaming_forwards_provider_response_h
|
|||
assert "".join([chunk.choices[0].delta.content or "" for chunk in collected]) == "hi"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"custom_llm_provider, enabled, expected",
|
||||
[("openai", True, True), ("openai", False, False), ("azure", True, False),
|
||||
("hosted_vllm", True, False), (None, True, False)],
|
||||
)
|
||||
def test_the_rust_responses_websocket_needs_openai_and_process_enablement(
|
||||
custom_llm_provider, enabled, expected, monkeypatch
|
||||
):
|
||||
@pytest.mark.parametrize("custom_llm_provider", ["openai", "azure", "hosted_vllm", None])
|
||||
def test_the_rust_responses_websocket_stays_on_python_with_the_switch_on(custom_llm_provider, monkeypatch):
|
||||
from litellm.rust_bridge import configuration
|
||||
|
||||
configuration.reset_rust_configuration()
|
||||
monkeypatch.setenv("LITELLM_RUST", "1" if enabled else "0")
|
||||
assert _rust_responses_websocket_enabled(custom_llm_provider) is expected
|
||||
monkeypatch.setenv("LITELLM_RUST", "1")
|
||||
assert _rust_responses_websocket_enabled(custom_llm_provider) is False
|
||||
|
||||
|
||||
def test_a_plain_callback_does_not_advertise_a_pre_call_deployment_hook(monkeypatch):
|
||||
|
|
|
|||
|
|
@ -24,23 +24,24 @@ def test_every_route_has_an_explicit_default_rule() -> None:
|
|||
(Context(Route.TRANSCRIPTION, provider="bedrock"), Rollout.RUST_REQUIRED),
|
||||
(Context(Route.TRANSCRIPTION, provider="openai"), Rollout.PYTHON_ONLY),
|
||||
(Context(Route.TRANSCRIPTION), Rollout.PYTHON_ONLY),
|
||||
(Context(Route.CHAT_COMPLETIONS, provider="anthropic"), Rollout.RUST_OPT_IN),
|
||||
(Context(Route.CHAT_COMPLETIONS, provider="bedrock"), Rollout.RUST_OPT_IN),
|
||||
(Context(Route.CHAT_COMPLETIONS, provider="anthropic", delivery=Delivery.STREAMING), Rollout.PYTHON_ONLY),
|
||||
(Context(Route.CHAT_COMPLETIONS, provider="openai"), Rollout.PYTHON_ONLY),
|
||||
(Context(Route.MESSAGES, provider="anthropic"), Rollout.RUST_OPT_IN),
|
||||
(Context(Route.MESSAGES, provider="azure_ai"), Rollout.RUST_OPT_IN),
|
||||
(Context(Route.MESSAGES, provider="bedrock"), Rollout.PYTHON_ONLY),
|
||||
(Context(Route.RESPONSES, provider="openai", delivery=Delivery.WEBSOCKET), Rollout.RUST_OPT_IN),
|
||||
(Context(Route.RESPONSES, provider="openai"), Rollout.PYTHON_ONLY),
|
||||
(Context(Route.RESPONSES, provider="azure", delivery=Delivery.WEBSOCKET), Rollout.PYTHON_ONLY),
|
||||
(Context(Route.EMBEDDING, provider="openai"), Rollout.PYTHON_ONLY),
|
||||
(Context(Route.CHAT_COMPLETIONS, provider="anthropic"), Rollout.PYTHON_ONLY),
|
||||
(Context(Route.CHAT_COMPLETIONS, provider="bedrock"), Rollout.PYTHON_ONLY),
|
||||
(Context(Route.MESSAGES, provider="anthropic"), Rollout.PYTHON_ONLY),
|
||||
(Context(Route.MESSAGES, provider="azure_ai"), Rollout.PYTHON_ONLY),
|
||||
(Context(Route.RESPONSES, provider="openai", delivery=Delivery.WEBSOCKET), Rollout.PYTHON_ONLY),
|
||||
),
|
||||
)
|
||||
def test_shipped_rules(context: Context, expected: Rollout) -> None:
|
||||
assert catalog.rollout(context) is expected
|
||||
|
||||
|
||||
def test_only_ocr_and_bedrock_transcription_can_reach_rust() -> None:
|
||||
rust_capable: Final = frozenset(
|
||||
(rule.route, rule.providers) for rule in catalog.RULES if rule.rollout is not Rollout.PYTHON_ONLY
|
||||
)
|
||||
assert rust_capable == frozenset({(Route.OCR, None), (Route.TRANSCRIPTION, frozenset({"bedrock"}))})
|
||||
|
||||
|
||||
def test_first_matching_rule_wins() -> None:
|
||||
rules: Final = (
|
||||
Rule(Route.EMBEDDING, Rollout.RUST_REQUIRED, providers=frozenset({"openai"}), models=frozenset({"m"})),
|
||||
|
|
|
|||
|
|
@ -9,7 +9,6 @@ from __future__ import annotations
|
|||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.rust_bridge import configuration
|
||||
from litellm.rust_bridge import chat_completions as bridge
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
|
@ -121,110 +120,16 @@ def _accepts(**overrides) -> bool:
|
|||
|
||||
|
||||
class TestGate:
|
||||
def test_declines_when_the_deployment_did_not_opt_in(self, monkeypatch):
|
||||
monkeypatch.delenv("LITELLM_RUST", raising=False)
|
||||
@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)
|
||||
assert _accepts(litellm_params={}) is False
|
||||
assert _accepts(litellm_params=None) is False
|
||||
assert gate.calls == [], "the gate must not be consulted before opt-in"
|
||||
|
||||
def test_accepts_when_the_deployment_opted_in_and_the_core_agrees(self, monkeypatch):
|
||||
monkeypatch.setenv("LITELLM_RUST", "1")
|
||||
gate = _RecordingDecline()
|
||||
bridge.set_rust_chat_completions(decline=gate)
|
||||
assert _accepts() is True
|
||||
assert gate.calls[0]["model"] == "claude-sonnet-4-5"
|
||||
assert gate.calls[0]["custom_llm_provider"] == "anthropic"
|
||||
|
||||
def test_process_enable_applies_without_request_override(self):
|
||||
bridge.set_rust_chat_completions(decline=_RecordingDecline())
|
||||
configuration.rust(True)
|
||||
|
||||
assert _accepts(litellm_params={}) is True
|
||||
|
||||
def test_the_env_var_opts_in_without_a_per_model_flag(self, monkeypatch):
|
||||
monkeypatch.setenv("LITELLM_RUST", "true")
|
||||
bridge.set_rust_chat_completions(decline=_RecordingDecline())
|
||||
assert _accepts(litellm_params={}) is True
|
||||
|
||||
def test_declines_streaming_and_providers_off_the_path(self, monkeypatch):
|
||||
monkeypatch.setenv("LITELLM_RUST", "1")
|
||||
gate = _RecordingDecline()
|
||||
bridge.set_rust_chat_completions(decline=gate)
|
||||
assert _accepts(stream=True) is False
|
||||
assert _accepts(custom_llm_provider="openai") is False
|
||||
assert _accepts(custom_llm_provider=None) is False
|
||||
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 test_declines_an_anthropic_request_carrying_a_litellm_metadata_user_id(self, monkeypatch):
|
||||
"""`AnthropicConfig.transform_request` copies a valid `user_id` into the Messages body.
|
||||
|
||||
It does that inside the function the Rust route replaces, and the core is
|
||||
handed `optional_params` only, so accepting here would send the request
|
||||
to Anthropic with the abuse-detection attribution silently missing.
|
||||
"""
|
||||
monkeypatch.setenv("LITELLM_RUST", "1")
|
||||
gate = _RecordingDecline()
|
||||
bridge.set_rust_chat_completions(decline=gate)
|
||||
assert _accepts(litellm_params={"metadata": {"user_id": "u-123"}}) is False
|
||||
assert gate.calls == [], "the core must not be consulted for a request it cannot see the key of"
|
||||
|
||||
# Bedrock's Converse transform reads no `user_id`, and an Anthropic request
|
||||
# whose metadata carries none is one Python would not attribute either.
|
||||
assert (
|
||||
_accepts(
|
||||
custom_llm_provider="bedrock",
|
||||
model="bedrock/us-east-1/anthropic.claude-v2",
|
||||
litellm_params={"metadata": {"user_id": "u-123"}},
|
||||
)
|
||||
is True
|
||||
)
|
||||
assert _accepts(litellm_params={"metadata": {"trace_id": "t-1"}}) is True
|
||||
assert _accepts(litellm_params={"metadata": {"user_id": None}}) is True
|
||||
assert _accepts(litellm_params={"metadata": None}) is True
|
||||
|
||||
def test_declines_a_bedrock_request_while_the_proxy_owns_request_metadata(self, monkeypatch):
|
||||
"""`AmazonConverseConfig` resolves proxy-owned `requestMetadata` onto the
|
||||
Converse body from `litellm_params`, and owning that field also means
|
||||
evicting a caller-supplied one. The core can do neither, so an operator
|
||||
who armed `bedrock_request_metadata_fields` keeps the Python path.
|
||||
"""
|
||||
monkeypatch.setenv("LITELLM_RUST", "1")
|
||||
gate = _RecordingDecline()
|
||||
bridge.set_rust_chat_completions(decline=gate)
|
||||
bedrock = {
|
||||
"custom_llm_provider": "bedrock",
|
||||
"model": "bedrock/us-east-1/anthropic.claude-v2",
|
||||
}
|
||||
|
||||
monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", ["user_api_key_team_id"])
|
||||
assert _accepts(**bedrock) is False
|
||||
assert gate.calls == [], "the core must not be consulted for a field it cannot write"
|
||||
assert _accepts() is True, "arming Bedrock attribution must not decline Anthropic"
|
||||
|
||||
monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", None)
|
||||
assert _accepts(**bedrock) is True, "the decline follows the operator's opt-in alone"
|
||||
|
||||
def test_declines_when_the_core_declines(self, monkeypatch):
|
||||
monkeypatch.setenv("LITELLM_RUST", "1")
|
||||
bridge.set_rust_chat_completions(decline=_RecordingDecline("streaming"))
|
||||
assert _accepts() is False
|
||||
|
||||
def test_declines_when_the_bridge_is_unavailable(self, monkeypatch):
|
||||
monkeypatch.setenv("LITELLM_RUST", "1")
|
||||
_hide_native_bridge(monkeypatch)
|
||||
assert _accepts() is False
|
||||
|
||||
def test_declines_when_the_gate_itself_raises(self, monkeypatch):
|
||||
monkeypatch.setenv("LITELLM_RUST", "1")
|
||||
|
||||
def exploding(**_kwargs):
|
||||
raise RuntimeError("boom")
|
||||
|
||||
bridge.set_rust_chat_completions(decline=exploding)
|
||||
assert _accepts() is False
|
||||
|
||||
|
||||
def _call_kwargs(model_response: ModelResponse) -> dict:
|
||||
return {
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue