From 103adf68e32aeef2fb58cde73a919212d1815e53 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=83=91=E6=B6=A6=E6=B3=BD?= <506166760@qq.com>
Date: Fri, 6 Nov 2026 03:41:34 +0800
Subject: [PATCH] fix(cohere): honor COHERE_API_BASE for embeddings
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The cohere embedding branch in the dispatch passed the caller's api_base
straight through with no env fallback, so COHERE_API_BASE was silently
ignored for embeddings — the request went to the hardcoded
api.cohere.ai/v2/embed host. The rerank path (rerank_api/main.py) already
reads the same variable, and the databricks/hosted_vllm/openrouter
embedding branches follow the same convention; this branch was the
asymmetric one (same family as voyage's #49411).

Adds two respx regression tests: the env-routed gateway URL is used when
the variable is set, and the hardcoded default still applies when it is
not.

Signed-off-by: 郑润泽 <506166760@qq.com>
---
 litellm/main.py                               |  4 ++
 .../embed/test_cohere_embedding_handler.py    | 39 +++++++++++++++++++
 2 files changed, 43 insertions(+)

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]