From c903b448498791480ea1ca9730cb5f2cf9f598a8 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Thu, 18 Apr 2024 21:27:22 -0700
Subject: [PATCH] (ci/cd) vertexai - skip 429 errors

---
 .../tests/test_amazing_vertex_completion.py   | 113 ++++++++++--------
 1 file changed, 64 insertions(+), 49 deletions(-)

diff --git a/litellm/tests/test_amazing_vertex_completion.py b/litellm/tests/test_amazing_vertex_completion.py
index a7296f03eb2..a27edd06f4d 100644
--- a/litellm/tests/test_amazing_vertex_completion.py
+++ b/litellm/tests/test_amazing_vertex_completion.py
@@ -145,30 +145,35 @@ def test_vertex_ai_anthropic():
 #     reason="Local test. Vertex AI Quota is low. Leads to rate limit errors on ci/cd."
 # )
 def test_vertex_ai_anthropic_streaming():
-    # load_vertex_ai_credentials()
+    try:
+        # load_vertex_ai_credentials()
 
-    # litellm.set_verbose = True
+        # litellm.set_verbose = True
 
-    model = "claude-3-sonnet@20240229"
+        model = "claude-3-sonnet@20240229"
 
-    vertex_ai_project = "adroit-crow-413218"
-    vertex_ai_location = "asia-southeast1"
-    json_obj = get_vertex_ai_creds_json()
-    vertex_credentials = json.dumps(json_obj)
+        vertex_ai_project = "adroit-crow-413218"
+        vertex_ai_location = "asia-southeast1"
+        json_obj = get_vertex_ai_creds_json()
+        vertex_credentials = json.dumps(json_obj)
 
-    response = completion(
-        model="vertex_ai/" + model,
-        messages=[{"role": "user", "content": "hi"}],
-        temperature=0.7,
-        vertex_ai_project=vertex_ai_project,
-        vertex_ai_location=vertex_ai_location,
-        stream=True,
-    )
-    # print("\nModel Response", response)
-    for chunk in response:
-        print(f"chunk: {chunk}")
+        response = completion(
+            model="vertex_ai/" + model,
+            messages=[{"role": "user", "content": "hi"}],
+            temperature=0.7,
+            vertex_ai_project=vertex_ai_project,
+            vertex_ai_location=vertex_ai_location,
+            stream=True,
+        )
+        # print("\nModel Response", response)
+        for chunk in response:
+            print(f"chunk: {chunk}")
 
     # raise Exception("it worked!")
+    except litellm.RateLimitError as e:
+        pass
+    except Exception as e:
+        pytest.fail(f"Error occurred: {e}")
 
 
 # test_vertex_ai_anthropic_streaming()
@@ -180,23 +185,28 @@ def test_vertex_ai_anthropic_streaming():
 @pytest.mark.asyncio
 async def test_vertex_ai_anthropic_async():
     # load_vertex_ai_credentials()
+    try:
 
-    model = "claude-3-sonnet@20240229"
+        model = "claude-3-sonnet@20240229"
 
-    vertex_ai_project = "adroit-crow-413218"
-    vertex_ai_location = "asia-southeast1"
-    json_obj = get_vertex_ai_creds_json()
-    vertex_credentials = json.dumps(json_obj)
+        vertex_ai_project = "adroit-crow-413218"
+        vertex_ai_location = "asia-southeast1"
+        json_obj = get_vertex_ai_creds_json()
+        vertex_credentials = json.dumps(json_obj)
 
-    response = await acompletion(
-        model="vertex_ai/" + model,
-        messages=[{"role": "user", "content": "hi"}],
-        temperature=0.7,
-        vertex_ai_project=vertex_ai_project,
-        vertex_ai_location=vertex_ai_location,
-        vertex_credentials=vertex_credentials,
-    )
-    print(f"Model Response: {response}")
+        response = await acompletion(
+            model="vertex_ai/" + model,
+            messages=[{"role": "user", "content": "hi"}],
+            temperature=0.7,
+            vertex_ai_project=vertex_ai_project,
+            vertex_ai_location=vertex_ai_location,
+            vertex_credentials=vertex_credentials,
+        )
+        print(f"Model Response: {response}")
+    except litellm.RateLimitError as e:
+        pass
+    except Exception as e:
+        pytest.fail(f"Error occurred: {e}")
 
 
 # asyncio.run(test_vertex_ai_anthropic_async())
@@ -208,26 +218,31 @@ async def test_vertex_ai_anthropic_async():
 @pytest.mark.asyncio
 async def test_vertex_ai_anthropic_async_streaming():
     # load_vertex_ai_credentials()
-    litellm.set_verbose = True
-    model = "claude-3-sonnet@20240229"
+    try:
+        litellm.set_verbose = True
+        model = "claude-3-sonnet@20240229"
 
-    vertex_ai_project = "adroit-crow-413218"
-    vertex_ai_location = "asia-southeast1"
-    json_obj = get_vertex_ai_creds_json()
-    vertex_credentials = json.dumps(json_obj)
+        vertex_ai_project = "adroit-crow-413218"
+        vertex_ai_location = "asia-southeast1"
+        json_obj = get_vertex_ai_creds_json()
+        vertex_credentials = json.dumps(json_obj)
 
-    response = await acompletion(
-        model="vertex_ai/" + model,
-        messages=[{"role": "user", "content": "hi"}],
-        temperature=0.7,
-        vertex_ai_project=vertex_ai_project,
-        vertex_ai_location=vertex_ai_location,
-        vertex_credentials=vertex_credentials,
-        stream=True,
-    )
+        response = await acompletion(
+            model="vertex_ai/" + model,
+            messages=[{"role": "user", "content": "hi"}],
+            temperature=0.7,
+            vertex_ai_project=vertex_ai_project,
+            vertex_ai_location=vertex_ai_location,
+            vertex_credentials=vertex_credentials,
+            stream=True,
+        )
 
-    async for chunk in response:
-        print(f"chunk: {chunk}")
+        async for chunk in response:
+            print(f"chunk: {chunk}")
+    except litellm.RateLimitError as e:
+        pass
+    except Exception as e:
+        pytest.fail(f"Error occurred: {e}")
 
 
 # asyncio.run(test_vertex_ai_anthropic_async_streaming())