diff --git a/litellm/main.py b/litellm/main.py index 8cb133bf76b..a45a77fb44f 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -7196,6 +7196,10 @@ def embedding( or get_secret_str("CO_API_KEY") or litellm.api_key ) + # Honor COHERE_API_BASE like the rerank path (rerank_api/main.py) and + # the databricks/hosted_vllm embedding branches do; without this the + # gateway URL was silently ignored for embeddings. + api_base = api_base or litellm.api_base or get_secret_str("COHERE_API_BASE") response = base_llm_http_handler.embedding( model=model, 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 f5cb874c6ef..e8dda46e119 100644 --- a/tests/unit/llms/cohere/embed/test_cohere_embedding_handler.py +++ b/tests/unit/llms/cohere/embed/test_cohere_embedding_handler.py @@ -33,3 +33,42 @@ 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" + + +async def test_cohere_embedding_honors_cohere_api_base_env( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> 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) + monkeypatch.setenv("COHERE_API_BASE", gateway) + route: Final = respx_mock.post(gateway).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" + ) + + assert route.calls.call_count == 1 + assert response.data[0]["embedding"] == [0.1, 0.2] + + +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.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" + ) + + assert route.calls.call_count == 1 + assert response.data[0]["embedding"] == [0.1, 0.2]