diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 7e573de261b..28d2c62f1f1 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -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 diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index cf3351c4ff8..181846fe289 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -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