Merge 60bc967239 into df9091deb6

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,

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()))
 
@@ -148,6 +150,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,7 +168,9 @@ 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)
 
@@ -253,6 +258,7 @@ class HuggingFaceEmbedding(BaseLLM):
             call_type="sync",
             optional_params=optional_params,
             embed_url=api_base,
+            timeout=timeout,
         )
 
         ## LOGGING
@@ -269,9 +275,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(
@@ -321,7 +328,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://")):
@@ -359,6 +368,7 @@ class HuggingFaceEmbedding(BaseLLM):
             call_type="sync",
             optional_params=optional_params,
             embed_url=embed_url,
+            timeout=timeout,
         )
 
         ## LOGGING
@@ -374,7 +384,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(

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

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]

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]]
 
@@ -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"
@@ -197,3 +197,33 @@ 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_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)
+@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})
+    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"],
+        "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]