From 8d4e2e00b8df4291ac11bc563cee2b64746d0b15 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 18 Jul 2026 18:44:59 +0000 Subject: [PATCH] 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> --- .../ai-gateway/src/messages/common_utils.rs | 2 + .../crates/ai-gateway/src/messages/tests.rs | 70 +++++++++++++++++-- litellm/llms/custom_httpx/llm_http_handler.py | 4 +- litellm/rust_bridge/messages.py | 20 +++++- litellm/rust_bridge/ocr.py | 17 ++--- .../test_litellm/rust_bridge/test_messages.py | 43 ++++++++++-- 6 files changed, 136 insertions(+), 20 deletions(-) diff --git a/litellm-rust/crates/ai-gateway/src/messages/common_utils.rs b/litellm-rust/crates/ai-gateway/src/messages/common_utils.rs index fe4ac4cf26f..33894d0ee64 100644 --- a/litellm-rust/crates/ai-gateway/src/messages/common_utils.rs +++ b/litellm-rust/crates/ai-gateway/src/messages/common_utils.rs @@ -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, } diff --git a/litellm-rust/crates/ai-gateway/src/messages/tests.rs b/litellm-rust/crates/ai-gateway/src/messages/tests.rs index 30f6642400e..680fd5b5d58 100644 --- a/litellm-rust/crates/ai-gateway/src/messages/tests.rs +++ b/litellm-rust/crates/ai-gateway/src/messages/tests.rs @@ -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")); } diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 3e6f9ee08ee..1c979ed7054 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -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 diff --git a/litellm/rust_bridge/messages.py b/litellm/rust_bridge/messages.py index 5abb21879d3..690faa4dbd6 100644 --- a/litellm/rust_bridge/messages.py +++ b/litellm/rust_bridge/messages.py @@ -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 diff --git a/litellm/rust_bridge/ocr.py b/litellm/rust_bridge/ocr.py index 35de2eb9727..386189700ed 100644 --- a/litellm/rust_bridge/ocr.py +++ b/litellm/rust_bridge/ocr.py @@ -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: diff --git a/tests/test_litellm/rust_bridge/test_messages.py b/tests/test_litellm/rust_bridge/test_messages.py index 26ca7d27210..7d1b10e6bc3 100644 --- a/tests/test_litellm/rust_bridge/test_messages.py +++ b/tests/test_litellm/rust_bridge/test_messages.py @@ -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