From e3efd65b067db5439064bfe80b2242e4d3e6257b Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 17 Jul 2026 03:23:04 +0000 Subject: [PATCH] feat(messages): route eligible Azure Anthropic streaming through Rust via buffered fake-stream Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> --- litellm/llms/custom_httpx/llm_http_handler.py | 33 +++++++++++++- .../test_litellm/rust_bridge/test_messages.py | 43 ++++++++++++++++++- 2 files changed, 72 insertions(+), 4 deletions(-) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index e58641a5d8d..133ce99818b 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -153,6 +153,9 @@ if TYPE_CHECKING: from aiohttp import ClientSession from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import ( + AnthropicMessagesStreamingResponse, + ) from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig from litellm.types.llms.openai_evals import ( CancelEvalResponse, @@ -2095,6 +2098,7 @@ class BaseLLMHTTPHandler: custom_llm_provider=custom_llm_provider, litellm_params=litellm_params, stream=stream or False, + rust_stream_eligible=bool(stream) and not self._has_agentic_completion_hook(logging_obj), model=model, api_key=api_key, api_base=api_base, @@ -2107,6 +2111,8 @@ class BaseLLMHTTPHandler: ), ) if rust_messages_response is not None: + if stream: + return self._rust_anthropic_messages_fake_stream(rust_messages_response) return await self._finalize_anthropic_messages_response( initial_response=rust_messages_response, model=model, @@ -2247,6 +2253,7 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params: GenericLiteLLMParams, stream: bool, + rust_stream_eligible: bool, model: str, api_key: str | None, api_base: str | None, @@ -2254,14 +2261,17 @@ class BaseLLMHTTPHandler: request_body: dict, timeout: float | httpx.Timeout | None, ) -> AnthropicMessagesResponse | None: - if stream or custom_llm_provider != "azure_ai" or litellm_params.get("rust") is not True: + if custom_llm_provider != "azure_ai" or litellm_params.get("rust") is not True: + return None + if stream and not rust_stream_eligible: return None from litellm.rust_bridge import messages as rust_messages_bridge + upstream_body = {key: value for key, value in request_body.items() if key != "stream"} rust_response = await rust_messages_bridge.amessages( model=model, - body=request_body, + body=upstream_body, api_key=api_key, api_base=api_base, custom_llm_provider=custom_llm_provider, @@ -2275,6 +2285,25 @@ class BaseLLMHTTPHandler: response_obj["_hidden_params"] = {"additional_headers": {"x-litellm-rust": "true"}} return response_obj + @staticmethod + def _rust_anthropic_messages_fake_stream( + rust_response: AnthropicMessagesResponse, + ) -> "AnthropicMessagesStreamingResponse": + from litellm.llms.anthropic.experimental_pass_through.messages.fake_stream_iterator import ( + FakeAnthropicMessagesStreamIterator, + ) + from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import ( + AnthropicMessagesStreamHiddenParams, + AnthropicMessagesStreamingResponse, + ) + + completion_stream = cast(AsyncIterator[bytes], FakeAnthropicMessagesStreamIterator(response=rust_response)) + hidden_params = AnthropicMessagesStreamHiddenParams(additional_headers={"x-litellm-rust": "true"}) + return AnthropicMessagesStreamingResponse( + completion_stream=completion_stream, + hidden_params=hidden_params, + ) + def anthropic_messages_handler( self, model: str, diff --git a/tests/test_litellm/rust_bridge/test_messages.py b/tests/test_litellm/rust_bridge/test_messages.py index f5be34ae0b6..f34b8c05819 100644 --- a/tests/test_litellm/rust_bridge/test_messages.py +++ b/tests/test_litellm/rust_bridge/test_messages.py @@ -1,12 +1,16 @@ """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") @@ -207,6 +211,7 @@ def _gate(**overrides): "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", @@ -271,16 +276,50 @@ async def test_gate_skips_rust_for_non_azure_provider(): @pytest.mark.asyncio -async def test_gate_skips_rust_when_streaming(): +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) + 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(