diff --git a/litellm/llms/huggingface_restapi.py b/litellm/llms/huggingface_restapi.py index 2910a644e59..80fd9592f2b 100644 --- a/litellm/llms/huggingface_restapi.py +++ b/litellm/llms/huggingface_restapi.py @@ -838,7 +838,12 @@ class Huggingface(BaseLLM): return {"inputs": input} # default to feature-extraction pipeline tag async def _async_transform_input( - self, model: str, task_type: Optional[str], embed_url: str, input: List + self, + model: str, + task_type: Optional[str], + embed_url: str, + input: List, + optional_params: dict, ) -> dict: hf_task = await async_get_hf_task_embedding_for_model( model=model, task_type=task_type, api_base=embed_url @@ -846,6 +851,9 @@ class Huggingface(BaseLLM): data = self._transform_input_on_pipeline_tag(input=input, pipeline_tag=hf_task) + if len(optional_params.keys()) > 0: + data["options"] = optional_params + return data def _transform_input( @@ -856,6 +864,7 @@ class Huggingface(BaseLLM): optional_params: dict, embed_url: str, ) -> dict: + data: Dict = {} ## TRANSFORMATION ## if "sentence-transformers" in model: if len(input) == 0: @@ -865,7 +874,7 @@ class Huggingface(BaseLLM): ) data = {"inputs": {"source_sentence": input[0], "sentences": input[1:]}} else: - data = {"inputs": input} # type: ignore + data = {"inputs": input} task_type = optional_params.pop("input_type", None) @@ -882,6 +891,9 @@ class Huggingface(BaseLLM): input=input, pipeline_tag=hf_task ) + if len(optional_params.keys()) > 0: + data["options"] = optional_params + return data def _process_embedding_response( diff --git a/litellm/tests/test_embedding.py b/litellm/tests/test_embedding.py index 62487b48865..2d2e125eb8a 100644 --- a/litellm/tests/test_embedding.py +++ b/litellm/tests/test_embedding.py @@ -1,3 +1,4 @@ +import json import os import sys import traceback @@ -11,7 +12,7 @@ load_dotenv() sys.path.insert( 0, os.path.abspath("../..") ) # Adds the parent directory to the system path -from unittest.mock import MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import litellm from litellm import completion, completion_cost, embedding @@ -740,3 +741,43 @@ async def test_databricks_embeddings(sync_mode): # print(response) # local_proxy_embeddings() + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_hf_embedddings_with_optional_params(sync_mode): + litellm.set_verbose = True + + if sync_mode: + client = HTTPHandler(concurrent_limit=1) + mock_obj = MagicMock() + else: + client = AsyncHTTPHandler(concurrent_limit=1) + mock_obj = AsyncMock() + + with patch.object(client, "post", new=mock_obj) as mock_client: + try: + if sync_mode: + response = embedding( + model="huggingface/jinaai/jina-embeddings-v2-small-en", + input=["good morning from litellm"], + wait_for_model=True, + client=client, + ) + else: + response = await litellm.aembedding( + model="huggingface/jinaai/jina-embeddings-v2-small-en", + input=["good morning from litellm"], + wait_for_model=True, + client=client, + ) + except Exception: + pass + + mock_client.assert_called_once() + + print(f"mock_client.call_args.kwargs: {mock_client.call_args.kwargs}") + assert "options" in mock_client.call_args.kwargs["data"] + json_data = json.loads(mock_client.call_args.kwargs["data"]) + assert "wait_for_model" in json_data["options"] + assert json_data["options"]["wait_for_model"] is True diff --git a/litellm/tests/test_optional_params.py b/litellm/tests/test_optional_params.py index d961190c29f..ad640fe295b 100644 --- a/litellm/tests/test_optional_params.py +++ b/litellm/tests/test_optional_params.py @@ -438,3 +438,11 @@ def test_get_optional_params_image_gen(): print(response) assert "aws_region_name" in response + + +def test_bedrock_optional_params_embeddings_provider_specific_params(): + optional_params = get_optional_params_embeddings( + custom_llm_provider="huggingface", + wait_for_model=True, + ) + assert len(optional_params) == 1