phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

(feat) proxy: cost tracking add track_cost_callback

Browse source
This commit is contained in:
ishaan-jaff 2023-11-24 05:53:23 -08:00
parent b03a7ae0cd
commit 0a23af7990
2 changed files with 11 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

11
litellm/proxy/proxy_server.py
View file

@ -267,6 +267,15 @@ def load_router_config(router: Optional[litellm.Router], config_file_path: str):
### CONNECT TO DATABASE ### ### CONNECT TO DATABASE ###
database_url = general_settings.get("database_url", None) database_url = general_settings.get("database_url", None)
prisma_setup(database_url=database_url) 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 ### ### START REDIS QUEUE ###
use_queue = general_settings.get("use_queue", False) use_queue = general_settings.get("use_queue", False)
celery_setup(use_queue=use_queue) 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
model_list = config.get('model_list', None) model_list = config.get('model_list', None)
if model_list: 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") print(f"\033[32mLiteLLM: Proxy initialized with Config, Set models:\033[0m")
for model in model_list: for model in model_list:
print(f"\033[32m {model.get('model_name', '')}\033[0m") print(f"\033[32m {model.get('model_name', '')}\033[0m")

2
litellm/proxy/utils.py
View file

@ -5,7 +5,7 @@ from typing import Optional
def track_cost_callback( def track_cost_callback(
kwargs, # kwargs to completion kwargs, # kwargs to completion
completion_response: ModelResponse, # response from completion completion_response: ModelResponse, # response from completion
start_time = None, start_time = None,
end_time = None, # start/end time for completion end_time = None, # start/end time for completion
): ):