phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-25 18:54:30 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

use DBSpendUpdateWriter class for managing DB spend updates

Browse source
This commit is contained in:
Ishaan Jaff 2025-03-27 16:31:23 -07:00
parent bc1f2fe3b2
commit c367a10c1c
5 changed files with 255 additions and 223 deletions
Download patch file Download diff file Expand all files Collapse all files

8
litellm/proxy/hooks/proxy_track_cost_callback.py
View file

@ -13,6 +13,7 @@ from litellm.litellm_core_utils.core_helpers import (
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.auth.auth_checks import log_db_metrics
from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter
from litellm.proxy.utils import ProxyUpdateSpend
from litellm.types.utils import (
StandardLoggingPayload,
@ -33,8 +34,6 @@ class _ProxyDBLogger(CustomLogger):
original_exception: Exception,
user_api_key_dict: UserAPIKeyAuth,
):
from litellm.proxy.proxy_server import update_database
if _ProxyDBLogger._should_track_errors_in_db() is False:
return
@ -67,7 +66,7 @@ class _ProxyDBLogger(CustomLogger):
request_data.get("proxy_server_request") or {}
)
request_data["litellm_params"]["metadata"] = existing_metadata
await update_database(
await DBSpendUpdateWriter.update_database(
token=user_api_key_dict.api_key,
response_cost=0.0,
user_id=user_api_key_dict.user_id,
@ -94,7 +93,6 @@ class _ProxyDBLogger(CustomLogger):
prisma_client,
proxy_logging_obj,
update_cache,
update_database,
)
verbose_proxy_logger.debug("INSIDE _PROXY_track_cost_callback")
@ -138,7 +136,7 @@ class _ProxyDBLogger(CustomLogger):
end_user_id=end_user_id,
):
## UPDATE DATABASE
await update_database(
await DBSpendUpdateWriter.update_database(
token=user_api_key,
response_cost=response_cost,
user_id=user_id,