diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index 1cba699abd..963121c4a7 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -7,4 +7,9 @@ model_list: general_settings: master_key: sk-1234 - proxy_batch_write_at: 5 \ No newline at end of file + proxy_batch_write_at: 5 # 👈 Frequency of batch writing logs to server (in seconds) + enable_jwt_auth: True + litellm_jwtauth: + admin_jwt_scope: "litellm_proxy_admin" + team_jwt_scope: "litellm_team" + public_key_ttl: 600 \ No newline at end of file diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index cfb2648a43..f2628cbe12 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -293,7 +293,7 @@ litellm_proxy_admin_name = "default_user_id" ui_access_mode: Literal["admin", "all"] = "all" proxy_budget_rescheduler_min_time = 597 proxy_budget_rescheduler_max_time = 605 -proxy_batch_write_at = 60 # in seconds +proxy_batch_write_at = 10 # in seconds litellm_master_key_hash = None disable_spend_logs = False jwt_handler = JWTHandler() @@ -1161,7 +1161,7 @@ async def _PROXY_track_cost_callback( verbose_proxy_logger.debug( f"user_api_key {user_api_key}, prisma_client: {prisma_client}, custom_db_client: {custom_db_client}" ) - if user_api_key is not None: + if user_api_key is not None or user_id is not None or team_id is not None: ## UPDATE DATABASE await update_database( token=user_api_key, @@ -1182,7 +1182,9 @@ async def _PROXY_track_cost_callback( response_cost=response_cost, ) else: - raise Exception("User API key missing from custom callback.") + raise Exception( + "User API key and team id and user id missing from custom callback." + ) else: if kwargs["stream"] != True or ( kwargs["stream"] == True and "complete_streaming_response" in kwargs @@ -1223,7 +1225,7 @@ async def update_database( verbose_proxy_logger.info( f"Enters prisma db call, response_cost: {response_cost}, token: {token}; user_id: {user_id}; team_id: {team_id}" ) - if isinstance(token, str) and token.startswith("sk-"): + if token is not None and isinstance(token, str) and token.startswith("sk-"): hashed_token = hash_token(token=token) else: hashed_token = token @@ -1337,6 +1339,8 @@ async def update_database( verbose_proxy_logger.debug( f"adding spend to key db. Response cost: {response_cost}. Token: {hashed_token}." ) + if hashed_token is None: + return if prisma_client is not None: prisma_client.key_list_transactons[hashed_token] = ( response_cost diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 9a35e7c016..f313a79852 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1901,6 +1901,7 @@ async def update_spend( spend_logs: list, """ n_retry_times = 3 + verbose_proxy_logger.debug("INSIDE UPDATE SPEND") ### UPDATE USER TABLE ### if len(prisma_client.user_list_transactons.keys()) > 0: for i in range(n_retry_times + 1): @@ -1989,6 +1990,11 @@ async def update_spend( raise e ### UPDATE TEAM TABLE ### + verbose_proxy_logger.debug( + "Team Spend transactions: {}".format( + len(prisma_client.team_list_transactons.keys()) + ) + ) if len(prisma_client.team_list_transactons.keys()) > 0: for i in range(n_retry_times + 1): try: @@ -2000,6 +2006,11 @@ async def update_spend( team_id, response_cost, ) in prisma_client.team_list_transactons.items(): + verbose_proxy_logger.debug( + "Updating spend for team id={} by {}".format( + team_id, response_cost + ) + ) batcher.litellm_teamtable.update_many( # 'update_many' prevents error from being raised if no row exists where={"team_id": team_id}, data={"spend": {"increment": response_cost}},