feat(messages): route native Anthropic /messages through Rust behind LITELLM_USE_RUST_MESSAGES

Extends the opt-in Rust Anthropic Messages path (previously azure_ai only,
per-deployment rust:true) to the native anthropic provider and switches
enablement to the LITELLM_USE_RUST_MESSAGES env var, mirroring the OCR bridge
(LITELLM_USE_RUST_OCR). When enabled, eligible providers route through Rust;
unsupported providers fall back to the Python path.

Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com>
This commit is contained in:
Devin AI 2026-07-18 18:44:59 +00:00
parent 1371a045ba
commit 8d4e2e00b8
6 changed files with 136 additions and 20 deletions

View file

@ -1,5 +1,6 @@
use litellm_core::error::{json_type_name, CoreError};
use litellm_core::messages::transformation::AnthropicMessagesProviderConfig;
use litellm_core::providers::anthropic::messages::transformation::ANTHROPIC_MESSAGES_CONFIG;
use litellm_core::providers::azure_ai::messages::transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG;
use litellm_core::CoreResult;
use serde_json::{Map, Value};
@ -18,6 +19,7 @@ pub(super) fn messages_provider_config(
provider: &str,
) -> Option<&'static dyn AnthropicMessagesProviderConfig> {
match provider {
"anthropic" => Some(&ANTHROPIC_MESSAGES_CONFIG),
"azure_ai" => Some(&AZURE_ANTHROPIC_MESSAGES_CONFIG),
_ => None,
}

View file

@ -8,6 +8,7 @@ use tokio::net::{TcpListener, TcpStream};
use super::common_utils::{
has_header, messages_provider_config, string_headers, truncate_error_body,
};
use super::prepare::prepare_messages_call;
use super::{messages, MessagesRequest};
async fn read_http_request(socket: &mut TcpStream) -> String {
@ -52,12 +53,73 @@ fn write_response(body: &str) -> String {
}
#[test]
fn provider_config_only_resolves_azure_ai() {
fn provider_config_resolves_supported_providers() {
assert!(messages_provider_config("azure_ai").is_some());
assert!(messages_provider_config("anthropic").is_none());
assert!(messages_provider_config("anthropic").is_some());
assert!(messages_provider_config("openai").is_none());
}
#[test]
fn prepare_messages_call_resolves_native_anthropic() {
let prepared = prepare_messages_call(MessagesRequest {
model: "claude-opus-4-8",
body: json!({
"model": "claude-opus-4-8",
"max_tokens": 16,
"messages": [{
"role": "user",
"content": [{
"type": "text",
"text": "hi",
"cache_control": {"type": "ephemeral", "scope": "global"}
}]
}]
}),
api_key: Some("sk-ant-test"),
api_base: None,
custom_llm_provider: Some("anthropic"),
extra_headers: None,
timeout: None,
})
.expect("native Anthropic provider resolves");
assert_eq!(prepared.url, "https://api.anthropic.com/v1/messages");
assert!(prepared
.upstream_headers
.iter()
.any(|(name, value)| name == "x-api-key" && value == "sk-ant-test"));
assert!(prepared
.upstream_headers
.iter()
.any(|(name, value)| name == "anthropic-version" && value == "2023-06-01"));
assert!(prepared
.upstream_headers
.iter()
.any(|(name, value)| name == "content-type" && value == "application/json"));
assert_eq!(
prepared.body["messages"][0]["content"][0]["cache_control"],
json!({"type": "ephemeral", "scope": "global"})
);
}
#[test]
fn prepare_messages_call_rejects_unknown_provider() {
let result = prepare_messages_call(MessagesRequest {
model: "some-model",
body: json!({"model": "some-model", "max_tokens": 8, "messages": []}),
api_key: Some("sk-test"),
api_base: None,
custom_llm_provider: Some("openai"),
extra_headers: None,
timeout: None,
});
assert!(matches!(
result,
Err(CoreError::InvalidProvider(provider)) if provider == "openai"
));
}
#[test]
fn truncate_error_body_caps_long_payloads() {
let body = "x".repeat(400);
@ -248,12 +310,12 @@ async fn messages_rejects_unsupported_provider() {
body: json!({"model": "claude-3-5-sonnet", "max_tokens": 8, "messages": []}),
api_key: Some("sk"),
api_base: Some("http://127.0.0.1:1"),
custom_llm_provider: Some("anthropic"),
custom_llm_provider: Some("openai"),
extra_headers: None,
timeout: Some(Duration::from_millis(50)),
})
.await
.expect_err("unsupported provider errors");
assert!(matches!(err, CoreError::InvalidProvider(provider) if provider == "anthropic"));
assert!(matches!(err, CoreError::InvalidProvider(provider) if provider == "openai"));
}

View file

@ -2261,7 +2261,9 @@ class BaseLLMHTTPHandler:
request_body: dict,
timeout: float | httpx.Timeout | None,
) -> AnthropicMessagesResponse | None:
if custom_llm_provider != "azure_ai" or litellm_params.get("rust") is not True:
from litellm.rust_bridge.messages import rust_messages_enabled
if not rust_messages_enabled():
return None
if stream and not rust_stream_eligible:
return None

View file

@ -2,6 +2,7 @@
from __future__ import annotations
import os
from dataclasses import dataclass
from typing import Awaitable, Final, Protocol, Union, cast
@ -47,24 +48,41 @@ _UNSET: Final[_Unset] = _Unset()
@dataclass(slots=True)
class _RustMessagesState:
enabled: bool = False
messages: RustMessages | None = None
amessages: RustAmessages | None = None
_STATE: Final[_RustMessagesState] = _RustMessagesState()
def _env_enables_rust_messages() -> bool:
return os.getenv("LITELLM_USE_RUST_MESSAGES", "").strip().lower() in {
"1",
"true",
"yes",
"on",
}
_STATE: Final[_RustMessagesState] = _RustMessagesState(enabled=_env_enables_rust_messages())
def set_rust_messages(
*,
enabled: bool | _Unset = _UNSET,
messages: RustMessages | None | _Unset = _UNSET,
amessages: RustAmessages | None | _Unset = _UNSET,
) -> None:
if not isinstance(enabled, _Unset):
_STATE.enabled = enabled
if not isinstance(messages, _Unset):
_STATE.messages = messages
if not isinstance(amessages, _Unset):
_STATE.amessages = amessages
def rust_messages_enabled() -> bool:
return _STATE.enabled
def load_rust_messages() -> RustMessages | None:
if _STATE.messages is not None:
return _STATE.messages

View file

@ -81,16 +81,17 @@ def use_litellm_rust(
_rust_ocr_impl = ocr
if not isinstance(aocr, _Unset):
_rust_aocr_impl = aocr
if not configuring_messages:
return
from litellm.rust_bridge.messages import set_rust_messages
if not isinstance(messages, _Unset) and not isinstance(amessages, _Unset):
set_rust_messages(messages=messages, amessages=amessages)
elif not isinstance(messages, _Unset):
set_rust_messages(messages=messages)
else:
set_rust_messages(amessages=amessages)
if configuring_messages or not configuring_ocr:
if not isinstance(messages, _Unset) and not isinstance(amessages, _Unset):
set_rust_messages(enabled=enabled, messages=messages, amessages=amessages)
elif not isinstance(messages, _Unset):
set_rust_messages(enabled=enabled, messages=messages)
elif not isinstance(amessages, _Unset):
set_rust_messages(enabled=enabled, amessages=amessages)
else:
set_rust_messages(enabled=enabled)
def rust_ocr_enabled() -> bool:

View file

@ -107,6 +107,15 @@ class RaisingAsyncMessages:
raise RuntimeError("upstream request failed with status 400: bad request")
class NoneAsyncMessages:
def __init__(self) -> None:
self.calls = 0
async def __call__(self, **kwargs: object) -> dict[str, object] | None:
self.calls += 1
return None
@pytest.fixture(autouse=True)
def _reset_rust_flag():
litellm.use_litellm_rust(False, messages=None, amessages=None)
@ -251,6 +260,28 @@ async def test_gate_invokes_rust_and_marks_response_header():
assert call["timeout_seconds"] == 30.0
@pytest.mark.asyncio
async def test_gate_invokes_rust_for_native_anthropic_provider():
bridge = RecordingAsyncMessages()
litellm.use_litellm_rust(True, amessages=bridge)
response = await _gate(
custom_llm_provider="anthropic",
api_key="sk-ant-test",
api_base="https://api.anthropic.com",
headers={"anthropic-version": "2023-06-01"},
litellm_params=GenericLiteLLMParams(api_key="sk-ant-test", rust=True),
)
assert response is not None
assert response["_hidden_params"]["additional_headers"] == {"x-litellm-rust": "true"}
call = bridge.calls[0]
assert call["custom_llm_provider"] == "anthropic"
assert call["api_key"] == "sk-ant-test"
assert call["api_base"] == "https://api.anthropic.com"
assert call["extra_headers"] == {"anthropic-version": "2023-06-01"}
@pytest.mark.asyncio
async def test_gate_falls_back_to_python_when_bridge_raises():
bridge = RaisingAsyncMessages()
@ -265,7 +296,7 @@ async def test_gate_falls_back_to_python_when_bridge_raises():
@pytest.mark.asyncio
async def test_gate_skips_rust_when_flag_absent():
bridge = ExplodingAsyncMessages()
litellm.use_litellm_rust(True, amessages=bridge)
litellm.use_litellm_rust(False, amessages=bridge)
response = await _gate(litellm_params=GenericLiteLLMParams(api_key="sk-azure"))
@ -276,7 +307,7 @@ async def test_gate_skips_rust_when_flag_absent():
@pytest.mark.asyncio
async def test_gate_skips_rust_when_flag_false():
bridge = ExplodingAsyncMessages()
litellm.use_litellm_rust(True, amessages=bridge)
litellm.use_litellm_rust(False, amessages=bridge)
response = await _gate(litellm_params=GenericLiteLLMParams(api_key="sk-azure", rust=False))
@ -285,14 +316,14 @@ async def test_gate_skips_rust_when_flag_false():
@pytest.mark.asyncio
async def test_gate_skips_rust_for_non_azure_provider():
bridge = ExplodingAsyncMessages()
async def test_gate_skips_rust_for_non_listed_provider():
bridge = NoneAsyncMessages()
litellm.use_litellm_rust(True, amessages=bridge)
response = await _gate(custom_llm_provider="anthropic")
response = await _gate(custom_llm_provider="openai")
assert response is None
assert bridge.calls == 0
assert bridge.calls == 1
@pytest.mark.asyncio