In Memory Store Persistance

dRAGon-Okinawa · Jun 1, 2024 · 3e65e96 · 3e65e96
1 parent 77712a1
commit 3e65e96
Show file tree

Hide file tree

Showing 14 changed files with 100 additions and 49 deletions.
diff --git a/backend/src/main/java/ai/dragon/component/DirectoryStructureComponent.java b/backend/src/main/java/ai/dragon/component/DirectoryStructureComponent.java
@@ -11,9 +11,8 @@
 import org.springframework.core.annotation.Order;
 import org.springframework.stereotype.Component;
 
-import ai.dragon.config.properties.DataProperties;
+import ai.dragon.properties.config.DataProperties;
 import jakarta.annotation.PostConstruct;
-import lombok.Getter;
 
 @Component
 @Order(Ordered.HIGHEST_PRECEDENCE)
@@ -23,53 +22,47 @@ public class DirectoryStructureComponent {
     @Autowired
     private DataProperties dataProperties;
 
-    @Getter
-    private File dataDirectory;
+    private File mainDataDirectory;
 
     @PostConstruct
     private void postConstruct() throws Exception {
-        dataDirectory = createDataDirectories();
-        createDatabaseDirectory();
+        mainDataDirectory = createMainDataDirectory();
     }
 
-    private File createDataDirectories() throws Exception {
+    public File directoryFor(String directoryName) {
+        if (mainDataDirectory == null || !mainDataDirectory.exists()) {
+            throw new RuntimeException("Data directory not found");
+        }
+        File directory = new File(mainDataDirectory, directoryName);
+        if (!directory.exists()) {
+            boolean creationStatus = directory.mkdirs();
+            if (creationStatus) {
+                logger.debug("Directory created successfully : " + directory);
+            } else {
+                logger.error("Failed to create directory");
+                throw new RuntimeException("Failed to create directory : " + directory);
+            }
+        } else {
+            logger.debug("Directory already exists : " + directory);
+        }
+        return directory;
+    }
+
+    private File createMainDataDirectory() throws Exception {
         String path = Optional.ofNullable(dataProperties.getPath()).orElse(":temp:");
         File dataDirectory = ":temp:".equals(path) ? Files.createTempDirectory("dRAGon").toFile() : new File(path);
         logger.info("Creating data directory at " + dataDirectory + " if doesn't exist");
-
         if (!dataDirectory.exists()) {
             boolean creationStatus = dataDirectory.mkdirs();
-
             if (creationStatus) {
                 logger.debug("Data directory created successfully : " + dataDirectory);
             } else {
                 logger.error("Failed to create data directory");
                 throw new RuntimeException("Failed to create data directory : " + dataDirectory);
             }
-
         } else {
             logger.debug("Data directory already exists : " + dataDirectory);
         }
-
         return dataDirectory;
     }
-
-    private void createDatabaseDirectory() {
-        logger.debug("Creating database directory if doesn't exist");
-        File databaseDirectory = new File(dataDirectory, "db");
-
-        if (!databaseDirectory.exists()) {
-            boolean creationStatus = databaseDirectory.mkdirs();
-
-            if (creationStatus) {
-                logger.debug("Database directory created successfully : " + databaseDirectory);
-            } else {
-                logger.error("Failed to create database directory");
-                throw new RuntimeException("Failed to create database directory : " + databaseDirectory);
-            }
-
-        } else {
-            logger.debug("Database directory already exists : " + databaseDirectory);
-        }
-    }
 }
diff --git a/backend/src/main/java/ai/dragon/config/PropertiesConfig.java b/backend/src/main/java/ai/dragon/config/PropertiesConfig.java
@@ -4,7 +4,7 @@
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 
-import ai.dragon.config.properties.DataProperties;
+import ai.dragon.properties.config.DataProperties;
 
 @Configuration
 public class PropertiesConfig {

diff --git a/...end/src/main/java/ai/dragon/job/silo/ingestor/dto/embedding/EmbeddingModelDefinition.java b/...end/src/main/java/ai/dragon/job/silo/ingestor/dto/embedding/EmbeddingModelDefinition.java
@@ -4,6 +4,7 @@
 import java.util.function.Function;
 
 import ai.dragon.enumeration.ProviderType;
+import ai.dragon.properties.embedding.EmbeddingModelSettings;
 import dev.langchain4j.model.embedding.EmbeddingModel;
 import lombok.Builder;
 import lombok.Data;

diff --git a/backend/src/main/java/ai/dragon/job/silo/ingestor/dto/embedding/EmbeddingModelSettings.java b/backend/src/main/java/ai/dragon/job/silo/ingestor/dto/embedding/EmbeddingModelSettings.java
diff --git a/backend/src/main/java/ai/dragon/job/silo/ingestor/handler/loader/SiloIngestorJobHandler.java b/backend/src/main/java/ai/dragon/job/silo/ingestor/handler/loader/SiloIngestorJobHandler.java
@@ -43,5 +43,6 @@ public void run(SiloIngestorJobRequest jobRequest) throws Exception {
                     jobContext().logger().info(ingestLogMessage.getMessage());
             }
         });
