tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-02 06:22:48 +00:00
Code Issues Packages Projects Releases Wiki Activity

Router - reduce p99 latency w/ redis enabled by 50% + OTEL - track pre_call hook latency (#13362)

Browse Source 
* 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
This commit is contained in:
Krish Dholakia
2025-08-09 16:09:51 -07:00
committed by GitHub
parent 60306d34a0
commit 1c8761111f
12 changed files with 627 additions and 180 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/caching/caching.py
+49 -25
View File
@@ -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)}")
litellm/caching/caching_handler.py
+58 -14
View File
@@ -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
litellm/caching/redis_cache.py
+71 -32
View File
@@ -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(
litellm/proxy/_experimental/mcp_server/db.py
+15 -7
View File
@@ -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
litellm/proxy/_new_secret_config.yaml
+22 -3
View File
@@ -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
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
litellm/proxy/hooks/parallel_request_limiter_v3.py
+51 -34
View File
@@ -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] = []
litellm/proxy/proxy_server.py
+14 -4
View File
@@ -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 (
litellm/proxy/utils.py
+60 -33
View File
@@ -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
),
)
litellm/router.py
+53 -9
View File
@@ -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(
tests/local_testing/test_amazing_vertex_completion.py
+1 -5
View File
@@ -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)
tests/local_testing/test_router_utils.py
+26 -14
View File
@@ -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
tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py
+207
View File
@@ -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"