feat: Implement user limits #82

sample/aidial.config.json

@@ -46,7 +46,12 @@
                     "endpoint": "http://localhost:7003",
                     "key": "modelKey3"
                 }
-            ]
+            ],
+            "userRoles": ["role1", "role2"],
+            "defaultUserLimit": {
+                "minute": "100000",
+                "day": "10000000"
+            }
         },
         "embedding-ada": {
             "type": "embedding",
@@ -57,7 +62,8 @@
                     "key": "modelKey4"
                 }
             ]
-        }
+        },
+        "userRoles": ["role3"]
     },
     "keys": {
         "proxyKey1": {
@@ -86,6 +92,22 @@
                 "search_assistant": {},
                 "app": {}
             }
+        },
+        "role1": {
+            "limits": {
+                "chat-gpt-35-turbo": {
+                    "minute": "200000",
+                    "day": "10000000"
+                }
+            }
+        },
+        "role2": {
+            "limits": {
+                "chat-gpt-35-turbo": {
+                    "minute": "100000",
+                    "day": "20000000"
+                }
+            }
         }
     }
 }

src/main/java/com/epam/aidial/core/config/Model.java

@@ -17,4 +17,5 @@ public class Model extends Deployment {
     private List<Upstream> upstreams = List.of();
     // if it's set then the model name is overridden with that name in the request body to the model adapter
     private String overrideName;
+    private Limit defaultUserLimit;
 }

src/main/java/com/epam/aidial/core/limiter/RateLimiter.java

@@ -4,6 +4,7 @@
 import com.epam.aidial.core.config.Deployment;
 import com.epam.aidial.core.config.Key;
 import com.epam.aidial.core.config.Limit;
+import com.epam.aidial.core.config.Model;
 import com.epam.aidial.core.config.Role;
 import com.epam.aidial.core.data.LimitStats;
 import com.epam.aidial.core.data.ResourceType;
