phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

fix(utils.py): fix streaming cost tracking

Browse source
This commit is contained in:
Krrish Dholakia 2024-01-23 14:39:35 -08:00
parent 7eb96e46a4
commit afada01ffc
4 changed files with 39 additions and 13 deletions
Download patch file Download diff file Expand all files Collapse all files

14
litellm/proxy/proxy_server.py
View file

@ -572,7 +572,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
@ -598,9 +598,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.get("complete_streaming_response", None) is not None
):
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)}")
@ -1514,7 +1518,7 @@ async def startup_event():
duration=None, models=[], aliases={}, config={}, spend=0, token=master_key
)
verbose_proxy_logger.debug(
f"custom_db_client client - Inserting master key {custom_db_client}. Master_key: {master_key}"
f"custom_db_client client {custom_db_client}. Master_key: {master_key}"
)
if custom_db_client is not None and master_key is not None:
# add master key to db