refactor(redis-store): login throttle store takes the source group and username, not keys

litellm/proxy/auth/login_throttle.py

@@ -298,7 +298,7 @@ class LoginThrottle:
         """Refuses a blocked key before any credential is looked at; otherwise hands back the attempt to settle."""
         if not self.enabled:
             return LoginAttempt(throttle=self, username=username)
-        block: Final = await self._active_block(self._keys(username))
+        block: Final = await self._active_block(username)
         if block is None:
             return LoginAttempt(throttle=self, username=username)
         verbose_proxy_logger.warning(
@@ -310,9 +310,9 @@ class LoginThrottle:
         )
         raise self.refused(block.retry_after)
 
-    async def _active_block(self, keys: LoginKeys) -> Block | None:
-        local: Final = self._local_block_ttls(keys)
-        shared: Final = await self._shared_block_ttls(keys)
+    async def _active_block(self, username: str) -> Block | None:
+        local: Final = self._local_block_ttls(self._keys(username))
+        shared: Final = await self._shared_block_ttls(username)
         user_ttl: Final = max(local[0], shared[0])
         source_ttl: Final = max(local[1], shared[1])
         if self.source_limit is not None and source_ttl > 0:
@@ -322,11 +322,11 @@ class LoginThrottle:
         return None
 
     @with_service_target(_LOGIN_THROTTLE_TARGET)
-    async def _shared_block_ttls(self, keys: LoginKeys) -> BlockTtls:
+    async def _shared_block_ttls(self, username: str) -> BlockTtls:
         store: Final = self._store()
         if store is None:
             return LOGIN_THROTTLE_NOT_BLOCKED
-        blocks: Final = await store.blocks(keys)
+        blocks: Final = await store.blocks(source_group(self.client_ip), username)
         return LOGIN_THROTTLE_NOT_BLOCKED if isinstance(blocks, RedisFailure) else blocks
 
     def _local_block_ttls(self, keys: LoginKeys) -> BlockTtls:
@@ -347,7 +347,8 @@ class LoginThrottle:
             None
             if store is None
             else await store.record_failure(
-                keys,
+                source_group(self.client_ip),
+                username,
                 user_limit=self.user_limit,
                 source_limit=source_limit,
                 window=timedelta(seconds=self.window_seconds),
@@ -376,7 +377,7 @@ class LoginThrottle:
         keys: Final = self._keys(username)
         store: Final = self._store()
         if store is not None:
-            _ = await store.clear(keys)
+            _ = await store.clear(source_group(self.client_ip), username)
         self.counters.delete_cache(keys.pair_counter)
 
     @staticmethod

litellm/redis/stores/login_throttle.py

@@ -50,30 +50,44 @@ def login_keys(source_group: str, username: str) -> LoginKeys:
 
 
 class LoginThrottleStore(Protocol):
-    async def blocks(self, keys: LoginKeys) -> BlockTtls | RedisFailure: ...
+    async def blocks(self, source_group: str, username: str) -> BlockTtls | RedisFailure: ...
 
     async def record_failure(
-        self, keys: LoginKeys, *, user_limit: int, source_limit: int, window: timedelta, block: timedelta
+        self,
+        source_group: str,
+        username: str,
+        *,
+        user_limit: int,
+        source_limit: int,
+        window: timedelta,
+        block: timedelta,
     ) -> BlockTtls | RedisFailure: ...
 
-    async def clear(self, keys: LoginKeys) -> None | RedisFailure: ...
+    async def clear(self, source_group: str, username: str) -> None | RedisFailure: ...
 
 
 @dataclass(frozen=True, slots=True)
 class TransportLoginThrottleStore:
     transport: RedisTransport
 
-    async def blocks(self, keys: LoginKeys) -> BlockTtls | RedisFailure:
-        return _ttls(await self.transport.eval(_BLOCK_TTLS_LUA, keys, ()))
+    async def blocks(self, source_group: str, username: str) -> BlockTtls | RedisFailure:
+        return _ttls(await self.transport.eval(_BLOCK_TTLS_LUA, login_keys(source_group, username), ()))
 
     async def record_failure(
-        self, keys: LoginKeys, *, user_limit: int, source_limit: int, window: timedelta, block: timedelta
+        self,
+        source_group: str,
+        username: str,
+        *,
+        user_limit: int,
+        source_limit: int,
+        window: timedelta,
+        block: timedelta,
     ) -> BlockTtls | RedisFailure:
         args: Final = (user_limit, source_limit, _whole_seconds(window), _whole_seconds(block))
-        return _ttls(await self.transport.eval(_RECORD_FAILURE_LUA, keys, args))
+        return _ttls(await self.transport.eval(_RECORD_FAILURE_LUA, login_keys(source_group, username), args))
 
-    async def clear(self, keys: LoginKeys) -> None | RedisFailure:
-        deleted: Final = await self.transport.delete(keys.pair_counter)
+    async def clear(self, source_group: str, username: str) -> None | RedisFailure:
+        deleted: Final = await self.transport.delete(login_keys(source_group, username).pair_counter)
         return deleted if isinstance(deleted, RedisFailure) else None