Enable tags in Advanced Custom Prompt (#296)

logancyang · Feb 16, 2024 · 01abcb7 · 01abcb7
1 parent 8ceadc2
commit 01abcb7
Show file tree

Hide file tree

Showing 7 changed files with 128 additions and 40 deletions.
diff --git a/.editorconfig b/.editorconfig
@@ -5,6 +5,6 @@ root = true
 charset = utf-8
 end_of_line = lf
 insert_final_newline = true
-indent_style = tab
-indent_size = 4
-tab_width = 4
+indent_style = space
+indent_size = 2
+tab_width = 2
diff --git a/src/components/AddPromptModal.tsx b/src/components/AddPromptModal.tsx
@@ -64,6 +64,11 @@ export class AddPromptModal extends Modal {
         { text: '- {FolderPath} represents a folder of notes. ' }
       );
       frag.createEl('br');
+      frag.createEl(
+        'strong',
+        { text: '- {#tag1, #tag2} represents ALL notes with ANY of the specified tags in their property (an OR operation). ' }
+      );
+      frag.createEl('br');
       frag.createEl('br');
       frag.appendText('Tip: turn on debug mode to show the processed prompt in the chat window.');
       frag.createEl('br');

diff --git a/src/components/AdhocPromptModal.tsx b/src/components/AdhocPromptModal.tsx
@@ -30,6 +30,12 @@ export class AdhocPromptModal extends Modal {
             { text: '- {FolderPath} represents a folder of notes. ' }
           );
           frag.createEl('br');
+          frag.createEl(
+            'strong',
+            { text: '- {#tag1, #tag2} represents ALL notes with ANY of the specified tags in their property (an OR operation). ' }
+          );
+          frag.createEl('br');
+          frag.createEl('br');
           frag.appendText('Tip: turn on debug mode to show the processed prompt in the chat window.');
           frag.createEl('br');
           frag.createEl('br');

diff --git a/src/customPromptProcessor.ts b/src/customPromptProcessor.ts
@@ -1,5 +1,11 @@
-import { getFileContent, getFileName, getNotesFromPath, processVariableName } from '@/utils';
-import { Notice, Vault } from 'obsidian';
+import {
+  getFileContent,
+  getFileName,
+  getNotesFromPath,
+  getNotesFromTags,
+  processVariableNameForNotePath,
+} from "@/utils";
+import { Notice, Vault } from "obsidian";
 
 export interface CustomPrompt {
   _id: string;
@@ -34,48 +40,75 @@ export class CustomPromptProcessor {
 
     while ((match = variableRegex.exec(customPrompt)) !== null) {
       const variableName = match[1].trim();
-      const processedVariableName = processVariableName(variableName);
-      const noteFiles = await getNotesFromPath(this.vault, processedVariableName);
       const notes = [];
 
-      for (const file of noteFiles) {
-        const content = await getFileContent(file, this.vault);
-        if (content) {
-          notes.push({ name: getFileName(file), content });
+      if (variableName.startsWith("#")) {
+        // Handle tag-based variable for multiple tags
+        const tagNames = variableName
+          .slice(1)
+          .split(",")
+          .map((tag) => tag.trim());
+        const noteFiles = await getNotesFromTags(this.vault, tagNames);
+        for (const file of noteFiles) {
+          const content = await getFileContent(file, this.vault);
+          if (content) {
+            notes.push({ name: getFileName(file), content });
+          }
+        }
+      } else {
+        const processedVariableName =
+          processVariableNameForNotePath(variableName);
+        const noteFiles = await getNotesFromPath(
+          this.vault,
+          processedVariableName
+        );
+        for (const file of noteFiles) {
+          const content = await getFileContent(file, this.vault);
+          if (content) {
+            notes.push({ name: getFileName(file), content });
+          }
         }
       }
 
       if (notes.length > 0) {
         variablesWithContent.push(JSON.stringify(notes));
       } else {
-        new Notice(`Warning: No valid notes found for the provided path '${variableName}'.`);
+        new Notice(
+          `Warning: No valid notes found for the provided path '${variableName}'.`
+        );
       }
     }
 
     return variablesWithContent;
   }
 
-  async processCustomPrompt(customPrompt: string, selectedText: string): Promise<string> {
-    const variablesWithContent = await this.extractVariablesFromPrompt(customPrompt);
+  async processCustomPrompt(
+    customPrompt: string,
+    selectedText: string
+  ): Promise<string> {
+    const variablesWithContent = await this.extractVariablesFromPrompt(
+      customPrompt
+    );
     let processedPrompt = customPrompt;
-    let index = 0; // Start with 0 for noteCollection0, noteCollection1, etc.
+    let index = 0; // Start with 0 for context0, context1, etc.
 
-    // Replace placeholders with noteCollectionX
+    // Replace placeholders with contextX
     processedPrompt = processedPrompt.replace(/\{([^}]+)\}/g, () => {
-      return `{noteCollection${index++}}`;
+      return `{context${index++}}`;
     });
 
-    let additionalInfo = '';
-    if (processedPrompt.includes('{}')) {
+    let additionalInfo = "";
+    if (processedPrompt.includes("{}")) {
       // Replace {} with {selectedText}
-      processedPrompt = processedPrompt.replace(/\{\}/g, '{selectedText}');
+      processedPrompt = processedPrompt.replace(/\{\}/g, "{selectedText}");
       additionalInfo += `selectedText:\n\n ${selectedText}`;
     }
 
     for (let i = 0; i < index; i++) {
-      additionalInfo += `\n\nnoteCollection${i}:\n\n${variablesWithContent[i]}`;
+      additionalInfo += `\n\ncontext${i}:\n\n${variablesWithContent[i]}`;
     }
 
-    return processedPrompt + '\n\n' + additionalInfo;
+    const endLine = "\nAvoid mentioning the variable names 'selectedText' or 'contextX' in the reply. ";
+    return processedPrompt + "\n\n" + additionalInfo + (index > 0 ? endLine : "");
   }
 }
diff --git a/src/utils.ts b/src/utils.ts
@@ -299,7 +299,7 @@ export function extractChatHistory(memoryVariables: MemoryVariables): [string, s
   return chatHistory;
 }
 
-export function processVariableName(variableName: string): string {
+export function processVariableNameForNotePath(variableName: string): string {
   variableName = variableName.trim();
   // Check if the variable name is enclosed in double brackets indicating it's a note
   if (variableName.startsWith('[[') && variableName.endsWith(']]')) {

diff --git a/tests/customPromptProcessor.test.ts b/tests/customPromptProcessor.test.ts
@@ -18,7 +18,7 @@ describe('CustomPromptProcessor', () => {
     processor = CustomPromptProcessor.getInstance(mockVault);
   });
 
-  it('should replace placeholders with 1 noteCollection and selectedText', async () => {
+  it('should replace placeholders with 1 context and selectedText', async () => {
     const doc: CustomPrompt = {
       _id: 'test-prompt',
       prompt: 'This is a {variable} and {}.'
@@ -30,12 +30,12 @@ describe('CustomPromptProcessor', () => {
 
     const result = await processor.processCustomPrompt(doc.prompt, selectedText);
 
-    expect(result).toContain('This is a {noteCollection0} and {selectedText}.');
+    expect(result).toContain('This is a {context0} and {selectedText}.');
     expect(result).toContain('here is some selected text 12345');
     expect(result).toContain('here is the note content for note0');
   });
 
-  it('should replace placeholders with 2 noteCollection and no selectedText', async () => {
+  it('should replace placeholders with 2 context and no selectedText', async () => {
     const doc: CustomPrompt = {
       _id: 'test-prompt',
       prompt: 'This is a {variable} and {var2}.'
@@ -50,12 +50,12 @@ describe('CustomPromptProcessor', () => {
 
     const result = await processor.processCustomPrompt(doc.prompt, selectedText);
 
-    expect(result).toContain('This is a {noteCollection0} and {noteCollection1}.');
+    expect(result).toContain('This is a {context0} and {context1}.');
     expect(result).toContain('here is the note content for note0');
     expect(result).toContain('note content for note1');
   });
 
-  it('should replace placeholders with 1 selectedText and no noteCollection', async () => {
+  it('should replace placeholders with 1 selectedText and no context', async () => {
     const doc: CustomPrompt = {
       _id: 'test-prompt',
       prompt: 'Rewrite the following text {}'
@@ -77,7 +77,7 @@ describe('CustomPromptProcessor', () => {
   });
 
   // This is not an expected use case but it's possible
-  it('should replace placeholders with 2 selectedText and no noteCollection', async () => {
+  it('should replace placeholders with 2 selectedText and no context', async () => {
     const doc: CustomPrompt = {
       _id: 'test-prompt',
       prompt: 'Rewrite the following text {} and {}'
@@ -112,4 +112,48 @@ describe('CustomPromptProcessor', () => {
 
     expect(result).toBe('This is a test prompt with no variables.\n\n');
   });
+
+  it("should process a single tag variable correctly", async () => {
+    const customPrompt = "Notes related to {#tag} are:";
+    const selectedText = "";
+
+    // Mock the extractVariablesFromPrompt method to simulate tag processing
+    jest
+      .spyOn(processor, "extractVariablesFromPrompt")
+      .mockResolvedValue([
+        '[{"name":"note","content":"Note content for #tag"}]',
+      ]);
+
+    const result = await processor.processCustomPrompt(
+      customPrompt,
+      selectedText
+    );
+
+    expect(result).toContain("Notes related to {context0} are:");
+    expect(result).toContain(
+      '[{"name":"note","content":"Note content for #tag"}]'
+    );
+  });
+
+  it("should process multiple tag variables correctly", async () => {
+    const customPrompt = "Notes related to {#tag1,#tag2,   #tag3} are:";
+    const selectedText = "";
+
+    // Mock the extractVariablesFromPrompt method to simulate processing of multiple tags
+    jest
+      .spyOn(processor, "extractVariablesFromPrompt")
+      .mockResolvedValue([
+        '[{"name":"note1","content":"Note content for #tag1"},{"name":"note2","content":"Note content for #tag2"}]',
+      ]);
+
+    const result = await processor.processCustomPrompt(
+      customPrompt,
+      selectedText
+    );
+
+    expect(result).toContain("Notes related to {context0} are:");
+    expect(result).toContain(
+      '[{"name":"note1","content":"Note content for #tag1"},{"name":"note2","content":"Note content for #tag2"}]'
+    );
+  });
 });
diff --git a/tests/utils.test.ts b/tests/utils.test.ts
@@ -4,7 +4,7 @@ import {
   getNotesFromPath,
   getNotesFromTags,
   isFolderMatch,
-  processVariableName
+  processVariableNameForNotePath,
 } from '../src/utils';
 
 describe('isFolderMatch', () => {
@@ -100,49 +100,49 @@ describe('getNotesFromPath', () => {
     expect(files).toEqual([]);
   });
 
-  describe('processVariableName', () => {
+  describe('processVariableNameForNotePath', () => {
     it('should return the note md filename', () => {
-      const variableName = processVariableName('[[test]]');
+      const variableName = processVariableNameForNotePath('[[test]]');
       expect(variableName).toEqual('test.md');
     });
 
     it('should return the note md filename with extra spaces 1', () => {
-      const variableName = processVariableName(' [[  test]]');
+      const variableName = processVariableNameForNotePath(' [[  test]]');
       expect(variableName).toEqual('test.md');
     });
 
     it('should return the note md filename with extra spaces 2', () => {
-      const variableName = processVariableName('[[ test   ]] ');
+      const variableName = processVariableNameForNotePath('[[ test   ]] ');
       expect(variableName).toEqual('test.md');
     });
 
     it('should return the note md filename with extra spaces 2', () => {
-      const variableName = processVariableName(' [[ test note   ]] ');
+      const variableName = processVariableNameForNotePath(' [[ test note   ]] ');
       expect(variableName).toEqual('test note.md');
     });
 
     it('should return the note md filename with extra spaces 2', () => {
-      const variableName = processVariableName(' [[    test_note note   ]] ');
+      const variableName = processVariableNameForNotePath(' [[    test_note note   ]] ');
       expect(variableName).toEqual('test_note note.md');
     });
 
     it('should return folder path with leading slash', () => {
-      const variableName = processVariableName('/testfolder');
+      const variableName = processVariableNameForNotePath('/testfolder');
       expect(variableName).toEqual('/testfolder');
     });
 
     it('should return folder path without slash', () => {
-      const variableName = processVariableName('testfolder');
+      const variableName = processVariableNameForNotePath('testfolder');
       expect(variableName).toEqual('testfolder');
     });
 
     it('should return folder path with trailing slash', () => {
-      const variableName = processVariableName('testfolder/');
+      const variableName = processVariableNameForNotePath('testfolder/');
       expect(variableName).toEqual('testfolder/');
     });
 
     it('should return folder path with leading spaces', () => {
-      const variableName = processVariableName('  testfolder ');
+      const variableName = processVariableNameForNotePath('  testfolder ');
       expect(variableName).toEqual('testfolder');
     });
   });