From 79b3dcdf39b94bfa5e87a0cbc6cc9a2c7bccae87 Mon Sep 17 00:00:00 2001 From: Wenzheng Jiang Date: Fri, 14 Feb 2025 16:32:25 +0900 Subject: [PATCH] Implement getFileTree intent (#1249) * Add FileTreeTools.ts * Output file tree in JSON format * Optimize file tree output --- src/LLMProviders/chainRunner.ts | 2 + src/LLMProviders/intentAnalyzer.ts | 29 ++++--- src/constants.ts | 1 + src/main.ts | 3 + src/tools/FileTreeTools.test.ts | 120 +++++++++++++++++++++++++++++ src/tools/FileTreeTools.ts | 55 +++++++++++++ 6 files changed, 200 insertions(+), 10 deletions(-) create mode 100644 src/tools/FileTreeTools.test.ts create mode 100644 src/tools/FileTreeTools.ts diff --git a/src/LLMProviders/chainRunner.ts b/src/LLMProviders/chainRunner.ts index 0a1f70da..05ab7d9c 100644 --- a/src/LLMProviders/chainRunner.ts +++ b/src/LLMProviders/chainRunner.ts @@ -632,6 +632,8 @@ class CopilotPlusChainRunner extends BaseChainRunner { updateLoadingMessage?.(LOADING_MESSAGES.READING_FILES); } else if (toolCall.tool.name === "webSearch") { updateLoadingMessage?.(LOADING_MESSAGES.SEARCHING_WEB); + } else if (toolCall.tool.name === "getFileTree") { + updateLoadingMessage?.(LOADING_MESSAGES.READING_FILE_TREE); } const output = await ToolManager.callTool(toolCall.tool, toolCall.args); toolOutputs.push({ tool: toolCall.tool.name, output }); diff --git a/src/LLMProviders/intentAnalyzer.ts b/src/LLMProviders/intentAnalyzer.ts index f39bcc9a..7aad5764 100644 --- a/src/LLMProviders/intentAnalyzer.ts +++ b/src/LLMProviders/intentAnalyzer.ts @@ -6,11 +6,13 @@ import { pomodoroTool, TimeInfo, } from "@/tools/TimeTools"; +import { createGetFileTreeTool } from "@/tools/FileTreeTools"; import { simpleYoutubeTranscriptionTool } from "@/tools/YoutubeTools"; import { ToolManager } from "@/tools/toolManager"; import { extractChatHistory, extractYoutubeUrl } from "@/utils"; import { BrevilabsClient } from "./brevilabsClient"; import MemoryManager from "./memoryManager"; +import { Vault } from "obsidian"; // TODO: Add @index with explicit pdf files in chat context menu export const COPILOT_TOOL_NAMES = ["@vault", "@web", "@youtube", "@pomodoro"]; @@ -21,16 +23,23 @@ type ToolCall = { }; export class IntentAnalyzer { - private static tools = [ - getCurrentTimeTool, - getTimeInfoByEpochTool, - getTimeRangeMsTool, - localSearchTool, - indexTool, - pomodoroTool, - webSearchTool, - simpleYoutubeTranscriptionTool, - ]; + private static tools: any[] = []; + + static initTools(vault: Vault) { + if (this.tools.length === 0) { + this.tools = [ + getCurrentTimeTool, + getTimeInfoByEpochTool, + getTimeRangeMsTool, + localSearchTool, + indexTool, + pomodoroTool, + webSearchTool, + simpleYoutubeTranscriptionTool, + createGetFileTreeTool(vault.getRoot()), + ]; + } + } static async analyzeIntent(originalMessage: string): Promise { try { diff --git a/src/constants.ts b/src/constants.ts index 6b16f96d..0f162d5e 100644 --- a/src/constants.ts +++ b/src/constants.ts @@ -29,6 +29,7 @@ export const LOADING_MESSAGES = { DEFAULT: "", READING_FILES: "Reading files", SEARCHING_WEB: "Searching the web", + READING_FILE_TREE: "Reading file tree", }; export const PLUS_UTM_MEDIUMS = { SETTINGS: "settings", diff --git a/src/main.ts b/src/main.ts index 19a9a741..d83cb9c1 100644 --- a/src/main.ts +++ b/src/main.ts @@ -31,6 +31,7 @@ import { TFolder, WorkspaceLeaf, } from "obsidian"; +import { IntentAnalyzer } from "./LLMProviders/intentAnalyzer"; export default class CopilotPlugin extends Plugin { // A chat history that stores the messages sent and received @@ -80,6 +81,8 @@ export default class CopilotPlugin extends Plugin { registerBuiltInCommands(this); + IntentAnalyzer.initTools(this.app.vault); + this.registerEvent( this.app.workspace.on("editor-menu", (menu: Menu, editor: Editor) => { const selectedText = editor.getSelection().trim(); diff --git a/src/tools/FileTreeTools.test.ts b/src/tools/FileTreeTools.test.ts new file mode 100644 index 00000000..f424b9a5 --- /dev/null +++ b/src/tools/FileTreeTools.test.ts @@ -0,0 +1,120 @@ +import { TFolder } from "obsidian"; +import { createGetFileTreeTool } from "./FileTreeTools"; + +// Mock TFile class +class MockTFile { + vault: any; + stat: { ctime: number; mtime: number; size: number }; + basename: string; + extension: string; + name: string; + parent: TFolder; + path: string; + + constructor(path: string, parent: TFolder) { + this.path = path; + this.name = path.split("/").pop() || ""; + this.basename = this.name.split(".")[0]; + this.extension = this.name.split(".")[1] || ""; + this.parent = parent; + this.vault = {}; + this.stat = { + ctime: Date.now(), + mtime: Date.now(), + size: 0, + }; + } +} + +// Mock TFolder class +class MockTFolder { + vault: any; + name: string; + path: string; + parent: TFolder | null; + children: Array; + + constructor(path: string, parent: TFolder | null) { + this.path = path; + this.name = path.split("/").pop() || ""; + this.parent = parent; + this.vault = {}; + this.children = []; + } + + isRoot(): boolean { + return this.parent === null; + } +} + +describe("FileTreeTools", () => { + let root: MockTFolder; + + beforeEach(() => { + root = new MockTFolder("", null); + + // Create a mock file structure + const docs = new MockTFolder("docs", root); + const projects = new MockTFolder("docs/projects", docs); + const notes = new MockTFolder("docs/notes", docs); + + docs.children = [projects, notes, new MockTFile("docs/readme.md", docs)]; + + projects.children = [ + new MockTFile("docs/projects/project1.md", projects), + new MockTFile("docs/projects/project2.md", projects), + ]; + + notes.children = [ + new MockTFile("docs/notes/note1.md", notes), + new MockTFile("docs/notes/note2.md", notes), + ]; + + root.children = [docs]; + }); + + it("should generate correct JSON file tree representation", async () => { + const tool = createGetFileTreeTool(root); + const result = await tool.invoke({}); + const parsedResult = JSON.parse(result); + + const expected = { + path: "", + children: [ + { + path: "docs", + children: [ + { + path: "docs/projects", + children: [ + { path: "docs/projects/project1.md" }, + { path: "docs/projects/project2.md" }, + ], + }, + { + path: "docs/notes", + children: [{ path: "docs/notes/note1.md" }, { path: "docs/notes/note2.md" }], + }, + { path: "docs/readme.md" }, + ], + }, + ], + }; + + expect(parsedResult).toEqual(expected); + }); + + it("should handle empty folder", async () => { + const emptyRoot = new MockTFolder("", null); + const tool = createGetFileTreeTool(emptyRoot); + const result = await tool.invoke({}); + const parsedResult = JSON.parse(result); + + const expected = { + path: "", + children: [], + }; + + expect(parsedResult).toEqual(expected); + }); +}); diff --git a/src/tools/FileTreeTools.ts b/src/tools/FileTreeTools.ts new file mode 100644 index 00000000..e752894b --- /dev/null +++ b/src/tools/FileTreeTools.ts @@ -0,0 +1,55 @@ +import { tool } from "@langchain/core/tools"; +import { TFile, TFolder } from "obsidian"; +import { z } from "zod"; + +interface FileTreeNode { + path: string; + children?: FileTreeNode[]; +} + +function isTFolder(item: any): item is TFolder { + return "children" in item && "path" in item; +} + +function isTFile(item: any): item is TFile { + return "path" in item && !("children" in item); +} + +function buildFileTree(folder: TFolder): FileTreeNode { + const node: FileTreeNode = { + path: folder.path, + children: [], + }; + + // Add folders first + for (const child of folder.children) { + if (isTFolder(child)) { + node.children?.push(buildFileTree(child)); + } + } + + // Then add files + for (const child of folder.children) { + if (isTFile(child)) { + node.children?.push({ + path: child.path, + }); + } + } + + return node; +} + +const createGetFileTreeTool = (root: TFolder) => + tool( + async () => { + return JSON.stringify(buildFileTree(root), null, 0); + }, + { + name: "getFileTree", + description: "Get the complete file tree structure of the folder as JSON", + schema: z.object({}), + } + ); + +export { createGetFileTreeTool, type FileTreeNode };