add import for co, anth

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:

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