From 8f1b6f8fea2bd7a321da461a06dea9aa6796303b Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Tue, 27 Oct 2026 10:03:45 -0700 Subject: [PATCH 1/3] 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 --- litellm/llms/cohere/embed/handler.py | 4 +-- litellm/llms/huggingface/embedding/handler.py | 5 +-- litellm/main.py | 1 + .../embed/test_cohere_embedding_handler.py | 35 +++++++++++++++++++ .../test_huggingface_embedding_handler.py | 33 +++++++++++++++-- 5 files changed, 71 insertions(+), 7 deletions(-) create mode 100644 tests/unit/llms/cohere/embed/test_cohere_embedding_handler.py diff --git a/litellm/llms/cohere/embed/handler.py b/litellm/llms/cohere/embed/handler.py index 2f60d94c08a..34fde444880 100644 --- a/litellm/llms/cohere/embed/handler.py +++ b/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, diff --git a/litellm/llms/huggingface/embedding/handler.py b/litellm/llms/huggingface/embedding/handler.py index a1d1a3f6a19..813615ff06f 100644 --- a/litellm/llms/huggingface/embedding/handler.py +++ b/litellm/llms/huggingface/embedding/handler.py @@ -269,9 +269,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( @@ -374,7 +375,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( diff --git a/litellm/main.py b/litellm/main.py index 23eac8ce0aa..d5d8d9d2d60 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -7118,6 +7118,7 @@ def embedding( aembedding=aembedding, litellm_params=litellm_params_dict, headers=headers, + timeout=timeout, ) elif custom_llm_provider == "bedrock": if isinstance(input, str): diff --git a/tests/unit/llms/cohere/embed/test_cohere_embedding_handler.py b/tests/unit/llms/cohere/embed/test_cohere_embedding_handler.py new file mode 100644 index 00000000000..f5cb874c6ef --- /dev/null +++ b/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] diff --git a/tests/unit/llms/huggingface/embedding/test_huggingface_embedding_handler.py b/tests/unit/llms/huggingface/embedding/test_huggingface_embedding_handler.py index ea1e04291ca..21a2b4519c2 100644 --- a/tests/unit/llms/huggingface/embedding/test_huggingface_embedding_handler.py +++ b/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]] @@ -197,3 +197,30 @@ 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_EMBED_URL: Final = "https://router.huggingface.co/hf-inference/pipeline/embed/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}) + 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"], + "input_type": "embed", + "api_key": "hf_test", + "timeout": 45, + } + + 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"] == MOCK_EMBEDDING_RESPONSE[0] From e95ae2f0652dc1ee4e72aa71613dcdc102a924db Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Tue, 27 Oct 2026 10:16:20 -0700 Subject: [PATCH 2/3] 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 --- litellm/llms/huggingface/embedding/handler.py | 32 ++++++++++++++----- .../test_huggingface_embedding_handler.py | 21 ++++++------ 2 files changed, 36 insertions(+), 17 deletions(-) diff --git a/litellm/llms/huggingface/embedding/handler.py b/litellm/llms/huggingface/embedding/handler.py index 813615ff06f..11fb74abcf2 100644 --- a/litellm/llms/huggingface/embedding/handler.py +++ b/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())) @@ -55,7 +57,9 @@ def get_hf_task_embedding_for_model(model: str, task_type: str | None, api_base: return pipeline_tag -async def async_get_hf_task_embedding_for_model(model: str, task_type: str | None, api_base: str) -> object: +async def async_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 @@ -65,7 +69,7 @@ async def async_get_hf_task_embedding_for_model(model: str, task_type: str | Non llm_provider=litellm.LlmProviders.HUGGINGFACE, ) - model_info: Final = await http_client.get(url=f"{api_base}/api/models/{model}") + model_info: Final = await 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())) @@ -107,8 +111,11 @@ class HuggingFaceEmbedding(BaseLLM): embed_url: str, input: list, optional_params: dict, + timeout: float | httpx.Timeout | None = None, ) -> dict: - hf_task = await async_get_hf_task_embedding_for_model(model=model, task_type=task_type, api_base=HF_HUB_URL) + hf_task = await async_get_hf_task_embedding_for_model( + model=model, task_type=task_type, api_base=HF_HUB_URL, timeout=timeout + ) data: Final = self._transform_input_on_pipeline_tag(input=input, pipeline_tag=hf_task) @@ -148,6 +155,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,9 +173,13 @@ 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) + return self._async_transform_input( + model=model, task_type=task_type, embed_url=embed_url, input=input, timeout=timeout + ) data = self._transform_input_on_pipeline_tag(input=input, pipeline_tag=hf_task) @@ -253,6 +265,7 @@ class HuggingFaceEmbedding(BaseLLM): call_type="sync", optional_params=optional_params, embed_url=api_base, + timeout=timeout, ) ## LOGGING @@ -322,7 +335,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://")): @@ -360,6 +375,7 @@ class HuggingFaceEmbedding(BaseLLM): call_type="sync", optional_params=optional_params, embed_url=embed_url, + timeout=timeout, ) ## LOGGING diff --git a/tests/unit/llms/huggingface/embedding/test_huggingface_embedding_handler.py b/tests/unit/llms/huggingface/embedding/test_huggingface_embedding_handler.py index 21a2b4519c2..7040fbe5fcc 100644 --- a/tests/unit/llms/huggingface/embedding/test_huggingface_embedding_handler.py +++ b/tests/unit/llms/huggingface/embedding/test_huggingface_embedding_handler.py @@ -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" @@ -199,7 +199,8 @@ async def test_hf_embedding_sentence_sim_sends_one_similarity_request( assert response.data == [{"object": "embedding", "index": 0, "embedding": 0.7708950042724609}] -HF_EMBED_URL: Final = "https://router.huggingface.co/hf-inference/pipeline/embed/BAAI/bge-m3" +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) @@ -209,18 +210,20 @@ async def test_huggingface_embedding_sends_the_caller_timeout( ) -> 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}) - route: Final = respx_mock.post(HF_EMBED_URL).mock( - return_value=httpx.Response(200, json=MOCK_EMBEDDING_RESPONSE) + 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"], - "input_type": "embed", "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] From 60bc9672399cd324d8818a87a1f0776f8d9d1d67 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Tue, 27 Oct 2026 10:40:15 -0700 Subject: [PATCH 3/3] 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/huggingface/embedding/handler.py | 15 ++++----------- 1 file changed, 4 insertions(+), 11 deletions(-) diff --git a/litellm/llms/huggingface/embedding/handler.py b/litellm/llms/huggingface/embedding/handler.py index 11fb74abcf2..5d997babc6e 100644 --- a/litellm/llms/huggingface/embedding/handler.py +++ b/litellm/llms/huggingface/embedding/handler.py @@ -57,9 +57,7 @@ def get_hf_task_embedding_for_model( return pipeline_tag -async def async_get_hf_task_embedding_for_model( - model: str, task_type: str | None, api_base: str, timeout: float | httpx.Timeout | None = None -) -> object: +async def async_get_hf_task_embedding_for_model(model: str, task_type: str | None, api_base: str) -> object: if task_type is not None: if task_type in get_args(hf_tasks_embeddings): return task_type @@ -69,7 +67,7 @@ async def async_get_hf_task_embedding_for_model( llm_provider=litellm.LlmProviders.HUGGINGFACE, ) - model_info: Final = await http_client.get(url=f"{api_base}/api/models/{model}", timeout=timeout) + model_info: Final = await http_client.get(url=f"{api_base}/api/models/{model}") model_info_dict: Final = decoded_json_object(_DECODED_JSON.validate_python(model_info.json())) @@ -111,11 +109,8 @@ class HuggingFaceEmbedding(BaseLLM): embed_url: str, input: list, optional_params: dict, - timeout: float | httpx.Timeout | None = None, ) -> dict: - hf_task = await async_get_hf_task_embedding_for_model( - model=model, task_type=task_type, api_base=HF_HUB_URL, timeout=timeout - ) + hf_task = await async_get_hf_task_embedding_for_model(model=model, task_type=task_type, api_base=HF_HUB_URL) data: Final = self._transform_input_on_pipeline_tag(input=input, pipeline_tag=hf_task) @@ -177,9 +172,7 @@ class HuggingFaceEmbedding(BaseLLM): 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, timeout=timeout - ) + return self._async_transform_input(model=model, task_type=task_type, embed_url=embed_url, input=input) data = self._transform_input_on_pipeline_tag(input=input, pipeline_tag=hf_task)