diff --git a/litellm/llms/bedrock.py b/litellm/llms/bedrock.py index 99550280a5..dd9c86200d 100644 --- a/litellm/llms/bedrock.py +++ b/litellm/llms/bedrock.py @@ -2,7 +2,7 @@ import json, copy, types import os from enum import Enum import time -from typing import Callable, Optional, Any +from typing import Callable, Optional, Any, Union import litellm from litellm.utils import ModelResponse, get_secret, Usage from .prompt_templates.factory import prompt_factory, custom_prompt @@ -714,7 +714,7 @@ def _embedding_func_single( def embedding( model: str, - input: list, + input: Union[list, str], api_key: Optional[str] = None, logging_obj=None, model_response=None, diff --git a/litellm/tests/test_embedding.py b/litellm/tests/test_embedding.py index 954a53e2a8..50dd6ee06e 100644 --- a/litellm/tests/test_embedding.py +++ b/litellm/tests/test_embedding.py @@ -189,10 +189,7 @@ def test_bedrock_embedding_titan(): litellm.set_verbose = True response = embedding( model="amazon.titan-embed-text-v1", - input=[ - "good morning from litellm, attempting to embed data", - "lets test a second string for good measure", - ], + input="good morning from litellm, attempting to embed data", ) print(f"response:", response) assert isinstance( @@ -206,7 +203,7 @@ def test_bedrock_embedding_titan(): pytest.fail(f"Error occurred: {e}") -# test_bedrock_embedding_titan() +test_bedrock_embedding_titan() def test_bedrock_embedding_cohere():