fix(rust): route agentic-completion-hook /messages requests to Python for all stream modes (#34126)

This commit is contained in:
Yassin Kortam 2026-07-21 17:36:47 -07:00 committed by GitHub
parent e4343eb148
commit 1ea7db2111
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 7 additions and 11 deletions

View file

@ -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

View file

@ -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,
)