File Tree: support exclusion and inclusion rules and simplify JSON st…

…ructure (#1284)

src/tools/FileTreeTools.test.ts

@@ -1,5 +1,12 @@
 import { TFolder } from "obsidian";
 import { createGetFileTreeTool } from "./FileTreeTools";
+import * as searchUtils from "@/search/searchUtils";
+
+// Mock the searchUtils functions
+jest.mock("@/search/searchUtils", () => ({
+  getMatchingPatterns: jest.fn(),
+  shouldIndexFile: jest.fn(),
+}));
 
 // Mock TFile class
 class MockTFile {
@@ -70,29 +77,67 @@ describe("FileTreeTools", () => {
       new MockTFile("docs/notes/note2.md", notes),
     ];
 
-    root.children = [docs];
+    root.children = [docs, new MockTFile("readme.md", root)];
+
+    // Reset mocks before each test
+    jest.clearAllMocks();
+
+    // Default mock implementations
+    (searchUtils.getMatchingPatterns as jest.Mock).mockReturnValue({
+      inclusions: null,
+      exclusions: null,
+    });
+    (searchUtils.shouldIndexFile as jest.Mock).mockReturnValue(true);
   });
 
-  it("should generate correct JSON file tree with filenames only", async () => {
+  it("should generate correct JSON file tree when no exclusions", async () => {
     const tool = createGetFileTreeTool(root);
     const result = await tool.invoke({});
-    const parsedResult = JSON.parse(result);
+    // Extract JSON part after the prompt
+    const jsonPart = result.substring(result.indexOf("{"));
+    const parsedResult = JSON.parse(jsonPart);
+
+    expect(searchUtils.getMatchingPatterns).toHaveBeenCalled();
+    expect(searchUtils.shouldIndexFile).toHaveBeenCalled();
 
     const expected = {
-      path: "",
-      children: [
+      vault: [
+        ["readme.md"],
         {
-          path: "docs",
-          children: [
+          docs: [
+            ["readme.md"],
             {
-              path: "projects",
-              children: [{ path: "project1.md" }, { path: "project2.md" }],
+              projects: ["project1.md", "project2.md"],
+              notes: ["note1.md", "note2.md"],
             },
+          ],
+        },
+      ],
+    };
+
+    expect(parsedResult).toEqual(expected);
+  });
+
+  it("should exclude files based on patterns", async () => {
+    // Mock shouldIndexFile to exclude all files in projects folder
+    (searchUtils.shouldIndexFile as jest.Mock).mockImplementation((file) => {
+      return !file.path.includes("projects");
+    });
+
+    const tool = createGetFileTreeTool(root);
+    const result = await tool.invoke({});
+    const jsonPart = result.substring(result.indexOf("{"));
+    const parsedResult = JSON.parse(jsonPart);
+
+    const expected = {
+      vault: [
+        ["readme.md"],
+        {
+          docs: [
+            ["readme.md"],
             {
-              path: "notes",
-              children: [{ path: "note1.md" }, { path: "note2.md" }],
+              notes: ["note1.md", "note2.md"],
             },
-            { path: "readme.md" },
           ],
         },
       ],
@@ -101,15 +146,35 @@ describe("FileTreeTools", () => {
     expect(parsedResult).toEqual(expected);
   });
 
-  it("should handle empty folder", async () => {
-    const emptyRoot = new MockTFolder("", null);
-    const tool = createGetFileTreeTool(emptyRoot);
+  it("should handle empty folder after exclusions", async () => {
+    // Mock shouldIndexFile to exclude all files
+    (searchUtils.shouldIndexFile as jest.Mock).mockReturnValue(false);
+
+    const tool = createGetFileTreeTool(root);
     const result = await tool.invoke({});
-    const parsedResult = JSON.parse(result);
+    const jsonPart = result.substring(result.indexOf("{"));
+    const parsedResult = JSON.parse(jsonPart);
+
+    expect(parsedResult).toEqual({});
+  });
+
+  it("should handle partial folder exclusions", async () => {
+    // Mock shouldIndexFile to only include files with "note" in the path
+    (searchUtils.shouldIndexFile as jest.Mock).mockImplementation((file) => {
+      return file.path.includes("note");
+    });
+
+    const tool = createGetFileTreeTool(root);
+    const result = await tool.invoke({});
+    const jsonPart = result.substring(result.indexOf("{"));
+    const parsedResult = JSON.parse(jsonPart);
 
     const expected = {
-      path: "",
-      children: [],
+      vault: {
+        docs: {
+          notes: ["note1.md", "note2.md"],
+        },
+      },
     };
 
     expect(parsedResult).toEqual(expected);

src/tools/FileTreeTools.ts

@@ -1,11 +1,9 @@
 import { tool } from "@langchain/core/tools";
 import { TFile, TFolder } from "obsidian";
 import { z } from "zod";
+import { getMatchingPatterns, shouldIndexFile } from "@/search/searchUtils";
 
-interface FileTreeNode {
-  path: string;
-  children?: FileTreeNode[];
-}
+type FileTreeNode = { [key: string]: string[] | FileTreeNode | (string[] | FileTreeNode)[] };
 
 function isTFolder(item: any): item is TFolder {
   return "children" in item && "path" in item;
@@ -16,38 +14,67 @@ function isTFile(item: any): item is TFile {
 }
 
 function buildFileTree(folder: TFolder): FileTreeNode {
-  const node: FileTreeNode = {
-    path: folder.name,
-    children: [],
-  };
+  const result: FileTreeNode = {};
+  const files: string[] = [];
+  const folders: FileTreeNode = {};
+
+  // Get exclusion patterns from settings
+  const { inclusions, exclusions } = getMatchingPatterns();
 
-  // Add folders first
+  // Separate files and folders
   for (const child of folder.children) {
-    if (isTFolder(child)) {
-      node.children?.push(buildFileTree(child));
+    if (isTFile(child)) {
+      // Only include file if it passes the pattern checks
+      if (shouldIndexFile(child, inclusions, exclusions)) {
+        files.push(child.name);
+      }
+    } else if (isTFolder(child)) {
+      const subTree = buildFileTree(child);
+      // Only include folder if it has any content after filtering
+      if (Object.keys(subTree).length > 0) {
+        Object.assign(folders, subTree);
+      }
     }
   }
 
-  // Then add files
-  for (const child of folder.children) {
-    if (isTFile(child)) {
-      node.children?.push({
-        path: child.name,
-      });
+  // If this is root folder, name it "vault" and return merged result
+  if (!folder.name) {
+    // Only return if there's content after filtering
+    if (files.length || Object.keys(folders).length) {
+      return {
+        vault: files.length ? (Object.keys(folders).length ? [files, folders] : files) : folders,
+      };
     }
+    return {};
+  }
+
+  // For named folders, nest everything under the folder name
+  // Only include folder if it has any content after filtering
+  if (files.length || Object.keys(folders).length) {
+    result[folder.name] = files.length
+      ? Object.keys(folders).length
+        ? [files, folders]
+        : files
+      : folders;
   }
 
-  return node;
+  return result;
 }
 
 const createGetFileTreeTool = (root: TFolder) =>
   tool(
     async () => {
-      return JSON.stringify(buildFileTree(root), null, 0);
+      const prompt = `A JSON represents the file tree as a nested structure:
+* The root object has a key "vault" which maps to an array with two items:
+ * An array of files at the current directory.
+ * An object of subdirectories, where each subdirectory follows the same structure as the root.
+
+`;
+      return prompt + JSON.stringify(buildFileTree(root));
     },
     {
       name: "getFileTree",
-      description: "Get the complete file tree structure of the folder as JSON",
+      description: "Get the file tree as JSON where folders are objects and files are arrays",
       schema: z.object({}),
     }
   );