diff --git a/litellm/llms/cohere/embed/transformation.py b/litellm/llms/cohere/embed/transformation.py index 1477e45997b..dd8e7015f16 100644 --- a/litellm/llms/cohere/embed/transformation.py +++ b/litellm/llms/cohere/embed/transformation.py @@ -21,6 +21,7 @@ from pydantic import BaseModel, ConfigDict, Field, TypeAdapter import litellm from litellm import COHERE_DEFAULT_EMBEDDING_INPUT_TYPE from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm import BaseEmbeddingConfig from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.types.llms.bedrock import ( @@ -205,7 +206,10 @@ class CohereEmbeddingConfig(BaseEmbeddingConfig): litellm_params: dict, stream: bool | None = None, ) -> str: - return api_base or "https://api.cohere.ai/v2/embed" + # ensure_api_route keeps the route-exists form working verbatim and + # joins the route onto a bare host, so a gateway host and a full + # endpoint are both usable as COHERE_API_BASE (see #50248). + return ensure_api_route(api_base or "https://api.cohere.ai/v2/embed", "v2/embed") def _transform_request( self, model: str, input: list[str], inference_params: dict diff --git a/tests/unit/llms/cohere/embed/test_cohere_embedding_handler.py b/tests/unit/llms/cohere/embed/test_cohere_embedding_handler.py index e8dda46e119..54e1c1170c1 100644 --- a/tests/unit/llms/cohere/embed/test_cohere_embedding_handler.py +++ b/tests/unit/llms/cohere/embed/test_cohere_embedding_handler.py @@ -34,6 +34,7 @@ async def test_bedrock_cohere_embedding_sends_the_caller_timeout( assert route.calls.last.request.extensions["timeout"]["read"] == 45 assert response.data[0]["embedding"] == [0.1, 0.2] + COHERE_DEFAULT_EMBED_URL: Final = "https://api.cohere.ai/v2/embed" @@ -42,16 +43,14 @@ async def test_cohere_embedding_honors_cohere_api_base_env( ) -> None: """`COHERE_API_BASE` must be honored for embeddings, like the rerank path and the databricks/hosted_vllm embedding branches already do (#49411 family).""" - gateway: Final = "https://gateway.example.com/v1/embed" - monkeypatch.setattr(litellm, 'disable_aiohttp_transport', True) + gateway: Final = "https://gateway.example.com" + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) monkeypatch.setenv("COHERE_API_BASE", gateway) - route: Final = respx_mock.post(gateway).mock( + route: Final = respx_mock.post(f"{gateway}/v2/embed").mock( return_value=httpx.Response(200, json=COHERE_EMBED_RESPONSE) ) - response: Final = await litellm.aembedding( - model="cohere/embed-english-v3.0", input=["hello"], api_key="mock-key" - ) + response: Final = await litellm.aembedding(model="cohere/embed-english-v3.0", input=["hello"], api_key="mock-key") assert route.calls.call_count == 1 assert response.data[0]["embedding"] == [0.1, 0.2] @@ -60,15 +59,13 @@ async def test_cohere_embedding_honors_cohere_api_base_env( async def test_cohere_embedding_keeps_the_default_url_without_the_env( respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch ) -> None: - monkeypatch.setattr(litellm, 'disable_aiohttp_transport', True) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) monkeypatch.delenv("COHERE_API_BASE", raising=False) route: Final = respx_mock.post(COHERE_DEFAULT_EMBED_URL).mock( return_value=httpx.Response(200, json=COHERE_EMBED_RESPONSE) ) - response: Final = await litellm.aembedding( - model="cohere/embed-english-v3.0", input=["hello"], api_key="mock-key" - ) + response: Final = await litellm.aembedding(model="cohere/embed-english-v3.0", input=["hello"], api_key="mock-key") assert route.calls.call_count == 1 assert response.data[0]["embedding"] == [0.1, 0.2] diff --git a/tests/unit/llms/cohere/embed/test_transformation.py b/tests/unit/llms/cohere/embed/test_transformation.py index abae30013fe..61060d8f6c3 100644 --- a/tests/unit/llms/cohere/embed/test_transformation.py +++ b/tests/unit/llms/cohere/embed/test_transformation.py @@ -443,3 +443,41 @@ async def test_cohere_embed_v4_image_embedding(respx_mock: respx.MockRouter, syn "input_type": "image", } assert response.data == [{"object": "embedding", "index": 0, "embedding": [0.125, -0.375]}] + + +# -------------------------------------------------------------------------------------- +# get_complete_url: COHERE_API_BASE routing semantics (#50248 follow-up) +# -------------------------------------------------------------------------------------- + + +def test_get_complete_url_appends_the_route_to_a_bare_host() -> None: + url: Final = CohereEmbeddingConfig().get_complete_url( + api_base="https://gateway.example.com", + api_key=None, + model="embed-english-v3.0", + optional_params={}, + litellm_params={}, + ) + assert url == "https://gateway.example.com/v2/embed" + + +def test_get_complete_url_keeps_a_full_endpoint_verbatim() -> None: + url: Final = CohereEmbeddingConfig().get_complete_url( + api_base="https://gateway.example.com/v2/embed", + api_key=None, + model="embed-english-v3.0", + optional_params={}, + litellm_params={}, + ) + assert url == "https://gateway.example.com/v2/embed" + + +def test_get_complete_url_default_is_unchanged() -> None: + url: Final = CohereEmbeddingConfig().get_complete_url( + api_base=None, + api_key=None, + model="embed-english-v3.0", + optional_params={}, + litellm_params={}, + ) + assert url == "https://api.cohere.ai/v2/embed"