+        jobContext().logger().info("End of Silo Ingestor Job.");
     }
 }
diff --git a/...src/main/java/ai/dragon/job/silo/ingestor/loader/filesystem/FileSystemIngestorLoader.java b/...src/main/java/ai/dragon/job/silo/ingestor/loader/filesystem/FileSystemIngestorLoader.java
@@ -7,6 +7,7 @@
 
 import ai.dragon.entity.SiloEntity;
 import ai.dragon.job.silo.ingestor.loader.ImplAbstractSiloIngestorLoader;
+import ai.dragon.properties.loader.FileSystemIngestorLoaderSettings;
 import ai.dragon.util.IniSettingUtil;
 import dev.langchain4j.data.document.Document;
 import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;

diff --git a/...gon/config/properties/DataProperties.java → ...gon/properties/config/DataProperties.java b/...gon/config/properties/DataProperties.java → ...gon/properties/config/DataProperties.java
@@ -1,4 +1,4 @@
-package ai.dragon.config.properties;
+package ai.dragon.properties.config;
 
 import lombok.Getter;
 import lombok.Setter;

diff --git a/backend/src/main/java/ai/dragon/properties/embedding/EmbeddingModelSettings.java b/backend/src/main/java/ai/dragon/properties/embedding/EmbeddingModelSettings.java
@@ -0,0 +1,11 @@
+package ai.dragon.properties.embedding;
+
+import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
+
+import lombok.Data;
+
+@Data
+@JsonIgnoreProperties(ignoreUnknown = true)
+public class EmbeddingModelSettings {
+    private String apiKey;
+}
diff --git a/...tem/FileSystemIngestorLoaderSettings.java → ...der/FileSystemIngestorLoaderSettings.java b/...tem/FileSystemIngestorLoaderSettings.java → ...der/FileSystemIngestorLoaderSettings.java
@@ -1,4 +1,4 @@
-package ai.dragon.job.silo.ingestor.loader.filesystem;
+package ai.dragon.properties.loader;
 
 import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
 

diff --git a/backend/src/main/java/ai/dragon/properties/store/InMemoryEmbeddingStoreSettings.java b/backend/src/main/java/ai/dragon/properties/store/InMemoryEmbeddingStoreSettings.java
@@ -0,0 +1,15 @@
+package ai.dragon.properties.store;
+
+import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
+
+import lombok.Data;
+
+@Data
+@JsonIgnoreProperties(ignoreUnknown = true)
+public class InMemoryEmbeddingStoreSettings {
+    private String persistance;
+
+    public InMemoryEmbeddingStoreSettings() {
+        persistance = ":memory:";
+    }
+}
diff --git a/backend/src/main/java/ai/dragon/service/DatabaseService.java b/backend/src/main/java/ai/dragon/service/DatabaseService.java
@@ -16,7 +16,7 @@
 import org.springframework.stereotype.Service;
 
 import ai.dragon.component.DirectoryStructureComponent;
-import ai.dragon.config.properties.DataProperties;
+import ai.dragon.properties.config.DataProperties;
 import jakarta.annotation.PostConstruct;
 import jakarta.annotation.PreDestroy;
 
