scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

fix(cohere): normalize rerank v2 api base

Browse source
This commit is contained in:
VectorPeak 2026-06-24 10:58:13 +08:00
parent fda08dd727
commit be8acac33f
2 changed files with 23 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
litellm/llms/cohere/rerank_v2/transformation.py
View file

@ -19,9 +19,10 @@ class CohereRerankV2Config(CohereRerankConfig):
optional_params: Optional[dict] = None,
) -> str:
if api_base:
# Remove trailing slashes and ensure clean base URL
api_base = api_base.rstrip("/")
if not api_base.endswith("/v2/rerank"):
if api_base.endswith("/v2"):
api_base = f"{api_base}/rerank"
elif not api_base.endswith("/v2/rerank"):
api_base = f"{api_base}/v2/rerank"
return api_base
return "https://api.cohere.ai/v2/rerank"

20
tests/test_litellm/llms/cohere/rerank/test_rerank_v2_transformation.py Normal file
View file

@ -0,0 +1,20 @@
import pytest
from litellm.llms.cohere.rerank_v2.transformation import CohereRerankV2Config
@pytest.mark.parametrize(
("api_base", "expected_url"),
[
("https://api.cohere.ai", "https://api.cohere.ai/v2/rerank"),
("https://api.cohere.ai/v2", "https://api.cohere.ai/v2/rerank"),
("https://api.cohere.ai/v2/rerank/", "https://api.cohere.ai/v2/rerank"),
],
)
def test_get_complete_url_normalizes_cohere_rerank_v2_api_base(
api_base: str, expected_url: str
) -> None:
assert (
CohereRerankV2Config().get_complete_url(api_base=api_base, model="rerank-v3.5")
== expected_url
)