scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-12 08:42:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 207

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

Browse source 
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
This commit is contained in:
ryan-crabbe-berri 2026-10-27 10:16:20 -07:00
parent 8f1b6f8fea
commit e95ae2f065
2 changed files with 36 additions and 17 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

32
litellm/llms/huggingface/embedding/handler.py
View file

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

21
tests/unit/llms/huggingface/embedding/test_huggingface_embedding_handler.py
View file

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