From a4bebe6815e28c97fab85b00587f010a3f253b88 Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 8 Aug 2023 10:46:20 -0700
Subject: [PATCH] add import for co, anth

---
 .DS_Store                        | Bin 6148 -> 6148 bytes
 litellm/main.py                  |   2 ++
 litellm/tests/test_completion.py |  16 +++++++++-------
 3 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/.DS_Store b/.DS_Store
index 17cee6a0acfa4d370d7d9b919baa2fb3a72b4c2f..22073074832468582ce431f5bcd4c2f80be08229 100644
GIT binary patch
delta 53
zcmZoMXfc@JFUrBdz`)4BAi&_6lb@WFlb;0S3v3qTSkBDIu{nT6m~rz`)+DBh4Q-p*
IIsWnk075|xYybcN

delta 182
zcmZoMXfc@JFUrNhz`)4BAi%(o!;s2Q!jQ?3&ycrSkYhQsD@aO+AsHx;%TRzU%@7Y%
zn95KLB=vwYm7Y2I$w@i+NkH{La~ReGvEF|$U|`tH!y?RRF9=kg%a8;lQ-Cm&A&(&)
Uzafm>D2B+erZR12=lIJH09xKHMgRZ+

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