diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index c48d75439a7..ec1301e5923 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -2107,8 +2107,7 @@ class BaseLLMHTTPHandler: rust_messages_response = await self._maybe_rust_anthropic_messages( 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), + has_agentic_hook=self._has_agentic_completion_hook(logging_obj), model=model, api_key=api_key, api_base=api_base, @@ -2266,8 +2265,7 @@ class BaseLLMHTTPHandler: *, custom_llm_provider: str, litellm_params: GenericLiteLLMParams, - stream: bool, - rust_stream_eligible: bool, + has_agentic_hook: bool, model: str, api_key: str | None, api_base: str | None, @@ -2279,7 +2277,7 @@ class BaseLLMHTTPHandler: return None if litellm_params.get("rust") is not True and not BaseLLMHTTPHandler._rust_env_enabled(): return None - if stream and not rust_stream_eligible: + if has_agentic_hook: return None from litellm.rust_bridge import messages as rust_messages_bridge 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 b745ca8eadf..fbd7e36e298 100644 --- a/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py +++ b/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py @@ -219,8 +219,7 @@ def _gate(**overrides): kwargs = { "custom_llm_provider": "azure_ai", "litellm_params": GenericLiteLLMParams(api_key="sk-azure", rust=True), - "stream": False, - "rust_stream_eligible": False, + "has_agentic_hook": False, "model": "claude-sonnet-4-5", "api_key": "sk-azure", "api_base": "https://resource.services.ai.azure.com/anthropic", @@ -345,11 +344,11 @@ async def test_gate_skips_rust_for_unsupported_provider(): @pytest.mark.asyncio -async def test_gate_skips_rust_when_streaming_but_not_eligible(): +async def test_gate_skips_rust_for_agentic_hook(): bridge = ExplodingAsyncMessages() litellm.use_litellm_rust(True, amessages=bridge) - response = await _gate(stream=True, rust_stream_eligible=False) + response = await _gate(has_agentic_hook=True) assert response is None assert bridge.calls == 0 @@ -362,8 +361,7 @@ async def test_gate_streams_through_rust_when_eligible_and_strips_stream_flag(): streaming_body = {**REQUEST_BODY, "stream": True} response = await _gate( - stream=True, - rust_stream_eligible=True, + has_agentic_hook=False, request_body=streaming_body, )