fix(cost_calculator.py): handle unexpected error in cost_calculator.py

This commit is contained in:
Krrish Dholakia 2024-06-28 14:53:00 -07:00
parent 224148d613
commit 0001683036
2 changed files with 25 additions and 2 deletions

View file

@ -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

View file

@ -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: