From a876748bf57f045061e547c965e4a32c983502ad Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Thu, 11 Jan 2024 22:56:18 +0530
Subject: [PATCH] v0

---
 litellm/llms/bedrock.py         | 4 ++--
 litellm/tests/test_embedding.py | 7 ++-----
 2 files changed, 4 insertions(+), 7 deletions(-)

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