This commit is contained in:
ishaan-jaff 2024-01-11 22:56:18 +05:30
parent d5a1458b8a
commit a876748bf5
2 changed files with 4 additions and 7 deletions

View file

@ -2,7 +2,7 @@ import json, copy, types
import os import os
from enum import Enum from enum import Enum
import time import time
from typing import Callable, Optional, Any from typing import Callable, Optional, Any, Union
import litellm import litellm
from litellm.utils import ModelResponse, get_secret, Usage from litellm.utils import ModelResponse, get_secret, Usage
from .prompt_templates.factory import prompt_factory, custom_prompt from .prompt_templates.factory import prompt_factory, custom_prompt
@ -714,7 +714,7 @@ def _embedding_func_single(
def embedding( def embedding(
model: str, model: str,
input: list, input: Union[list, str],
api_key: Optional[str] = None, api_key: Optional[str] = None,
logging_obj=None, logging_obj=None,
model_response=None, model_response=None,

View file

@ -189,10 +189,7 @@ def test_bedrock_embedding_titan():
litellm.set_verbose = True litellm.set_verbose = True
response = embedding( response = embedding(
model="amazon.titan-embed-text-v1", model="amazon.titan-embed-text-v1",
input=[ input="good morning from litellm, attempting to embed data",
"good morning from litellm, attempting to embed data",
"lets test a second string for good measure",
],
) )
print(f"response:", response) print(f"response:", response)
assert isinstance( assert isinstance(
@ -206,7 +203,7 @@ def test_bedrock_embedding_titan():
pytest.fail(f"Error occurred: {e}") pytest.fail(f"Error occurred: {e}")
# test_bedrock_embedding_titan() test_bedrock_embedding_titan()
def test_bedrock_embedding_cohere(): def test_bedrock_embedding_cohere():