@@ -59,7 +59,7 @@ public void openDatabase() {
                 .loadModule(new JacksonMapperModule());
 
         if (!isDatabaseInMemory()) {
-            File databaseFile = new File(directoryStructureComponent.getDataDirectory(), "db/" + getDatabaseFilename());
+            File databaseFile = new File(directoryStructureComponent.directoryFor("db"), getDatabaseFilename());
             logger.debug("Will use database file: " + databaseFile);
 
             MVStoreModule storeModule = MVStoreModule.withConfig()

diff --git a/backend/src/main/java/ai/dragon/service/EmbeddingModelService.java b/backend/src/main/java/ai/dragon/service/EmbeddingModelService.java
@@ -3,7 +3,7 @@
 import org.springframework.stereotype.Service;
 
 import ai.dragon.entity.SiloEntity;
-import ai.dragon.job.silo.ingestor.dto.embedding.EmbeddingModelSettings;
+import ai.dragon.properties.embedding.EmbeddingModelSettings;
 import ai.dragon.util.IniSettingUtil;
 import dev.langchain4j.model.embedding.EmbeddingModel;
 

diff --git a/backend/src/main/java/ai/dragon/service/EmbeddingStoreService.java b/backend/src/main/java/ai/dragon/service/EmbeddingStoreService.java
@@ -1,5 +1,7 @@
 package ai.dragon.service;
 
+import java.io.File;
+import java.nio.file.Path;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.UUID;
@@ -8,9 +10,12 @@
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.stereotype.Service;
 
+import ai.dragon.component.DirectoryStructureComponent;
 import ai.dragon.entity.SiloEntity;
 import ai.dragon.listener.EntityChangeListener;
+import ai.dragon.properties.store.InMemoryEmbeddingStoreSettings;
 import ai.dragon.repository.SiloRepository;
+import ai.dragon.util.IniSettingUtil;
 import dev.langchain4j.data.embedding.Embedding;
 import dev.langchain4j.data.segment.TextSegment;
 import dev.langchain4j.model.embedding.EmbeddingModel;
@@ -32,7 +37,11 @@ public class EmbeddingStoreService {
     @Autowired
     private EmbeddingModelService embeddingModelService;
 
+    @Autowired
+    private DirectoryStructureComponent directoryStructureComponent;
+
     private EntityChangeListener<SiloEntity> entityChangeListener;
+    private Map<UUID, Path> inMemoryEmbededdingStorePersistPaths = new HashMap<>();
 
     @PostConstruct
     private void init() {
@@ -57,7 +66,7 @@ private void destroy() {
         closeAllEmbeddingStores();
     }
 
-    public EmbeddingStore<TextSegment> retrieveEmbeddingStore(UUID siloUuid) {
+    public EmbeddingStore<TextSegment> retrieveEmbeddingStore(UUID siloUuid) throws Exception {
         if (embeddingStores.containsKey(siloUuid)) {
             return embeddingStores.get(siloUuid);
         }
@@ -68,11 +77,24 @@ public EmbeddingStore<TextSegment> retrieveEmbeddingStore(UUID siloUuid) {
     }
 
     public void closeEmbeddingStore(UUID siloUuid) {
-        if (embeddingStores.containsKey(siloUuid)) {
+        EmbeddingStore<TextSegment> embeddingStore = embeddingStores.get(siloUuid);
+        if (embeddingStore != null) {
+            persistEmbeddingStore(siloUuid);
             embeddingStores.remove(siloUuid);
         }
     }
 
+    public void persistEmbeddingStore(UUID siloUuid) {
+        EmbeddingStore<TextSegment> embeddingStore = embeddingStores.get(siloUuid);
+        if (embeddingStore instanceof InMemoryEmbeddingStore) {
+            InMemoryEmbeddingStore<TextSegment> inMemoryEmbeddingStore = (InMemoryEmbeddingStore<TextSegment>) embeddingStore;
+            Path vectorFile = inMemoryEmbededdingStorePersistPaths.get(siloUuid);
+            if (vectorFile != null) {
+                inMemoryEmbeddingStore.serializeToFile(vectorFile);
+            }
+        }
+    }
+
     public void closeAllEmbeddingStores() {
         for (UUID siloUuid : embeddingStores.keySet()) {
             closeEmbeddingStore(siloUuid);
@@ -94,6 +116,7 @@ public void query(UUID siloUuid, String query) throws Exception {
         EmbeddingSearchRequest embeddingSearchRequest1 = EmbeddingSearchRequest.builder()
                 .queryEmbedding(queryEmbedding)
                 // .filter(onlyForUser1)
+                .maxResults(10)
                 .build();
         EmbeddingSearchResult<TextSegment> embeddingSearchResult1 = embeddingStore.search(embeddingSearchRequest1);
         for (EmbeddingMatch<TextSegment> embeddingMatch : embeddingSearchResult1.matches()) {
@@ -104,12 +127,24 @@ public void query(UUID siloUuid, String query) throws Exception {
         }
     }
 
-    private EmbeddingStore<TextSegment> buildEmbeddingStore(SiloEntity siloEntity) {
+    private EmbeddingStore<TextSegment> buildEmbeddingStore(SiloEntity siloEntity) throws Exception {
         EmbeddingStore<TextSegment> embeddingStore = null;
 
         switch (siloEntity.getVectorStoreType()) {
             case InMemoryEmbeddingStore:
+                InMemoryEmbeddingStoreSettings storeSettings = IniSettingUtil.convertIniSettingsToObject(
+                        siloEntity.getVectorStoreSettings(), InMemoryEmbeddingStoreSettings.class);
+                String persistance = storeSettings.getPersistance();
                 embeddingStore = new InMemoryEmbeddingStore<>();
+                if (!":memory:".equals(persistance)) {
+                    File vectorFile = new File(directoryStructureComponent.directoryFor("vector"),
+                            siloEntity.getUuid().toString() + ".json");
+                    Path vectorPath = vectorFile.toPath();
+                    inMemoryEmbededdingStorePersistPaths.put(siloEntity.getUuid(), vectorPath);
+                    if (vectorFile.exists()) {
+                        embeddingStore = InMemoryEmbeddingStore.fromFile(vectorPath);
+                    }
+                }
                 break;
             default:
                 throw new UnsupportedOperationException(

diff --git a/backend/src/main/java/ai/dragon/service/IngestorService.java b/backend/src/main/java/ai/dragon/service/IngestorService.java
@@ -26,7 +26,6 @@ public class IngestorService {
     @Autowired
     private EmbeddingModelService embeddingModelService;
 
-
     public void runSiloIngestion(SiloEntity siloEntity, Consumer<Integer> progressCallback,
             Consumer<SiloIngestLoaderLogMessage> logCallback)
             throws Exception {
@@ -38,6 +37,9 @@ public void runSiloIngestion(SiloEntity siloEntity, Consumer<Integer> progressCa
         logCallback.accept(SiloIngestLoaderLogMessage.builder()
                 .message(String.format("Will ingest %d documents to Silo...", documents.size())).build());
         ingestDocumentsToSilo(documents, siloEntity, progressCallback, logCallback);
+        logCallback.accept(SiloIngestLoaderLogMessage.builder()
+                .message(String.format("Persisting the Embedding Store...", documents.size())).build());
+        embeddingStoreService.persistEmbeddingStore(siloEntity.getUuid());
     }
 
     private void ingestDocumentsToSilo(List<Document> documents, SiloEntity siloEntity,
@@ -62,7 +64,7 @@ private void ingestDocumentsToSilo(List<Document> documents, SiloEntity siloEnti
             progressCallback.accept(progress);
         }
         logCallback.accept(SiloIngestLoaderLogMessage.builder()
-                .message("End.").build());
+                .message("End of ingestion.").build());
         progressCallback.accept(100);
     }
-Original file line number
+Diff line change
@@ Expand Up @@
                         jobContext().logger().info(ingestLogMessage.getMessage());
                 }
             });
+            jobContext().logger().info("End of Silo Ingestor Job.");
         }
     }