From 08784061e4581c6658321ed6ebdaa2e7d37cdaba Mon Sep 17 00:00:00 2001 From: narmstrong-berry Date: Thu, 5 Nov 2026 23:59:02 +0000 Subject: [PATCH] fix(router): keep async_increment_count inside the basedpyright gate --- litellm/caching/dual_cache.py | 10 +++------- litellm/caching/redis_cache.py | 35 +++++++++++++++++++++++----------- litellm/router_rust/core.py | 3 +-- 3 files changed, 28 insertions(+), 20 deletions(-) diff --git a/litellm/caching/dual_cache.py b/litellm/caching/dual_cache.py index 532b06a5ea7..e338e0aa1ab 100644 --- a/litellm/caching/dual_cache.py +++ b/litellm/caching/dual_cache.py @@ -587,14 +587,10 @@ class DualCache(BaseCache): local_only: bool = False, ttl: float | None = None, parent_otel_span: Span | None = None, - ) -> int | None: + ) -> int: """``async_increment_cache`` for whole-number counters: Redis gets INCRBY, so the key holds an - integer. Returns the new count, the in-memory one when Redis fails, or None with no backend.""" - in_memory: Final = ( - None - if self.in_memory_cache is None - else int(await self.in_memory_cache.async_increment(key, value, ttl=ttl)) - ) + integer. Returns the new count, or the in-memory one when Redis fails.""" + in_memory: Final = int(self.in_memory_cache.increment_cache(key=key, value=value, ttl=ttl)) # pyright: ignore[reportUnknownMemberType] # InMemoryCache.increment_cache takes untyped **kwargs if self.redis_cache is None or local_only: return in_memory try: diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 68acc26877a..a3db5820c5f 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -1528,6 +1528,22 @@ class RedisCache(BaseCache): ) raise e + @staticmethod + async def _incrby_with_ttl( + _redis_client: "Redis", # pyright: ignore[reportUnknownParameterType, reportMissingTypeArgument] # redis-py's Redis is generic over an untyped response + key: str, + value: int, + ttl: int | None, + ) -> int: + """INCRBY and the key's TTL in one round trip; an EXPIRE only when the key has none.""" + async with _redis_client.pipeline(transaction=False) as pipe: # pyright: ignore[reportUnknownVariableType] # redis-py's Pipeline is generic over an untyped response + pipe.incrby(name=key, amount=value) + pipe.ttl(key) + result, current_ttl = await pipe.execute() # pyright: ignore[reportUnknownVariableType, reportUnknownMemberType] # see above + if ttl is not None and current_ttl == -1: + await _redis_client.expire(key, ttl) + return int(result) # pyright: ignore[reportUnknownArgumentType] # INCRBY replies with an int + @_redis_circuit_breaker_guard async def async_increment_count( self, @@ -1539,20 +1555,17 @@ class RedisCache(BaseCache): """INCRBY, for counters that hold whole numbers (the router's allowed-fails counts), so the key stays an integer the way the sync ``increment_cache`` leaves it. The TTL is set only when the key has none.""" - _redis_client: Final = self.init_async_client() + from redis.asyncio import Redis + + _redis_client: Final[Redis] = self.init_async_client() # pyright: ignore[reportUnknownVariableType, reportUnknownMemberType, reportAssignmentType, reportMissingTypeArgument] # as in async_increment: the cluster client serves the same commands start_time: Final = time.time() - _used_ttl: Final = self.get_ttl(ttl=ttl) + _used_ttl: Final[int | None] = self.get_ttl(ttl=ttl) # pyright: ignore[reportUnknownMemberType] # BaseCache.get_ttl takes untyped **kwargs namespaced_key: Final = self.check_and_fix_namespace(key=key) try: - async with _redis_client.pipeline(transaction=False) as pipe: - pipe.incrby(name=namespaced_key, amount=value) - pipe.ttl(namespaced_key) - result, current_ttl = await pipe.execute() - if _used_ttl is not None and current_ttl == -1: - await _redis_client.expire(namespaced_key, _used_ttl) + result: Final = await self._incrby_with_ttl(_redis_client, namespaced_key, value, _used_ttl) # pyright: ignore[reportUnknownMemberType] # see above end_time: Final = time.time() asyncio.create_task( - self.service_logger_obj.async_service_success_hook( + self.service_logger_obj.async_service_success_hook( # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # ServiceLogging's hooks are untyped, as in async_increment service=ServiceTypes.REDIS, duration=end_time - start_time, call_type="async_increment_count", @@ -1562,11 +1575,10 @@ class RedisCache(BaseCache): parent_otel_span=parent_otel_span, ) ) - return int(result) except Exception as e: failed_at: Final = time.time() asyncio.create_task( - self.service_logger_obj.async_service_failure_hook( + self.service_logger_obj.async_service_failure_hook( # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # see above service=ServiceTypes.REDIS, duration=failed_at - start_time, error=e, @@ -1584,6 +1596,7 @@ class RedisCache(BaseCache): e, ) raise + return result @_redis_circuit_breaker_guard async def async_set_max( diff --git a/litellm/router_rust/core.py b/litellm/router_rust/core.py index f0e7b4af1cd..756d3789f39 100644 --- a/litellm/router_rust/core.py +++ b/litellm/router_rust/core.py @@ -3573,8 +3573,7 @@ class RustCore( # The allowed-fails counters: INCRBY, so the key holds an integer as the Python # router's increment_cache left it. count: Final = _count(step.get("value")) - counted: Final = await self._store_for(key).async_increment_count(key, count, local_only=local, ttl=ttl) - return count if counted is None else counted + return await self._store_for(key).async_increment_count(key, count, local_only=local, ttl=ttl) case "increment_usage": await self._increment_usage(_as_list(step.get("keys")), _as_list(step.get("value")), ttl) return None