chore:Auto correct context; support extract image paths from selections (#27)

Author: googlicius

src/constants.ts

@@ -12,24 +12,43 @@ export const STW_EMBEDDED_CONVERSATION_VIEW_CONFIG = {
   icon: 'message-square',
 };
 
-// Pattern only, without flags, to avoid global regex state issues
-// Captures image path (group 1) which may be followed by size parameter after |
+/**
+ * Pattern only, without flags, to avoid global regex state issues
+ * Captures image path (group 1) which may be followed by size parameter after |
+ */
 export const IMAGE_LINK_PATTERN = '!\\[\\[(.*?\\.(jpg|jpeg|png|webp|svg))(?:\\|.*?)?\\]\\]';
-// Stw-selected pattern constants for reuse across the application
-// Pattern to match any stw-selected block (with capture group for splitting)
+
+/**
+ * Stw-selected pattern constants for reuse across the application
+ * Pattern to match any stw-selected block (with capture group for splitting)
+ */
 export const STW_SELECTED_PATTERN = '(\\{\\{stw-selected.*?\\}\\})';
 
-// Pattern to match {{stw-squeezed [[<path>]] }}
+/**
+ * The placeholder for the stw-selected blocks in the original query
+ * This helps to reduce the complexity for the planner to just put the placeholder rather than extract the stw-selected blocks
+ */
+export const STW_SELECTED_PLACEHOLDER = '<stwSelected>';
+
+/**
+ * Pattern to match {{stw-squeezed [[<path>]] }}
+ */
 export const STW_SQUEEZED_PATTERN = '\\{\\{stw-squeezed \\[\\[([^\\]]+)\\]\\] \\}\\}';
 
-// Pattern to match any wikilink
+/**
+ * Pattern to match any wikilink
+ */
 export const WIKI_LINK_PATTERN = '\\[\\[([^\\]]+)\\]\\]';
 
-// Pattern to extract metadata from stw-selected blocks
+/**
+ * Pattern to extract metadata from stw-selected blocks
+ */
 export const STW_SELECTED_METADATA_PATTERN =
   '\\{\\{stw-selected from:(\\d+),to:(\\d+),selection:(.+?),path:(.+?)\\}\\}';
 
-// Supported command prefixes
+/**
+ * Supported command prefixes
+ */
 export const COMMAND_PREFIXES = [
   '/ ',
   '/search',

src/lib/modelfusion/extractions/contentUpdateExtraction.ts

@@ -3,7 +3,7 @@ import { contentUpdatePrompt } from '../prompts/contentUpdatePrompt';
 import { userLanguagePrompt } from '../prompts/languagePrompt';
 import { logger } from 'src/utils/logger';
 import { AbortService } from 'src/services/AbortService';
-import { prepareUserMessage } from '../utils/userMessageUtils';
+import { prepareMessage } from '../utils/messageUtils';
 import { App } from 'obsidian';
 import { LLMService } from 'src/services/LLMService';
 import { z } from 'zod';
@@ -64,7 +64,7 @@ export async function extractContentUpdate(params: {
       generateType: 'object',
     });
 
-    const userMessage = await prepareUserMessage(command.query, app);
+    const userMessage = await prepareMessage(command.query, app);
 
     const { object } = await generateObject({
       ...llmConfig,

src/lib/modelfusion/extractions/noteCreationExtraction.ts

@@ -2,13 +2,13 @@ import { generateObject } from 'ai';
 import { noteCreationPrompt } from '../prompts/noteCreationPrompt';
 import { userLanguagePrompt } from '../prompts/languagePrompt';
 import { AbortService } from 'src/services/AbortService';
-import { prepareUserMessage } from '../utils/userMessageUtils';
 import { App } from 'obsidian';
 import { LLMService } from 'src/services/LLMService';
 import { z } from 'zod';
 import { CommandIntent } from 'src/types/types';
 import { explanationFragment, confidenceFragment } from '../prompts/fragments';
 import { logger } from 'src/utils/logger';
+import { prepareMessage } from '../utils/messageUtils';
 
 const abortService = AbortService.getInstance();
 
@@ -56,7 +56,7 @@ export async function extractNoteCreation(params: {
     });
 
     // Prepare user message with potential image content
-    const userMessage = await prepareUserMessage(command.query, app);
+    const userMessage = await prepareMessage(command.query, app);
 
     const { object } = await generateObject({
       ...llmConfig,

src/lib/modelfusion/index.ts

@@ -8,4 +8,4 @@ export * from './validators';
 export { getClassifier } from './classifiers/getClassifier';
 
 // Export image utils
-export { prepareUserMessage } from './utils/userMessageUtils';
+export { prepareMessage } from './utils/messageUtils';

src/lib/modelfusion/prompts/__snapshots__/commands.test.ts.snap

@@ -26,16 +26,20 @@ exports[`commands formatCommandsForPrompt should return formatted commands for p
   Use when: Delete notes from the artifact
 - "build_search_index" (aliases: index, build-index, search-index): Build or rebuild the search index for all markdown files in the vault
   Use when: Build or rebuild the search index for all markdown files in the vault
-- "generate": Generate content with the LLM help (either in a new note or in the conversation). You also can "generate" from the provided content in the user's query without reading the note. Example: "Help me update this list to the numbered list:\\n- Item 1\\n- Item 2" -> ["generate"]. The list is already in the query.
+- "generate": Generate content with the LLM help (either in a new note or in the conversation). You also can "generate" from the provided content in the user's query without reading the note. Example: "Help me update this list to the numbered list:
+  - Item 1
+  - Item 2" -> ["generate"]. The list is already in the query.
   Use when: Ask or generate content with your help
 - "read": Read content from the current note or specific position: "above", "below". Use this when you don't know the content and need to retrieve it before proceeding
+  Can read any content type, including code blocks, tables, lists, paragraphs, images, and more.
   Use when: Read or Find content based on a specific pattern in their current note"
 `;
 
 exports[`commands formatCommandsForPrompt should return formatted commands for prompt with given command names 1`] = `
 "- "search": Find files using the search engine to search files locally and store the result as an artifact
   Use when: Search for files (and doesn't mention existing search results)
 - "read": Read content from the current note or specific position: "above", "below". Use this when you don't know the content and need to retrieve it before proceeding
+  Can read any content type, including code blocks, tables, lists, paragraphs, images, and more.
   Use when: Read or Find content based on a specific pattern in their current note"
 `;
 

src/lib/modelfusion/prompts/commands.ts

@@ -181,8 +181,9 @@ export const COMMAND_DEFINITIONS: CommandDefinition[] = [
   },
   {
     commandType: 'generate',
-    description:
-      'Generate content with the LLM help (either in a new note or in the conversation). You also can "generate" from the provided content in the user\'s query without reading the note. Example: "Help me update this list to the numbered list:\\n- Item 1\\n- Item 2" -> ["generate"]. The list is already in the query.',
+    description: `Generate content with the LLM help (either in a new note or in the conversation). You also can "generate" from the provided content in the user's query without reading the note. Example: "Help me update this list to the numbered list:
+  - Item 1
+  - Item 2" -> ["generate"]. The list is already in the query.`,
     category: 'intent-based',
     queryTemplate: `Extract the query for the generate command follows this format: <query_in_natural_language>; [note name: <noteName>]
 - <query_in_natural_language>: Tailored query for the generate command.
@@ -192,8 +193,8 @@ export const COMMAND_DEFINITIONS: CommandDefinition[] = [
   },
   {
     commandType: 'read',
-    description:
-      'Read content from the current note or specific position: "above", "below". Use this when you don\'t know the content and need to retrieve it before proceeding',
+    description: `Read content from the current note or specific position: "above", "below". Use this when you don't know the content and need to retrieve it before proceeding
+  Can read any content type, including code blocks, tables, lists, paragraphs, images, and more.`,
     category: 'intent-based',
     queryTemplate: `Extract a specific query for a read command:
 1. Extract the query for the read command follows this format: <query_in_natural_language>; read type: <readType>[; note name: <noteName>]

src/lib/modelfusion/utils/userMessageUtils.ts renamed to src/lib/modelfusion/utils/messageUtils.ts

@@ -6,29 +6,28 @@ import { NoteContentService } from 'src/services/NoteContentService';
 import { resizeImageWithCanvas } from 'src/utils/resizeImageWithCanvas';
 import { logger } from 'src/utils/logger';
 
-export function getTextContentWithoutImages(userInput: string): string {
+export function getTextContentWithoutImages(input: string): string {
   // Create a new RegExp instance with flags each time to avoid stateful issues
   const imageRegex = new RegExp(IMAGE_LINK_PATTERN, 'gi');
-  return userInput.replace(imageRegex, '').trim();
+  return input.replace(imageRegex, '').trim();
 }
 
 /**
- * Prepares user message content with images and wikilinks's content for OpenAI's Vision API
- * @param userInput Original user input text
+ * Prepares message content with images and wikilinks's content
+ * @param input Original input text
  * @param app Obsidian App instance for accessing vault
- * @returns An array of content items for OpenAI's ChatMessage.user
  */
-export async function prepareUserMessage(
-  userInput: string,
+export async function prepareMessage(
+  input: string,
   app: App
 ): Promise<Array<TextPart | ImagePart>> {
   const noteContentService = NoteContentService.getInstance(app);
-  const imagePaths = noteContentService.extractImageLinks(userInput);
-  const wikilinks = noteContentService.extractWikilinks(userInput);
+  const imagePaths = noteContentService.extractImageLinks(input);
+  const wikilinks = noteContentService.extractWikilinks(input);
   const messageContent: Array<TextPart | ImagePart> = [];
 
   // Add the original user input first
-  messageContent.push({ type: 'text', text: userInput });
+  messageContent.push({ type: 'text', text: input });
 
   const mediaTools = MediaTools.getInstance(app);
 

src/services/CommandProcessorService.ts

@@ -24,9 +24,9 @@ import {
   SummaryCommandHandler,
   ContextAugmentationHandler,
 } from '../solutions/commands/handlers';
-import { getTextContentWithoutImages } from 'src/lib/modelfusion/utils/userMessageUtils';
 
 import type StewardPlugin from '../main';
+import { getTextContentWithoutImages } from 'src/lib/modelfusion/utils/messageUtils';
 
 export class CommandProcessorService {
   private readonly commandProcessor: CommandProcessor;
@@ -103,7 +103,7 @@ export class CommandProcessorService {
     this.commandProcessor.registerHandler('read', readHandler);
 
     // Register the generate command handler
-    const generateHandler = new GenerateCommandHandler(this.plugin);
+    const generateHandler = new GenerateCommandHandler(this.plugin, this.commandProcessor);
     this.commandProcessor.registerHandler('generate', generateHandler);
 
     // Register the stop command handler

src/services/NoteContentService.test.ts

@@ -1,3 +1,4 @@
+import { MarkdownUtil } from 'src/utils/markdownUtils';
 import { NoteContentService } from './NoteContentService';
 import { App, TFile } from 'obsidian';
 
@@ -53,6 +54,13 @@ Is the image above a lake, pond, reservoir, or sea?`;
       const imageLinks = noteContentService.extractImageLinks(content);
       expect(imageLinks).toEqual(['Pasted image 20250222171626.png']);
     });
+
+    it('should extract images from stw-selected blocks', () => {
+      const content = `Here is a selected block with an image:
+{{stw-selected from:1,to:5,selection: ${new MarkdownUtil('This contains an image: ![[image2.jpg]] and some text').escape().getText()},path:test.md}}`;
+      const imageLinks = noteContentService.extractImageLinks(content);
+      expect(imageLinks).toEqual(['image2.jpg']);
+    });
   });
 
   describe('extractWikilinks', () => {

src/services/NoteContentService.ts

@@ -1,6 +1,12 @@
-import { IMAGE_LINK_PATTERN, WIKI_LINK_PATTERN } from 'src/constants';
+import {
+  IMAGE_LINK_PATTERN,
+  WIKI_LINK_PATTERN,
+  STW_SELECTED_PATTERN,
+  STW_SELECTED_METADATA_PATTERN,
+} from 'src/constants';
 import { App } from 'obsidian';
 import { logger } from 'src/utils/logger';
+import { MarkdownUtil } from 'src/utils/markdownUtils';
 
 export class NoteContentService {
   private static instance: NoteContentService;
@@ -59,17 +65,44 @@ export class NoteContentService {
    * @returns Array of image paths extracted from the content
    */
   public extractImageLinks(content: string): string[] {
-    // Create a new RegExp instance with flags each time to avoid stateful issues
+    const imagePaths: string[] = [];
+
+    // Extract images from regular image links
     const imageRegex = new RegExp(IMAGE_LINK_PATTERN, 'gi');
     const matches = content.matchAll(imageRegex);
-    const imagePaths: string[] = [];
 
     for (const match of matches) {
       if (match[1]) {
         imagePaths.push(match[1]);
       }
     }
 
+    // Early check if content has stw-selected blocks
+    if (!content.includes('{{stw-selected')) {
+      return imagePaths;
+    }
+
+    // Extract images from stw-selected blocks
+    const stwSelectedMatches = content.matchAll(new RegExp(STW_SELECTED_PATTERN, 'g'));
+
+    for (const stwMatch of stwSelectedMatches) {
+      if (stwMatch[1]) {
+        const stwBlock = stwMatch[1];
+        const metadataMatch = stwBlock.match(new RegExp(STW_SELECTED_METADATA_PATTERN));
+
+        if (metadataMatch) {
+          const [, , , escapedSelection] = metadataMatch;
+          // Unescape the selection content
+          const unescapedSelection = new MarkdownUtil(escapedSelection)
+            .unescape()
+            .decodeURI()
+            .getText();
+          const selectionImagePaths = this.extractImageLinks(unescapedSelection);
+          imagePaths.push(...selectionImagePaths);
+        }
+      }
+    }
+
     return imagePaths;
   }