diff --git a/litellm/llms/cohere/embed/handler.py b/litellm/llms/cohere/embed/handler.py index 2f60d94c08a..34fde444880 100644 --- a/litellm/llms/cohere/embed/handler.py +++ b/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, diff --git a/litellm/llms/huggingface/embedding/handler.py b/litellm/llms/huggingface/embedding/handler.py index a1d1a3f6a19..813615ff06f 100644 --- a/litellm/llms/huggingface/embedding/handler.py +++ b/litellm/llms/huggingface/embedding/handler.py @@ -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( diff --git a/litellm/main.py b/litellm/main.py index 23eac8ce0aa..d5d8d9d2d60 100644 --- a/litellm/main.py +++ b/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): diff --git a/tests/unit/llms/cohere/embed/test_cohere_embedding_handler.py b/tests/unit/llms/cohere/embed/test_cohere_embedding_handler.py new file mode 100644 index 00000000000..f5cb874c6ef --- /dev/null +++ b/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] 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 ea1e04291ca..21a2b4519c2 100644 --- a/tests/unit/llms/huggingface/embedding/test_huggingface_embedding_handler.py +++ b/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]] @@ -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]