Refactor note reference and fix dupe title issue (#1275)

* Refactor note reference handling with improved file resolution * Improve note processing with deduplication and enhanced context formatting * Refactor active note type-in and time filtering to use just modified time * Update tests and remove deprecated getNoteFileFromTitle
logancyang · Feb 20, 2025 · 5a63112 · 5a63112
1 parent d2c8f42
commit 5a63112
Show file tree

Hide file tree

Showing 8 changed files with 244 additions and 213 deletions.
diff --git a/src/components/chat-components/ChatInput.tsx b/src/components/chat-components/ChatInput.tsx
@@ -1,14 +1,23 @@
 import { useChainType, useModelKey } from "@/aiParams";
 import { ChainType } from "@/chainFactory";
+import { AddContextNoteModal } from "@/components/modals/AddContextNoteModal";
 import { AddImageModal } from "@/components/modals/AddImageModal";
 import { ListPromptModal } from "@/components/modals/ListPromptModal";
+import { Button } from "@/components/ui/button";
+import {
+  DropdownMenu,
+  DropdownMenuContent,
+  DropdownMenuItem,
+  DropdownMenuTrigger,
+} from "@/components/ui/dropdown-menu";
+import { ModelDisplay } from "@/components/ui/model-display";
 import { ContextProcessor } from "@/contextProcessor";
 import { CustomPromptProcessor } from "@/customPromptProcessor";
 import { COPILOT_TOOL_NAMES } from "@/LLMProviders/intentAnalyzer";
 import { Mention } from "@/mentions/Mention";
 import { getModelKeyFromModel, useSettingsValue } from "@/settings/model";
 import { getToolDescription } from "@/tools/toolManager";
-import { err2String, extractNoteTitles, checkModelApiKey } from "@/utils";
+import { checkModelApiKey, err2String, extractNoteFiles, isNoteTitleUnique } from "@/utils";
 import {
   ArrowBigUp,
   ChevronDown,
@@ -18,28 +27,18 @@ import {
   StopCircle,
   X,
 } from "lucide-react";
-import { App, Platform, TFile } from "obsidian";
+import { App, Notice, Platform, TFile } from "obsidian";
 import React, {
   forwardRef,
+  useCallback,
   useEffect,
   useImperativeHandle,
   useMemo,
   useRef,
   useState,
-  useCallback,
 } from "react";
-import ContextControl from "./ContextControl";
-import {
-  DropdownMenu,
-  DropdownMenuContent,
-  DropdownMenuItem,
-  DropdownMenuTrigger,
-} from "@/components/ui/dropdown-menu";
-import { Button } from "@/components/ui/button";
 import { useDropzone } from "react-dropzone";
-import { AddContextNoteModal } from "@/components/modals/AddContextNoteModal";
-import { ModelDisplay } from "@/components/ui/model-display";
-import { Notice } from "obsidian";
+import ContextControl from "./ContextControl";
 
 interface ChatInputProps {
   inputMessage: string;
@@ -165,7 +164,12 @@ const ChatInput = forwardRef<{ focus: () => void }, ChatInputProps>(
           onNoteSelect: async (note: TFile) => {
             const before = inputMessage.slice(0, cursorPos - 2);
             const after = inputMessage.slice(cursorPos - 1);
-            const newInputMessage = `${before}[[${note.basename}]]${after}`;
+
+            // Check if this note title has duplicates
+            const isUnique = isNoteTitleUnique(note.basename, app.vault);
+            // If the title is unique, just show the title, otherwise show the full path
+            const noteRef = isUnique ? note.basename : note.path;
+            const newInputMessage = `${before}[[${noteRef}]]${after}`;
             setInputMessage(newInputMessage);
 
             const activeNote = app.workspace.getActiveFile();
@@ -183,7 +187,7 @@ const ChatInput = forwardRef<{ focus: () => void }, ChatInputProps>(
             // Add a delay to ensure the cursor is set after inputMessage is updated
             setTimeout(() => {
               if (textAreaRef.current) {
-                const newCursorPos = cursorPos + note.basename.length + 2;
+                const newCursorPos = cursorPos + noteRef.length + 2;
                 textAreaRef.current.setSelectionRange(newCursorPos, newCursorPos);
               }
             }, 0);
@@ -325,7 +329,7 @@ const ChatInput = forwardRef<{ focus: () => void }, ChatInputProps>(
 
     useEffect(() => {
       // Get all note titles that are referenced using [[note]] syntax in the input
-      const currentTitles = new Set(extractNoteTitles(inputMessage));
+      const currentFiles = new Set(extractNoteFiles(inputMessage, app.vault));
       // Get all URLs mentioned in the input
       const currentUrls = mention.extractAllUrls(inputMessage);
 
@@ -344,22 +348,19 @@ const ChatInput = forwardRef<{ focus: () => void }, ChatInputProps>(
           if (note.path === currentActiveNote?.path) {
             if (wasAddedViaReference) {
               // Case 1: Active note was added by typing [[note]]
-              // Keep it only if its title is still in the input
-              // This ensures it's removed when you delete the [[note]] reference
-              return currentTitles.has(note.basename);
+              // Keep it only if its file is still in the input
+              return currentFiles.has(note);
             } else {
               // Case 2: Active note was NOT added by [[note]], but by the includeActiveNote toggle
               // Keep it only if includeActiveNote is true
-              // This handles the "Include active note" toggle in the UI
               return includeActiveNote;
             }
           } else {
             // Handling for all other notes (not the active note)
             if (wasAddedViaReference) {
               // Case 3: Other note was added by typing [[note]]
-              // Keep it only if its title is still in the input
-              // This ensures it's removed when you delete the [[note]] reference
-              return currentTitles.has(note.basename);
+              // Keep it only if its file is still in the input
+              return currentFiles.has(note);
             } else {
               // Case 4: Other note was added via "Add Note to Context" button
               // Always keep these notes as they were manually added
@@ -371,7 +372,7 @@ const ChatInput = forwardRef<{ focus: () => void }, ChatInputProps>(
 
       // Remove any URLs that are no longer present in the input
       setContextUrls((prev) => prev.filter((url) => currentUrls.includes(url)));
-    }, [inputMessage, includeActiveNote, currentActiveNote, mention, setContextNotes]);
+    }, [inputMessage, includeActiveNote, currentActiveNote, mention, setContextNotes, app.vault]);
 
     // Update the current active note whenever it changes
     useEffect(() => {

diff --git a/src/components/modals/OramaSearchModal.tsx b/src/components/modals/OramaSearchModal.tsx
@@ -1,5 +1,5 @@
 import CopilotPlugin from "@/main";
-import { extractNoteTitles } from "@/utils";
+import { extractNoteFiles } from "@/utils";
 import { App, Modal, Notice, TFile } from "obsidian";
 
 export class OramaSearchModal extends Modal {
@@ -36,7 +36,7 @@ export class OramaSearchModal extends Modal {
 
     searchButton.addEventListener("click", async () => {
       const input = this.searchInput.value;
-      const notePaths = extractNoteTitles(input);
+      const notePaths = extractNoteFiles(input, this.app.vault).map((file) => file.path);
 
       if (notePaths.length === 0) {
         new Notice("No valid note paths found. Use format: - [[Note Name]]");

diff --git a/src/contextProcessor.ts b/src/contextProcessor.ts
@@ -57,6 +57,12 @@ export class ContextProcessor {
 
     const processNote = async (note: TFile) => {
       try {
+        // Skip if this note was already processed by processCustomPrompt
+        const noteRef = `[[${note.basename}]]`;
+        if (processedVars.has(noteRef)) {
+          return;
+        }
+
         if (currentChain !== ChainType.COPILOT_PLUS_CHAIN && note.extension !== "md") {
           if (!fileParserManager.supportsExtension(note.extension)) {
             console.warn(`Unsupported file type: ${note.extension}`);
@@ -77,10 +83,10 @@ export class ContextProcessor {
           content = await this.processEmbeddedPDFs(content, vault, fileParserManager);
         }
 
-        additionalContext += `\n\n[[${note.basename}.${note.extension}]]:\n\n${content}`;
+        additionalContext += `\n\nTitle: [[${note.basename}]]\nPath: ${note.path}\n\n${content}`;
       } catch (error) {
         console.error(`Error processing file ${note.path}:`, error);
-        additionalContext += `\n\n[[${note.basename}]]: [Error: Could not process file]`;
+        additionalContext += `\n\nTitle: [[${note.basename}]]\nPath: ${note.path}\n\n[Error: Could not process file]`;
       }
     };
 
@@ -114,30 +120,27 @@ export class ContextProcessor {
     setContextNotes: (notes: TFile[] | ((prev: TFile[]) => TFile[])) => void,
     setIncludeActiveNote: (include: boolean) => void
   ): Promise<void> {
-    // First check if this note can be added
-    if (
-      contextNotes.some((existing) => existing.path === note.path) ||
-      (activeNote && note.path === activeNote.path)
-    ) {
+    // Only check if the note exists in contextNotes
+    if (contextNotes.some((existing) => existing.path === note.path)) {
       return; // Note already exists in context
     }
 
     // Read the note content
     const content = await vault.read(note);
     const hasEmbeddedPDFs = await this.hasEmbeddedPDFs(content);
 
-    // If it's the active note, set includeActiveNote to true
+    // Set includeActiveNote if it's the active note
     if (activeNote && note.path === activeNote.path) {
       setIncludeActiveNote(true);
-    } else {
-      // Otherwise add it to contextNotes
-      setContextNotes((prev: TFile[]) => [
-        ...prev,
-        Object.assign(note, {
-          wasAddedViaReference: true,
-          hasEmbeddedPDFs,
-        }),
-      ]);
     }
+
+    // Add to contextNotes with wasAddedViaReference flag
+    setContextNotes((prev: TFile[]) => [
+      ...prev,
+      Object.assign(note, {
+        wasAddedViaReference: true,
+        hasEmbeddedPDFs,
+      }),
+    ]);
   }
 }
diff --git a/src/customPromptProcessor.test.ts b/src/customPromptProcessor.test.ts
@@ -1,5 +1,5 @@
 import { CustomPrompt, CustomPromptProcessor } from "@/customPromptProcessor";
-import { extractNoteTitles, getFileContent, getNoteFileFromTitle, getNotesFromPath } from "@/utils";
+import { extractNoteFiles, getFileContent, getNotesFromPath } from "@/utils";
 import { Notice, TFile, Vault } from "obsidian";
 
 // Mock Obsidian
@@ -11,8 +11,7 @@ jest.mock("obsidian", () => ({
 
 // Mock the utility functions
 jest.mock("@/utils", () => ({
-  extractNoteTitles: jest.fn().mockReturnValue([]),
-  getNoteFileFromTitle: jest.fn(),
+  extractNoteFiles: jest.fn().mockReturnValue([]),
   getFileContent: jest.fn(),
   getFileName: jest.fn(),
   getNotesFromPath: jest.fn(),
@@ -176,14 +175,15 @@ describe("CustomPromptProcessor", () => {
     const selectedText = "";
 
     // Mock the necessary functions
-    (extractNoteTitles as jest.Mock).mockReturnValue(["Test Note"]);
-    (getNoteFileFromTitle as jest.Mock).mockResolvedValue({} as TFile);
+    (extractNoteFiles as jest.Mock).mockReturnValue([
+      { basename: "Test Note", path: "Test Note.md" },
+    ]);
     (getFileContent as jest.Mock).mockResolvedValue("Test note content");
 
     const result = await processor.processCustomPrompt(customPrompt, selectedText, mockActiveNote);
 
     expect(result).toContain("Content of [[Test Note]] is important.");
-    expect(result).toContain("[[Test Note]]:\n\nTest note content");
+    expect(result).toContain("Title: [[Test Note]]\nPath: Test Note.md\n\nTest note content");
   });
 
   it("should process {[[note title]]} syntax correctly without duplication", async () => {
@@ -194,21 +194,17 @@ describe("CustomPromptProcessor", () => {
     jest
       .spyOn(processor, "extractVariablesFromPrompt")
       .mockResolvedValue([JSON.stringify([{ name: "Test Note", content: "Test note content" }])]);
-    (extractNoteTitles as jest.Mock).mockReturnValue(["Test Note"]);
+    (extractNoteFiles as jest.Mock).mockReturnValue([
+      { basename: "Test Note", path: "Test Note.md" },
+    ]);
 
     const result = await processor.processCustomPrompt(customPrompt, selectedText, mockActiveNote);
 
     expect(result).toContain("Content of {[[Test Note]]} is important.");
     expect(result).toContain(
       '[[Test Note]]:\n\n[{"name":"Test Note","content":"Test note content"}]'
     );
-
-    // Ensure the content is not duplicated
-    const occurrences = (result.match(/Test note content/g) || []).length;
-    expect(occurrences).toBe(1);
-
-    // Verify that getNoteFileFromTitle was not called
-    expect(getNoteFileFromTitle).not.toHaveBeenCalled();
+    expect((result.match(/Test note content/g) || []).length).toBe(1);
   });
 
   it("should process both {[[note title]]} and [[note title]] syntax correctly", async () => {
@@ -219,32 +215,29 @@ describe("CustomPromptProcessor", () => {
     jest
       .spyOn(processor, "extractVariablesFromPrompt")
       .mockResolvedValue([JSON.stringify([{ name: "Note1", content: "Note1 content" }])]);
-    (extractNoteTitles as jest.Mock).mockReturnValue(["Note1", "Note2"]);
-    (getNoteFileFromTitle as jest.Mock).mockResolvedValue({} as TFile);
+    (extractNoteFiles as jest.Mock).mockReturnValue([
+      { basename: "Note1", path: "Note1.md" },
+      { basename: "Note2", path: "Note2.md" },
+    ]);
     (getFileContent as jest.Mock).mockResolvedValue("Note2 content");
 
     const result = await processor.processCustomPrompt(customPrompt, selectedText, mockActiveNote);
 
     expect(result).toContain("{[[Note1]]} content and [[Note2]] are both important.");
     expect(result).toContain('[[Note1]]:\n\n[{"name":"Note1","content":"Note1 content"}]');
-    expect(result).toContain("[[Note2]]:\n\nNote2 content");
-
-    // Ensure Note1 content is not duplicated
-    const note1Occurrences = (result.match(/Note1 content/g) || []).length;
-    expect(note1Occurrences).toBe(1);
-
-    // Verify that getNoteFileFromTitle was called only for Note2
-    expect(getNoteFileFromTitle).toHaveBeenCalledTimes(1);
-    expect(getNoteFileFromTitle).toHaveBeenCalledWith(expect.anything(), "Note2");
+    expect(result).toContain("Title: [[Note2]]\nPath: Note2.md\n\nNote2 content");
   });
 
   it("should handle multiple occurrences of [[note title]] syntax", async () => {
     const customPrompt = "[[Note1]] is related to [[Note2]] and [[Note3]].";
     const selectedText = "";
 
     // Mock the necessary functions
-    (extractNoteTitles as jest.Mock).mockReturnValue(["Note1", "Note2", "Note3"]);
-    (getNoteFileFromTitle as jest.Mock).mockResolvedValue({} as TFile);
+    (extractNoteFiles as jest.Mock).mockReturnValue([
+      { basename: "Note1", path: "Note1.md" },
+      { basename: "Note2", path: "Note2.md" },
+      { basename: "Note3", path: "Note3.md" },
+    ]);
     (getFileContent as jest.Mock)
       .mockResolvedValueOnce("Note1 content")
       .mockResolvedValueOnce("Note2 content")
@@ -253,18 +246,17 @@ describe("CustomPromptProcessor", () => {
     const result = await processor.processCustomPrompt(customPrompt, selectedText, mockActiveNote);
 
     expect(result).toContain("[[Note1]] is related to [[Note2]] and [[Note3]].");
-    expect(result).toContain("[[Note1]]:\n\nNote1 content");
-    expect(result).toContain("[[Note2]]:\n\nNote2 content");
-    expect(result).toContain("[[Note3]]:\n\nNote3 content");
+    expect(result).toContain("Title: [[Note1]]\nPath: Note1.md\n\nNote1 content");
+    expect(result).toContain("Title: [[Note2]]\nPath: Note2.md\n\nNote2 content");
+    expect(result).toContain("Title: [[Note3]]\nPath: Note3.md\n\nNote3 content");
   });
 
   it("should handle non-existent note titles gracefully", async () => {
     const customPrompt = "[[Non-existent Note]] should not cause errors.";
     const selectedText = "";
 
     // Mock the necessary functions
-    (extractNoteTitles as jest.Mock).mockReturnValue(["Non-existent Note"]);
-    (getNoteFileFromTitle as jest.Mock).mockResolvedValue(null);
+    (extractNoteFiles as jest.Mock).mockReturnValue(["Non-existent Note"]);
 
     const result = await processor.processCustomPrompt(customPrompt, selectedText, mockActiveNote);
 
@@ -329,20 +321,6 @@ describe("CustomPromptProcessor", () => {
     expect(result).toContain("Summarize this: {selectedText}. Additional info: {activeNote}");
     expect(result).toContain("selectedText (entire active note):\n\n Content of the active note");
     expect(result).not.toContain("activeNote:");
-    expect((result.match(/Content of the active note/g) || []).length).toBe(1);
-  });
-
-  it("should handle {} when no active note or selected text is available", async () => {
-    const doc: CustomPrompt = {
-      title: "test-prompt",
-      content: "Process this: {}",
-    };
-    const selectedText = "";
-
-    const result = await processor.processCustomPrompt(doc.content, selectedText, undefined);
-
-    expect(result).toContain("Process this: {selectedText}");
-    expect(result).toContain("selectedText:\n\n (No selected text or active note available)");
   });
 
   it("should prioritize selected text over active note when both are available", async () => {
@@ -358,6 +336,6 @@ describe("CustomPromptProcessor", () => {
 
     expect(result).toContain("Analyze this: {selectedText}");
     expect(result).toContain("selectedText:\n\n This is the selected text");
-    expect(result).not.toContain("Content of the active note");
+    expect(result).not.toContain("selectedText (entire active note):");
   });
 });