From cda3d8bd33131ae96841ad3304cb406ede1e986c Mon Sep 17 00:00:00 2001 From: Vineeth Sai Date: Mon, 24 Aug 2026 17:27:23 -0700 Subject: [PATCH] fix(get_api_base): read stream from the parsed params, not the raw dict get_api_base accepts either a plain dict or a LiteLLM_Params, and parses the dict into _optional_params at the top. The stream lookup then read the original argument instead of the parsed one: stream: Final[bool] = getattr(optional_params, "stream", False) getattr on a dict does not see its keys, so a dict caller always resolved stream=False and got the non-streaming Gemini and Vertex URL back even when streaming was requested. A LiteLLM_Params caller was unaffected, which is why this survived. Read it from _optional_params so both accepted shapes agree. Review feedback: annotated the parametrized test parameters, per the fully typed convention in CLAUDE.md. --- .../llm_response_utils/get_api_base.py | 2 +- .../llm_response_utils/test_get_api_base.py | 46 +++++++++++++++++++ 2 files changed, 47 insertions(+), 1 deletion(-) create mode 100644 tests/test_litellm/litellm_core_utils/llm_response_utils/test_get_api_base.py diff --git a/litellm/litellm_core_utils/llm_response_utils/get_api_base.py b/litellm/litellm_core_utils/llm_response_utils/get_api_base.py index 26e79fa0ea8..d26d8fc538e 100644 --- a/litellm/litellm_core_utils/llm_response_utils/get_api_base.py +++ b/litellm/litellm_core_utils/llm_response_utils/get_api_base.py @@ -62,7 +62,7 @@ def get_api_base(model: str, optional_params: dict | LiteLLM_Params) -> str | No if dynamic_api_base is not None: return dynamic_api_base - stream: Final[bool] = getattr(optional_params, "stream", False) + stream: Final[bool] = getattr(_optional_params, "stream", False) if _optional_params.vertex_location is not None and _optional_params.vertex_project is not None: from litellm.llms.vertex_ai.vertex_llm_base import VertexBase diff --git a/tests/test_litellm/litellm_core_utils/llm_response_utils/test_get_api_base.py b/tests/test_litellm/litellm_core_utils/llm_response_utils/test_get_api_base.py new file mode 100644 index 00000000000..40ea80f29b0 --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/llm_response_utils/test_get_api_base.py @@ -0,0 +1,46 @@ +""" +Tests for litellm.litellm_core_utils.llm_response_utils.get_api_base + +`optional_params` is documented as accepting either a plain dict or a +LiteLLM_Params object, so both shapes must resolve `stream` the same way and +report the same streaming vs non-streaming Gemini / Vertex URL. +""" + +import pytest + +from litellm.litellm_core_utils.llm_response_utils.get_api_base import get_api_base +from litellm.types.router import LiteLLM_Params + +GEMINI_STREAM = "https://generativelanguage.googleapis.com/v1beta/models/gemini-pro:streamGenerateContent" +GEMINI_NON_STREAM = "https://generativelanguage.googleapis.com/v1beta/models/gemini-pro:generateContent" + + +@pytest.mark.parametrize( + "optional_params, expected_api_base", + [ + ({"stream": True}, GEMINI_STREAM), + ({"stream": False}, GEMINI_NON_STREAM), + ({}, GEMINI_NON_STREAM), + (LiteLLM_Params(model="gemini/gemini-pro", stream=True), GEMINI_STREAM), + (LiteLLM_Params(model="gemini/gemini-pro"), GEMINI_NON_STREAM), + ], +) +def test_get_api_base_gemini_stream( + optional_params: dict[str, bool] | LiteLLM_Params, expected_api_base: str +) -> None: + assert get_api_base(model="gemini/gemini-pro", optional_params=optional_params) == expected_api_base + + +@pytest.mark.parametrize("stream, suffix", [(True, "streamGenerateContent"), (False, "generateContent")]) +def test_get_api_base_vertex_stream_from_dict(stream: bool, suffix: str) -> None: + api_base = get_api_base( + model="gemini-pro", + optional_params={ + "stream": stream, + "vertex_location": "us-central1", + "vertex_project": "my-project", + }, + ) + + assert api_base is not None + assert api_base.endswith(f":{suffix}")