feat: Add natural language suggestion mode (non-tag based prompts)

tests/js/autocomplete.test.js

@@ -7,6 +7,9 @@ import {
     ModelTagSource,
     autoCompleteData
 } from "../../web/js/data.js";
+import {
+    settingValues
+} from "../../web/js/settings.js";
 import {
     createFlexSearchDocument,
     createFlexSearchDocumentForModel
@@ -24,13 +27,16 @@ const {
 
 // Helper function to create mock textarea element
 function createMockTextarea(value, selectionStart, selectionEnd = -1) {
-    return {
+    let mock = {
         value: value,
         selectionStart: selectionStart,
         selectionEnd: selectionEnd > selectionStart ? selectionEnd : selectionStart,
         nodeName: 'TEXTAREA',
         focus() { document.setFocus(this); },
-        setSelectionRange: () => { },
+        setSelectionRange: (start, end) => {
+            mock.selectionStart = start;
+            mock.selectionEnd = end;
+        },
         getBoundingClientRect: () => ({
             top: 0,
             left: 0,
@@ -45,6 +51,7 @@ function createMockTextarea(value, selectionStart, selectionEnd = -1) {
         },
         dispatchEvent: () => { }
     };
+    return mock;
 }
 
 // Test data setup function
@@ -75,6 +82,7 @@ function setupTestData() {
         new TagData('1girl', 0, 1000000, ['one_girl'], TagSource.Danbooru),
         new TagData('blue_hair', 0, 500000, ['blue hair'], TagSource.Danbooru),
         new TagData('long_hair', 0, 800000, [], TagSource.Danbooru),
+        new TagData('smile', 0, 700000, [], TagSource.Danbooru),
         new TagData('__wildcard__', 0, 100, [], TagSource.Danbooru),
         new TagData(':d', 0, 50000, [], TagSource.Danbooru),
         new TagData('test_tag', 0, 1000, ['test'], TagSource.E621),
@@ -101,62 +109,85 @@ function setupTestData() {
         autoCompleteData[source].sortedTags.sort((a, b) => b.count - a.count);
 
         const isModelSrc = Object.values(ModelTagSource).includes(source);
-        const doc = isModelSrc ? createFlexSearchDocumentForModel() : createFlexSearchDocument();
+        let doc;
+        if (isModelSrc) {
+            doc = createFlexSearchDocumentForModel();
+        } else {
+            doc = createFlexSearchDocument();
+        }
         autoCompleteData[source].sortedTags.forEach((tagData, i) => {
-            doc.add(i, tagData);
+            doc.add({
+                id: i,
+                tag: tagData.tag,
+                alias: tagData.alias.join(',')
+            });
         });
 
         autoCompleteData[source].flexSearchDocument = doc;
     });
+
+
+    // Mock global document
+    global.document = {
+        focusedElement: null,
+        setFocus: function (element) {
+            this.focusedElement = element;
+        },
+        execCommand: (commandId, showUI, value) => {
+            switch (commandId) {
+                case "insertText":
+                    if (document.focusedElement) {
+                        const el = document.focusedElement;
+                        const start = el.selectionStart;
+                        const end = el.selectionEnd;
+                        el.value = el.value.substring(0, start) + value + el.value.substring(end);
+                        el.selectionStart = el.selectionEnd = start + value.length;
+                    }
+                    return true;
+                default:
+                    throw new Error('Not implemented!');
+            }
+        },
+        createElement: () => ({
+            id: '',
+            style: {},
+            innerHTML: '',
+            appendChild: () => { },
+            getBoundingClientRect: () => ({
+                top: 0,
+                left: 0,
+                width: 300,
+                height: 100
+            })
+        }),
+        setFocus: (element) => {
+            document.focusedElement = element;
+        },
+        body: {
+            appendChild: () => { },
+            removeChild: () => { }
+        },
+        focusedElement: null
+    };
+
+    // Mock global window
+    global.window = {
+        getComputedStyle: () => ({
+            lineHeight: '20px',
+            fontSize: '14px',
+            fontFamily: 'Arial'
+        })
+    };
 }
 
 describe('Autocomplete Functions', () => {
     beforeEach(() => {
+        // Set tag mode as default
+        settingValues.useFastSearch = false;
+        settingValues.naturalLanguageMode = false;
+        settingValues.replaceUnderscoreWithSpace = true;
+
         setupTestData();
-
-        // Mock global document
-        global.document = {
-            execCommand: (commandId, showUI, value) => {
-                switch (commandId) {
-                    case "insertText":
-                        if (document.focusedElement) {
-                            document.focusedElement.value += value;
-                        }
-                        break;
-                    default:
-                        throw new Error('Not implemented!');
-                }
-            },
-            createElement: () => ({
-                id: '',
-                style: {},
-                innerHTML: '',
-                appendChild: () => { },
-                getBoundingClientRect: () => ({
-                    top: 0,
-                    left: 0,
-                    width: 300,
-                    height: 100
-                })
-            }),
-            setFocus: (element) => {
-                document.focusedElement = element;
-            },
-            body: {
-                appendChild: () => { },
-                removeChild: () => { }
-            },
-            focusedElement: null
-        };
-
-        // Mock global window
-        global.window = {
-            getComputedStyle: () => ({
-                lineHeight: '20px',
-                fontSize: '14px',
-                fontFamily: 'Arial'
-            })
-        };
     });
 
     describe('matchWord', () => {
@@ -375,4 +406,99 @@ describe('Autocomplete Functions', () => {
         });
 
     });
+});
+
+describe('Natural Language Mode', () => {
+    beforeEach(() => {
+        // Enable NL mode for these tests
+        settingValues.useFastSearch = true;
+        settingValues.naturalLanguageMode = true;
+        settingValues.replaceUnderscoreWithSpace = false; // Explicitly set for NL tests
+        // Re-setup data with NL mode enabled
+        setupTestData();
+    });
+
+    afterEach(() => {
+        // Disable NL mode after tests
+        settingValues.useFastSearch = true;
+        settingValues.naturalLanguageMode = false;
+        settingValues.replaceUnderscoreWithSpace = true; // Reset to default
+    });
+
+    describe('getCurrentPartialTag (NL Mode)', () => {
+        test('should extract multi-word phrase for tag matching', () => {
+            const textarea = createMockTextarea('a beautiful girl with long ha', 29);
+            const result = getCurrentPartialTag(textarea);
+            expect(result).toBe('a beautiful girl with long ha');
+        });
+
+        test('should extract context from the start of the segment', () => {
+            const textarea = createMockTextarea('1girl, a girl smil', 18);
+            const result = getCurrentPartialTag(textarea);
+            expect(result).toBe('a girl smil');
+        });
+
+        test('should handle single partial word', () => {
+            const textarea = createMockTextarea('smil', 4);
+            const result = getCurrentPartialTag(textarea);
+            expect(result).toBe('smil');
+        });
+
+        test('should handle two-word phrase', () => {
+            const textarea = createMockTextarea('long ha', 7);
+            const result = getCurrentPartialTag(textarea);
+            expect(result).toBe('long ha');
+        });
+    });
+
+    describe('matchWord (NL Mode)', () => {
+        test('should treat space and underscore as equivalent for matching', () => {
+            const queries = new Set(['long hair']);
+            const result = matchWord('long_hair', queries);
+            expect(result.matched).toBe(true);
+        });
+
+        test('should perform partial match with space/underscore normalization', () => {
+            const queries = new Set(['long h']);
+            const result = matchWord('long_hair', queries);
+            expect(result.matched).toBe(true);
+            expect(result.isExactMatch).toBe(false);
+        });
+    });
+
+    describe('searchCompletionCandidates (NL Mode)', () => {
+        test('should find tags when typing with spaces', () => {
+            const textarea = createMockTextarea('a girl with long ha', 19);
+            const results = searchCompletionCandidates(textarea);
+            expect(results.map(t => t.tag)).toContain('long_hair');
+        });
+
+        test('should find tags with underscores when typing with spaces', () => {
+            const textarea = createMockTextarea('blue hai', 8);
+            const results = searchCompletionCandidates(textarea);
+            expect(results.map(t => t.tag)).toContain('blue_hair');
+        });
+    });
+
+    describe('insertTagToTextArea (NL Mode)', () => {
+        test('should insert tag with spaces instead of underscores', () => {
+            const textarea = createMockTextarea('a girl with smil', 16);
+            console.log('[DEBUG] original:', textarea.value);
+            textarea.focus();
+            textarea.setSelectionRange(16, 16);
+            const tagData = { tag: 'smile', source: 'danbooru' };
+            insertTagToTextArea(textarea, tagData);
+            console.log('[DEBUG] result:', textarea.value);
+            expect(textarea.value).toBe('a girl with smile ');
+        });
+
+        test('should insert tag with a single space suffix', () => {
+            const textarea = createMockTextarea('blue hai', 8);
+            textarea.focus();
+            textarea.setSelectionRange(8, 8);
+            const tagData = { tag: 'blue_hair', source: 'danbooru' };
+            insertTagToTextArea(textarea, tagData);
+            expect(textarea.value).toBe('blue hair ');
+        });
+    });
 });

web/js/autocomplete.js

@@ -76,6 +76,40 @@ function matchWord(target, queries) {
     return { matched, isExactMatch };
 }
 
+/**
+ * Make query variations for searching.
+ * In natural language mode, adds space/underscore variants.
+ * In tag mode, adds CJK character variants.
+ * @param {string} partialTag - The partial tag/text to generate variations for
+ * @returns {Set<string>} Set of query variations
+ */
+function makeQueryVariations(partialTag) {
+    const variations = new Set([
+        partialTag.toLowerCase(),
+        normalizeTagToSearch(partialTag).toLowerCase()
+    ]);
+
+    if (settingValues.naturalLanguageMode) {
+        // Natural language mode: add underscore/space variants
+        // "long hair" → add "long_hair"
+        variations.add(partialTag.replace(/\s+/g, '_').toLowerCase());
+        // "long_hair" → add "long hair"
+        variations.add(partialTag.replace(/_+/g, ' ').toLowerCase());
+    }
+
+    // CJK support (existing functionality for both modes)
+    const kataQuery = hiraToKata(partialTag);
+    if (kataQuery !== partialTag) {
+        variations.add(kataQuery);
+    }
+    const hiraQuery = kataToHira(partialTag);
+    if (hiraQuery !== partialTag) {
+        variations.add(hiraQuery);
+    }
+
+    return variations;
+}
+
 /**
  * Search tag completion candidates based on the current input and cursor position in the textarea.
  * @param {HTMLTextAreaElement} textareaElement The partial tag input.
@@ -91,16 +125,7 @@ function searchCompletionCandidates(textareaElement) {
         return []; // No valid input for autocomplete
     }
 
-    // Generate Hiragana/Katakana variations if applicable
+    const queryVariations = makeQueryVariations(partialTag);
-    const queryVariations = new Set([partialTag.toLowerCase(), normalizeTagToSearch(partialTag).toLowerCase()]);
-    const kataQuery = hiraToKata(partialTag);
-    if (kataQuery !== partialTag) {
-        queryVariations.add(kataQuery);
-    }
-    const hiraQuery = kataToHira(partialTag);
-    if (hiraQuery !== partialTag) {
-        queryVariations.add(hiraQuery);
-    }
 
     if (settingValues.useFastSearch) {
         return searchWithFlexSearch(partialTag, queryVariations);
@@ -199,49 +224,63 @@ function sequentialSearch(partialTag, queryVariations) {
 function searchWithFlexSearch(partialTag, queryVariations) {
     const startTime = performance.now();
 
+    const createSubstrings = (input) => {
+        const words = input.split(' ');
+        return words.map((_, i) => words.slice(i).join(' '));
+    };
+
     let mergedResult = [];
     let totalSearchCount = 0;
 
     const sources = getEnabledTagSourceInPriorityOrder();
     for (const source of sources) {
         if (!autoCompleteData[source].flexSearchDocument) continue;
-        if (mergedResult.length >= settingValues.maxSuggestions) break;
 
-        // Use the FlexSearch Document to search
+        const substrings = settingValues.naturalLanguageMode ? createSubstrings(partialTag) : [partialTag];
-        // NOTE: The limit param is reflected separately for "tag" and "alias".
-        let searchResult = autoCompleteData[source].flexSearchDocument.search(partialTag, {
-            field: ["tag", "alias"],
-            limit: Math.min(settingValues.maxSuggestions * 10, 500),
-            merge: true,
-            suggest: false,
-            cache: true,
-        });
 
-        if (!searchResult || searchResult.length <= 0) continue;
+        // Search for each substring and merge results
+        for (const element of substrings) {
+            if (mergedResult.length >= settingValues.maxSuggestions) break;
+            if (element.length <= 1) break;
 
-        // Sort results based on exact matches and counts
+            // Use the FlexSearch Document to search
-        searchResult = searchResult
+            // NOTE: The limit param is reflected separately for "tag" and "alias".
-            .map(r => autoCompleteData[source].sortedTags[r.id])
+            let searchResult = autoCompleteData[source].flexSearchDocument.search(element, {
-            .sort((aTag, bTag) => {
+                field: ["tag", "alias"],
-                if (matchWord(bTag.tag, queryVariations).isExactMatch) {
+                limit: Math.min(settingValues.maxSuggestions * 10, 500) / substrings.length,
-                    return 999999999999;
+                merge: true,
-                }
+                suggest: false,
-                if (matchWord(aTag.tag, queryVariations).isExactMatch) {
+                cache: true,
-                    return -999999999999;
-                }
-                if (bTag.alias && bTag.alias.some(alias => matchWord(alias, queryVariations).isExactMatch)) {
-                    return 999999999999;
-                }
-                if (aTag.alias && aTag.alias.some(alias => matchWord(alias, queryVariations).isExactMatch)) {
-                    return -999999999999;
-                }
-                return bTag.count - aTag.count;
             });
 
-        // Merge results into the final array
+            console.debug(`[Autocomplete-Plus] Fast Search for "${element}" in substrings, Found ${searchResult.length}.`);
-        mergedResult = mergedResult.concat(searchResult.slice(0, settingValues.maxSuggestions - mergedResult.length));
 
-        totalSearchCount += searchResult.length;
+            if (!searchResult || searchResult.length <= 0) continue;
+
+            // Sort results based on exact matches and counts
+            searchResult = searchResult
+                .map(r => autoCompleteData[source].sortedTags[r.id])
+                .sort((aTag, bTag) => {
+                    if (matchWord(bTag.tag, queryVariations).isExactMatch) {
+                        return 999999999999;
+                    }
+                    if (matchWord(aTag.tag, queryVariations).isExactMatch) {
+                        return -999999999999;
+                    }
+                    if (bTag.alias && bTag.alias.some(alias => matchWord(alias, queryVariations).isExactMatch)) {
+                        return 999999999999;
+                    }
+                    if (aTag.alias && aTag.alias.some(alias => matchWord(alias, queryVariations).isExactMatch)) {
+                        return -999999999999;
+                    }
+                    return bTag.count - aTag.count;
+                });
+
+            // Merge results into the final array
+            mergedResult = mergedResult.concat(searchResult.slice(0, settingValues.maxSuggestions - mergedResult.length));
+
+            totalSearchCount += searchResult.length;
+        }
     }
 
     if (settingValues._logprocessingTime) {
@@ -255,14 +294,10 @@ function searchWithFlexSearch(partialTag, queryVariations) {
 
 /**
  * Extracts the current tag being typed before the cursor.
- * @param {HTMLTextAreaElement} inputElement
+ * @param {HTMLTextAreaElement} inputElement 
- * @returns {string} The current partial tag.
+ * @returns 
  */
-function getCurrentPartialTag(inputElement) {
+function extractTagSegment(inputElement) {
-    if (!inputElement) {
-        return "";
-    }
-    
     const text = inputElement.value;
     const cursorPos = inputElement.selectionStart;
 
@@ -301,6 +336,63 @@ function getCurrentPartialTag(inputElement) {
     return normalizeTagToSearch(partial);
 }
 
+/**
+ * Extracts natural language context from the input for autocomplete.
+ * In natural language mode, returns the last partial phrase for tag matching.
+ * @param {HTMLTextAreaElement} inputElement - The textarea element
+ * @param {number} extact maximum number of words to extract (default 10)
+ * @returns {string} The partial phrase for matching
+ */
+function extractNaturalLanguageContext(inputElement, extract = 10) {
+    if (!inputElement) {
+        return "";
+    }
+
+    const text = inputElement.value;
+    const cursorPos = inputElement.selectionStart;
+
+    // Step 1: Find segment boundary (comma or newline before cursor)
+    const lastNewLine = text.lastIndexOf('\n', cursorPos - 1);
+    const lastComma = text.lastIndexOf(',', cursorPos - 1);
+    const segmentStart = Math.max(lastNewLine, lastComma) + 1;
+
+    // Step 2: Extract segment text from boundary to cursor
+    const segment = text.substring(segmentStart, cursorPos).trimStart();
+
+    // Step 3: Split into words
+    const words = segment.split(/\s+/).filter(Boolean);
+
+    if (words.length === 0) {
+        return "";
+    } else if (words.length === 1) {
+        // Single word
+        return words[0];
+    } else {
+        // Multiple words: check if last 2 words could be a multi-word tag
+        // For "a girl with long h", extract "long h"
+        // For "blue hai", extract "blue hai"
+        const lastTwoWords = words.slice(-extract).join(' ');
+        return lastTwoWords;
+    }
+}
+
+/**
+ * 
+ * @param {HTMLTextAreaElement} inputElement
+ * @returns {string} The current partial tag or natural language context.
+ */
+function getCurrentPartialTag(inputElement, wordLimitForNL) {
+    if (!inputElement) {
+        return "";
+    }
+
+    if (settingValues.naturalLanguageMode) {
+        return extractNaturalLanguageContext(inputElement, wordLimitForNL);
+    } else {
+        return extractTagSegment(inputElement);
+    }
+}
+
 /**
  * Inserts the selected tag into the textarea, replacing the partial tag,
  * making the change undoable.
@@ -311,13 +403,13 @@ function insertTagToTextArea(inputElement, tagDataToInsert) {
     if (!inputElement || !tagDataToInsert) {
         return;
     }
-    
+
     const text = inputElement.value;
     const cursorPos = inputElement.selectionStart;
 
     const tagRange = getCurrentTagRange(text, cursorPos);
     let tagStart, tagEnd, currentTag;
-    
+
     if (!tagRange) {
         // Fallback: insert at cursor position
         tagStart = cursorPos;
@@ -326,7 +418,7 @@ function insertTagToTextArea(inputElement, tagDataToInsert) {
     } else {
         ({ start: tagStart, end: tagEnd, tag: currentTag } = tagRange);
     }
-    
+
     const replaceStart = Math.min(cursorPos, tagStart);
     let replaceEnd = cursorPos;
 
@@ -550,7 +642,7 @@ class AutocompleteUI {
         tagCount.className = `autocomplete-plus-tag-count`;
         tagCount.textContent = formatCountHumanReadable(tagData.count);
 
-         // Create tooltip with more info
+        // Create tooltip with more info
         let tooltipText = `Count: ${tagData.count}\nCategory: ${categoryText}`;
         if (aliasText.length > 0) {
             tooltipText += `\nAlias: ${aliasText}`;

web/js/main.js

@@ -284,6 +284,17 @@ app.registerExtension({
         },
 
         // --- Autocomplete Settings ---
+        {
+            id: id + ".natural_language_mode",
+            name: "Natural Language Mode",
+            type: "boolean",
+            defaultValue: false,
+            tooltip: "Enable autocomplete for natural language prompts. Tags will be suggested as space-separated words. ⚠️ Requires browser reload to apply.",
+            category: [name, "Autocompletion", "Natural Language"],
+            onChange: (value) => {
+                settingValues.naturalLanguageMode = value;
+            }
+        },
         {
             id: id + ".use_fast_search",
             name: "Use Fast Search",

web/js/settings.js

@@ -10,6 +10,7 @@ export const settingValues = {
     enableModels: true, // Enable Lora and Embedding suggestions
     useFastSearch: false,
     replaceUnderscoreWithSpace: true, // Replace underscores with spaces in tag insertion
+    naturalLanguageMode: false,  // Enable/disable natural language mode
 
     // Related tags feature settings
     enableRelatedTags: true,