From b377655d8ff6033f13b1fc8c0f96f4b28c1b75c5 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Sat, 9 Sep 2023 18:36:12 -0700 Subject: [PATCH] try/except completion_cost + custom logger func --- litellm/integrations/custom_logger.py | 36 +++++++++++++ litellm/utils.py | 74 ++++++++++++++++----------- 2 files changed, 80 insertions(+), 30 deletions(-) create mode 100644 litellm/integrations/custom_logger.py diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py new file mode 100644 index 0000000000..1f936ad27f --- /dev/null +++ b/litellm/integrations/custom_logger.py @@ -0,0 +1,36 @@ +#### What this does #### +# On success, logs events to Promptlayer +import dotenv, os +import requests +import requests + +dotenv.load_dotenv() # Loading env variables using dotenv +import traceback + + +class CustomLogger: + # Class variables or attributes + def __init__(self, callback_func): + # Instance variables + self.callback_func = callback_func + + def log_event(self, kwargs, response_obj, start_time, end_time, print_verbose): + # Method definition + try: + print_verbose( + f"Custom Logger - Enters logging function for model {kwargs}" + ) + self.callback_func( + kwargs, # kwargs to func + response_obj, + start_time, + end_time, + print_verbose + ) + print_verbose( + f"Custom Logger - final response object: {response_obj}" + ) + except: + # traceback.print_exc() + print_verbose(f"Custom Logger Error - {traceback.format_exc()}") + pass diff --git a/litellm/utils.py b/litellm/utils.py index fb97037ba1..614d824fed 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -17,6 +17,7 @@ from .integrations.berrispend import BerriSpendLogger from .integrations.supabase import Supabase from .integrations.llmonitor import LLMonitorLogger from .integrations.prompt_layer import PromptLayerLogger +from .integrations.custom_logger import CustomLogger from .integrations.langfuse import LangFuseLogger from .integrations.litedebugger import LiteDebugger from openai.error import OpenAIError as OriginalError @@ -46,6 +47,7 @@ slack_app = None alerts_channel = None heliconeLogger = None promptLayerLogger = None +customLogger = None langFuseLogger = None llmonitorLogger = None aispendLogger = None @@ -677,35 +679,37 @@ def completion_cost( completion="", total_time=0.0, # used for replicate ): - - # Handle Inputs to completion_cost - prompt_tokens = 0 - completion_tokens = 0 - if completion_response != None: - # get input/output tokens from completion_response - prompt_tokens = completion_response['usage']['prompt_tokens'] - completion_tokens = completion_response['usage']['completion_tokens'] - model = completion_response['model'] # get model from completion_response - else: - prompt_tokens = token_counter(model=model, text=prompt) - completion_tokens = token_counter(model=model, text=completion) - - # Calculate cost based on prompt_tokens, completion_tokens - if "togethercomputer" in model: - # together ai prices based on size of llm - # get_model_params_and_category takes a model name and returns the category of LLM size it is in model_prices_and_context_window.json - model = get_model_params_and_category(model) - # replicate llms are calculate based on time for request running - # see https://replicate.com/pricing - elif ( - model in litellm.replicate_models or - "replicate" in model - ): - return get_replicate_completion_pricing(completion_response, total_time) - prompt_tokens_cost_usd_dollar, completion_tokens_cost_usd_dollar = cost_per_token( - model=model, prompt_tokens=prompt_tokens, completion_tokens=completion_tokens - ) - return prompt_tokens_cost_usd_dollar + completion_tokens_cost_usd_dollar + try: + # Handle Inputs to completion_cost + prompt_tokens = 0 + completion_tokens = 0 + if completion_response != None: + # get input/output tokens from completion_response + prompt_tokens = completion_response['usage']['prompt_tokens'] + completion_tokens = completion_response['usage']['completion_tokens'] + model = completion_response['model'] # get model from completion_response + else: + prompt_tokens = token_counter(model=model, text=prompt) + completion_tokens = token_counter(model=model, text=completion) + + # Calculate cost based on prompt_tokens, completion_tokens + if "togethercomputer" in model: + # together ai prices based on size of llm + # get_model_params_and_category takes a model name and returns the category of LLM size it is in model_prices_and_context_window.json + model = get_model_params_and_category(model) + # replicate llms are calculate based on time for request running + # see https://replicate.com/pricing + elif ( + model in litellm.replicate_models or + "replicate" in model + ): + return get_replicate_completion_pricing(completion_response, total_time) + prompt_tokens_cost_usd_dollar, completion_tokens_cost_usd_dollar = cost_per_token( + model=model, prompt_tokens=prompt_tokens, completion_tokens=completion_tokens + ) + return prompt_tokens_cost_usd_dollar + completion_tokens_cost_usd_dollar + except: + return 0.0 # this should not block a users execution path ####### HELPER FUNCTIONS ################ def get_litellm_params( @@ -993,7 +997,7 @@ def validate_environment(): return api_key def set_callbacks(callback_list, function_id=None): - global sentry_sdk_instance, capture_exception, add_breadcrumb, posthog, slack_app, alerts_channel, traceloopLogger, heliconeLogger, aispendLogger, berrispendLogger, supabaseClient, liteDebuggerClient, llmonitorLogger, promptLayerLogger, langFuseLogger + global sentry_sdk_instance, capture_exception, add_breadcrumb, posthog, slack_app, alerts_channel, traceloopLogger, heliconeLogger, aispendLogger, berrispendLogger, supabaseClient, liteDebuggerClient, llmonitorLogger, promptLayerLogger, langFuseLogger, customLogger try: for callback in callback_list: print_verbose(f"callback: {callback}") @@ -1073,6 +1077,8 @@ def set_callbacks(callback_list, function_id=None): liteDebuggerClient = LiteDebugger(email=litellm.email) else: liteDebuggerClient = LiteDebugger(email=str(uuid.uuid4())) + elif callable(callback): + customLogger = CustomLogger() except Exception as e: raise e @@ -1366,6 +1372,14 @@ def handle_success(args, kwargs, result, start_time, end_time): litellm_call_id=kwargs["litellm_call_id"], print_verbose=print_verbose, ) + elif callable(callback): # custom logger functions + customLogger.log_event( + kwargs=kwargs, + response_obj=result, + start_time=start_time, + end_time=end_time, + print_verbose=print_verbose, + ) except Exception as e: # LOGGING exception_logging(logger_fn=user_logger_fn, exception=e)