fix(utils.py): fix proxy streaming spend tracking

This commit is contained in:
Krrish Dholakia 2024-01-23 15:59:03 -08:00
parent 01a2514b98
commit f8870fb48e
4 changed files with 130 additions and 14 deletions

View file

@ -570,7 +570,7 @@ async def track_cost_callback(
litellm_params = kwargs.get("litellm_params", {}) or {}
proxy_server_request = litellm_params.get("proxy_server_request") or {}
user_id = proxy_server_request.get("body", {}).get("user", None)
if "response_cost" in kwargs:
if kwargs.get("response_cost", None) is not None:
response_cost = kwargs["response_cost"]
user_api_key = kwargs["litellm_params"]["metadata"].get(
"user_api_key", None
@ -596,9 +596,13 @@ async def track_cost_callback(
end_time=end_time,
)
else:
raise Exception(
f"Model not in litellm model cost map. Add custom pricing - https://docs.litellm.ai/docs/proxy/custom_pricing"
)
if kwargs["stream"] != True or (
kwargs["stream"] == True
and kwargs.get("complete_streaming_response") in kwargs
):
raise Exception(
f"Model not in litellm model cost map. Add custom pricing - https://docs.litellm.ai/docs/proxy/custom_pricing"
)
except Exception as e:
verbose_proxy_logger.debug(f"error in tracking cost callback - {str(e)}")