add import for co, anth

This commit is contained in:
ishaan-jaff 2023-08-08 10:46:20 -07:00
parent b6f601bb86
commit a4bebe6815
3 changed files with 11 additions and 7 deletions

BIN
.DS_Store vendored

Binary file not shown.

View file

@ -156,6 +156,7 @@ def completion(
elif "replicate" in model:
# import replicate/if it fails then pip install replicate
install_and_import("replicate")
import replicate
# replicate defaults to os.environ.get("REPLICATE_API_TOKEN")
# checking in case user set it to REPLICATE_API_KEY instead
if not get_secret("REPLICATE_API_TOKEN") and get_secret("REPLICATE_API_KEY"):
@ -246,6 +247,7 @@ def completion(
elif model in litellm.cohere_models:
# import cohere/if it fails then pip install cohere
install_and_import("cohere")
import cohere
if api_key:
cohere_key = api_key
elif litellm.cohere_key:

View file

@ -26,6 +26,14 @@ def test_completion_claude():
except Exception as e:
pytest.fail(f"Error occurred: {e}")
def test_completion_cohere():
try:
response = completion(model="command-nightly", messages=messages, max_tokens=500)
# Add any assertions here to check the response
print(response)
except Exception as e:
pytest.fail(f"Error occurred: {e}")
def test_completion_openai():
try:
response = completion(model="gpt-3.5-turbo", messages=messages)
@ -94,13 +102,7 @@ def test_completion_azure():
except Exception as e:
pytest.fail(f"Error occurred: {e}")
def test_completion_cohere():
try:
response = completion(model="command-nightly", messages=messages, max_tokens=500)
# Add any assertions here to check the response
print(response)
except Exception as e:
pytest.fail(f"Error occurred: {e}")
# Replicate API endpoints are unstable -> throw random CUDA errors -> this means our tests can fail even if our tests weren't incorrect.
# [TODO] improve our try-except block to handle for these