mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-10 22:41:41 +00:00
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:
parent
998a372417
commit
970ea2949e
2 changed files with 183 additions and 8 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue