mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
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:
parent
cfe0100e1f
commit
8f1b6f8fea
5 changed files with 71 additions and 7 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
@ -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]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue