diff --git a/.DS_Store b/.DS_Store index 17cee6a0a..220730748 100644 Binary files a/.DS_Store and b/.DS_Store differ diff --git a/litellm/main.py b/litellm/main.py index ea4d43a63..33f7676d3 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -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: diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py index 925483f32..52c9373cc 100644 --- a/litellm/tests/test_completion.py +++ b/litellm/tests/test_completion.py @@ -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