From 8da3e5ad66ff4d55bf964ea120b20fb1c7aae875 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Wed, 19 Jun 2024 19:25:27 -0700
Subject: [PATCH] test(test_amazing_vertex_completion.py): fix test

---
 litellm/tests/test_amazing_vertex_completion.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/litellm/tests/test_amazing_vertex_completion.py b/litellm/tests/test_amazing_vertex_completion.py
index 742e038801..6433b06410 100644
--- a/litellm/tests/test_amazing_vertex_completion.py
+++ b/litellm/tests/test_amazing_vertex_completion.py
@@ -548,6 +548,7 @@ async def test_gemini_pro_vision(provider, sync_mode):
             resp = await litellm.acompletion(
                 model="{}/gemini-1.5-flash-preview-0514".format(provider),
                 messages=[
+                    {"role": "system", "content": "Be a good bot"},
                     {
                         "role": "user",
                         "content": [
@@ -559,7 +560,7 @@ async def test_gemini_pro_vision(provider, sync_mode):
                                 },
                             },
                         ],
-                    }
+                    },
                 ],
             )
         print(resp)
@@ -568,7 +569,7 @@ async def test_gemini_pro_vision(provider, sync_mode):
 
         # DO Not DELETE this ASSERT
         # Google counts the prompt tokens for us, we should ensure we use the tokens from the orignal response
-        assert prompt_tokens == 263  # the gemini api returns 263 to us
+        assert prompt_tokens == 267  # the gemini api returns 267 to us
 
     except litellm.RateLimitError as e:
         pass