scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-08 22:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 78

Merge cda3d8bd33 into b03e913ccf

Browse source
This commit is contained in:
Vineeth Sai Varikuntla 2026-08-26 14:32:56 -04:00 • committed by GitHub
commit 1da3b6cce5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 47 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm/litellm_core_utils/llm_response_utils/get_api_base.py
View file

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

46
tests/test_litellm/litellm_core_utils/llm_response_utils/test_get_api_base.py Normal file
View file

@ -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}")