From f7537f2cdfbc6041650de447702994f38c24561c Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Fri, 5 Apr 2024 14:02:26 -0700
Subject: [PATCH] fix(test_amazing_vertex_completion.py): handle vertex ai rate
 limit error

---
 .../tests/test_amazing_vertex_completion.py   | 56 +++++++++++++++----
 1 file changed, 44 insertions(+), 12 deletions(-)

diff --git a/litellm/tests/test_amazing_vertex_completion.py b/litellm/tests/test_amazing_vertex_completion.py
index 7641badfb..05782b51d 100644
--- a/litellm/tests/test_amazing_vertex_completion.py
+++ b/litellm/tests/test_amazing_vertex_completion.py
@@ -487,6 +487,35 @@ def test_gemini_pro_vision_base64():
 
 
 def test_gemini_pro_function_calling():
+    load_vertex_ai_credentials()
+    tools = [
+        {
+            "type": "function",
+            "function": {
+                "name": "get_current_weather",
+                "description": "Get the current weather in a given location",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "location": {
+                            "type": "string",
+                            "description": "The city and state, e.g. San Francisco, CA",
+                        },
+                        "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},
+                    },
+                    "required": ["location"],
+                },
+            },
+        }
+    ]
+
+    messages = [{"role": "user", "content": "What's the weather like in Boston today?"}]
+    completion = litellm.completion(
+        model="gemini-pro", messages=messages, tools=tools, tool_choice="auto"
+    )
+    print(f"completion: {completion}")
+    assert completion.choices[0].message.content is None
+    assert len(completion.choices[0].message.tool_calls) == 1
     try:
         load_vertex_ai_credentials()
         tools = [
@@ -557,18 +586,21 @@ def test_gemini_pro_function_calling_streaming():
         }
     ]
     messages = [{"role": "user", "content": "What's the weather like in Boston today?"}]
-    completion = litellm.completion(
-        model="gemini-pro",
-        messages=messages,
-        tools=tools,
-        tool_choice="auto",
-        stream=True,
-    )
-    print(f"completion: {completion}")
-    # assert completion.choices[0].message.content is None
-    # assert len(completion.choices[0].message.tool_calls) == 1
-    for chunk in completion:
-        print(f"chunk: {chunk}")
+    try:
+        completion = litellm.completion(
+            model="gemini-pro",
+            messages=messages,
+            tools=tools,
+            tool_choice="auto",
+            stream=True,
+        )
+        print(f"completion: {completion}")
+        # assert completion.choices[0].message.content is None
+        # assert len(completion.choices[0].message.tool_calls) == 1
+        for chunk in completion:
+            print(f"chunk: {chunk}")
+    except litellm.RateLimitError as e:
+        pass
 
 
 @pytest.mark.asyncio