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]