mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-16 23:41:43 +00:00
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:
parent
1371a045ba
commit
8d4e2e00b8
6 changed files with 136 additions and 20 deletions
|
|
@ -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,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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"));
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue