From 00016830363fded6516c95f1576655c362af1e35 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Fri, 28 Jun 2024 14:53:00 -0700
Subject: [PATCH] fix(cost_calculator.py): handle unexpected error in
 cost_calculator.py

---
 litellm/cost_calculator.py      |  8 ++++++++
 litellm/tests/test_streaming.py | 19 +++++++++++++++++--
 2 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py
index 9a7df7ebefb..062e98be97b 100644
--- a/litellm/cost_calculator.py
+++ b/litellm/cost_calculator.py
@@ -1,6 +1,7 @@
 # What is this?
 ## File for 'response_cost' calculation in Logging
 import time
+import traceback
 from typing import List, Literal, Optional, Tuple, Union
 
 import litellm
@@ -668,3 +669,10 @@ def response_cost_calculator(
             f"Model={model} for LLM Provider={custom_llm_provider} not found in completion cost map."
         )
         return None
+    except Exception as e:
+        verbose_logger.error(
+            "litellm.cost_calculator.py::response_cost_calculator - Exception occurred - {}/n{}".format(
+                str(e), traceback.format_exc()
+            )
+        )
+        return None
diff --git a/litellm/tests/test_streaming.py b/litellm/tests/test_streaming.py
index 5cd0e35a9e3..eec6929f2e3 100644
--- a/litellm/tests/test_streaming.py
+++ b/litellm/tests/test_streaming.py
@@ -1040,14 +1040,27 @@ def test_vertex_ai_stream(provider):
     litellm.vertex_project = "adroit-crow-413218"
     import random
 
-    test_models = ["gemini-1.0-pro"]
+    test_models = ["gemini-1.5-pro"]
     for model in test_models:
         try:
             print("making request", model)
             response = completion(
                 model="{}/{}".format(provider, model),
                 messages=[
-                    {"role": "user", "content": "write 10 line code code for saying hi"}
+                    {"role": "user", "content": "Hey, how's it going?"},
+                    {
+                        "role": "assistant",
+                        "content": "I'm doing well. Would like to hear the rest of the story?",
+                    },
+                    {"role": "user", "content": "Na"},
+                    {
+                        "role": "assistant",
+                        "content": "No problem, is there anything else i can help you with today?",
+                    },
+                    {
+                        "role": "user",
+                        "content": "I think you're getting cut off sometimes",
+                    },
                 ],
                 stream=True,
             )
@@ -1064,6 +1077,8 @@ def test_vertex_ai_stream(provider):
                 raise Exception("Empty response received")
             print(f"completion_response: {complete_response}")
             assert is_finished == True
+
+            assert False
         except litellm.RateLimitError as e:
             pass
         except Exception as e: