Merge pull request HamaWhiteGG#39 from HamaWhiteGG/dev

Add DirectoryLoader
zhangxiaojiawow · Jul 14, 2023 · 107d080 · 107d080
2 parents c4a291d + efe2f6c
commit 107d080
Show file tree

Hide file tree

Showing 3 changed files with 228 additions and 0 deletions.
diff --git a/...chain-core/src/main/java/com/hw/langchain/document/loaders/directory/DirectoryLoader.java b/...chain-core/src/main/java/com/hw/langchain/document/loaders/directory/DirectoryLoader.java
@@ -0,0 +1,140 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.hw.langchain.document.loaders.directory;
+
+import com.hw.langchain.document.loaders.base.BaseLoader;
+import com.hw.langchain.document.loaders.text.TextLoader;
+import com.hw.langchain.exception.LangChainException;
+import com.hw.langchain.schema.Document;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import lombok.Builder;
+
+import java.io.IOException;
+import java.nio.file.FileSystems;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutorService;
+import java.util.stream.Stream;
+
+import static com.google.common.base.Preconditions.checkArgument;
+import static com.hw.langchain.utils.ThreadPoolUtil.createThreadPool;
+
+/**
+ * Load documents from a directory.
+ *
+ * @author HamaWhite
+ */
+@Builder
+public class DirectoryLoader extends BaseLoader {
+
+    private static final Logger LOG = LoggerFactory.getLogger(DirectoryLoader.class);
+
+    /**
+     * Path to directory.
+     */
+    private Path path;
+
+    /**
+     * Whether to load hidden files.
+     */
+    private boolean loadHidden;
+
+    /**
+     * Support UnstructuredFileLoader, TextLoader and BSHTMLLoader
+     */
+    @Builder.Default
+    private Class<? extends BaseLoader> loaderCls = TextLoader.class;
+
+    /**
+     * Whether to recursively search for files.
+     */
+    private boolean recursive;
+
+    /**
+     * Whether to use multithreading.
+     */
+    private boolean useMultithreading;
+
+    /**
+     * The maximum number of threads to use.
+     */
+    @Builder.Default
+    private int maxConcurrency = 4;
+
+    /**
+     * check if a path is visible.
+     */
+    private boolean isVisible(Path p) {
+        String separator = FileSystems.getDefault().getSeparator();
+        String[] parts = p.toAbsolutePath().toString().split(separator);
+        for (String part : parts) {
+            if (part.startsWith(".")) {
+                return false;
+            }
+        }
+        return true;
+    }
+
+    private void loadFile(Path item, Path path, List<Document> docs) {
+        if (isVisible(path.relativize(item)) || loadHidden) {
+            try {
+                LOG.info("Loading file: {}", item);
+                // Create an instance of the loader class and load sub docs
+                var subDocs = loaderCls.getConstructor(String.class)
+                        .newInstance(item.toString())
+                        .load();
+                docs.addAll(subDocs);
+            } catch (Exception e) {
+                throw new LangChainException(errorMessage(path.toString()), e);
+            }
+        }
+    }
+
+    @Override
+    public List<Document> load() {
+        checkArgument(path.toFile().exists(), "Directory not found: '%s'", path.toString());
+        checkArgument(path.toFile().isDirectory(), "Expected directory, got file: '%s'", path.toString());
+        try {
+            List<Path> items;
+            try (Stream<Path> stream = recursive ? Files.walk(path) : Files.list(path)) {
+                items = stream.filter(p -> !p.toFile().isDirectory()).toList();
+            }
+            List<Document> docs = new ArrayList<>();
+            if (useMultithreading) {
+                ExecutorService executor = createThreadPool(maxConcurrency);
+                var futures = items.stream()
+                        .map(item -> CompletableFuture.runAsync(() -> loadFile(item, path, docs), executor))
+                        .toList();
+                CompletableFuture.allOf(futures.toArray(CompletableFuture[]::new)).join();
+                executor.shutdown();
+            } else {
+                items.forEach(item -> loadFile(item, path, docs));
+            }
+            return docs;
+        } catch (IOException e) {
+            throw new LangChainException(errorMessage(path.toString()), e);
+        }
+    }
+}
diff --git a/langchain-core/src/main/java/com/hw/langchain/utils/ThreadPoolUtil.java b/langchain-core/src/main/java/com/hw/langchain/utils/ThreadPoolUtil.java
@@ -0,0 +1,42 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.hw.langchain.utils;
+
+import java.util.concurrent.*;
+
+/**
+ * @author HamaWhite
+ */
+public class ThreadPoolUtil {
+
+    private ThreadPoolUtil() {
+    }
+
+    public static ExecutorService createThreadPool(int maxConcurrency) {
+        int corePoolSize = Math.min(maxConcurrency, Runtime.getRuntime().availableProcessors());
+        long keepAliveTime = 0L;
+        TimeUnit unit = TimeUnit.MILLISECONDS;
+        BlockingQueue<Runnable> workQueue = new LinkedBlockingQueue<>();
+        ThreadFactory threadFactory = Executors.defaultThreadFactory();
+        RejectedExecutionHandler handler = new ThreadPoolExecutor.AbortPolicy();
+
+        return new ThreadPoolExecutor(corePoolSize, maxConcurrency, keepAliveTime, unit, workQueue, threadFactory,
+                handler);
+    }
+}
diff --git a/...n-core/src/test/java/com/hw/langchain/document/loaders/directory/DirectoryLoaderTest.java b/...n-core/src/test/java/com/hw/langchain/document/loaders/directory/DirectoryLoaderTest.java
@@ -0,0 +1,46 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.hw.langchain.document.loaders.directory;
+
+import com.hw.langchain.document.loaders.text.TextLoader;
+
+import org.junit.jupiter.api.Test;
+
+import java.nio.file.Path;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+/**
+ * @author HamaWhite
+ */
+class DirectoryLoaderTest {
+
+    @Test
+    void testLoad() {
+        String path = "../docs/extras/modules/";
+
+        var loader = DirectoryLoader.builder()
+                .path(Path.of(path))
+                .loaderCls(TextLoader.class)
+                .build();
+
+        var docs = loader.load();
+        assertEquals(1, docs.size());
+    }
+}