@@ -20,10 +21,16 @@
 import lombok.SneakyThrows;
 import lombok.extern.slf4j.Slf4j;
 
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+
 @Slf4j
 @RequiredArgsConstructor
 public class RateLimiter {
 
+    private static final Limit DEFAULT_LIMIT = new Limit();
+
     private final Vertx vertx;
 
     private final ResourceService resourceService;
@@ -34,10 +41,7 @@ public Future<Void> increase(ProxyContext context) {
             if (resourceService == null) {
                 return Future.succeededFuture();
             }
-            Key key = context.getKey();
-            if (key == null) {
-                return Future.succeededFuture();
-            }
+
             Deployment deployment = context.getDeployment();
             TokenUsage usage = context.getTokenUsage();
 
@@ -62,8 +66,7 @@ public Future<RateLimitResult> limit(ProxyContext context) {
             Deployment deployment = context.getDeployment();
             Limit limit;
             if (key == null) {
-                // don't support user limits yet
-                return Future.succeededFuture(RateLimitResult.SUCCESS);
+                limit = getLimitByUser(context);
             } else {
                 limit = getLimitByApiKey(context, deployment.getName());
             }
@@ -176,8 +179,38 @@ private Limit getLimitByApiKey(ProxyContext context, String deploymentName) {
         return role.getLimits().get(deploymentName);
     }
 
+    private Limit getLimitByUser(ProxyContext context) {
+        List<String> userRoles = context.getUserRoles();
+        Limit defaultUserLimit = getDefaultUserLimit(context.getDeployment());
+        if (userRoles.isEmpty()) {
+            return defaultUserLimit;
+        }
+        String deploymentName = context.getDeployment().getName();
+        Map<String, Role> userRoleToDeploymentLimits = context.getConfig().getRoles();
+        long minuteLimit = 0;
+        long dayLimit = 0;
+        for (String userRole : userRoles) {
+            Limit limit = Optional.ofNullable(userRoleToDeploymentLimits.get(userRole))
+                    .map(role -> role.getLimits().get(deploymentName))
+                    .orElse(defaultUserLimit);
+            minuteLimit = Math.max(minuteLimit, limit.getMinute());
+            dayLimit = Math.max(dayLimit, limit.getDay());
+        }
+        Limit limit = new Limit();
+        limit.setMinute(minuteLimit);
+        limit.setDay(dayLimit);
+        return limit;
+    }
+
     private static String getPath(String deploymentName) {
         return String.format("%s/tokens", deploymentName);
     }
 
+    private static Limit getDefaultUserLimit(Deployment deployment) {
+        if (deployment instanceof Model model) {
+            return model.getDefaultUserLimit() == null ? DEFAULT_LIMIT : model.getDefaultUserLimit();
+        }
+        return DEFAULT_LIMIT;
+    }
+
 }

src/test/java/com/epam/aidial/core/limiter/RateLimiterTest.java

@@ -15,7 +15,6 @@
 import com.epam.aidial.core.storage.BlobStorage;
 import com.epam.aidial.core.token.TokenUsage;
 import com.epam.aidial.core.util.HttpStatus;
-import com.epam.aidial.core.util.ProxyUtil;
 import io.vertx.core.Future;
 import io.vertx.core.Vertx;
 import io.vertx.core.http.HttpServerRequest;
@@ -34,7 +33,7 @@
 import redis.embedded.RedisServer;
 
 import java.io.IOException;
-import java.util.Collections;
+import java.util.List;
 import java.util.Map;
 import java.util.concurrent.Callable;
 
@@ -128,17 +127,6 @@ public void testLimit_EntityNotFound() {
         assertEquals(HttpStatus.FORBIDDEN, result.result().status());
     }
 
-    @Test
-    public void testLimit_SuccessUser() {
-        ProxyContext proxyContext = new ProxyContext(new Config(), request, new ApiKeyData(), new ExtractedClaims("sub", Collections.emptyList(), "hash"), "trace-id", "span-id");
-
-        Future<RateLimitResult> result = rateLimiter.limit(proxyContext);
-
-        assertNotNull(result);
-        assertNotNull(result.result());
-        assertEquals(HttpStatus.OK, result.result().status());
-    }
-
     @Test
     public void testLimit_ApiKeyLimitNotFound() {
         Key key = new Key();
@@ -340,4 +328,101 @@ public void testGetLimitStats_ApiKey() {
 
     }
 
+    @Test
+    public void testLimit_User_LimitFound() {
+        Config config = new Config();
+
+        Role role1 = new Role();
+        Limit limit = new Limit();
+        limit.setDay(10000);
+        limit.setMinute(100);
+        role1.setLimits(Map.of("model", limit));
+
+        Role role2 = new Role();
+        limit = new Limit();
+        limit.setDay(20000);
+        limit.setMinute(200);
+        role2.setLimits(Map.of("model", limit));
+
+        config.getRoles().put("role1", role1);
+        config.getRoles().put("role2", role2);
+
+        ApiKeyData apiKeyData = new ApiKeyData();
+        ProxyContext proxyContext = new ProxyContext(config, request, apiKeyData, new ExtractedClaims("sub", List.of("role1", "role2"), "user-hash"), "trace-id", "span-id");
+        Model model = new Model();
+        model.setName("model");
+        proxyContext.setDeployment(model);
+
+        when(vertx.executeBlocking(any(Callable.class))).thenAnswer(invocation -> {
+            Callable<?> callable = invocation.getArgument(0);
+            return Future.succeededFuture(callable.call());
+        });
+
+        TokenUsage tokenUsage = new TokenUsage();
+        tokenUsage.setTotalTokens(150);
+        proxyContext.setTokenUsage(tokenUsage);
+
+        Future<Void> increaseLimitFuture = rateLimiter.increase(proxyContext);
+        assertNotNull(increaseLimitFuture);
+        assertNull(increaseLimitFuture.cause());
+
+        Future<RateLimitResult> checkLimitFuture = rateLimiter.limit(proxyContext);
+
+        assertNotNull(checkLimitFuture);
+        assertNotNull(checkLimitFuture.result());
+        assertEquals(HttpStatus.OK, checkLimitFuture.result().status());
+
+        increaseLimitFuture = rateLimiter.increase(proxyContext);
+        assertNotNull(increaseLimitFuture);
+        assertNull(increaseLimitFuture.cause());
+
+        checkLimitFuture = rateLimiter.limit(proxyContext);
+
+        assertNotNull(checkLimitFuture);
+        assertNotNull(checkLimitFuture.result());
+        assertEquals(HttpStatus.TOO_MANY_REQUESTS, checkLimitFuture.result().status());
+
+    }
+
+    @Test
+    public void testLimit_User_LimitNotFound() {
+        Config config = new Config();
+
+        ApiKeyData apiKeyData = new ApiKeyData();
+        ProxyContext proxyContext = new ProxyContext(config, request, apiKeyData, new ExtractedClaims("sub", List.of("role1"), "user-hash"), "trace-id", "span-id");
+        Model model = new Model();
+        model.setName("model");
+        proxyContext.setDeployment(model);
+
+        when(vertx.executeBlocking(any(Callable.class))).thenAnswer(invocation -> {
+            Callable<?> callable = invocation.getArgument(0);
+            return Future.succeededFuture(callable.call());
+        });
+
+        TokenUsage tokenUsage = new TokenUsage();
+        tokenUsage.setTotalTokens(90);
+        proxyContext.setTokenUsage(tokenUsage);
+
+        Future<Void> increaseLimitFuture = rateLimiter.increase(proxyContext);
+        assertNotNull(increaseLimitFuture);
+        assertNull(increaseLimitFuture.cause());
+
+        Future<RateLimitResult> checkLimitFuture = rateLimiter.limit(proxyContext);
+
+        assertNotNull(checkLimitFuture);
+        assertNotNull(checkLimitFuture.result());
+        assertEquals(HttpStatus.OK, checkLimitFuture.result().status());
+
+        increaseLimitFuture = rateLimiter.increase(proxyContext);
+        assertNotNull(increaseLimitFuture);
+        assertNull(increaseLimitFuture.cause());
+
+        checkLimitFuture = rateLimiter.limit(proxyContext);
+
+        assertNotNull(checkLimitFuture);
+        assertNotNull(checkLimitFuture.result());
+        assertEquals(HttpStatus.OK, checkLimitFuture.result().status());
+
+    }
+
 }