From 0a23af7990ce68666913996d903464640ca292e2 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Fri, 24 Nov 2023 05:53:23 -0800 Subject: [PATCH] (feat) proxy: cost tracking add track_cost_callback --- litellm/proxy/proxy_server.py | 11 ++++++++++- litellm/proxy/utils.py | 2 +- 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 1299318e25..fbeca71112 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -267,6 +267,15 @@ def load_router_config(router: Optional[litellm.Router], config_file_path: str): ### CONNECT TO DATABASE ### database_url = general_settings.get("database_url", None) prisma_setup(database_url=database_url) + ## Cost Tracking for master key + auth setup ## + if master_key is not None: + if isinstance(litellm.success_callback, list): + import utils + print("setting litellm success callback to track cost") + if (utils.track_cost_callback) not in litellm.success_callback: + litellm.success_callback.append(utils.track_cost_callback) + else: + litellm.success_callback = utils.track_cost_callback ### START REDIS QUEUE ### use_queue = general_settings.get("use_queue", False) celery_setup(use_queue=use_queue) @@ -303,7 +312,7 @@ def load_router_config(router: Optional[litellm.Router], config_file_path: str): ## MODEL LIST model_list = config.get('model_list', None) if model_list: - router = litellm.Router(model_list=model_list) + router = litellm.Router(model_list=model_list, num_retries=2) print(f"\033[32mLiteLLM: Proxy initialized with Config, Set models:\033[0m") for model in model_list: print(f"\033[32m {model.get('model_name', '')}\033[0m") diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 855fc6b1c4..c61bcfa85d 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -5,7 +5,7 @@ from typing import Optional def track_cost_callback( kwargs, # kwargs to completion - completion_response: ModelResponse, # response from completion + completion_response: ModelResponse, # response from completion start_time = None, end_time = None, # start/end time for completion ):