phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

fix(bedrock.py): adding support for cohere embeddings

Browse source
This commit is contained in:
Krrish Dholakia 2023-12-06 13:24:49 -08:00
parent cf6ecc03a5
commit d962d5d4c0
4 changed files with 85 additions and 38 deletions
Download patch file Download diff file Expand all files Collapse all files

15
litellm/tests/test_embedding.py
View file

@ -151,7 +151,7 @@ def test_cohere_embedding3():
# test_cohere_embedding3()
def test_bedrock_embedding():
def test_bedrock_embedding_titan():
try:
response = embedding(
model="amazon.titan-embed-text-v1", input=["good morning from litellm, attempting to embed data",
@ -162,6 +162,19 @@ def test_bedrock_embedding():
pytest.fail(f"Error occurred: {e}")
# test_bedrock_embedding()
def test_bedrock_embedding_cohere():
try:
# litellm.set_verbose=True
response = embedding(
model="cohere.embed-multilingual-v3", input=["good morning from litellm, attempting to embed data", "lets test a second string for good measure"],
aws_region_name="os.environ/AWS_REGION_NAME_2"
)
# print(f"response:", response)
except Exception as e:
pytest.fail(f"Error occurred: {e}")
test_bedrock_embedding_cohere()
# comment out hf tests - since hf endpoints are unstable
def test_hf_embedding():
try: