diff --git a/litellm/llms/cohere/rerank_v2/transformation.py b/litellm/llms/cohere/rerank_v2/transformation.py index 4c800d6455d..494f5e6fc8f 100644 --- a/litellm/llms/cohere/rerank_v2/transformation.py +++ b/litellm/llms/cohere/rerank_v2/transformation.py @@ -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" diff --git a/tests/test_litellm/llms/cohere/rerank/test_rerank_v2_transformation.py b/tests/test_litellm/llms/cohere/rerank/test_rerank_v2_transformation.py new file mode 100644 index 00000000000..4caf1354078 --- /dev/null +++ b/tests/test_litellm/llms/cohere/rerank/test_rerank_v2_transformation.py @@ -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 + )