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

feat(proxy/utils.py): enable updating db in a separate server

Browse source
This commit is contained in:
Krrish Dholakia 2024-03-27 16:02:36 -07:00
parent 9b7383ac67
commit 1e856443e1
6 changed files with 89 additions and 57 deletions
Download patch file Download diff file Expand all files Collapse all files

22
litellm/proxy/proxy_server.py
View file

@ -97,7 +97,6 @@ from litellm.proxy.utils import (
_is_projected_spend_over_limit,
_get_projected_spend_over_limit,
update_spend,
monitor_spend_list,
)
from litellm.proxy.secret_managers.google_kms import load_google_kms
from litellm.proxy.secret_managers.aws_secret_manager import load_aws_secret_manager
@ -118,6 +117,7 @@ from litellm.proxy.auth.auth_checks import (
allowed_routes_check,
get_actual_routes,
)
from litellm.llms.custom_httpx.httpx_handler import HTTPHandler
try:
from litellm._version import version
@ -305,6 +305,8 @@ proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache)
async_result = None
celery_app_conn = None
celery_fn = None # Redis Queue for handling requests
### DB WRITER ###
db_writer_client: Optional[HTTPHandler] = None
### logger ###
@ -1363,7 +1365,15 @@ async def update_database(
payload["spend"] = response_cost
if prisma_client is not None:
await prisma_client.insert_data(data=payload, table_name="spend")
prisma_client.spend_log_transactons.append(payload)
# if db_writer_client is not None:
# print("Tries to make call")
# response = await db_writer_client.post(
# url="http://0.0.0.0:3000/spend/update",
# data=json.dumps(payload),
# headers={"Content-Type": "application/json"},
# )
# print(f"response: {response}")
except Exception as e:
verbose_proxy_logger.debug(
f"Update Spend Logs DB failed to execute - {str(e)}\n{traceback.format_exc()}"
@ -2693,7 +2703,7 @@ def on_backoff(details):
@router.on_event("startup")
async def startup_event():
global prisma_client, master_key, use_background_health_checks, llm_router, llm_model_list, general_settings, proxy_budget_rescheduler_min_time, proxy_budget_rescheduler_max_time, litellm_proxy_admin_name
global prisma_client, master_key, use_background_health_checks, llm_router, llm_model_list, general_settings, proxy_budget_rescheduler_min_time, proxy_budget_rescheduler_max_time, litellm_proxy_admin_name, db_writer_client
import json
### LOAD MASTER KEY ###
@ -2726,6 +2736,8 @@ async def startup_event():
## COST TRACKING ##
cost_tracking()
db_writer_client = HTTPHandler()
proxy_logging_obj._init_litellm_callbacks() # INITIALIZE LITELLM CALLBACKS ON SERVER STARTUP <- do this to catch any logging errors on startup, not when calls are being made
## JWT AUTH ##
@ -2836,7 +2848,7 @@ async def startup_event():
update_spend,
"interval",
seconds=batch_writing_interval,
args=[prisma_client],
args=[prisma_client, db_writer_client],
)
scheduler.start()
@ -7985,6 +7997,8 @@ async def shutdown_event():
await jwt_handler.close()
if db_writer_client is not None:
await db_writer_client.close()
## RESET CUSTOM VARIABLES ##
cleanup_router_config_variables()