From 9d4f9067a618be776537fe81bb8c4cdde7bfc3b7 Mon Sep 17 00:00:00 2001
From: Arnaud Mengus <amengus@immonot.com>
Date: Thu, 30 May 2024 09:38:30 +0000
Subject: [PATCH] refactor: Update SiloJobService to use SiloEntity in job
 methods

---
 .../java/ai/dragon/entity/SiloEntity.java     | 12 ++++++---
 .../ai/dragon/service/SiloJobService.java     | 26 ++++++++++++-------
 2 files changed, 25 insertions(+), 13 deletions(-)

diff --git a/backend/src/main/java/ai/dragon/entity/SiloEntity.java b/backend/src/main/java/ai/dragon/entity/SiloEntity.java
index d817fa5c..cbc8e486 100644
--- 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
index d8c8952c..4bc1dde6 100644
--- 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);
+        }
     }
 }