fix(vertex): decide rawPredict passthrough streaming from the request body

Vertex passthrough classified any target URL containing "stream" as a streaming
request. `:streamRawPredict` carries that substring, so a unary Claude-on-Vertex
call whose body omits `stream` was routed through the streaming logging path.
That path never consults the response content-type, so a complete
`"type": "message"` JSON body was handed to the Anthropic SSE chunk parser,
which recognises none of it; the spend log recorded 0 prompt tokens,
0 completion tokens and zero cost

Streaming for the rawPredict family now comes from the request body, which is
what the Anthropic Messages contract uses for those endpoints. The
generateContent family keeps its URL signal because the Gemini REST body has no
`stream` field, and `?alt=sse` is still appended for every request that is
classified as streaming, so Gemini framing and its usage parsing are unchanged

Both passthrough streaming predicates read `.get("stream")` off a body that is
only annotated as a dict; `_read_request_body` returns whatever the JSON parser
produced, so an array body raised AttributeError. The two predicates are now one
owner that answers False for any non-object body, which covers the vertex,
mistral, anthropic, vllm and azure passthrough routes
This commit is contained in:
Tin Chi Lo 2026-07-25 17:13:23 -07:00
parent 998a372417
commit 970ea2949e
2 changed files with 183 additions and 8 deletions

View file

@ -86,11 +86,16 @@ def is_passthrough_request_using_router_model(request_body: dict, llm_router: Op
return False
def is_passthrough_request_streaming(request_body: dict) -> bool:
def is_passthrough_request_streaming(request_body: object) -> bool:
"""
Returns True if the request is streaming
Returns True if the request is streaming.
A JSON body need not be an object, so a list or scalar can reach here; it
carries no streaming flag.
"""
return request_body.get("stream", False)
if not isinstance(request_body, dict):
return False
return bool(request_body.get("stream", False))
async def llm_passthrough_factory_proxy_route(
@ -551,8 +556,7 @@ async def is_streaming_request_fn(request: Request) -> bool:
_request_body = await get_form_data(request)
else:
_request_body = await _read_request_body(request)
if _request_body.get("stream"):
return True
return is_passthrough_request_streaming(_request_body)
return False
@ -1755,9 +1759,11 @@ async def _base_vertex_proxy_route(
## check for streaming
target = str(updated_url)
is_streaming_request = False
if "stream" in str(updated_url):
is_streaming_request = True
if ":rawPredict" in target or ":streamRawPredict" in target:
is_streaming_request = await is_streaming_request_fn(request)
else:
is_streaming_request = "stream" in target
if is_streaming_request:
target += "?alt=sse"
## CREATE PASS-THROUGH

View file

@ -3022,3 +3022,172 @@ class TestCursorProxyRoute:
assert call_args["target"] == "https://api.cursor.com/v0/agents"
assert result["id"] == "bc_abc123"
assert result["status"] == "CREATING"
class TestVertexRawPredictStreamingClassification:
"""
Regression coverage for LIT-4761.
`_base_vertex_proxy_route` classified any target URL containing "stream" as a
streaming request. `:streamRawPredict` carries that substring, so a unary
Anthropic-on-Vertex call (no `stream` field in the body) was sent with
`?alt=sse` and logged through the streaming chunk collector, which parses
Anthropic SSE deltas and finds no usage in a complete `"type": "message"`
body; the spend log recorded 0 tokens and $0 cost.
Streaming for the rawPredict family is decided by the request body, per the
Anthropic Messages contract. The Gemini generateContent family stays
URL-signalled because the Gemini REST body has no `stream` field.
"""
RAW_PREDICT_ENDPOINT = (
"v1/projects/test-project/locations/us-east5/publishers/anthropic/models/"
"claude-sonnet-4-6:streamRawPredict"
)
GENERATE_CONTENT_ENDPOINT = (
"v1/projects/test-project/locations/us-east5/publishers/google/models/"
"gemini-2.5-flash:streamGenerateContent"
)
async def _capture_passthrough_kwargs(self, endpoint: str, body: object) -> dict:
raw_body = json.dumps(body).encode("utf-8")
async def receive():
return {"type": "http.request", "body": raw_body, "more_body": False}
request = Request(
{
"type": "http",
"method": "POST",
"path": f"/vertex_ai/{endpoint}",
"headers": [(b"content-type", b"application/json")],
"query_string": b"",
},
receive=receive,
)
captured: dict = {}
def fake_create_pass_through_route(**kwargs):
captured.update(kwargs)
return AsyncMock(return_value={"status": "success"})
mock_credentials = Mock()
mock_credentials.token = "test-token"
base_url = "https://us-east5-aiplatform.googleapis.com/"
mock_handler = Mock()
mock_handler.get_default_base_target_url.return_value = base_url
mock_handler.update_base_target_url_with_credential_location = Mock(return_value=base_url)
module = "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints"
with (
mock.patch(
"litellm.llms.vertex_ai.vertex_llm_base.VertexBase.load_auth",
return_value=(mock_credentials, "test-project"),
),
mock.patch(f"{module}.create_pass_through_route", side_effect=fake_create_pass_through_route),
mock.patch(f"{module}.get_litellm_virtual_key", return_value="Bearer test-key"),
mock.patch(f"{module}.user_api_key_auth", new=AsyncMock(return_value={"api_key": "test-key"})),
mock.patch(f"{module}.get_vertex_pass_through_handler", return_value=mock_handler),
):
await vertex_proxy_route(
endpoint=endpoint,
request=request,
fastapi_response=Response(),
user_api_key_dict=UserAPIKeyAuth(token="test-key"),
)
assert captured, "create_pass_through_route was never called"
return captured
@pytest.mark.asyncio
async def test_raw_predict_without_stream_field_is_not_streaming(self):
captured = await self._capture_passthrough_kwargs(
endpoint=self.RAW_PREDICT_ENDPOINT,
body={
"anthropic_version": "vertex-2023-10-16",
"messages": [{"role": "user", "content": "Explain MLOps"}],
"max_tokens": 5000,
},
)
assert captured["is_streaming_request"] is False
assert "alt=sse" not in captured["target"]
@pytest.mark.asyncio
async def test_raw_predict_with_stream_false_is_not_streaming(self):
captured = await self._capture_passthrough_kwargs(
endpoint=self.RAW_PREDICT_ENDPOINT,
body={
"anthropic_version": "vertex-2023-10-16",
"stream": False,
"messages": [{"role": "user", "content": "Explain MLOps"}],
"max_tokens": 5000,
},
)
assert captured["is_streaming_request"] is False
assert "alt=sse" not in captured["target"]
@pytest.mark.asyncio
async def test_raw_predict_with_stream_true_still_streams(self):
captured = await self._capture_passthrough_kwargs(
endpoint=self.RAW_PREDICT_ENDPOINT,
body={
"anthropic_version": "vertex-2023-10-16",
"stream": True,
"messages": [{"role": "user", "content": "Explain MLOps"}],
"max_tokens": 5000,
},
)
assert captured["is_streaming_request"] is True
assert captured["target"].endswith("?alt=sse")
@pytest.mark.asyncio
async def test_gemini_stream_generate_content_stays_url_signalled(self):
captured = await self._capture_passthrough_kwargs(
endpoint=self.GENERATE_CONTENT_ENDPOINT,
body={"contents": [{"role": "user", "parts": [{"text": "Explain MLOps"}]}]},
)
assert captured["is_streaming_request"] is True
assert captured["target"].endswith("?alt=sse")
@pytest.mark.asyncio
async def test_raw_predict_with_non_object_body_is_not_streaming(self):
captured = await self._capture_passthrough_kwargs(
endpoint=self.RAW_PREDICT_ENDPOINT,
body=[{"role": "user", "content": "Explain MLOps"}],
)
assert captured["is_streaming_request"] is False
assert "alt=sse" not in captured["target"]
@pytest.mark.parametrize(
"request_body, expected",
[
({"stream": True}, True),
({"stream": "true"}, True),
({"stream": False}, False),
({}, False),
([{"role": "user"}], False),
([], False),
("stream", False),
(7, False),
(None, False),
],
)
def test_is_passthrough_request_streaming_tolerates_non_object_bodies(request_body, expected):
"""
A JSON request body is not required to be an object. Every passthrough
streaming decision funnels through this predicate, so a list or scalar body
must answer False instead of raising AttributeError.
"""
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
is_passthrough_request_streaming,
)
assert is_passthrough_request_streaming(request_body) is expected