feat: initial application refactoring (#498)

epam · Oct 2, 2024 · df98208 · df98208
1 parent 4aef750
commit df98208
Show file tree

Hide file tree

Showing 27 changed files with 663 additions and 544 deletions.
diff --git a/src/main/java/com/epam/aidial/core/AiDial.java b/src/main/java/com/epam/aidial/core/AiDial.java
@@ -11,7 +11,7 @@
 import com.epam.aidial.core.security.AccessTokenValidator;
 import com.epam.aidial.core.security.ApiKeyStore;
 import com.epam.aidial.core.security.EncryptionService;
-import com.epam.aidial.core.service.CustomApplicationService;
+import com.epam.aidial.core.service.ApplicationService;
 import com.epam.aidial.core.service.HeartbeatService;
 import com.epam.aidial.core.service.InvitationService;
 import com.epam.aidial.core.service.LockService;
@@ -113,7 +113,6 @@ void start() throws Exception {
             resourceService = new ResourceService(vertx, redis, storage, lockService, settings("resources"), storage.getPrefix());
             InvitationService invitationService = new InvitationService(resourceService, encryptionService, settings("invitations"));
             ShareService shareService = new ShareService(resourceService, invitationService, encryptionService);
-            ResourceOperationService resourceOperationService = new ResourceOperationService(resourceService, invitationService, shareService);
             RuleService ruleService = new RuleService(resourceService);
             AccessService accessService = new AccessService(encryptionService, shareService, ruleService, settings("access"));
             NotificationService notificationService = new NotificationService(resourceService, encryptionService);
@@ -124,18 +123,19 @@ void start() throws Exception {
             ApiKeyStore apiKeyStore = new ApiKeyStore(resourceService, vertx);
             ConfigStore configStore = new FileConfigStore(vertx, settings("config"), apiKeyStore, upstreamRouteProvider);
 
-            CustomApplicationService customApplicationService = new CustomApplicationService(encryptionService,
-                    resourceService, shareService, accessService, settings("applications"));
+            ApplicationService applicationService = new ApplicationService(encryptionService, resourceService,
+                    settings("applications"));
 
             TokenStatsTracker tokenStatsTracker = new TokenStatsTracker(vertx, resourceService);
+            ResourceOperationService resourceOperationService = new ResourceOperationService(applicationService, resourceService, invitationService, shareService);
 
             HeartbeatService heartbeatService = new HeartbeatService(
                     vertx, settings("resources").getLong("heartbeatPeriod"));
             proxy = new Proxy(vertx, client, configStore, logStore,
                     rateLimiter, upstreamRouteProvider, accessTokenValidator,
                     storage, encryptionService, apiKeyStore, tokenStatsTracker, resourceService, invitationService,
                     shareService, publicationService, accessService, lockService, resourceOperationService, ruleService,
-                    notificationService, customApplicationService, heartbeatService, version());
+                    notificationService, applicationService, heartbeatService, version());
 
             server = vertx.createHttpServer(new HttpServerOptions(settings("server"))).requestHandler(proxy);
             open(server, HttpServer::listen);

diff --git a/src/main/java/com/epam/aidial/core/Proxy.java b/src/main/java/com/epam/aidial/core/Proxy.java
@@ -12,7 +12,7 @@
 import com.epam.aidial.core.security.ApiKeyStore;
 import com.epam.aidial.core.security.EncryptionService;
 import com.epam.aidial.core.security.ExtractedClaims;
-import com.epam.aidial.core.service.CustomApplicationService;
+import com.epam.aidial.core.service.ApplicationService;
 import com.epam.aidial.core.service.HeartbeatService;
 import com.epam.aidial.core.service.InvitationService;
 import com.epam.aidial.core.service.LockService;
@@ -92,7 +92,7 @@ public class Proxy implements Handler<HttpServerRequest> {
     private final ResourceOperationService resourceOperationService;
     private final RuleService ruleService;
     private final NotificationService notificationService;
-    private final CustomApplicationService customApplicationService;
+    private final ApplicationService applicationService;
     private final HeartbeatService heartbeatService;
     private final String version;
 
@@ -268,7 +268,7 @@ private Future<?> processAuthorizationResult(ExtractedClaims extractedClaims, Co
                                                  HttpServerRequest request, ApiKeyData apiKeyData, String traceId, String spanId) {
         Future<?> future;
         try {
-            ProxyContext context = new ProxyContext(config, request, apiKeyData, extractedClaims, traceId, spanId);
+            ProxyContext context = new ProxyContext(this, config, request, apiKeyData, extractedClaims, traceId, spanId);
             Controller controller = ControllerSelector.select(this, context);
             future = controller.handle();
         } catch (Exception t) {

diff --git a/src/main/java/com/epam/aidial/core/ProxyContext.java b/src/main/java/com/epam/aidial/core/ProxyContext.java
@@ -46,6 +46,7 @@ public class ProxyContext {
             .map(header -> header.toString().toLowerCase())
             .collect(Collectors.toUnmodifiableSet());
 
+    private final Proxy proxy;
     private final Config config;
     // API key of root requester
     private final Key key;
@@ -89,7 +90,8 @@ public class ProxyContext {
     private List<String> interceptors;
     private boolean isStreamingRequest;
 
-    public ProxyContext(Config config, HttpServerRequest request, ApiKeyData apiKeyData, ExtractedClaims extractedClaims, String traceId, String spanId) {
+    public ProxyContext(Proxy proxy, Config config, HttpServerRequest request, ApiKeyData apiKeyData, ExtractedClaims extractedClaims, String traceId, String spanId) {
+        this.proxy = proxy;
         this.config = config;
         this.apiKeyData = apiKeyData;
         this.request = request;

diff --git a/src/main/java/com/epam/aidial/core/config/Application.java b/src/main/java/com/epam/aidial/core/config/Application.java
@@ -10,4 +10,41 @@
 @EqualsAndHashCode(callSuper = true)
 @JsonInclude(JsonInclude.Include.NON_NULL)
 public class Application extends Deployment {
+
+    private Function function;
+
+    @Data
+    @Accessors(chain = true)
+    @JsonInclude(JsonInclude.Include.NON_NULL)
+    public static class Function {
+
+        private String sourceFolder;
+        private String targetFolder;
+        private Status status;
+        private State state;
+
+        public enum Status {
+            CREATED, STARTING, STOPPING, STARTED, STOPPED, FAILED
+        }
+
+        @Data
+        @Accessors(chain = true)
+        @JsonInclude(JsonInclude.Include.NON_NULL)
+        public static class State {
+            private Image image;
+            private Deployment deployment;
+        }
+
+        @Data
+        @Accessors(chain = true)
+        @JsonInclude(JsonInclude.Include.NON_NULL)
+        public static class Image {
+        }
+
+        @Data
+        @Accessors(chain = true)
+        @JsonInclude(JsonInclude.Include.NON_NULL)
+        public static class Deployment {
+        }
+    }
 }
diff --git a/src/main/java/com/epam/aidial/core/controller/ApplicationController.java b/src/main/java/com/epam/aidial/core/controller/ApplicationController.java
@@ -6,7 +6,7 @@
 import com.epam.aidial.core.config.Config;
 import com.epam.aidial.core.data.ApplicationData;
 import com.epam.aidial.core.data.ListData;
-import com.epam.aidial.core.service.CustomApplicationService;
+import com.epam.aidial.core.service.ApplicationService;
 import com.epam.aidial.core.service.PermissionDeniedException;
 import com.epam.aidial.core.service.ResourceNotFoundException;
 import com.epam.aidial.core.util.HttpStatus;
@@ -22,98 +22,67 @@ public class ApplicationController {
 
     private final ProxyContext context;
     private final Vertx vertx;
-    private final CustomApplicationService customApplicationService;
-    private final boolean includeCustomApplications;
+    private final ApplicationService applicationService;
 
     public ApplicationController(ProxyContext context, Proxy proxy) {
         this.context = context;
         this.vertx = proxy.getVertx();
-        this.customApplicationService = proxy.getCustomApplicationService();
-        this.includeCustomApplications = customApplicationService.includeCustomApplications();
+        this.applicationService = proxy.getApplicationService();
     }
 
     public Future<?> getApplication(String applicationId) {
-        Config config = context.getConfig();
-        Application application = config.getApplications().get(applicationId);
-
-        Future<Application> applicationFuture;
-        if (application != null) {
-            if (!DeploymentController.hasAccess(context, application)) {
-                return context.respond(HttpStatus.FORBIDDEN);
-            }
-            applicationFuture = Future.succeededFuture(application);
-        } else {
-            applicationFuture = vertx.executeBlocking(() -> customApplicationService.getCustomApplication(applicationId, context), false);
-        }
-
-        applicationFuture.map(app -> {
-            if (app == null) {
-                throw new ResourceNotFoundException(applicationId);
-            }
-
-            ApplicationData data = ApplicationUtil.mapApplication(app);
-            context.respond(HttpStatus.OK, data);
-            return null;
-        }).onFailure(error -> handleRequestError(applicationId, error));
+        DeploymentController.selectDeployment(context, applicationId)
+                .map(deployment -> {
+                    if (deployment instanceof Application application) {
+                        return application;
+                    }
+
+                    throw new ResourceNotFoundException("Application is not found: " + applicationId);
+                })
+                .map(ApplicationUtil::mapApplication)
+                .onSuccess(data -> context.respond(HttpStatus.OK, data))
+                .onFailure(this::handleRequestError);
 
         return Future.succeededFuture();
     }
 
-    public Future<?> getApplications() {
+    public Future<?> getApplicationService() {
         Config config = context.getConfig();
-        List<ApplicationData> applications = new ArrayList<>();
+        List<ApplicationData> list = new ArrayList<>();
 
         for (Application application : config.getApplications().values()) {
             if (DeploymentController.hasAccess(context, application)) {
                 ApplicationData data = ApplicationUtil.mapApplication(application);
-                applications.add(data);
+                list.add(data);
             }
         }
 
-        ListData<ApplicationData> list = new ListData<>();
-        list.setData(applications);
+        Future<List<ApplicationData>> future = Future.succeededFuture(list);
 
-        if (includeCustomApplications) {
-            vertx.executeBlocking(() -> {
-                List<Application> ownCustomApplications = customApplicationService.getOwnCustomApplications(context);
-                for (Application application : ownCustomApplications) {
-                    ApplicationData data = ApplicationUtil.mapApplication(application);
-                    applications.add(data);
-                }
-                List<Application> sharedApplications = customApplicationService.getSharedApplications(context);
-                for (Application application : sharedApplications) {
-                    ApplicationData data = ApplicationUtil.mapApplication(application);
-                    applications.add(data);
-                }
-                List<Application> publicApplications = customApplicationService.getPublicApplications(context);
-                for (Application application : publicApplications) {
-                    ApplicationData data = ApplicationUtil.mapApplication(application);
-                    applications.add(data);
-                }
-                return null;
-            }, false)
-                    .onSuccess(ignore -> context.respond(HttpStatus.OK, list)
-                    .onFailure(error -> {
-                        log.error("Can't fetch custom applications", error);
-                        context.respond(HttpStatus.INTERNAL_SERVER_ERROR, error.getMessage());
-                    }));
-        } else {
-            context.respond(HttpStatus.OK, list);
+        if (applicationService.isIncludeCustomApps()) {
+            future = vertx.executeBlocking(() -> applicationService.getAllApplications(context), false)
+                    .map(apps -> {
+                        apps.forEach(app -> list.add(ApplicationUtil.mapApplication(app)));
+                        return list;
+                    });
         }
 
+        future.onSuccess(apps -> context.respond(HttpStatus.OK, new ListData<>(apps)))
+                .onFailure(this::handleRequestError);
+
         return Future.succeededFuture();
     }
 
-    private void handleRequestError(String applicationId, Throwable error) {
-        if (error instanceof PermissionDeniedException) {
-            log.error("Forbidden application {}. Key: {}. User sub: {}", applicationId, context.getProject(), context.getUserSub());
+    private void handleRequestError(Throwable error) {
+        if (error instanceof IllegalArgumentException) {
+            context.respond(HttpStatus.BAD_REQUEST, error.getMessage());
+        } else if (error instanceof PermissionDeniedException) {
             context.respond(HttpStatus.FORBIDDEN, error.getMessage());
         } else if (error instanceof ResourceNotFoundException) {
-            log.error("Application not found {}", applicationId, error);
             context.respond(HttpStatus.NOT_FOUND, error.getMessage());
         } else {
-            log.error("Failed to load application {}", applicationId, error);
-            context.respond(HttpStatus.INTERNAL_SERVER_ERROR, "Failed to load application: " + applicationId);
+            log.error("Failed to handle application request", error);
+            context.respond(error, "Internal error");
         }
     }
 }
diff --git a/src/main/java/com/epam/aidial/core/controller/BucketController.java b/src/main/java/com/epam/aidial/core/controller/BucketController.java
@@ -26,7 +26,7 @@ public Future<?> getBucket() {
             appDataLocation = null;
         } else {
             String encryptedAppDataBucket = encryptionService.encrypt(appDataBucket);
-            String encodedSourceDeployment = UrlUtil.encodePath(context.getSourceDeployment());
+            String encodedSourceDeployment = UrlUtil.encodePath(context.getSourceDeployment()); // bucket/my-app
             appDataLocation = encryptedAppDataBucket + BlobStorageUtil.PATH_SEPARATOR + BlobStorageUtil.APPDATA_PATTERN.formatted(encodedSourceDeployment);
         }
         return context.respond(HttpStatus.OK, new Bucket(encryptedBucket, appDataLocation));

diff --git a/src/main/java/com/epam/aidial/core/controller/ControllerSelector.java b/src/main/java/com/epam/aidial/core/controller/ControllerSelector.java
@@ -17,16 +17,16 @@
 public class ControllerSelector {
 
     private static final Pattern PATTERN_POST_DEPLOYMENT = Pattern.compile("^/+openai/deployments/(.+?)/(completions|chat/completions|embeddings)$");
-    private static final Pattern PATTERN_DEPLOYMENT = Pattern.compile("^/+openai/deployments/([^/]+)$");
+    private static final Pattern PATTERN_DEPLOYMENT = Pattern.compile("^/+openai/deployments/(.+?)$");
     private static final Pattern PATTERN_DEPLOYMENTS = Pattern.compile("^/+openai/deployments$");
 
-    private static final Pattern PATTERN_MODEL = Pattern.compile("^/+openai/models/([^/]+)$");
+    private static final Pattern PATTERN_MODEL = Pattern.compile("^/+openai/models/(.+?)$");
     private static final Pattern PATTERN_MODELS = Pattern.compile("^/+openai/models$");
 
-    private static final Pattern PATTERN_ADDON = Pattern.compile("^/+openai/addons/([^/]+)$");
+    private static final Pattern PATTERN_ADDON = Pattern.compile("^/+openai/addons/(.+?)$");
     private static final Pattern PATTERN_ADDONS = Pattern.compile("^/+openai/addons$");
 
-    private static final Pattern PATTERN_ASSISTANT = Pattern.compile("^/+openai/assistants/([^/]+)$");
+    private static final Pattern PATTERN_ASSISTANT = Pattern.compile("^/+openai/assistants/(.+?)$");
     private static final Pattern PATTERN_ASSISTANTS = Pattern.compile("^/+openai/assistants$");
 
     private static final Pattern PATTERN_APPLICATION = Pattern.compile("^/+openai/applications/(.+?)$");
@@ -56,7 +56,7 @@ public class ControllerSelector {
 
     private static final Pattern RESOURCE_OPERATIONS = Pattern.compile("^/v1/ops/resource/(move|subscribe)$");
 
-    private static final Pattern DEPLOYMENT_LIMITS = Pattern.compile("^/v1/deployments/([^/]+)/limits$");
+    private static final Pattern DEPLOYMENT_LIMITS = Pattern.compile("^/v1/deployments/(.+?)/limits$");
 
     private static final Pattern NOTIFICATIONS = Pattern.compile("^/v1/ops/notification/(list|delete)$");
 
@@ -143,7 +143,7 @@ private static Controller selectGet(Proxy proxy, ProxyContext context, String pa
         match = match(PATTERN_APPLICATIONS, path);
         if (match != null) {
             ApplicationController controller = new ApplicationController(context, proxy);
-            return controller::getApplications;
+            return controller::getApplicationService;
         }
 
         match = match(PATTERN_FILES_METADATA, path);

diff --git a/src/main/java/com/epam/aidial/core/controller/DeploymentController.java b/src/main/java/com/epam/aidial/core/controller/DeploymentController.java
@@ -1,17 +1,20 @@
 package com.epam.aidial.core.controller;
 
+import com.epam.aidial.core.Proxy;
 import com.epam.aidial.core.ProxyContext;
 import com.epam.aidial.core.config.Config;
 import com.epam.aidial.core.config.Deployment;
 import com.epam.aidial.core.config.Features;
-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.DeploymentData;
 import com.epam.aidial.core.data.FeaturesData;
 import com.epam.aidial.core.data.ListData;
+import com.epam.aidial.core.data.ResourceType;
+import com.epam.aidial.core.service.PermissionDeniedException;
+import com.epam.aidial.core.service.ResourceNotFoundException;
+import com.epam.aidial.core.storage.ResourceDescription;
 import com.epam.aidial.core.util.HttpStatus;
+import com.epam.aidial.core.util.UrlUtil;
 import io.vertx.core.Future;
 import lombok.RequiredArgsConstructor;
 
@@ -57,6 +60,41 @@ public Future<?> getDeployments() {
         return context.respond(HttpStatus.OK, list);
     }
 
+    public static Future<Deployment> selectDeployment(ProxyContext context, String id) {
+        Deployment deployment = context.getConfig().selectDeployment(id);
+
+        if (deployment != null) {
+            if (!DeploymentController.hasAccess(context, deployment)) {
+                return Future.failedFuture(new PermissionDeniedException("Forbidden deployment: " + id));
+            } else {
+                return Future.succeededFuture(deployment);
+            }
+        }
+
+        Proxy proxy = context.getProxy();
+        return proxy.getVertx().executeBlocking(() -> {
+            String url;
+            ResourceDescription resource;
+
+            try {
+                url = UrlUtil.encodePath(id);
+                resource = ResourceDescription.fromAnyUrl(url, proxy.getEncryptionService());
+            } catch (Throwable ignore) {
+                throw new ResourceNotFoundException("Unknown application: " + id);
+            }
+
+            if (resource.isFolder() || resource.getType() != ResourceType.APPLICATION) {
+                throw new ResourceNotFoundException("Invalid application url: " + url);
+            }
+
+            if (!proxy.getAccessService().hasReadAccess(resource, context)) {
+                throw new PermissionDeniedException();
+            }
+
+            return proxy.getApplicationService().getApplication(resource).getValue();
+        }, false);
+    }
+
     public static boolean hasAccess(ProxyContext context, Deployment deployment) {
         Set<String> expectedUserRoles = deployment.getUserRoles();
         List<String> actualUserRoles = context.getUserRoles();