fix(embeddings): honor the caller's timeout on Cohere sync and HuggingFace embedding requests (#48593)

* fix(embeddings): honor the caller's timeout on Cohere sync and HuggingFace embedding requests

Bedrock Cohere embeddings ran the sync request with the client default instead of the timeout
passed to litellm.embedding, and HuggingFace embeddings dropped it on both the sync and async paths.
Pass the timeout through to the HTTP request on each path.

Fixes #47301
Fixes #47299

* fix(huggingface): pass the caller's timeout to the embedding model lookup too

Without input_type the handler first fetches the model's pipeline tag from huggingface.co, and that GET still used the default timeout. Thread the caller's timeout through both lookup helpers and the transform step so every request the embedding makes honors it.

Refs #47299

* fix(huggingface): keep the unreachable async lookup path untouched

The async model lookup helper and the call_type="async" transform branch are never reached, since the async embedding path transforms through the sync lookup. Leave them as they were so the change only touches requests that actually run.

litellm/llms/cohere/embed/handler.py

@@ -82,7 +82,7 @@ async def async_embedding(
         )
 
     try:
-        response: Final = await client.post(api_base, headers=headers, data=json.dumps(data))
+        response: Final = await client.post(api_base, headers=headers, data=json.dumps(data), timeout=timeout)
     except httpx.HTTPStatusError as e:
         ## LOGGING
         logging_obj.post_call(
@@ -166,7 +166,7 @@ def embedding(
     if client is None or not isinstance(client, HTTPHandler):
         client = HTTPHandler(concurrent_limit=1)
 
-    response: Final = client.post(embed_url, headers=headers, data=json.dumps(data))
+    response: Final = client.post(embed_url, headers=headers, data=json.dumps(data), timeout=timeout)
 
     return CohereEmbeddingConfig().transform_response(
         response=response,

litellm/llms/huggingface/embedding/handler.py

@@ -38,7 +38,9 @@ class _SupportsTokenEncode(Protocol):
     def encode(self, text: str) -> list[int]: ...
 
 
-def get_hf_task_embedding_for_model(model: str, task_type: str | None, api_base: str) -> object:
+def get_hf_task_embedding_for_model(
+    model: str, task_type: str | None, api_base: str, timeout: float | httpx.Timeout | None = None
+) -> object:
     if task_type is not None:
         if task_type in get_args(hf_tasks_embeddings):
             return task_type
@@ -46,7 +48,7 @@ def get_hf_task_embedding_for_model(model: str, task_type: str | None, api_base:
             raise Exception(f"Invalid task_type={task_type}. Expected one of={hf_tasks_embeddings}")
     http_client: Final = HTTPHandler(concurrent_limit=1)
 
-    model_info: Final = http_client.get(url=f"{api_base}/api/models/{model}")
+    model_info: Final = http_client.get(url=f"{api_base}/api/models/{model}", timeout=timeout)
 
     model_info_dict: Final = decoded_json_object(_DECODED_JSON.validate_python(model_info.json()))
 
@@ -148,6 +150,7 @@ class HuggingFaceEmbedding(BaseLLM):
         call_type: Literal["sync", "async"],
         optional_params: dict,
         embed_url: str,
+        timeout: float | httpx.Timeout | None = None,
     ) -> dict:
         data: dict = {}
 
@@ -165,7 +168,9 @@ class HuggingFaceEmbedding(BaseLLM):
             task_type: Final = optional_params.pop("input_type", None)
 
             if call_type == "sync":
-                hf_task: Final = get_hf_task_embedding_for_model(model=model, task_type=task_type, api_base=HF_HUB_URL)
+                hf_task: Final = get_hf_task_embedding_for_model(
+                    model=model, task_type=task_type, api_base=HF_HUB_URL, timeout=timeout
+                )
             elif call_type == "async":
                 return self._async_transform_input(model=model, task_type=task_type, embed_url=embed_url, input=input)
 
@@ -253,6 +258,7 @@ class HuggingFaceEmbedding(BaseLLM):
             call_type="sync",
             optional_params=optional_params,
             embed_url=api_base,
+            timeout=timeout,
         )
 
         ## LOGGING
@@ -269,9 +275,10 @@ class HuggingFaceEmbedding(BaseLLM):
         if client is None:
             client = get_async_httpx_client(
                 llm_provider=litellm.LlmProviders.HUGGINGFACE,
+                params={"timeout": timeout},
             )
 
-        response: Final = await client.post(api_base, headers=headers, data=json.dumps(data))
+        response: Final = await client.post(api_base, headers=headers, data=json.dumps(data), timeout=timeout)
 
         ## LOGGING
         logging_obj.post_call(
@@ -321,7 +328,9 @@ class HuggingFaceEmbedding(BaseLLM):
             litellm_params=litellm_params,
         )
         task_type: Final = optional_params.get("input_type", None)
-        task: Final = get_hf_task_embedding_for_model(model=model, task_type=task_type, api_base=HF_HUB_URL)
+        task: Final = get_hf_task_embedding_for_model(
+            model=model, task_type=task_type, api_base=HF_HUB_URL, timeout=timeout
+        )
         # print_verbose(f"{model}, {task}")
         embed_url = ""
         if model.startswith(("http://", "https://")):
@@ -359,6 +368,7 @@ class HuggingFaceEmbedding(BaseLLM):
             call_type="sync",
             optional_params=optional_params,
             embed_url=embed_url,
+            timeout=timeout,
         )
 
         ## LOGGING
@@ -374,7 +384,7 @@ class HuggingFaceEmbedding(BaseLLM):
         ## COMPLETION CALL
         if client is None or not isinstance(client, HTTPHandler):
             client = HTTPHandler(concurrent_limit=1)
-        response: Final = client.post(embed_url, headers=headers, data=json.dumps(data))
+        response: Final = client.post(embed_url, headers=headers, data=json.dumps(data), timeout=timeout)
 
         ## LOGGING
         logging_obj.post_call(

litellm/main.py

@@ -7400,6 +7400,7 @@ def embedding(
                 aembedding=aembedding,
                 litellm_params=litellm_params_dict,
                 headers=headers,
+                timeout=timeout,
             )
         elif custom_llm_provider == "bedrock":
             if isinstance(input, str):

tests/unit/llms/cohere/embed/test_cohere_embedding_handler.py

@@ -0,0 +1,35 @@
+from typing import Final
+
+import httpx
+import pytest
+import respx
+
+import litellm
+
+BEDROCK_COHERE_INVOKE_URL: Final = (
+    "https://bedrock-runtime.us-east-1.amazonaws.com/model/cohere.embed-english-v3/invoke"
+)
+COHERE_EMBED_RESPONSE: Final = {"embeddings": [[0.1, 0.2]], "meta": {"billed_units": {"input_tokens": 1}}}
+
+
+@pytest.mark.parametrize("sync_mode", [True, False])
+async def test_bedrock_cohere_embedding_sends_the_caller_timeout(
+    respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, sync_mode: bool
+) -> None:
+    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+    route: Final = respx_mock.post(BEDROCK_COHERE_INVOKE_URL).mock(
+        return_value=httpx.Response(200, json=COHERE_EMBED_RESPONSE)
+    )
+    request: Final = {
+        "model": "bedrock/cohere.embed-english-v3",
+        "input": ["hello"],
+        "timeout": 45,
+        "aws_region_name": "us-east-1",
+        "aws_access_key_id": "mockaws_access_key_id",
+        "aws_secret_access_key": "mockaws_secret_access_key",
+    }
+
+    response: Final = litellm.embedding(**request) if sync_mode else await litellm.aembedding(**request)
+
+    assert route.calls.last.request.extensions["timeout"]["read"] == 45
+    assert response.data[0]["embedding"] == [0.1, 0.2]

tests/unit/llms/huggingface/embedding/test_huggingface_embedding_handler.py

@@ -1,12 +1,12 @@
 import json
 from typing import Final
-from unittest.mock import patch, MagicMock, AsyncMock
+from unittest.mock import AsyncMock, MagicMock, patch
 
-
-import litellm
+import httpx
 import pytest
 import respx
 
+import litellm
 
 MOCK_EMBEDDING_RESPONSE = [[0.1, 0.2, 0.3, 0.4, 0.5]]
 
@@ -52,12 +52,10 @@ def mock_embedding_async_http_handler():
 class TestHuggingFaceEmbedding:
     @pytest.fixture(autouse=True)
     def setup(self, mock_embedding_http_handler, mock_embedding_async_http_handler, mock_hf_config_fetch):
-        self.mock_get_task_patcher = patch(
-            "litellm.llms.huggingface.embedding.handler.get_hf_task_embedding_for_model"
-        )
+        self.mock_get_task_patcher = patch("litellm.llms.huggingface.embedding.handler.get_hf_task_embedding_for_model")
         self.mock_get_task = self.mock_get_task_patcher.start()
 
-        def mock_get_task_side_effect(model, task_type, api_base):
+        def mock_get_task_side_effect(model, task_type, api_base, timeout=None):
             if task_type is not None:
                 return task_type
             return "sentence-similarity"
@@ -180,7 +178,9 @@ async def test_hf_embedding_sentence_sim_sends_one_similarity_request(
     monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
     monkeypatch.setenv("HUGGINGFACE_API_KEY", "hf-test-token")
     with respx.mock(assert_all_called=False) as hf:
-        hf.get(url__regex=r"https://huggingface\.co/api/models/.*").respond(json={"pipeline_tag": "sentence-similarity"})
+        hf.get(url__regex=r"https://huggingface\.co/api/models/.*").respond(
+            json={"pipeline_tag": "sentence-similarity"}
+        )
         hf.get(url__regex=r"https://huggingface\.co/.*/config\.json").respond(json={"max_position_embeddings": 512})
         inference: Final = hf.post(
             "https://router.huggingface.co/hf-inference/pipeline/sentence-similarity/sentence-transformers/TaylorAI/bge-micro-v2"
@@ -197,3 +197,33 @@ async def test_hf_embedding_sentence_sim_sends_one_similarity_request(
     }
     assert inference.calls.last.request.headers["Authorization"] == "Bearer hf-test-token"
     assert response.data == [{"object": "embedding", "index": 0, "embedding": 0.7708950042724609}]
+
+
+HF_MODEL_LOOKUP_URL: Final = "https://huggingface.co/api/models/BAAI/bge-m3"
+HF_EMBED_URL: Final = "https://router.huggingface.co/hf-inference/pipeline/feature-extraction/BAAI/bge-m3"
+
+
+@pytest.mark.respx(assert_all_called=False)
+@pytest.mark.parametrize("sync_mode", [True, False])
+async def test_huggingface_embedding_sends_the_caller_timeout(
+    respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, sync_mode: bool
+) -> None:
+    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+    respx_mock.get(url__regex=r"https://huggingface\.co/.*/config\.json").respond(json={"max_position_embeddings": 512})
+    lookup: Final = respx_mock.get(HF_MODEL_LOOKUP_URL).mock(
+        return_value=httpx.Response(200, json={"pipeline_tag": "feature-extraction"})
+    )
+    route: Final = respx_mock.post(HF_EMBED_URL).mock(return_value=httpx.Response(200, json=MOCK_EMBEDDING_RESPONSE))
+    request: Final = {
+        "model": "huggingface/BAAI/bge-m3",
+        "input": ["hello"],
+        "api_key": "hf_test",
+        "timeout": 45,
+    }
+
+    response: Final = litellm.embedding(**request) if sync_mode else await litellm.aembedding(**request)
+
+    assert all(call.request.extensions["timeout"]["read"] == 45 for call in lookup.calls)
+    assert lookup.called
+    assert route.calls.last.request.extensions["timeout"]["read"] == 45
+    assert response.data[0]["embedding"] == MOCK_EMBEDDING_RESPONSE[0]