fix(test_amazing_vertex_completion.py): handle vertex ai rate limit error

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(
+    try:
-        model="gemini-pro",
+        completion = litellm.completion(
-        messages=messages,
+            model="gemini-pro",
-        tools=tools,
+            messages=messages,
-        tool_choice="auto",
+            tools=tools,
-        stream=True,
+            tool_choice="auto",
-    )
+            stream=True,
-    print(f"completion: {completion}")
+        )
-    # assert completion.choices[0].message.content is None
+        print(f"completion: {completion}")
-    # assert len(completion.choices[0].message.tool_calls) == 1
+        # assert completion.choices[0].message.content is None
-    for chunk in completion:
+        # assert len(completion.choices[0].message.tool_calls) == 1
-        print(f"chunk: {chunk}")
+        for chunk in completion:
+            print(f"chunk: {chunk}")
+    except litellm.RateLimitError as e:
+        pass
 
 
 @pytest.mark.asyncio