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

add bedrock tests

Browse source
This commit is contained in:
ishaan-jaff 2023-09-14 14:15:32 -07:00
parent 2c0694bb2f
commit d121029a56
1 changed files with 30 additions and 14 deletions
Download patch file Download diff file Expand all files Collapse all files

44
litellm/tests/test_completion.py
View file

@ -634,20 +634,36 @@ def test_completion_sagemaker():
pytest.fail(f"Error occurred: {e}") pytest.fail(f"Error occurred: {e}")
# def test_completion_bedrock(): def test_completion_bedrock_titan():
# try: try:
# response = completion( response = completion(
# model="bedrock/amazon.titan-tg1-large", model="bedrock/amazon.titan-tg1-large",
# messages=messages, messages=messages,
# temperature=0.2, temperature=0.2,
# max_tokens=80, max_tokens=20,
# logger_fn=logger_fn logger_fn=logger_fn
# ) )
# # Add any assertions here to check the response # Add any assertions here to check the response
# print(response) print(response)
# except Exception as e: except Exception as e:
# pytest.fail(f"Error occurred: {e}") pytest.fail(f"Error occurred: {e}")
# test_completion_bedrock() test_completion_bedrock_titan()
def test_completion_bedrock_ai21():
try:
response = completion(
model="bedrock/ai21.j2-mid",
messages=messages,
temperature=0.2,
max_tokens=20,
logger_fn=logger_fn
)
# Add any assertions here to check the response
print(response)
except Exception as e:
pytest.fail(f"Error occurred: {e}")
test_completion_bedrock_ai21()
# test_completion_sagemaker() # test_completion_sagemaker()