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 fe4ac4cf26f0..33894d0ee642 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 30f6642400ea..680fd5b5d58e 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 3e6f9ee08eeb..1c979ed7054b 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -2261,7 +2261,9 @@ async def _maybe_rust_anthropic_messages( 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 5abb21879d3e..690faa4dbd6e 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 @@ class _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 35de2eb9727a..386189700eda 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/documentation_tests/test_env_keys.py b/tests/documentation_tests/test_env_keys.py index 60fbd505d674..fb6b7f639f2d 100644 --- a/tests/documentation_tests/test_env_keys.py +++ b/tests/documentation_tests/test_env_keys.py @@ -27,6 +27,7 @@ # Temporary/internal rollout flags are intentionally not added to the public # environment settings docs until the feature is ready for broad use. EXCLUDED_ROLLOUT_FLAGS = { + "LITELLM_USE_RUST_MESSAGES", "LITELLM_USE_RUST_OCR", } diff --git a/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py b/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py index 26ca7d272105..7d1b10e6bc38 100644 --- a/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py +++ b/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py @@ -107,6 +107,15 @@ async def __call__(self, **kwargs: object) -> dict[str, object]: 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 diff --git a/tests/test_litellm/rust_bridge/__init__.py b/tests/test_litellm/rust_bridge/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/test_litellm/rust_bridge/test_messages.py b/tests/test_litellm/rust_bridge/test_messages.py new file mode 100644 index 000000000000..7d1b10e6bc38 --- /dev/null +++ b/tests/test_litellm/rust_bridge/test_messages.py @@ -0,0 +1,385 @@ +"""Tests for the optional Rust-backed Anthropic Messages path.""" + +import importlib +from typing import cast + +import httpx +import pytest + +import litellm +from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler +from litellm.types.llms.anthropic_messages.anthropic_response import ( + AnthropicMessagesResponse, +) +from litellm.types.router import GenericLiteLLMParams + +rust_messages = importlib.import_module("litellm.rust_bridge.messages") +rust_bridge_loader = importlib.import_module("litellm.rust_bridge.loader") + +FAKE_MESSAGES_RESPONSE: dict[str, object] = { + "id": "msg_123", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5-20250929", + "content": [{"type": "text", "text": "hello world"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 5, "output_tokens": 3}, +} + +REQUEST_BODY: dict[str, object] = { + "model": "claude-sonnet-4-5", + "max_tokens": 64, + "messages": [{"role": "user", "content": "hi"}], +} + + +class RecordingMessages: + def __init__(self) -> None: + self.calls: list[dict[str, object]] = [] + + def __call__( + self, + model: str, + body: dict[str, object], + api_key: str | None, + api_base: str | None, + custom_llm_provider: str | None, + extra_headers: dict[str, object] | None, + timeout_seconds: float | None, + ) -> dict[str, object]: + self.calls.append( + { + "model": model, + "body": body, + "api_key": api_key, + "api_base": api_base, + "custom_llm_provider": custom_llm_provider, + "extra_headers": extra_headers, + "timeout_seconds": timeout_seconds, + } + ) + return dict(FAKE_MESSAGES_RESPONSE) + + +class RecordingAsyncMessages: + def __init__(self) -> None: + self.calls: list[dict[str, object]] = [] + + async def __call__( + self, + model: str, + body: dict[str, object], + api_key: str | None, + api_base: str | None, + custom_llm_provider: str | None, + extra_headers: dict[str, object] | None, + timeout_seconds: float | None, + ) -> dict[str, object]: + self.calls.append( + { + "model": model, + "body": body, + "api_key": api_key, + "api_base": api_base, + "custom_llm_provider": custom_llm_provider, + "extra_headers": extra_headers, + "timeout_seconds": timeout_seconds, + } + ) + return dict(FAKE_MESSAGES_RESPONSE) + + +class ExplodingAsyncMessages: + def __init__(self) -> None: + self.calls = 0 + + async def __call__(self, **kwargs: object) -> dict[str, object]: + self.calls += 1 + 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") + + +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) + rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL + yield + litellm.use_litellm_rust(False, messages=None, amessages=None) + rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL + + +def test_load_rust_messages_returns_injected_impl(): + bridge = RecordingMessages() + litellm.use_litellm_rust(True, messages=bridge) + assert rust_messages.load_rust_messages() is bridge + + +def test_configuring_messages_does_not_enable_ocr(): + from litellm.rust_bridge.ocr import rust_ocr_enabled + + litellm.use_litellm_rust(False) + assert rust_ocr_enabled() is False + + litellm.use_litellm_rust(True, messages=RecordingMessages()) + + assert rust_ocr_enabled() is False + + +def test_bare_use_litellm_rust_still_toggles_ocr(): + from litellm.rust_bridge.ocr import rust_ocr_enabled + + litellm.use_litellm_rust(True) + assert rust_ocr_enabled() is True + + litellm.use_litellm_rust(False) + assert rust_ocr_enabled() is False + + +def test_load_rust_amessages_returns_injected_impl(): + bridge = RecordingAsyncMessages() + litellm.use_litellm_rust(True, amessages=bridge) + assert rust_messages.load_rust_amessages() is bridge + + +def test_messages_wrapper_returns_none_when_bridge_absent(monkeypatch): + monkeypatch.setattr( + importlib.import_module("litellm.rust_bridge"), + "get_native_bridge", + lambda: None, + ) + litellm.use_litellm_rust(True) + assert rust_messages.load_rust_messages() is None + result = rust_messages.messages( + model="claude", + body=REQUEST_BODY, + api_key="k", + api_base="b", + custom_llm_provider="azure_ai", + extra_headers={}, + timeout=30.0, + ) + assert result is None + + +def test_messages_wrapper_forwards_args_and_converts_timeout(): + bridge = RecordingMessages() + litellm.use_litellm_rust(True, messages=bridge) + + response = rust_messages.messages( + model="claude-sonnet-4-5", + body=REQUEST_BODY, + api_key="sk-azure", + api_base="https://resource.services.ai.azure.com/anthropic", + custom_llm_provider="azure_ai", + extra_headers={"anthropic-beta": "token-efficient-tools-2025-02-19"}, + timeout=httpx.Timeout(600.0, read=42.0), + ) + + assert response == FAKE_MESSAGES_RESPONSE + assert bridge.calls[0] == { + "model": "claude-sonnet-4-5", + "body": REQUEST_BODY, + "api_key": "sk-azure", + "api_base": "https://resource.services.ai.azure.com/anthropic", + "custom_llm_provider": "azure_ai", + "extra_headers": {"anthropic-beta": "token-efficient-tools-2025-02-19"}, + "timeout_seconds": 42.0, + } + + +@pytest.mark.asyncio +async def test_amessages_wrapper_forwards_args(): + bridge = RecordingAsyncMessages() + litellm.use_litellm_rust(True, amessages=bridge) + + response = await rust_messages.amessages( + model="claude-sonnet-4-5", + body=REQUEST_BODY, + api_key="sk-azure", + api_base="https://resource.services.ai.azure.com/anthropic", + custom_llm_provider="azure_ai", + extra_headers=None, + timeout=12.5, + ) + + assert response == FAKE_MESSAGES_RESPONSE + assert bridge.calls[0]["model"] == "claude-sonnet-4-5" + assert bridge.calls[0]["timeout_seconds"] == 12.5 + + +def _gate(**overrides): + kwargs = { + "custom_llm_provider": "azure_ai", + "litellm_params": GenericLiteLLMParams(api_key="sk-azure", rust=True), + "stream": False, + "rust_stream_eligible": False, + "model": "claude-sonnet-4-5", + "api_key": "sk-azure", + "api_base": "https://resource.services.ai.azure.com/anthropic", + "headers": {"x-api-key": "sk-azure", "anthropic-version": "2023-06-01"}, + "request_body": dict(REQUEST_BODY), + "timeout": 30.0, + } + kwargs.update(overrides) + return BaseLLMHTTPHandler._maybe_rust_anthropic_messages(**kwargs) + + +@pytest.mark.asyncio +async def test_gate_invokes_rust_and_marks_response_header(): + bridge = RecordingAsyncMessages() + litellm.use_litellm_rust(True, 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_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() + litellm.use_litellm_rust(True, 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() + litellm.use_litellm_rust(False, 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_skips_rust_when_flag_false(): + bridge = ExplodingAsyncMessages() + litellm.use_litellm_rust(False, amessages=bridge) + + response = await _gate(litellm_params=GenericLiteLLMParams(api_key="sk-azure", rust=False)) + + assert response is None + assert bridge.calls == 0 + + +@pytest.mark.asyncio +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="openai") + + assert response is None + assert bridge.calls == 1 + + +@pytest.mark.asyncio +async def test_gate_skips_rust_when_streaming_but_not_eligible(): + bridge = ExplodingAsyncMessages() + litellm.use_litellm_rust(True, amessages=bridge) + + response = await _gate(stream=True, rust_stream_eligible=False) + + 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.use_litellm_rust(True, amessages=bridge) + + streaming_body = {**REQUEST_BODY, "stream": True} + response = await _gate( + stream=True, + rust_stream_eligible=True, + 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)) + stream = BaseLLMHTTPHandler._rust_anthropic_messages_fake_stream(response) + + assert stream._hidden_params["additional_headers"] == {"x-litellm-rust": "true"} + + chunks = [chunk async for chunk in stream] + joined = b"".join(chunks) + + assert b"event: message_start" in joined + 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.use_litellm_rust(True) + + response = await _gate() + + assert response is None