expose flag to disable_spend_updates

This commit is contained in:
Ishaan Jaff 2025-03-17 20:45:49 -07:00
parent 7c556a008e
commit e6975a56da
2 changed files with 41 additions and 8 deletions

View file

@ -23,6 +23,7 @@ from typing import (
get_origin,
get_type_hints,
)
from litellm.types.utils import (
ModelResponse,
ModelResponseStream,
@ -254,6 +255,7 @@ from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import (
from litellm.proxy.utils import (
PrismaClient,
ProxyLogging,
ProxyUpdateSpend,
_cache_user_row,
_get_docs_url,
_get_projected_spend_over_limit,
@ -924,6 +926,8 @@ async def update_database( # noqa: PLR0915
verbose_proxy_logger.debug(
f"Enters prisma db call, response_cost: {response_cost}, token: {token}; user_id: {user_id}; team_id: {team_id}"
)
if ProxyUpdateSpend.disable_spend_updates() is True:
return
if token is not None and isinstance(token, str) and token.startswith("sk-"):
hashed_token = hash_token(token=token)
else:
@ -3047,7 +3051,11 @@ async def async_data_generator(
):
verbose_proxy_logger.debug("inside generator")
try:
async for chunk in proxy_logging_obj.async_post_call_streaming_iterator_hook(user_api_key_dict=user_api_key_dict, response=response, request_data=request_data):
async for chunk in proxy_logging_obj.async_post_call_streaming_iterator_hook(
user_api_key_dict=user_api_key_dict,
response=response,
request_data=request_data,
):
verbose_proxy_logger.debug(
"async_data_generator: received streaming chunk - {}".format(chunk)
)