From 1c8761111fdbb5ee11cb09593dbae8fff11f1407 Mon Sep 17 00:00:00 2001 From: Krish Dholakia Date: Sat, 9 Aug 2025 16:09:51 -0700 Subject: [PATCH] Router - reduce p99 latency w/ redis enabled by 50% + OTEL - track pre_call hook latency (#13362) * feat(proxy/utils.py): track pre-call hooks in OTEL some pre call hooks can cause latency in high traffic - make sure this is tracked * fix(router.py): move redis call on deployment_callback_on_success to pipeline operation reduces p99 latency by half when redis is enabled * fix(parallel_request_limiter_v3.py): only run check if any item has rate limits set Prevents unnecessary latency added by rate limit checks * test: add unit tests * Latency Improvements: only track tpm/rpm usage when set on deployment+ LLM Caching - use an in-memory cache to reduce redis calls + OTEL - track time spent on LLM caching (#13472) * fix(router.py): only track usage for deployments with tpm/rpm set ensures additional latency avoided for non-tpm/rpm models * fix(caching_handler.py): log time spent on request get cache to OTEL enables easy debugging of call latency * fix(caching_handler.py): use dual cache object for in-memory caching + trace redis call within caching handler * fix(caching_handler.py): working in-memory cache for redis calls ensures dual cache works when redis cache setup for llm calls makes calls quicker by only checking redis when in-memory cache missed for llm api call * test: remove redundant test * test: add unit tests --- litellm/caching/caching.py | 74 ++++--- litellm/caching/caching_handler.py | 72 ++++-- litellm/caching/redis_cache.py | 103 ++++++--- litellm/proxy/_experimental/mcp_server/db.py | 22 +- litellm/proxy/_new_secret_config.yaml | 25 ++- .../hooks/parallel_request_limiter_v3.py | 85 ++++--- litellm/proxy/proxy_server.py | 18 +- litellm/proxy/utils.py | 93 +++++--- litellm/router.py | 62 +++++- .../test_amazing_vertex_completion.py | 6 +- tests/local_testing/test_router_utils.py | 40 ++-- .../hooks/test_parallel_request_limiter_v3.py | 207 ++++++++++++++++++ 12 files changed, 627 insertions(+), 180 deletions(-) diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index 1455e011bc..175eaa112a 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -28,13 +28,13 @@ from .azure_blob_cache import AzureBlobCache from .base_cache import BaseCache from .disk_cache import DiskCache from .dual_cache import DualCache # noqa +from .gcs_cache import GCSCache from .in_memory_cache import InMemoryCache from .qdrant_semantic_cache import QdrantSemanticCache from .redis_cache import RedisCache from .redis_cluster_cache import RedisClusterCache from .redis_semantic_cache import RedisSemanticCache from .s3_cache import S3Cache -from .gcs_cache import GCSCache def print_verbose(print_statement): @@ -177,7 +177,7 @@ class Cache: cluster_kwargs["gcp_service_account"] = gcp_service_account if gcp_ssl_ca_certs is not None: cluster_kwargs["gcp_ssl_ca_certs"] = gcp_ssl_ca_certs - + self.cache: BaseCache = RedisClusterCache(**cluster_kwargs) else: self.cache = RedisCache( @@ -481,7 +481,7 @@ class Cache: return cached_response return cached_result - def get_cache(self, **kwargs): + def get_cache(self, dynamic_cache_object: Optional[BaseCache] = None, **kwargs): """ Retrieves the cached result for the given arguments. @@ -507,8 +507,12 @@ class Cache: or cache_control_args.get("s-max-age") or float("inf") ) - cached_result = self.cache.get_cache(cache_key, messages=messages) - cached_result = self.cache.get_cache(cache_key, messages=messages) + if dynamic_cache_object is not None: + cached_result = dynamic_cache_object.get_cache( + cache_key, messages=messages + ) + else: + cached_result = self.cache.get_cache(cache_key, messages=messages) return self._get_cache_logic( cached_result=cached_result, max_age=max_age ) @@ -516,7 +520,9 @@ class Cache: print_verbose(f"An exception occurred: {traceback.format_exc()}") return None - async def async_get_cache(self, **kwargs): + async def async_get_cache( + self, dynamic_cache_object: Optional[BaseCache] = None, **kwargs + ): """ Async get cache implementation. @@ -537,7 +543,14 @@ class Cache: max_age = cache_control_args.get( "s-max-age", cache_control_args.get("s-maxage", float("inf")) ) - cached_result = await self.cache.async_get_cache(cache_key, **kwargs) + if dynamic_cache_object is not None: + cached_result = await dynamic_cache_object.async_get_cache( + cache_key, **kwargs + ) + else: + cached_result = await self.cache.async_get_cache( + cache_key, **kwargs + ) return self._get_cache_logic( cached_result=cached_result, max_age=max_age ) @@ -596,7 +609,9 @@ class Cache: except Exception as e: verbose_logger.exception(f"LiteLLM Cache: Excepton add_cache: {str(e)}") - async def async_add_cache(self, result, **kwargs): + async def async_add_cache( + self, result, dynamic_cache_object: Optional[BaseCache], **kwargs + ): """ Async implementation of add_cache """ @@ -610,12 +625,18 @@ class Cache: cache_key, cached_data, kwargs = self._add_cache_logic( result=result, **kwargs ) - - await self.cache.async_set_cache(cache_key, cached_data, **kwargs) + if dynamic_cache_object is not None: + await dynamic_cache_object.async_set_cache( + cache_key, cached_data, **kwargs + ) + else: + await self.cache.async_set_cache(cache_key, cached_data, **kwargs) except Exception as e: verbose_logger.exception(f"LiteLLM Cache: Excepton add_cache: {str(e)}") - def _convert_to_cached_embedding(self, embedding_response: Any, model: Optional[str]) -> CachedEmbedding: + def _convert_to_cached_embedding( + self, embedding_response: Any, model: Optional[str] + ) -> CachedEmbedding: """ Convert any embedding response into the standardized CachedEmbedding TypedDict format. """ @@ -627,7 +648,7 @@ class Cache: "object": embedding_response.get("object"), "model": model, } - elif hasattr(embedding_response, 'model_dump'): + elif hasattr(embedding_response, "model_dump"): data = embedding_response.model_dump() return { "embedding": data.get("embedding"), @@ -646,7 +667,6 @@ class Cache: except KeyError as e: raise ValueError(f"Missing expected key in embedding response: {e}") - def add_embedding_response_to_cache( self, result: EmbeddingResponse, @@ -657,18 +677,22 @@ class Cache: preset_cache_key = self.get_cache_key(**{**kwargs, "input": input}) kwargs["cache_key"] = preset_cache_key embedding_response = result.data[idx_in_result_data] - + # Always convert to properly typed CachedEmbedding model_name = result.model - embedding_dict: CachedEmbedding = self._convert_to_cached_embedding(embedding_response, model_name) - + embedding_dict: CachedEmbedding = self._convert_to_cached_embedding( + embedding_response, model_name + ) + cache_key, cached_data, kwargs = self._add_cache_logic( result=embedding_dict, **kwargs, ) return cache_key, cached_data, kwargs - async def async_add_cache_pipeline(self, result, **kwargs): + async def async_add_cache_pipeline( + self, result, dynamic_cache_object: Optional[BaseCache], **kwargs + ): """ Async implementation of add_cache for Embedding calls @@ -697,14 +721,14 @@ class Cache: ) cache_list.append((cache_key, cached_data)) - await self.cache.async_set_cache_pipeline(cache_list=cache_list, **kwargs) - # if async_set_cache_pipeline: - # await async_set_cache_pipeline(cache_list=cache_list, **kwargs) - # else: - # tasks = [] - # for val in cache_list: - # tasks.append(self.cache.async_set_cache(val[0], val[1], **kwargs)) - # await asyncio.gather(*tasks) + if dynamic_cache_object is not None: + await dynamic_cache_object.async_set_cache_pipeline( + cache_list=cache_list, **kwargs + ) + else: + await self.cache.async_set_cache_pipeline( + cache_list=cache_list, **kwargs + ) except Exception as e: verbose_logger.exception(f"LiteLLM Cache: Excepton add_cache: {str(e)}") diff --git a/litellm/caching/caching_handler.py b/litellm/caching/caching_handler.py index dcc59b2071..f41b745bb1 100644 --- a/litellm/caching/caching_handler.py +++ b/litellm/caching/caching_handler.py @@ -1,5 +1,5 @@ """ -This contains LLMCachingHandler +This contains LLMCachingHandler This exposes two methods: - async_get_cache @@ -18,6 +18,7 @@ import asyncio import datetime import inspect import threading +from functools import lru_cache, wraps from typing import ( TYPE_CHECKING, Any, @@ -35,11 +36,13 @@ from pydantic import BaseModel import litellm from litellm._logging import print_verbose, verbose_logger +from litellm._service_logger import ServiceLogging +from litellm.caching import InMemoryCache from litellm.caching.caching import S3Cache -from litellm.types.caching import CachedEmbedding from litellm.litellm_core_utils.logging_utils import ( _assemble_complete_response_from_streaming_chunks, ) +from litellm.types.caching import CachedEmbedding from litellm.types.rerank import RerankResponse from litellm.types.utils import ( CallTypes, @@ -68,7 +71,12 @@ class CachingHandlerResponse(BaseModel): cached_result: Optional[Any] = None final_embedding_cached_response: Optional[EmbeddingResponse] = None - embedding_all_elements_cache_hit: bool = False # this is set to True when all elements in the list have a cache hit in the embedding cache, if true return the final_embedding_cached_response no need to make an API call + embedding_all_elements_cache_hit: bool = ( + False # this is set to True when all elements in the list have a cache hit in the embedding cache, if true return the final_embedding_cached_response no need to make an API call + ) + + +in_memory_cache_obj = InMemoryCache() class LLMCachingHandler: @@ -78,11 +86,20 @@ class LLMCachingHandler: request_kwargs: Dict[str, Any], start_time: datetime.datetime, ): + from litellm.caching import DualCache, RedisCache + self.async_streaming_chunks: List[ModelResponse] = [] self.sync_streaming_chunks: List[ModelResponse] = [] self.request_kwargs = request_kwargs self.original_function = original_function self.start_time = start_time + if litellm.cache is not None and isinstance(litellm.cache.cache, RedisCache): + self.dual_cache: Optional[DualCache] = DualCache( + redis_cache=litellm.cache.cache, + in_memory_cache=in_memory_cache_obj, + ) + else: + self.dual_cache = None pass async def _async_get_cache( @@ -115,10 +132,16 @@ class LLMCachingHandler: Raises: None """ + from litellm.litellm_core_utils.core_helpers import ( + _get_parent_otel_span_from_kwargs, + ) from litellm.utils import CustomStreamWrapper + kwargs = kwargs.copy() args = args or () + parent_otel_span = _get_parent_otel_span_from_kwargs(kwargs) + kwargs["parent_otel_span"] = parent_otel_span final_embedding_cached_response: Optional[EmbeddingResponse] = None embedding_all_elements_cache_hit: bool = False cached_result: Optional[Any] = None @@ -306,13 +329,15 @@ class LLMCachingHandler: else: raise ValueError("input must be a string or a list") - def _extract_model_from_cached_results(self, non_null_list: List[Tuple[int, CachedEmbedding]]) -> Optional[str]: + def _extract_model_from_cached_results( + self, non_null_list: List[Tuple[int, CachedEmbedding]] + ) -> Optional[str]: """ Helper method to extract the model name from cached results. - + Args: non_null_list: List of (idx, cr) tuples where cr is the cached result dict - + Returns: Optional[str]: The model name if found, None otherwise """ @@ -558,7 +583,12 @@ class LLMCachingHandler: preset_cache_key = litellm.cache.get_cache_key( **{**new_kwargs, "input": i} ) - tasks.append(litellm.cache.async_get_cache(cache_key=preset_cache_key)) + tasks.append( + litellm.cache.async_get_cache( + cache_key=preset_cache_key, + dynamic_cache_object=self.dual_cache, + ) + ) cached_result = await asyncio.gather(*tasks) ## check if cached result is None ## if cached_result is not None and isinstance(cached_result, list): @@ -567,9 +597,14 @@ class LLMCachingHandler: cached_result = None else: if litellm.cache._supports_async() is True: - cached_result = await litellm.cache.async_get_cache(**new_kwargs) + ## check if dual cache is supported ## + cached_result = await litellm.cache.async_get_cache( + dynamic_cache_object=self.dual_cache, **new_kwargs + ) else: # for s3 caching. [NOT RECOMMENDED IN PROD - this will slow down responses since boto3 is sync] - cached_result = litellm.cache.get_cache(**new_kwargs) + cached_result = litellm.cache.get_cache( + dynamic_cache_object=self.dual_cache, **new_kwargs + ) return cached_result def _convert_cached_result_to_model_response( @@ -735,6 +770,9 @@ class LLMCachingHandler: Raises: None """ + from litellm.litellm_core_utils.core_helpers import ( + _get_parent_otel_span_from_kwargs, + ) if litellm.cache is None: return @@ -746,6 +784,8 @@ class LLMCachingHandler: args, ) ) + parent_otel_span = _get_parent_otel_span_from_kwargs(new_kwargs) + new_kwargs["parent_otel_span"] = parent_otel_span # [OPTIONAL] ADD TO CACHE if self._should_store_result_in_cache( original_function=original_function, kwargs=new_kwargs @@ -764,7 +804,9 @@ class LLMCachingHandler: ) # s3 doesn't support bulk writing. Exclude. ): asyncio.create_task( - litellm.cache.async_add_cache_pipeline(result, **new_kwargs) + litellm.cache.async_add_cache_pipeline( + result, dynamic_cache_object=self.dual_cache, **new_kwargs + ) ) elif isinstance(litellm.cache.cache, S3Cache): threading.Thread( @@ -775,7 +817,9 @@ class LLMCachingHandler: else: asyncio.create_task( litellm.cache.async_add_cache( - result.model_dump_json(), **new_kwargs + result.model_dump_json(), + dynamic_cache_object=self.dual_cache, + **new_kwargs, ) ) else: @@ -933,9 +977,9 @@ class LLMCachingHandler: } if litellm.cache is not None: - litellm_params[ - "preset_cache_key" - ] = litellm.cache._get_preset_cache_key_from_kwargs(**kwargs) + litellm_params["preset_cache_key"] = ( + litellm.cache._get_preset_cache_key_from_kwargs(**kwargs) + ) else: litellm_params["preset_cache_key"] = None diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index b8091187bf..47bc0222ed 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -43,6 +43,45 @@ else: Span = Any +def _get_call_stack_info(num_frames: int = 2) -> str: + """ + Get the function names from the previous 1-2 functions in the call stack. + + Args: + num_frames: Number of previous frames to include (default: 2) + + Returns: + A string with format "current_function <- caller_function [<- grandparent_function]" + """ + try: + current_frame = inspect.currentframe() + if current_frame is None: + return "unknown" + + # Skip this function and the immediate caller (which sets call_type) + f_back = current_frame.f_back + if f_back is None: + return "unknown" + frame = f_back.f_back + if frame is None: + return "unknown" + function_names = [] + + for _ in range(num_frames): + if frame is None: + break + func_name = frame.f_code.co_name + function_names.append(func_name) + frame = frame.f_back + + if not function_names: + return "unknown" + + return " <- ".join(function_names) + except Exception: + return "unknown" + + class RedisCache(BaseCache): # if users don't provider one, use the default litellm cache @@ -181,7 +220,7 @@ class RedisCache(BaseCache): self.service_logger_obj.service_success_hook( service=ServiceTypes.REDIS, duration=_duration, - call_type="set_cache", + call_type=f"set_cache <- {_get_call_stack_info()}", start_time=start_time, end_time=end_time, ) @@ -205,7 +244,7 @@ class RedisCache(BaseCache): self.service_logger_obj.service_success_hook( service=ServiceTypes.REDIS, duration=_duration, - call_type="increment_cache", + call_type=f"increment_cache <- {_get_call_stack_info()}", start_time=start_time, end_time=end_time, ) @@ -219,7 +258,7 @@ class RedisCache(BaseCache): self.service_logger_obj.service_success_hook( service=ServiceTypes.REDIS, duration=_duration, - call_type="increment_cache_ttl", + call_type=f"increment_cache_ttl <- {_get_call_stack_info()}", start_time=start_time, end_time=end_time, ) @@ -232,7 +271,7 @@ class RedisCache(BaseCache): self.service_logger_obj.service_success_hook( service=ServiceTypes.REDIS, duration=_duration, - call_type="increment_cache_expire", + call_type=f"increment_cache_expire <- {_get_call_stack_info()}", start_time=start_time, end_time=end_time, ) @@ -271,7 +310,7 @@ class RedisCache(BaseCache): self.service_logger_obj.async_service_success_hook( service=ServiceTypes.REDIS, duration=_duration, - call_type="async_scan_iter", + call_type=f"async_scan_iter <- {_get_call_stack_info()}", start_time=start_time, end_time=end_time, ) @@ -287,7 +326,7 @@ class RedisCache(BaseCache): service=ServiceTypes.REDIS, duration=_duration, error=e, - call_type="async_scan_iter", + call_type=f"async_scan_iter <- {_get_call_stack_info()}", start_time=start_time, end_time=end_time, ) @@ -341,7 +380,7 @@ class RedisCache(BaseCache): start_time=start_time, end_time=end_time, parent_otel_span=_get_parent_otel_span_from_kwargs(kwargs), - call_type="async_set_cache", + call_type=f"async_set_cache <- {_get_call_stack_info()}", ) ) verbose_logger.error( @@ -374,7 +413,7 @@ class RedisCache(BaseCache): self.service_logger_obj.async_service_success_hook( service=ServiceTypes.REDIS, duration=_duration, - call_type="async_set_cache", + call_type=f"async_set_cache <- {_get_call_stack_info()}", start_time=start_time, end_time=end_time, parent_otel_span=_get_parent_otel_span_from_kwargs(kwargs), @@ -390,7 +429,7 @@ class RedisCache(BaseCache): service=ServiceTypes.REDIS, duration=_duration, error=e, - call_type="async_set_cache", + call_type=f"async_set_cache <- {_get_call_stack_info()}", start_time=start_time, end_time=end_time, parent_otel_span=_get_parent_otel_span_from_kwargs(kwargs), @@ -463,7 +502,7 @@ class RedisCache(BaseCache): self.service_logger_obj.async_service_success_hook( service=ServiceTypes.REDIS, duration=_duration, - call_type="async_set_cache_pipeline", + call_type=f"async_set_cache_pipeline <- {_get_call_stack_info()}", start_time=start_time, end_time=end_time, parent_otel_span=_get_parent_otel_span_from_kwargs(kwargs), @@ -479,7 +518,7 @@ class RedisCache(BaseCache): service=ServiceTypes.REDIS, duration=_duration, error=e, - call_type="async_set_cache_pipeline", + call_type=f"async_set_cache_pipeline <- {_get_call_stack_info()}", start_time=start_time, end_time=end_time, parent_otel_span=_get_parent_otel_span_from_kwargs(kwargs), @@ -528,7 +567,7 @@ class RedisCache(BaseCache): start_time=start_time, end_time=end_time, parent_otel_span=_get_parent_otel_span_from_kwargs(kwargs), - call_type="async_set_cache_sadd", + call_type=f"async_set_cache_sadd <- {_get_call_stack_info()}", ) ) # NON blocking - notify users Redis is throwing an exception @@ -554,7 +593,7 @@ class RedisCache(BaseCache): self.service_logger_obj.async_service_success_hook( service=ServiceTypes.REDIS, duration=_duration, - call_type="async_set_cache_sadd", + call_type=f"async_set_cache_sadd <- {_get_call_stack_info()}", start_time=start_time, end_time=end_time, parent_otel_span=_get_parent_otel_span_from_kwargs(kwargs), @@ -568,7 +607,7 @@ class RedisCache(BaseCache): service=ServiceTypes.REDIS, duration=_duration, error=e, - call_type="async_set_cache_sadd", + call_type=f"async_set_cache_sadd <- {_get_call_stack_info()}", start_time=start_time, end_time=end_time, parent_otel_span=_get_parent_otel_span_from_kwargs(kwargs), @@ -620,7 +659,7 @@ class RedisCache(BaseCache): self.service_logger_obj.async_service_success_hook( service=ServiceTypes.REDIS, duration=_duration, - call_type="async_increment", + call_type=f"async_increment <- {_get_call_stack_info()}", start_time=start_time, end_time=end_time, parent_otel_span=parent_otel_span, @@ -636,7 +675,7 @@ class RedisCache(BaseCache): service=ServiceTypes.REDIS, duration=_duration, error=e, - call_type="async_increment", + call_type=f"async_increment <- {_get_call_stack_info()}", start_time=start_time, end_time=end_time, parent_otel_span=parent_otel_span, @@ -683,7 +722,7 @@ class RedisCache(BaseCache): self.service_logger_obj.service_success_hook( service=ServiceTypes.REDIS, duration=_duration, - call_type="get_cache", + call_type=f"get_cache <- {_get_call_stack_info()}", start_time=start_time, end_time=end_time, parent_otel_span=parent_otel_span, @@ -745,7 +784,7 @@ class RedisCache(BaseCache): self.service_logger_obj.service_success_hook( service=ServiceTypes.REDIS, duration=_duration, - call_type="batch_get_cache", + call_type=f"batch_get_cache <- {_get_call_stack_info()}", start_time=start_time, end_time=end_time, parent_otel_span=parent_otel_span, @@ -790,7 +829,7 @@ class RedisCache(BaseCache): self.service_logger_obj.async_service_success_hook( service=ServiceTypes.REDIS, duration=_duration, - call_type="async_get_cache", + call_type=f"async_get_cache <- {_get_call_stack_info()}", start_time=start_time, end_time=end_time, parent_otel_span=parent_otel_span, @@ -806,7 +845,7 @@ class RedisCache(BaseCache): service=ServiceTypes.REDIS, duration=_duration, error=e, - call_type="async_get_cache", + call_type=f"async_get_cache <- {_get_call_stack_info()}", start_time=start_time, end_time=end_time, parent_otel_span=parent_otel_span, @@ -851,7 +890,7 @@ class RedisCache(BaseCache): self.service_logger_obj.async_service_success_hook( service=ServiceTypes.REDIS, duration=_duration, - call_type="async_batch_get_cache", + call_type=f"async_batch_get_cache <- {_get_call_stack_info()}", start_time=start_time, end_time=end_time, parent_otel_span=parent_otel_span, @@ -879,7 +918,7 @@ class RedisCache(BaseCache): service=ServiceTypes.REDIS, duration=_duration, error=e, - call_type="async_batch_get_cache", + call_type=f"async_batch_get_cache <- {_get_call_stack_info()}", start_time=start_time, end_time=end_time, parent_otel_span=parent_otel_span, @@ -903,7 +942,7 @@ class RedisCache(BaseCache): self.service_logger_obj.service_success_hook( service=ServiceTypes.REDIS, duration=_duration, - call_type="sync_ping", + call_type=f"sync_ping <- {_get_call_stack_info()}", start_time=start_time, end_time=end_time, ) @@ -917,7 +956,7 @@ class RedisCache(BaseCache): service=ServiceTypes.REDIS, duration=_duration, error=e, - call_type="sync_ping", + call_type=f"sync_ping <- {_get_call_stack_info()}", ) verbose_logger.error( f"LiteLLM Redis Cache PING: - Got exception from REDIS : {str(e)}" @@ -938,7 +977,7 @@ class RedisCache(BaseCache): self.service_logger_obj.async_service_success_hook( service=ServiceTypes.REDIS, duration=_duration, - call_type="async_ping", + call_type=f"async_ping <- {_get_call_stack_info()}", ) ) return response @@ -952,7 +991,7 @@ class RedisCache(BaseCache): service=ServiceTypes.REDIS, duration=_duration, error=e, - call_type="async_ping", + call_type=f"async_ping <- {_get_call_stack_info()}", ) ) verbose_logger.error( @@ -1051,7 +1090,7 @@ class RedisCache(BaseCache): self.service_logger_obj.async_service_success_hook( service=ServiceTypes.REDIS, duration=_duration, - call_type="async_increment_pipeline", + call_type=f"async_increment_pipeline <- {_get_call_stack_info()}", start_time=start_time, end_time=end_time, parent_otel_span=_get_parent_otel_span_from_kwargs(kwargs), @@ -1067,7 +1106,7 @@ class RedisCache(BaseCache): service=ServiceTypes.REDIS, duration=_duration, error=e, - call_type="async_increment_pipeline", + call_type=f"async_increment_pipeline <- {_get_call_stack_info()}", start_time=start_time, end_time=end_time, parent_otel_span=_get_parent_otel_span_from_kwargs(kwargs), @@ -1131,7 +1170,7 @@ class RedisCache(BaseCache): self.service_logger_obj.async_service_success_hook( service=ServiceTypes.REDIS, duration=_duration, - call_type="async_rpush", + call_type=f"async_rpush <- {_get_call_stack_info()}", ) ) return response @@ -1145,7 +1184,7 @@ class RedisCache(BaseCache): service=ServiceTypes.REDIS, duration=_duration, error=e, - call_type="async_rpush", + call_type=f"async_rpush <- {_get_call_stack_info()}", ) ) verbose_logger.error( @@ -1202,7 +1241,7 @@ class RedisCache(BaseCache): self.service_logger_obj.async_service_success_hook( service=ServiceTypes.REDIS, duration=_duration, - call_type="async_lpop", + call_type=f"async_lpop <- {_get_call_stack_info()}", ) ) @@ -1230,7 +1269,7 @@ class RedisCache(BaseCache): service=ServiceTypes.REDIS, duration=_duration, error=e, - call_type="async_lpop", + call_type=f"async_lpop <- {_get_call_stack_info()}", ) ) verbose_logger.error( diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index 3d90c99eee..d5d9f97890 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -1,6 +1,7 @@ import uuid from typing import Any, Dict, Iterable, List, Optional, Set, Union +from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ( LiteLLM_MCPServerTable, LiteLLM_ObjectPermissionTable, @@ -53,11 +54,20 @@ async def get_all_mcp_servers( """ Returns all of the mcp servers from the db """ - mcp_servers = await prisma_client.db.litellm_mcpservertable.find_many() + try: + mcp_servers = await prisma_client.db.litellm_mcpservertable.find_many() - return [ - LiteLLM_MCPServerTable(**mcp_server.model_dump()) for mcp_server in mcp_servers - ] + return [ + LiteLLM_MCPServerTable(**mcp_server.model_dump()) + for mcp_server in mcp_servers + ] + except Exception as e: + verbose_proxy_logger.debug( + "litellm.proxy._experimental.mcp_server.db.py::get_all_mcp_servers - {}".format( + str(e) + ) + ) + return [] async def get_mcp_server( @@ -91,9 +101,7 @@ async def get_mcp_servers( ) final_mcp_servers: List[LiteLLM_MCPServerTable] = [] for _mcp_server in _mcp_servers: - final_mcp_servers.append( - LiteLLM_MCPServerTable(**_mcp_server.model_dump()) - ) + final_mcp_servers.append(LiteLLM_MCPServerTable(**_mcp_server.model_dump())) return final_mcp_servers diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index b587913a08..cb41dfbd75 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -1,5 +1,24 @@ +general_settings: + store_model_in_db: true + database_connection_pool_limit: 20 + store_prompts_in_spend_logs: true + maximum_spend_logs_retention_period: "14d" + model_list: - - model_name: openai-test + - model_name: fake-openai-endpoint litellm_params: - model: gpt-3.5-turbo - api_key: os.environ/OPENAI_API_KEY \ No newline at end of file + model: openai/my-fake-model + api_key: my-fake-key + api_base: https://exampleopenaiendpoint-production.up.railway.app/ + - model_name: bedrock-thinking-us.anthropic.claude-3-7-sonnet-20250219-v1:0 + litellm_params: + model: bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0 + thinking: {"type": "enabled", "budget_tokens": 1024} + max_tokens: 1080 + merge_reasoning_content_in_choices: true + +litellm_settings: + return_response_headers: true + store_audit_logs: true + callbacks: ["prometheus", "resend_email", "otel"] + cache: true \ No newline at end of file diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index b5fd819eb2..dde0542c7d 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -1,8 +1,9 @@ """ -This is a rate limiter implementation based on a similar one by Envoy proxy. +This is a rate limiter implementation based on a similar one by Envoy proxy. This is currently in development and not yet ready for production. """ + import os from datetime import datetime from typing import ( @@ -309,13 +310,15 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): continue key_metadata[window_key] = { - "requests_limit": int(requests_limit) - if requests_limit is not None - else None, + "requests_limit": ( + int(requests_limit) if requests_limit is not None else None + ), "tokens_limit": int(tokens_limit) if tokens_limit is not None else None, - "max_parallel_requests_limit": int(max_parallel_requests_limit) - if max_parallel_requests_limit is not None - else None, + "max_parallel_requests_limit": ( + int(max_parallel_requests_limit) + if max_parallel_requests_limit is not None + else None + ), "window_size": int(window_size), "descriptor_key": descriptor_key, } @@ -394,7 +397,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): descriptors = [] # API Key rate limits - if user_api_key_dict.api_key: + if user_api_key_dict.api_key and ( + user_api_key_dict.rpm_limit is not None + or user_api_key_dict.tpm_limit is not None + or user_api_key_dict.max_parallel_requests is not None + ): descriptors.append( RateLimitDescriptor( key="api_key", @@ -409,7 +416,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) # User rate limits - if user_api_key_dict.user_id: + if user_api_key_dict.user_id and ( + user_api_key_dict.user_rpm_limit is not None + or user_api_key_dict.user_tpm_limit is not None + ): descriptors.append( RateLimitDescriptor( key="user", @@ -423,7 +433,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) # Team rate limits - if user_api_key_dict.team_id: + if user_api_key_dict.team_id and ( + user_api_key_dict.team_rpm_limit is not None + or user_api_key_dict.team_tpm_limit is not None + ): descriptors.append( RateLimitDescriptor( key="team", @@ -437,7 +450,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) # End user rate limits - if user_api_key_dict.end_user_id: + if user_api_key_dict.end_user_id and ( + user_api_key_dict.end_user_rpm_limit is not None + or user_api_key_dict.end_user_tpm_limit is not None + ): descriptors.append( RateLimitDescriptor( key="end_user", @@ -483,28 +499,29 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) ) - # Check rate limits - response = await self.should_rate_limit( - descriptors=descriptors, - parent_otel_span=user_api_key_dict.parent_otel_span, - ) + # Only check rate limits if we have descriptors with actual limits + if descriptors: + response = await self.should_rate_limit( + descriptors=descriptors, + parent_otel_span=user_api_key_dict.parent_otel_span, + ) - if response["overall_code"] == "OVER_LIMIT": - # Find which descriptor hit the limit - for i, status in enumerate(response["statuses"]): - if status["code"] == "OVER_LIMIT": - descriptor = descriptors[i] - raise HTTPException( - status_code=429, - detail=f"Rate limit exceeded for {descriptor['key']}: {descriptor['value']}. Remaining: {status['limit_remaining']}", - headers={ - "retry-after": str(self.window_size) - }, # Retry after 1 minute - ) + if response["overall_code"] == "OVER_LIMIT": + # Find which descriptor hit the limit + for i, status in enumerate(response["statuses"]): + if status["code"] == "OVER_LIMIT": + descriptor = descriptors[i] + raise HTTPException( + status_code=429, + detail=f"Rate limit exceeded for {descriptor['key']}: {descriptor['value']}. Remaining: {status['limit_remaining']}", + headers={ + "retry-after": str(self.window_size) + }, # Retry after 1 minute + ) - else: - # add descriptors to request headers - data["litellm_proxy_rate_limit_response"] = response + else: + # add descriptors to request headers + data["litellm_proxy_rate_limit_response"] = response def _create_pipeline_operations( self, @@ -690,9 +707,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): from litellm.types.caching import RedisPipelineIncrementOperation try: - litellm_parent_otel_span: Union[ - Span, None - ] = _get_parent_otel_span_from_kwargs(kwargs) + litellm_parent_otel_span: Union[Span, None] = ( + _get_parent_otel_span_from_kwargs(kwargs) + ) user_api_key = kwargs["litellm_params"]["metadata"].get("user_api_key") pipeline_operations: List[RedisPipelineIncrementOperation] = [] diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f29ab7c588..6b33e92c48 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2720,7 +2720,10 @@ class ProxyConfig: alert_types=general_settings["alert_types"], llm_router=llm_router ) - if _general_settings is not None and "alert_to_webhook_url" in _general_settings: + if ( + _general_settings is not None + and "alert_to_webhook_url" in _general_settings + ): general_settings["alert_to_webhook_url"] = _general_settings[ "alert_to_webhook_url" ] @@ -2940,9 +2943,16 @@ class ProxyConfig: async def _init_prompts_in_db(self, prisma_client: PrismaClient): from litellm.proxy.prompts.prompt_registry import IN_MEMORY_PROMPT_REGISTRY - prompts_in_db = await prisma_client.db.litellm_prompttable.find_many() - for prompt in prompts_in_db: - IN_MEMORY_PROMPT_REGISTRY.initialize_prompt(prompt=prompt) + try: + prompts_in_db = await prisma_client.db.litellm_prompttable.find_many() + for prompt in prompts_in_db: + IN_MEMORY_PROMPT_REGISTRY.initialize_prompt(prompt=prompt) + except Exception as e: + verbose_proxy_logger.debug( + "litellm.proxy.proxy_server.py::ProxyConfig:_init_prompts_in_db - {}".format( + str(e) + ) + ) async def _init_guardrails_in_db(self, prisma_client: PrismaClient): from litellm.proxy.guardrails.guardrail_registry import ( diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 0e6b242bdb..f3affa707a 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -22,7 +22,7 @@ from typing import ( overload, ) -from litellm.constants import MAX_TEAM_LIST_LIMIT, DEFAULT_MODEL_CREATED_AT_TIME +from litellm.constants import DEFAULT_MODEL_CREATED_AT_TIME, MAX_TEAM_LIST_LIMIT from litellm.proxy._types import ( DB_CONNECTION_ERROR_TYPES, CommonProxyErrors, @@ -448,7 +448,6 @@ class ProxyLogging: litellm_parent_otel_span=None, ) - def _convert_user_api_key_auth_to_dict(self, user_api_key_auth_obj): """ Helper function to convert UserAPIKeyAuth object to dictionary. @@ -728,15 +727,19 @@ class ProxyLogging: } return result - def _create_mcp_request_object_from_kwargs(self, kwargs: dict) -> "MCPPreCallRequestObject": + def _create_mcp_request_object_from_kwargs( + self, kwargs: dict + ) -> "MCPPreCallRequestObject": """ Helper function to create MCPPreCallRequestObject from kwargs for standard pre_call_hook. """ from litellm.types.llms.base import HiddenParams from litellm.types.mcp import MCPPreCallRequestObject - user_api_key_auth_dict = self._convert_user_api_key_auth_to_dict(kwargs.get("user_api_key_auth")) - + user_api_key_auth_dict = self._convert_user_api_key_auth_to_dict( + kwargs.get("user_api_key_auth") + ) + return MCPPreCallRequestObject( tool_name=kwargs.get("name", ""), arguments=kwargs.get("arguments", {}), @@ -745,22 +748,23 @@ class ProxyLogging: hidden_params=HiddenParams(), ) - def _convert_mcp_hook_response_to_kwargs(self, response_data: Optional[dict], original_kwargs: dict) -> dict: + def _convert_mcp_hook_response_to_kwargs( + self, response_data: Optional[dict], original_kwargs: dict + ) -> dict: """ Helper function to convert pre_call_hook response back to kwargs for MCP usage. """ if not response_data: return original_kwargs - + # Apply any argument modifications from the hook response modified_kwargs = original_kwargs.copy() - + # If the response contains modified arguments, apply them if response_data.get("modified_arguments"): modified_kwargs["arguments"] = response_data["modified_arguments"] - - return modified_kwargs + return modified_kwargs async def process_pre_call_hook_response(self, response, data, call_type): if isinstance(response, Exception): @@ -894,6 +898,7 @@ class ProxyLogging: try: for callback in litellm.callbacks: + start_time = time.time() _callback = None if isinstance(callback, str): _callback = litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class( @@ -903,15 +908,13 @@ class ProxyLogging: _callback = callback # type: ignore if _callback is not None and isinstance(_callback, CustomGuardrail): from litellm.types.guardrails import GuardrailEventHooks - + event_type = GuardrailEventHooks.pre_call if call_type == "mcp_call": event_type = GuardrailEventHooks.pre_mcp_call - + if ( - _callback.should_run_guardrail( - data=data, event_type=event_type - ) + _callback.should_run_guardrail(data=data, event_type=event_type) is not True ): continue @@ -936,7 +939,7 @@ class ProxyLogging: ): if call_type == "mcp_call" and user_api_key_dict is None: continue - + response = await _callback.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=self.call_details["user_api_key_cache"], @@ -948,6 +951,19 @@ class ProxyLogging: response=response, data=data, call_type=call_type ) + end_time = time.time() + duration = end_time - start_time + if ( + hasattr(self, "service_logging_obj") and duration > 0.01 + ): # only if duration is non-negligible - don't spam the logs + await self.service_logging_obj.async_service_success_hook( + service=ServiceTypes.PROXY_PRE_CALL, + duration=duration, + call_type=f"{_callback.__class__.__name__}", + parent_otel_span=user_api_key_dict.parent_otel_span, + start_time=start_time, + end_time=end_time, + ) return data except Exception as e: raise e @@ -999,7 +1015,9 @@ class ProxyLogging: continue # Convert user_api_key_dict to proper format for async_moderation_hook if call_type == "mcp_call": - user_api_key_auth_dict = self._convert_user_api_key_auth_to_dict(user_api_key_dict) + user_api_key_auth_dict = ( + self._convert_user_api_key_auth_to_dict(user_api_key_dict) + ) else: user_api_key_auth_dict = user_api_key_dict @@ -3668,9 +3686,10 @@ def construct_database_url_from_env_vars() -> Optional[str]: database_url = f"postgresql://{database_username_enc}@{database_host}/{database_name_enc}" return database_url - + return None + async def count_tokens_with_anthropic_api( model_to_use: str, messages: Optional[List[Dict[str, Any]]], @@ -3691,9 +3710,10 @@ async def count_tokens_with_anthropic_api( return None try: - import anthropic import os + import anthropic + # Get Anthropic API key from deployment config anthropic_api_key = None if deployment is not None: @@ -3712,7 +3732,7 @@ async def count_tokens_with_anthropic_api( response = client.beta.messages.count_tokens( model=model_to_use, messages=messages, # type: ignore - betas=["token-counting-2024-11-01"] + betas=["token-counting-2024-11-01"], ) total_tokens = response.input_tokens tokenizer_used = "anthropic_api" @@ -3723,11 +3743,16 @@ async def count_tokens_with_anthropic_api( } except ImportError: - verbose_proxy_logger.warning("Anthropic library not available, falling back to LiteLLM tokenizer") + verbose_proxy_logger.warning( + "Anthropic library not available, falling back to LiteLLM tokenizer" + ) except Exception as e: - verbose_proxy_logger.warning(f"Error calling Anthropic API: {e}, falling back to LiteLLM tokenizer") + verbose_proxy_logger.warning( + f"Error calling Anthropic API: {e}, falling back to LiteLLM tokenizer" + ) return None + async def get_available_models_for_user( user_api_key_dict: "UserAPIKeyAuth", llm_router: Optional["Router"], @@ -3743,7 +3768,7 @@ async def get_available_models_for_user( ) -> List[str]: """ Get the list of models available to a user based on their API key and team permissions. - + Args: user_api_key_dict: User API key authentication object llm_router: LiteLLM router instance @@ -3755,18 +3780,18 @@ async def get_available_models_for_user( include_model_access_groups: Whether to include model access groups only_model_access_groups: Whether to only return model access groups return_wildcard_routes: Whether to return wildcard routes - + Returns: List of model names available to the user """ + from litellm.proxy.auth.auth_checks import get_team_object from litellm.proxy.auth.model_checks import ( + get_complete_model_list, get_key_models, get_team_models, - get_complete_model_list, ) - from litellm.proxy.auth.auth_checks import get_team_object from litellm.proxy.management_endpoints.team_endpoints import validate_membership - + # Get proxy model list and access groups if llm_router is None: proxy_model_list = [] @@ -3831,19 +3856,19 @@ def create_model_info_response( ) -> dict: """ Create a standardized model info response. - + Args: model_id: The model ID provider: The model provider include_metadata: Whether to include metadata fallback_type: Type of fallbacks to include llm_router: LiteLLM router instance - + Returns: Dictionary containing model information """ from litellm.proxy.auth.model_checks import get_all_fallbacks - + model_info = { "id": model_id, "object": "model", @@ -3886,16 +3911,18 @@ def validate_model_access( ) -> None: """ Validate that a model is accessible to the user. - + Args: model_id: The model ID to validate available_models: List of models available to the user - + Raises: HTTPException: If the model is not accessible """ if model_id not in available_models: raise HTTPException( status_code=404, - detail="The model `{}` does not exist or is not accessible".format(model_id) + detail="The model `{}` does not exist or is not accessible".format( + model_id + ), ) diff --git a/litellm/router.py b/litellm/router.py index 3be88596b1..dbd5b4c6b2 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -4308,6 +4308,8 @@ class Router: """ Track remaining tpm/rpm quota for model in model_list """ + from litellm.types.caching import RedisPipelineIncrementOperation + try: standard_logging_object: Optional[StandardLoggingPayload] = kwargs.get( "standard_logging_object", None @@ -4327,6 +4329,39 @@ class Router: elif isinstance(id, int): id = str(id) + ## get deployment info + deployment_info = self.get_deployment(model_id=id) + + if deployment_info is None: + return + else: + deployment_model_info = self.get_router_model_info( + deployment=deployment_info.model_dump(), + received_model_name=model_group, + ) + # get tpm/rpm from deployment info + tpm = deployment_info.get("tpm", None) + rpm = deployment_info.get("rpm", None) + + ## check tpm/rpm in litellm_params + tpm_litellm_params = deployment_info.litellm_params.tpm + rpm_litellm_params = deployment_info.litellm_params.rpm + + ## check tpm/rpm in model_info + tpm_model_info = deployment_model_info.get("tpm", None) + rpm_model_info = deployment_model_info.get("rpm", None) + + ## if all are none, return - no need to track current tpm/rpm usage for models with no tpm/rpm set + if ( + tpm is None + and rpm is None + and tpm_litellm_params is None + and rpm_litellm_params is None + and tpm_model_info is None + and rpm_model_info is None + ): + return + parent_otel_span = _get_parent_otel_span_from_kwargs(kwargs) total_tokens: float = standard_logging_object.get("total_tokens", 0) @@ -4346,23 +4381,32 @@ class Router: # ------------ # update cache + pipeline_operations: List[RedisPipelineIncrementOperation] = [] + ## TPM - await self.cache.async_increment_cache( - key=tpm_key, - value=total_tokens, - parent_otel_span=parent_otel_span, - ttl=RoutingArgs.ttl.value, + pipeline_operations.append( + RedisPipelineIncrementOperation( + key=tpm_key, + increment_value=total_tokens, + ttl=RoutingArgs.ttl.value, + ) ) ## RPM rpm_key = RouterCacheEnum.RPM.value.format( id=id, current_minute=current_minute, model=deployment_name ) - await self.cache.async_increment_cache( - key=rpm_key, - value=1, + pipeline_operations.append( + RedisPipelineIncrementOperation( + key=rpm_key, + increment_value=1, + ttl=RoutingArgs.ttl.value, + ) + ) + + await self.cache.async_increment_cache_pipeline( + increment_list=pipeline_operations, parent_otel_span=parent_otel_span, - ttl=RoutingArgs.ttl.value, ) increment_deployment_successes_for_current_minute( diff --git a/tests/local_testing/test_amazing_vertex_completion.py b/tests/local_testing/test_amazing_vertex_completion.py index 398a57e340..23502a1434 100644 --- a/tests/local_testing/test_amazing_vertex_completion.py +++ b/tests/local_testing/test_amazing_vertex_completion.py @@ -167,7 +167,6 @@ async def test_get_response(): pytest.fail(f"An error occurred - {str(e)}") - @pytest.mark.skip( reason="Local test. Vertex AI Quota is low. Leads to rate limit errors on ci/cd." ) @@ -547,7 +546,6 @@ def test_completion_function_plus_pdf(load_pdf): except Exception as e: pytest.fail("Got={}".format(str(e))) - def encode_image(image_path): import base64 @@ -765,9 +763,7 @@ def test_gemini_pro_grounding(value_in_dict): # @pytest.mark.skip(reason="exhausted vertex quota. need to refactor to mock the call") -@pytest.mark.parametrize( - "model", ["vertex_ai_beta/gemini-1.5-pro"] -) # "vertex_ai", +@pytest.mark.parametrize("model", ["vertex_ai_beta/gemini-1.5-pro"]) # "vertex_ai", @pytest.mark.parametrize("sync_mode", [True]) # "vertex_ai", @pytest.mark.asyncio @pytest.mark.flaky(retries=3, delay=1) diff --git a/tests/local_testing/test_router_utils.py b/tests/local_testing/test_router_utils.py index f3b0238be5..e5ead7a701 100644 --- a/tests/local_testing/test_router_utils.py +++ b/tests/local_testing/test_router_utils.py @@ -240,7 +240,7 @@ async def test_call_router_callbacks_on_success(): ) with patch.object( - router.cache, "async_increment_cache", new=AsyncMock() + router.cache, "async_increment_cache_pipeline", new=AsyncMock() ) as mock_callback: await router.acompletion( model="gemini/gemini-1.5-flash", @@ -248,18 +248,22 @@ async def test_call_router_callbacks_on_success(): mock_response="Hello, I'm good.", ) await asyncio.sleep(1) - assert mock_callback.call_count == 2 + assert mock_callback.call_count == 1 - assert ( - mock_callback.call_args_list[0] - .kwargs["key"] - .startswith("global_router:1:gemini/gemini-1.5-flash:tpm") - ) - assert ( - mock_callback.call_args_list[1] - .kwargs["key"] - .startswith("global_router:1:gemini/gemini-1.5-flash:rpm") - ) + increment_list = mock_callback.call_args_list[0].kwargs["increment_list"] + assert len(increment_list) == 2 + + for increment in increment_list: + if "tpm" in increment["key"]: + assert increment["key"].startswith( + "global_router:1:gemini/gemini-1.5-flash:tpm" + ) + assert increment["increment_value"] == 30 + elif "rpm" in increment["key"]: + assert increment["key"].startswith( + "global_router:1:gemini/gemini-1.5-flash:rpm" + ) + assert increment["increment_value"] == 1 @pytest.mark.asyncio @@ -456,7 +460,15 @@ def test_router_get_deployment_credentials(): def test_router_get_deployment_model_info(): router = Router( - model_list=[{"model_name": "gemini/*", "litellm_params": {"model": "gemini/*"}, "model_info": {"id": "1"}}] + model_list=[ + { + "model_name": "gemini/*", + "litellm_params": {"model": "gemini/*"}, + "model_info": {"id": "1"}, + } + ] + ) + model_info = router.get_deployment_model_info( + model_id="1", model_name="gemini/gemini-1.5-flash" ) - model_info = router.get_deployment_model_info(model_id="1", model_name="gemini/gemini-1.5-flash") assert model_info is not None diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py index 2ae1fea59d..f76bc225e5 100644 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py @@ -1,6 +1,7 @@ """ Unit Tests for the max parallel request limiter v3 for the proxy """ + import asyncio import os import sys @@ -515,3 +516,209 @@ async def test_async_log_failure_event_v3(): assert op["key"] == f"{{api_key:{_api_key}}}:max_parallel_requests" assert op["increment_value"] == -1 assert op["ttl"] == 60 # default window size + + +@pytest.mark.asyncio +async def test_should_rate_limit_only_called_when_limits_exist_v3(): + """ + Test that should_rate_limit is only called when actual rate limits are configured. + This verifies the optimization that avoids unnecessary rate limit checks. + """ + _api_key = "sk-12345" + _api_key = hash_token(_api_key) + local_cache = DualCache() + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + + # Mock should_rate_limit to track if it's called + should_rate_limit_called = False + + async def mock_should_rate_limit(*args, **kwargs): + nonlocal should_rate_limit_called + should_rate_limit_called = True + return {"overall_code": "OK", "statuses": []} + + parallel_request_handler.should_rate_limit = mock_should_rate_limit + + # Test 1: No rate limits configured - should_rate_limit should NOT be called + should_rate_limit_called = False + user_api_key_dict_no_limits = UserAPIKeyAuth( + api_key=_api_key, + user_id="test_user", + team_id="test_team", + end_user_id="test_end_user", + # No rpm_limit, tpm_limit, max_parallel_requests, etc. + ) + + await parallel_request_handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict_no_limits, + cache=local_cache, + data={"model": "gpt-3.5-turbo"}, + call_type="", + ) + + assert ( + not should_rate_limit_called + ), "should_rate_limit should not be called when no rate limits are configured" + + # Test 2: API key rate limits configured - should_rate_limit SHOULD be called + should_rate_limit_called = False + user_api_key_dict_with_api_limits = UserAPIKeyAuth( + api_key=_api_key, + rpm_limit=100, # Rate limit configured + ) + + await parallel_request_handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict_with_api_limits, + cache=local_cache, + data={"model": "gpt-3.5-turbo"}, + call_type="", + ) + + assert ( + should_rate_limit_called + ), "should_rate_limit should be called when API key rate limits are configured" + + # Test 3: User rate limits configured - should_rate_limit SHOULD be called + should_rate_limit_called = False + user_api_key_dict_with_user_limits = UserAPIKeyAuth( + api_key=_api_key, + user_id="test_user", + user_tpm_limit=1000, # User rate limit configured + ) + + await parallel_request_handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict_with_user_limits, + cache=local_cache, + data={"model": "gpt-3.5-turbo"}, + call_type="", + ) + + assert ( + should_rate_limit_called + ), "should_rate_limit should be called when user rate limits are configured" + + # Test 4: Team rate limits configured - should_rate_limit SHOULD be called + should_rate_limit_called = False + user_api_key_dict_with_team_limits = UserAPIKeyAuth( + api_key=_api_key, + team_id="test_team", + team_rpm_limit=500, # Team rate limit configured + ) + + await parallel_request_handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict_with_team_limits, + cache=local_cache, + data={"model": "gpt-3.5-turbo"}, + call_type="", + ) + + assert ( + should_rate_limit_called + ), "should_rate_limit should be called when team rate limits are configured" + + # Test 5: End user rate limits configured - should_rate_limit SHOULD be called + should_rate_limit_called = False + user_api_key_dict_with_end_user_limits = UserAPIKeyAuth( + api_key=_api_key, + end_user_id="test_end_user", + end_user_rpm_limit=200, # End user rate limit configured + ) + + await parallel_request_handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict_with_end_user_limits, + cache=local_cache, + data={"model": "gpt-3.5-turbo"}, + call_type="", + ) + + assert ( + should_rate_limit_called + ), "should_rate_limit should be called when end user rate limits are configured" + + # Test 6: Max parallel requests configured - should_rate_limit SHOULD be called + should_rate_limit_called = False + user_api_key_dict_with_parallel_limits = UserAPIKeyAuth( + api_key=_api_key, + max_parallel_requests=5, # Max parallel requests configured + ) + + await parallel_request_handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict_with_parallel_limits, + cache=local_cache, + data={"model": "gpt-3.5-turbo"}, + call_type="", + ) + + assert ( + should_rate_limit_called + ), "should_rate_limit should be called when max parallel requests are configured" + + +@pytest.mark.asyncio +async def test_model_specific_rate_limits_only_called_when_configured_v3(): + """ + Test that model-specific rate limits only trigger should_rate_limit when actually configured for the requested model. + """ + from litellm.proxy.auth.auth_utils import ( + get_key_model_rpm_limit, + get_key_model_tpm_limit, + ) + + _api_key = "sk-12345" + _api_key = hash_token(_api_key) + local_cache = DualCache() + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + + # Mock should_rate_limit to track if it's called + should_rate_limit_called = False + + async def mock_should_rate_limit(*args, **kwargs): + nonlocal should_rate_limit_called + should_rate_limit_called = True + return {"overall_code": "OK", "statuses": []} + + parallel_request_handler.should_rate_limit = mock_should_rate_limit + + # Test 1: Model-specific rate limits configured but for different model - should NOT be called + should_rate_limit_called = False + user_api_key_dict_with_model_limits = UserAPIKeyAuth( + api_key=_api_key, + metadata={ + "model_tpm_limit": {"gpt-4": 1000} + }, # Rate limit for gpt-4, not gpt-3.5-turbo + ) + + await parallel_request_handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict_with_model_limits, + cache=local_cache, + data={"model": "gpt-3.5-turbo"}, # Requesting different model + call_type="", + ) + + assert ( + not should_rate_limit_called + ), "should_rate_limit should not be called when model-specific limits don't match requested model" + + # Test 2: Model-specific rate limits configured for requested model - SHOULD be called + should_rate_limit_called = False + user_api_key_dict_with_matching_model_limits = UserAPIKeyAuth( + api_key=_api_key, + metadata={ + "model_tpm_limit": {"gpt-3.5-turbo": 1000} + }, # Rate limit for requested model + ) + + await parallel_request_handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict_with_matching_model_limits, + cache=local_cache, + data={"model": "gpt-3.5-turbo"}, # Requesting same model + call_type="", + ) + + assert ( + should_rate_limit_called + ), "should_rate_limit should be called when model-specific limits match requested model"