#### What this does #### # picks based on response time (for streaming, this is time to first token) from datetime import datetime, timedelta from typing import Dict, List, Optional, Union import litellm from litellm import ModelResponse, token_counter, verbose_logger from litellm._logging import verbose_router_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger class LowestCostLoggingHandler(CustomLogger): test_flag: bool = False logged_success: int = 0 logged_failure: int = 0 def __init__( self, router_cache: DualCache, model_list: list, routing_args: dict = {} ): self.router_cache = router_cache self.model_list = model_list def log_success_event(self, kwargs, response_obj, start_time, end_time): try: """ Update usage on success """ if kwargs["litellm_params"].get("metadata") is None: pass else: model_group = kwargs["litellm_params"]["metadata"].get( "model_group", None ) id = kwargs["litellm_params"].get("model_info", {}).get("id", None) if model_group is None or id is None: return elif isinstance(id, int): id = str(id) # ------------ # Setup values # ------------ """ { {model_group}_map: { id: { f"{date:hour:minute}" : {"tpm": 34, "rpm": 3} } } } """ current_date = datetime.now().strftime("%Y-%m-%d") current_hour = datetime.now().strftime("%H") current_minute = datetime.now().strftime("%M") precise_minute = f"{current_date}-{current_hour}-{current_minute}" cost_key = f"{model_group}_map" response_ms: timedelta = end_time - start_time total_tokens = 0 if isinstance(response_obj, ModelResponse): _usage = getattr(response_obj, "usage", None) if _usage is not None and isinstance(_usage, litellm.Usage): completion_tokens = _usage.completion_tokens total_tokens = _usage.total_tokens float(response_ms.total_seconds() / completion_tokens) # ------------ # Update usage # ------------ request_count_dict = self.router_cache.get_cache(key=cost_key) or {} # check local result first if id not in request_count_dict: request_count_dict[id] = {} if precise_minute not in request_count_dict[id]: request_count_dict[id][precise_minute] = {} ## TPM request_count_dict[id][precise_minute]["tpm"] = ( request_count_dict[id][precise_minute].get("tpm", 0) + total_tokens ) ## RPM request_count_dict[id][precise_minute]["rpm"] = ( request_count_dict[id][precise_minute].get("rpm", 0) + 1 ) self.router_cache.set_cache(key=cost_key, value=request_count_dict) ### TESTING ### if self.test_flag: self.logged_success += 1 except Exception as e: verbose_logger.exception( "litellm.router_strategy.lowest_cost.py::log_success_event(): Exception occured - {}".format( str(e) ) ) pass async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): try: """ Update cost usage on success """ if kwargs["litellm_params"].get("metadata") is None: pass else: model_group = kwargs["litellm_params"]["metadata"].get( "model_group", None ) id = kwargs["litellm_params"].get("model_info", {}).get("id", None) if model_group is None or id is None: return elif isinstance(id, int): id = str(id) # ------------ # Setup values # ------------ """ { {model_group}_map: { id: { "cost": [..] f"{date:hour:minute}" : {"tpm": 34, "rpm": 3} } } } """ cost_key = f"{model_group}_map" current_date = datetime.now().strftime("%Y-%m-%d") current_hour = datetime.now().strftime("%H") current_minute = datetime.now().strftime("%M") precise_minute = f"{current_date}-{current_hour}-{current_minute}" response_ms: timedelta = end_time - start_time total_tokens = 0 if isinstance(response_obj, ModelResponse): _usage = getattr(response_obj, "usage", None) if _usage is not None and isinstance(_usage, litellm.Usage): completion_tokens = _usage.completion_tokens total_tokens = _usage.total_tokens float(response_ms.total_seconds() / completion_tokens) # ------------ # Update usage # ------------ request_count_dict = ( await self.router_cache.async_get_cache(key=cost_key) or {} ) if id not in request_count_dict: request_count_dict[id] = {} if precise_minute not in request_count_dict[id]: request_count_dict[id][precise_minute] = {} ## TPM request_count_dict[id][precise_minute]["tpm"] = ( request_count_dict[id][precise_minute].get("tpm", 0) + total_tokens ) ## RPM request_count_dict[id][precise_minute]["rpm"] = ( request_count_dict[id][precise_minute].get("rpm", 0) + 1 ) await self.router_cache.async_set_cache( key=cost_key, value=request_count_dict ) # reset map within window ### TESTING ### if self.test_flag: self.logged_success += 1 except Exception as e: verbose_logger.exception( "litellm.proxy.hooks.prompt_injection_detection.py::async_pre_call_hook(): Exception occured - {}".format( str(e) ) ) pass async def async_get_available_deployments( # noqa: PLR0915 self, model_group: str, healthy_deployments: list, messages: Optional[List[Dict[str, str]]] = None, input: Optional[Union[str, List]] = None, request_kwargs: Optional[Dict] = None, ): """ Returns a deployment with the lowest cost """ cost_key = f"{model_group}_map" request_count_dict = await self.router_cache.async_get_cache(key=cost_key) or {} # ----------------------- # Find lowest used model # ---------------------- float("inf") current_date = datetime.now().strftime("%Y-%m-%d") current_hour = datetime.now().strftime("%H") current_minute = datetime.now().strftime("%M") precise_minute = f"{current_date}-{current_hour}-{current_minute}" if request_count_dict is None: # base case return all_deployments = request_count_dict for d in healthy_deployments: ## if healthy deployment not yet used if d["model_info"]["id"] not in all_deployments: all_deployments[d["model_info"]["id"]] = { precise_minute: {"tpm": 0, "rpm": 0}, } try: input_tokens = token_counter(messages=messages, text=input) except Exception: input_tokens = 0 # randomly sample from all_deployments, incase all deployments have latency=0.0 _items = all_deployments.items() ### GET AVAILABLE DEPLOYMENTS ### filter out any deployments > tpm/rpm limits potential_deployments = [] _cost_per_deployment = {} for item, item_map in all_deployments.items(): ## get the item from model list _deployment = None for m in healthy_deployments: if item == m["model_info"]["id"]: _deployment = m if _deployment is None: continue # skip to next one _deployment_tpm = ( _deployment.get("tpm", None) or _deployment.get("litellm_params", {}).get("tpm", None) or _deployment.get("model_info", {}).get("tpm", None) or float("inf") ) _deployment_rpm = ( _deployment.get("rpm", None) or _deployment.get("litellm_params", {}).get("rpm", None) or _deployment.get("model_info", {}).get("rpm", None) or float("inf") ) item_litellm_model_name = _deployment.get("litellm_params", {}).get("model") item_litellm_model_cost_map = litellm.model_cost.get( item_litellm_model_name, {} ) # check if user provided input_cost_per_token and output_cost_per_token in litellm_params item_input_cost = None item_output_cost = None if _deployment.get("litellm_params", {}).get("input_cost_per_token", None): item_input_cost = _deployment.get("litellm_params", {}).get( "input_cost_per_token" ) if _deployment.get("litellm_params", {}).get("output_cost_per_token", None): item_output_cost = _deployment.get("litellm_params", {}).get( "output_cost_per_token" ) if item_input_cost is None: item_input_cost = item_litellm_model_cost_map.get( "input_cost_per_token", 5.0 ) if item_output_cost is None: item_output_cost = item_litellm_model_cost_map.get( "output_cost_per_token", 5.0 ) # if litellm["model"] is not in model_cost map -> use item_cost = $10 item_cost = item_input_cost + item_output_cost item_rpm = item_map.get(precise_minute, {}).get("rpm", 0) item_tpm = item_map.get(precise_minute, {}).get("tpm", 0) verbose_router_logger.debug( f"item_cost: {item_cost}, item_tpm: {item_tpm}, item_rpm: {item_rpm}, model_id: {_deployment.get('model_info', {}).get('id')}" ) # -------------- # # Debugging Logic # -------------- # # We use _cost_per_deployment to log to langfuse, slack - this is not used to make a decision on routing # this helps a user to debug why the router picked a specfic deployment # _deployment_api_base = _deployment.get("litellm_params", {}).get( "api_base", "" ) if _deployment_api_base is not None: _cost_per_deployment[_deployment_api_base] = item_cost # -------------- # # End of Debugging Logic # -------------- # if ( item_tpm + input_tokens > _deployment_tpm or item_rpm + 1 > _deployment_rpm ): # if user passed in tpm / rpm in the model_list continue else: potential_deployments.append((_deployment, item_cost)) if len(potential_deployments) == 0: return None potential_deployments = sorted(potential_deployments, key=lambda x: x[1]) selected_deployment = potential_deployments[0][0] return selected_deployment