feat: restrict use of paid features

l3s-learnweb · Nov 7, 2024 · 219e729 · 219e729
1 parent 47940ea
commit 219e729
Show file tree

Hide file tree

Showing 5 changed files with 25 additions and 4 deletions.
diff --git a/interweb-server/src/main/java/de/l3s/interweb/server/features/chat/ChatResource.java b/interweb-server/src/main/java/de/l3s/interweb/server/features/chat/ChatResource.java
@@ -21,6 +21,7 @@
 import de.l3s.interweb.core.chat.Message;
 import de.l3s.interweb.core.util.StringUtils;
 import de.l3s.interweb.server.Roles;
+import de.l3s.interweb.server.features.user.User;
 import de.l3s.interweb.server.features.api.ApiChatRequest;
 import de.l3s.interweb.server.features.api.ApiKey;
 
@@ -41,8 +42,9 @@ public class ChatResource {
     @POST
     @Path("/completions")
     public Uni<CompletionsResults> completions(@Valid CompletionsQuery query) {
+        User user = (User) securityIdentity.getPrincipal();
         ApiKey apikey = securityIdentity.getCredential(ApiKey.class);
-        return chatService.completions(query).chain(results -> {
+        return chatService.completions(query, user.allowPaidModels).chain(results -> {
             results.setChatId(null); // reset chatId if it was set
             bus.send("api-request-chat", ApiChatRequest.of(results, apikey));
 

diff --git a/interweb-server/src/main/java/de/l3s/interweb/server/features/chat/ChatService.java b/interweb-server/src/main/java/de/l3s/interweb/server/features/chat/ChatService.java
@@ -17,10 +17,13 @@ public class ChatService {
     @Inject
     ModelsService modelsService;
 
-    public Uni<CompletionsResults> completions(CompletionsQuery query) {
+    public Uni<CompletionsResults> completions(CompletionsQuery query, boolean allowPaidModels) {
         return modelsService.getModel(query.getModel()).flatMap(model -> {
             ModelsConnector connector = modelsService.getConnector(model.getProvider());
             if (connector instanceof ChatConnector chatConnector) {
+                if (model.getPrice().getInput() > 0 && !allowPaidModels) {
+                    return Uni.createFrom().failure(new ConnectorException("Please contact the administrator to enable this model"));
+                }
                 return completions(query, model, chatConnector);
             }
 
@@ -56,6 +59,6 @@ public Uni<String> generateTitle(final Chat chat) {
             Summarize the conversation in 5 words or less, in a way that sounds like a book title.
             Don't use any formatting. You can use emojis. Only print the title, nothing else.
             """.formatted(sb), Role.user);
-        return completions(query).map(results -> results.getLastMessage().getContent().trim());
+        return completions(query, false).map(results -> results.getLastMessage().getContent().trim());
     }
 }
diff --git a/interweb-server/src/main/java/de/l3s/interweb/server/features/openai/OpenaiV1Resource.java b/interweb-server/src/main/java/de/l3s/interweb/server/features/openai/OpenaiV1Resource.java
@@ -9,6 +9,8 @@
 import de.l3s.interweb.server.features.api.ApiKey;
 import de.l3s.interweb.server.features.chat.ChatService;
 import de.l3s.interweb.server.features.models.ModelsService;
+import de.l3s.interweb.server.features.user.User;
+
 import io.quarkus.security.identity.SecurityIdentity;
 import io.smallrye.mutiny.Uni;
 import io.vertx.core.eventbus.EventBus;
@@ -50,8 +52,9 @@ public Uni<CompletionsResults> completions() {
     @POST
     @Path("/chat/completions")
     public Uni<CompletionsResults> chatCompletions(@Valid CompletionsQuery query) {
+        User user = (User) securityIdentity.getPrincipal();
         ApiKey apikey = securityIdentity.getCredential(ApiKey.class);
-        return chatService.completions(query).chain(results -> {
+        return chatService.completions(query, user.allowPaidModels).chain(results -> {
             results.setChatId(null); // reset chatId if it was set
             bus.send("api-request-chat", ApiChatRequest.of(results, apikey));
             return Uni.createFrom().item(results);

diff --git a/interweb-server/src/main/java/de/l3s/interweb/server/features/search/SearchResource.java b/interweb-server/src/main/java/de/l3s/interweb/server/features/search/SearchResource.java
@@ -18,6 +18,7 @@
 import jakarta.ws.rs.Path;
 import jakarta.ws.rs.core.Context;
 
+import io.quarkus.security.ForbiddenException;
 import io.quarkus.security.identity.SecurityIdentity;
 import io.smallrye.mutiny.Uni;
 import org.eclipse.microprofile.openapi.annotations.parameters.Parameter;
@@ -27,6 +28,7 @@
 import de.l3s.interweb.core.search.*;
 import de.l3s.interweb.core.util.StringUtils;
 import de.l3s.interweb.server.Roles;
+import de.l3s.interweb.server.features.user.User;
 
 @Tag(name = "Search", description = "Search internet by query")
 @Path("/search")
@@ -70,6 +72,11 @@ public Uni<SearchResults> search(@Parameter(description = "The search query", ex
 
     @POST
     public Uni<SearchResults> search(@NotNull @Valid SearchQuery query, @HeaderParam("Cache-Control") String cacheControl) {
+        final User user = (User) securityIdentity.getPrincipal();
+        if (!user.allowSearch) {
+            return Uni.createFrom().failure(new ForbiddenException("Please contact the administrator to enable search"));
+        }
+
         long start = System.currentTimeMillis();
         if (NO_CACHE.equals(cacheControl)) {
             query.setIgnoreCache(true);

diff --git a/interweb-server/src/main/java/de/l3s/interweb/server/features/user/User.java b/interweb-server/src/main/java/de/l3s/interweb/server/features/user/User.java
@@ -39,6 +39,12 @@ public class User extends PanacheEntityBase implements Principal {
     @NotNull
     public boolean approved = false;
 
+    @NotNull
+    public boolean allowSearch = false;
+
+    @NotNull
+    public boolean allowPaidModels = false;
+
     @UpdateTimestamp
     public Instant updated;