diff --git a/litellm/llms/bedrock.py b/litellm/llms/bedrock.py index 99550280a57..617964a749b 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, @@ -737,18 +737,28 @@ def embedding( aws_region_name=aws_region_name, aws_bedrock_runtime_endpoint=aws_bedrock_runtime_endpoint, ) - - ## Embedding Call - embeddings = [ - _embedding_func_single( - model, - i, - optional_params=optional_params, - client=client, - logging_obj=logging_obj, - ) - for i in input - ] # [TODO]: make these parallel calls + if type(input) == str: + embeddings = [ + _embedding_func_single( + model, + input, + optional_params=optional_params, + client=client, + logging_obj=logging_obj, + ) + ] + else: + ## Embedding Call + embeddings = [ + _embedding_func_single( + model, + i, + optional_params=optional_params, + client=client, + logging_obj=logging_obj, + ) + for i in input + ] # [TODO]: make these parallel calls ## Populate OpenAI compliant dictionary embedding_response = [] diff --git a/litellm/tests/test_embedding.py b/litellm/tests/test_embedding.py index 954a53e2a83..6505d432dc1 100644 --- a/litellm/tests/test_embedding.py +++ b/litellm/tests/test_embedding.py @@ -186,13 +186,16 @@ def test_cohere_embedding3(): def test_bedrock_embedding_titan(): try: + # this tests if we support str input for bedrock embedding litellm.set_verbose = True + litellm.enable_cache() + import time + + current_time = str(time.time()) + # DO NOT MAKE THE INPUT A LIST in this test 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", - ], + model="bedrock/amazon.titan-embed-text-v1", + input=f"good morning from litellm, attempting to embed data {current_time}", # input should always be a string in this test ) print(f"response:", response) assert isinstance( @@ -202,11 +205,28 @@ def test_bedrock_embedding_titan(): assert all( isinstance(x, float) for x in response["data"][0]["embedding"] ), "Expected response to be a list of floats" + + # this also tests if we can return a cache response for this scenario + import time + + start_time = time.time() + + response = embedding( + model="bedrock/amazon.titan-embed-text-v1", + input=f"good morning from litellm, attempting to embed data {current_time}", # input should always be a string in this test + ) + print(response) + + end_time = time.time() + print(f"Embedding 2 response time: {end_time - start_time} seconds") + + assert end_time - start_time < 0.1 + litellm.disable_cache() except Exception as e: pytest.fail(f"Error occurred: {e}") -# test_bedrock_embedding_titan() +test_bedrock_embedding_titan() def test_bedrock_embedding_cohere(): @@ -280,7 +300,7 @@ def test_aembedding(): pytest.fail(f"Error occurred: {e}") -test_aembedding() +# test_aembedding() def test_aembedding_azure():