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
This commit is contained in:
ryan-crabbe-berri 2026-10-27 10:03:45 -07:00
parent cfe0100e1f
commit 8f1b6f8fea
5 changed files with 71 additions and 7 deletions

View file

@ -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,

View file

@ -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(

View file

@ -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):

View file

@ -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]

View file

@ -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]