refactor: Update SiloJobService to use SiloEntity in job methods

dRAGon-Okinawa · May 30, 2024 · 9d4f906 · 9d4f906
1 parent ca1a16e
commit 9d4f906
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 13 deletions.
diff --git a/backend/src/main/java/ai/dragon/entity/SiloEntity.java b/backend/src/main/java/ai/dragon/entity/SiloEntity.java
@@ -23,6 +23,8 @@
 @Getter
 @Setter
 public class SiloEntity implements AbstractEntity {
+    public static final String DEFAULT_CRON_EXPRESSION = "*/15 * * * *";
+
     @Id
     @NotNull
     @Schema(description = "Identifier of the Silo")
@@ -32,9 +34,6 @@ public class SiloEntity implements AbstractEntity {
     @Schema(description = "Name of the Silo. Must be unique.")
     private String name;
 
-    @Schema(description = "Type of the Silo")
-    private SiloType type;
-
     @NotNull
     @Schema(description = "Java Class to be used for the Vector Store", example = "InMemoryEmbeddingStore")
     private VectorStoreType vectorStoreType;
@@ -43,6 +42,12 @@ public class SiloEntity implements AbstractEntity {
     @Schema(description = "Type to be used for the Embedding Model", example = "BgeSmallEnV15QuantizedEmbeddingModel")
     private EmbeddingModelType embeddingModelType;
 
+    @Schema(description = "Type of the Silo's Ingestor Type")
+    private SiloType ingestorType;
+
+    @Schema(description = "Cron Expression for the Silo's Ingestor Job", example = "Launch the Silo ingestor every 15 minutes : */15 * * * *")
+    private String ingestorSchedule;
+
     @Schema(description = "Settings to be linked to the Silo (if applicable) in the form of key-value pairs.")
     private Map<String, String> settings;
 
@@ -51,5 +56,6 @@ public SiloEntity() {
         this.name = String.format("Silo %s", this.uuid.toString());
         this.vectorStoreType = VectorStoreType.InMemoryEmbeddingStore;
         this.embeddingModelType = EmbeddingModelType.BgeSmallEnV15QuantizedEmbeddingModel;
+        this.ingestorSchedule = DEFAULT_CRON_EXPRESSION;
     }
 }
diff --git a/backend/src/main/java/ai/dragon/service/SiloJobService.java b/backend/src/main/java/ai/dragon/service/SiloJobService.java
@@ -7,6 +7,8 @@
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.stereotype.Service;
 
+import java.util.Optional;
+
 import ai.dragon.entity.SiloEntity;
 import ai.dragon.job.silo.SiloIngestorJobHandler;
 import ai.dragon.job.silo.SiloIngestorJobRequest;
@@ -78,15 +80,19 @@ public void removeSiloIngestorJob(SiloEntity siloEntity) {
 
     public void createSiloIngestorJob(SiloEntity siloEntity) {
         logger.info(String.format("Creating Silo Ingestor Job : %s -> %s", siloEntity.getUuid(), siloEntity.getName()));
-        SiloIngestorJobRequest jobRequest = SiloIngestorJobRequest
-                .create()
-                .uuid(siloEntity.getUuid());
-        jobRequestScheduler.scheduleRecurrently(
-                siloEntity.getUuid().toString(),
-                "* * * * * ", // TODO Move to SiloEntity settings
-                java.time.ZoneId.of("UTC"), // TODO Move to configuration
-                jobRequest);
-        jobService.getRecurringJob(siloEntity.getUuid().toString())
-                .setJobName(String.format("%s : %s", SiloIngestorJobHandler.JOB_NAME, siloEntity.getName()));
+        try {
+            SiloIngestorJobRequest jobRequest = SiloIngestorJobRequest
+                    .create()
+                    .uuid(siloEntity.getUuid());
+            jobRequestScheduler.scheduleRecurrently(
+                    siloEntity.getUuid().toString(),
+                    Optional.ofNullable(siloEntity.getIngestorSchedule()).orElse(SiloEntity.DEFAULT_CRON_EXPRESSION),
+                    java.time.ZoneId.of("UTC"), // TODO Move to configuration
+                    jobRequest);
+            jobService.getRecurringJob(siloEntity.getUuid().toString())
+                    .setJobName(String.format("%s : %s", SiloIngestorJobHandler.JOB_NAME, siloEntity.getName()));
+        } catch (Exception ex) {
+            logger.error("Error creating Silo Ingestor Job", ex);
+        }
     }
 }