From 527886e2b8f7726a9216ddf34f219e2af270ff32 Mon Sep 17 00:00:00 2001 From: Harsh16gupta Date: Sun, 26 Jul 2026 04:38:24 +0530 Subject: [PATCH 1/2] feat: add 2-step AI-powered cluster naming using Joplin native AI --- src/pipeline/clustering/aiNamingService.ts | 310 ++++++++++++++++++ src/pipeline/runPipeline.ts | 7 + src/types/joplinAi.d.ts | 11 + .../clustering/aiNamingService.test.ts | 243 ++++++++++++++ 4 files changed, 571 insertions(+) create mode 100644 src/pipeline/clustering/aiNamingService.ts create mode 100644 test/pipeline/clustering/aiNamingService.test.ts diff --git a/src/pipeline/clustering/aiNamingService.ts b/src/pipeline/clustering/aiNamingService.ts new file mode 100644 index 0000000..cf69f31 --- /dev/null +++ b/src/pipeline/clustering/aiNamingService.ts @@ -0,0 +1,310 @@ +import joplin from 'api'; +import { JoplinAi, ChatMessage, ChatResult } from '../../types/joplinAi'; +import { BenchmarkResult } from '../../types/cluster'; +import { DocumentText } from './tfidf'; +import { log, logErr } from '../../utils/logger'; + +/** Maximum time (ms) to wait for the AI chat response before falling back. */ +const AI_TIMEOUT_MS = 45_000; + +/** Maximum number of representative note titles to include per cluster in the prompt. */ +const MAX_TITLES_PER_CLUSTER = 6; + +/** Maximum number of TF-IDF keywords to include per cluster in the prompt. */ +const MAX_KEYWORDS_PER_CLUSTER = 3; + +interface ClusterSummary { + clusterId: number; + noteCount: number; + sampleTitles: string[]; + topKeywords: string[]; +} + +/** + * Sanitizes a raw AI-generated cluster name by removing common LLM formatting artifacts. + * Strips surrounding quotes, backticks, markdown headers, labels like "Title:" or + * "Cluster Name:", and truncates to a maximum of 5 words. + */ +export function sanitizeAiName(raw: string): string { + let name = raw.trim(); + + // Strip surrounding quotes (single, double, or backticks) + name = name.replace(/^["'`]+|["'`]+$/g, ''); + + // Strip markdown header prefixes (e.g. "## ", "### ") + name = name.replace(/^#{1,6}\s*/, ''); + + // Strip common LLM preamble labels + name = name.replace(/^(?:title|cluster(?:\s+name)?|name|category|group|topic)\s*:\s*/i, ''); + + // Strip trailing periods + name = name.replace(/\.+$/, ''); + + // Re-trim after all replacements + name = name.trim(); + + // Truncate to 5 words max + const words = name.split(/\s+/); + if (words.length > 5) { + name = words.slice(0, 5).join(' '); + } + + return name; +} + +/** + * Builds the system and user prompt messages for the batched AI naming call. + * All clusters are included in a single prompt to minimize API round-trips. + */ +export function buildNamingPrompt(clusterSummaries: ClusterSummary[]): ChatMessage[] { + const systemMessage: ChatMessage = { + role: 'system', + content: + 'You are a concise note organizer. For each numbered group of notes below, generate a short descriptive category title (2 to 4 words). ' + + 'Reply with ONLY a valid JSON object mapping group numbers to titles. ' + + 'Example: {"0": "Web Development", "1": "Travel Plans", "2": "Fitness Routines"}. ' + + 'Do NOT include any other text, explanation, or markdown formatting in your response.', + }; + + let userContent = ''; + for (const summary of clusterSummaries) { + const titlesStr = summary.sampleTitles.map((t) => `"${t}"`).join(', '); + userContent += `Group ${summary.clusterId} (${summary.noteCount} notes): ${titlesStr}\n`; + if (summary.topKeywords.length > 0) { + userContent += `Keywords: ${summary.topKeywords.join(', ')}\n`; + } + userContent += '\n'; + } + + const userMessage: ChatMessage = { + role: 'user', + content: userContent.trim(), + }; + + return [systemMessage, userMessage]; +} + +/** + * Parses the AI response text as JSON and extracts cluster names. + * Returns null if parsing fails or the response doesn't contain valid mappings. + */ +export function parseAiNamesResponse( + responseText: string, + clusterIds: number[], +): { [clusterId: number]: string } | null { + // Try to extract JSON from the response (handle cases where LLM wraps in ```json ... ```) + let jsonStr = responseText.trim(); + const jsonBlockMatch = jsonStr.match(/```(?:json)?\s*([\s\S]*?)```/); + if (jsonBlockMatch) { + jsonStr = jsonBlockMatch[1].trim(); + } + + // Try to find a JSON object in the response + const jsonObjectMatch = jsonStr.match(/\{[\s\S]*\}/); + if (!jsonObjectMatch) { + return null; + } + + let parsed: Record; + try { + parsed = JSON.parse(jsonObjectMatch[0]); + } catch { + return null; + } + + if (typeof parsed !== 'object' || parsed === null || Array.isArray(parsed)) { + return null; + } + + const result: { [clusterId: number]: string } = {}; + let validCount = 0; + + for (const id of clusterIds) { + const raw = parsed[String(id)]; + if (typeof raw === 'string' && raw.trim().length > 0) { + const sanitized = sanitizeAiName(raw); + if (sanitized.length > 0) { + result[id] = sanitized; + validCount++; + } + } + } + + // Only accept if we got names for at least half the clusters + if (validCount < Math.ceil(clusterIds.length / 2)) { + return null; + } + + return result; +} + +/** + * Wraps a promise with a timeout. Rejects with a timeout error if the promise + * does not resolve within the specified duration. + */ +function withTimeout(promise: Promise, ms: number): Promise { + return new Promise((resolve, reject) => { + const timer = setTimeout(() => reject(new Error(`AI naming timed out after ${ms}ms`)), ms); + promise.then( + (value) => { + clearTimeout(timer); + resolve(value); + }, + (err) => { + clearTimeout(timer); + reject(err); + }, + ); + }); +} + +/** + * Attempts to upgrade cluster names in all BenchmarkResults using Joplin's AI Chat API. + * + * This function is designed to be called AFTER enrichResultsWithTags() has already + * populated result.clusterNames with TF-IDF-generated names. If AI naming succeeds, + * it overwrites those names. If it fails for any reason, the existing TF-IDF names + * remain untouched. + * + * JOPLIN PLUGIN SANDBOX PROXY: + * We invoke `(joplin as unknown as { ai: JoplinAi }).ai.chat(messages)` directly + * in a single unchained expression as a defensive practice to prevent proxy path + * state accumulation across Joplin sandbox runtime versions. + */ +export async function upgradeClusterNamesWithAi(results: BenchmarkResult[], documents: DocumentText[]): Promise { + await Promise.all( + results.map(async (result) => { + if (!result.clusterNames || Object.keys(result.clusterNames).length === 0) { + return; + } + + try { + // Build cluster summaries from the assignments and documents + const clusterIndices: { [clusterId: number]: number[] } = {}; + result.assignments.forEach((clusterId, noteIdx) => { + if (clusterId !== -1) { + if (!clusterIndices[clusterId]) { + clusterIndices[clusterId] = []; + } + clusterIndices[clusterId].push(noteIdx); + } + }); + + const clusterIds = Object.keys(clusterIndices).map(Number); + if (clusterIds.length === 0) return; + + const summaries: ClusterSummary[] = clusterIds.map((clusterId) => { + const indices = clusterIndices[clusterId]; + const clusterDocs = indices.map((idx) => documents[idx]); + + // Get representative titles (skip empty/generic titles) + const sampleTitles = clusterDocs + .map((doc) => doc.title) + .filter((t) => t && t.trim().length > 2) + .slice(0, MAX_TITLES_PER_CLUSTER); + + // Get top TF-IDF keywords from the existing tags + const topKeywords = (result.tags?.[clusterId] || []).slice(0, MAX_KEYWORDS_PER_CLUSTER); + + return { + clusterId, + noteCount: indices.length, + sampleTitles, + topKeywords, + }; + }); + + // Skip if no cluster has meaningful titles + const hasAnyTitles = summaries.some((s) => s.sampleTitles.length > 0); + if (!hasAnyTitles) { + log('AI naming skipped for strategy: no meaningful note titles found'); + return; + } + + const messages = buildNamingPrompt(summaries); + log(`AI naming: sending prompt with ${summaries.length} clusters to joplin.ai.chat (45s timeout)...`); + + let chatResult: ChatResult | string | undefined; + try { + // Single direct chained expression — DO NOT store `const joplinAi = joplin.ai` + chatResult = await withTimeout( + (joplin as unknown as { ai: JoplinAi }).ai.chat(messages), + AI_TIMEOUT_MS, + ); + } catch (chatErr) { + const errMsg = chatErr instanceof Error ? chatErr.message : String(chatErr); + log('AI naming skipped (keeping TF-IDF names):', errMsg); + return; + } + + // Joplin Plugin API returns { text: string } (or { message: { content: string } } in legacy mock) + const responseText = + typeof chatResult === 'string' + ? chatResult + : chatResult?.text || chatResult?.message?.content || null; + + if (!responseText) { + log('AI naming: empty response from chat API, keeping TF-IDF names'); + return; + } + + log(`AI naming: received response (${responseText.length} chars)`); + + const aiNames = parseAiNamesResponse(responseText, clusterIds); + if (!aiNames) { + log('AI naming: failed to parse response, keeping TF-IDF names'); + return; + } + + // Overwrite TF-IDF names with AI-generated names + let upgradedCount = 0; + for (const clusterId of clusterIds) { + if (aiNames[clusterId]) { + result.clusterNames![clusterId] = aiNames[clusterId]; + upgradedCount++; + } + // If AI didn't provide a name for this cluster, keep the TF-IDF name + } + + log(`AI naming: upgraded ${upgradedCount}/${clusterIds.length} cluster names`); + + // Resolve name collisions (same logic pattern as postProcess.ts) + resolveNameCollisions(result.clusterNames!); + } catch (err) { + const message = err instanceof Error ? err.message : String(err); + logErr('AI naming failed, keeping TF-IDF names:', message); + // result.clusterNames already has TF-IDF names, so we just continue + } + }), + ); +} + +/** + * Resolves duplicate cluster names by appending numeric suffixes. + * Mutates the clusterNames object in-place. + */ +function resolveNameCollisions(clusterNames: { [clusterId: number]: string }): void { + const nameCounts: { [name: string]: number } = {}; + for (const id of Object.keys(clusterNames)) { + const name = clusterNames[Number(id)]; + nameCounts[name] = (nameCounts[name] || 0) + 1; + } + + const usedNames = new Set(); + for (const id of Object.keys(clusterNames)) { + const clusterId = Number(id); + const name = clusterNames[clusterId]; + if (nameCounts[name] > 1) { + let resolved = name; + if (usedNames.has(resolved)) { + let suffix = 2; + while (usedNames.has(`${name} ${suffix}`)) suffix++; + resolved = `${name} ${suffix}`; + } + clusterNames[clusterId] = resolved; + usedNames.add(resolved); + } else { + usedNames.add(name); + } + } +} diff --git a/src/pipeline/runPipeline.ts b/src/pipeline/runPipeline.ts index bd226a5..1a1a1f1 100644 --- a/src/pipeline/runPipeline.ts +++ b/src/pipeline/runPipeline.ts @@ -7,6 +7,7 @@ import { VectorCache } from './vectorCache'; import { isNativeAiReady, fetchNativeEmbeddings } from './nativeEmbeddingPipeline'; import { DEFAULT_CONFIG, isValidEmbeddingVector } from './pipelineConfig'; import { enrichResultsWithTags } from './clustering/postProcess'; +import { upgradeClusterNamesWithAi } from './clustering/aiNamingService'; import { EmbeddingWorkerOrchestrator } from './EmbeddingWorkerOrchestrator'; export interface PipelineCallbacks { @@ -89,6 +90,9 @@ export const runPipeline = async (installDir: string, callbacks: PipelineCallbac enrichResultsWithTags(results, allPipelineDocuments); + callbacks.onStatus('Generating AI cluster names...'); + await upgradeClusterNamesWithAi(results, allPipelineDocuments); + const panelNotes: PanelNote[] = validNotes.map((n) => ({ noteId: n.id, title: n.title, @@ -164,6 +168,9 @@ export const runPipeline = async (installDir: string, callbacks: PipelineCallbac enrichResultsWithTags(results, allPipelineDocuments); + callbacks.onStatus('Generating AI cluster names...'); + await upgradeClusterNamesWithAi(results, allPipelineDocuments); + const panelNotes: PanelNote[] = noteVectors.map((nv) => ({ noteId: nv.noteId, title: nv.title, diff --git a/src/types/joplinAi.d.ts b/src/types/joplinAi.d.ts index b3b7fc3..b099518 100644 --- a/src/types/joplinAi.d.ts +++ b/src/types/joplinAi.d.ts @@ -29,7 +29,18 @@ export interface AiGetEmbeddingsOptions { limit?: number; } +export interface ChatMessage { + role: 'system' | 'user' | 'assistant'; + content: string; +} + +export interface ChatResult { + message?: ChatMessage; + text?: string; +} + export interface JoplinAi { getIndexStatus(): Promise; getEmbeddings(options: AiGetEmbeddingsOptions): Promise; + chat(messages: ChatMessage[]): Promise; } diff --git a/test/pipeline/clustering/aiNamingService.test.ts b/test/pipeline/clustering/aiNamingService.test.ts new file mode 100644 index 0000000..6efe025 --- /dev/null +++ b/test/pipeline/clustering/aiNamingService.test.ts @@ -0,0 +1,243 @@ +jest.mock('api', () => ({}), { virtual: true }); + +import { + sanitizeAiName, + buildNamingPrompt, + parseAiNamesResponse, +} from '../../../src/pipeline/clustering/aiNamingService'; + +describe('sanitizeAiName', () => { + it('trims whitespace', () => { + expect(sanitizeAiName(' Machine Learning ')).toBe('Machine Learning'); + }); + + it('strips surrounding double quotes', () => { + expect(sanitizeAiName('"Machine Learning"')).toBe('Machine Learning'); + }); + + it('strips surrounding single quotes', () => { + expect(sanitizeAiName("'Machine Learning'")).toBe('Machine Learning'); + }); + + it('strips surrounding backticks', () => { + expect(sanitizeAiName('`Machine Learning`')).toBe('Machine Learning'); + }); + + it('strips markdown header prefixes', () => { + expect(sanitizeAiName('## Machine Learning')).toBe('Machine Learning'); + expect(sanitizeAiName('### Deep Learning Notes')).toBe('Deep Learning Notes'); + }); + + it('strips "Title:" preamble', () => { + expect(sanitizeAiName('Title: Machine Learning')).toBe('Machine Learning'); + }); + + it('strips "Cluster Name:" preamble (case insensitive)', () => { + expect(sanitizeAiName('Cluster Name: Web Dev')).toBe('Web Dev'); + expect(sanitizeAiName('cluster name: Web Dev')).toBe('Web Dev'); + }); + + it('strips "Category:" preamble', () => { + expect(sanitizeAiName('Category: Travel Plans')).toBe('Travel Plans'); + }); + + it('strips "Name:" preamble', () => { + expect(sanitizeAiName('Name: Fitness')).toBe('Fitness'); + }); + + it('strips "Group:" preamble', () => { + expect(sanitizeAiName('Group: Recipes')).toBe('Recipes'); + }); + + it('strips "Topic:" preamble', () => { + expect(sanitizeAiName('Topic: Finance')).toBe('Finance'); + }); + + it('strips trailing periods', () => { + expect(sanitizeAiName('Machine Learning.')).toBe('Machine Learning'); + expect(sanitizeAiName('Deep Learning...')).toBe('Deep Learning'); + }); + + it('truncates to 5 words maximum', () => { + expect(sanitizeAiName('This Is A Very Long Category Name')).toBe('This Is A Very Long'); + }); + + it('handles combined artifacts (quotes + label + period)', () => { + expect(sanitizeAiName('"Title: Machine Learning."')).toBe('Machine Learning'); + }); + + it('returns empty string for whitespace-only input', () => { + expect(sanitizeAiName(' ')).toBe(''); + }); + + it('preserves valid short names', () => { + expect(sanitizeAiName('Recipes')).toBe('Recipes'); + }); +}); + +describe('buildNamingPrompt', () => { + it('returns system and user messages', () => { + const messages = buildNamingPrompt([ + { + clusterId: 0, + noteCount: 3, + sampleTitles: ['React Hooks Guide'], + topKeywords: ['react hooks'], + }, + ]); + + expect(messages).toHaveLength(2); + expect(messages[0].role).toBe('system'); + expect(messages[1].role).toBe('user'); + }); + + it('includes cluster ID and note count in user message', () => { + const messages = buildNamingPrompt([ + { + clusterId: 2, + noteCount: 5, + sampleTitles: ['Note A'], + topKeywords: [], + }, + ]); + + expect(messages[1].content).toContain('Group 2'); + expect(messages[1].content).toContain('5 notes'); + }); + + it('includes sample titles in quotes', () => { + const messages = buildNamingPrompt([ + { + clusterId: 0, + noteCount: 2, + sampleTitles: ['React Hooks', 'Redux State'], + topKeywords: [], + }, + ]); + + expect(messages[1].content).toContain('"React Hooks"'); + expect(messages[1].content).toContain('"Redux State"'); + }); + + it('includes keywords when provided', () => { + const messages = buildNamingPrompt([ + { + clusterId: 0, + noteCount: 2, + sampleTitles: ['Note A'], + topKeywords: ['react', 'typescript'], + }, + ]); + + expect(messages[1].content).toContain('Keywords: react, typescript'); + }); + + it('omits keywords line when empty', () => { + const messages = buildNamingPrompt([ + { + clusterId: 0, + noteCount: 2, + sampleTitles: ['Note A'], + topKeywords: [], + }, + ]); + + expect(messages[1].content).not.toContain('Keywords:'); + }); + + it('handles multiple clusters', () => { + const messages = buildNamingPrompt([ + { clusterId: 0, noteCount: 3, sampleTitles: ['A'], topKeywords: [] }, + { clusterId: 1, noteCount: 5, sampleTitles: ['B'], topKeywords: [] }, + { clusterId: 2, noteCount: 2, sampleTitles: ['C'], topKeywords: [] }, + ]); + + expect(messages[1].content).toContain('Group 0'); + expect(messages[1].content).toContain('Group 1'); + expect(messages[1].content).toContain('Group 2'); + }); + + it('system message asks for JSON output', () => { + const messages = buildNamingPrompt([{ clusterId: 0, noteCount: 1, sampleTitles: ['X'], topKeywords: [] }]); + + expect(messages[0].content).toContain('JSON'); + }); +}); + +describe('parseAiNamesResponse', () => { + it('parses valid JSON response', () => { + const response = '{"0": "Machine Learning", "1": "Travel Plans"}'; + const result = parseAiNamesResponse(response, [0, 1]); + expect(result).toEqual({ 0: 'Machine Learning', 1: 'Travel Plans' }); + }); + + it('parses JSON wrapped in markdown code block', () => { + const response = '```json\n{"0": "ML", "1": "Travel"}\n```'; + const result = parseAiNamesResponse(response, [0, 1]); + expect(result).toEqual({ 0: 'ML', 1: 'Travel' }); + }); + + it('parses JSON wrapped in plain code block', () => { + const response = '```\n{"0": "ML", "1": "Travel"}\n```'; + const result = parseAiNamesResponse(response, [0, 1]); + expect(result).toEqual({ 0: 'ML', 1: 'Travel' }); + }); + + it('extracts JSON embedded in surrounding text', () => { + const response = 'Here are the names:\n{"0": "ML", "1": "Travel"}\nHope this helps!'; + const result = parseAiNamesResponse(response, [0, 1]); + expect(result).toEqual({ 0: 'ML', 1: 'Travel' }); + }); + + it('sanitizes names in the response', () => { + const response = '{"0": "Title: Machine Learning.", "1": "\\"Travel Plans\\""}'; + const result = parseAiNamesResponse(response, [0, 1]); + expect(result).not.toBeNull(); + expect(result![0]).toBe('Machine Learning'); + expect(result![1]).toBe('Travel Plans'); + }); + + it('returns null for non-JSON response', () => { + const response = 'I think the clusters should be named Machine Learning and Travel.'; + const result = parseAiNamesResponse(response, [0, 1]); + expect(result).toBeNull(); + }); + + it('returns null for empty response', () => { + expect(parseAiNamesResponse('', [0, 1])).toBeNull(); + }); + + it('returns null for malformed JSON', () => { + expect(parseAiNamesResponse('{invalid json}', [0, 1])).toBeNull(); + }); + + it('returns null if fewer than half the clusters have names', () => { + const response = '{"0": "ML"}'; // only 1 of 3 clusters + const result = parseAiNamesResponse(response, [0, 1, 2]); + expect(result).toBeNull(); + }); + + it('accepts partial results if at least half the clusters have names', () => { + const response = '{"0": "ML", "1": "Travel"}'; // 2 of 3 clusters + const result = parseAiNamesResponse(response, [0, 1, 2]); + expect(result).not.toBeNull(); + expect(result![0]).toBe('ML'); + expect(result![1]).toBe('Travel'); + expect(result![2]).toBeUndefined(); + }); + + it('skips entries with empty string values', () => { + const response = '{"0": "ML", "1": " "}'; + const result = parseAiNamesResponse(response, [0, 1]); + // "1" is whitespace-only, sanitized to empty, so only "0" is valid + // 1 out of 2 = 50%, which meets the threshold (ceil(2/2) = 1) + expect(result).not.toBeNull(); + expect(result![0]).toBe('ML'); + }); + + it('handles numeric string cluster IDs', () => { + const response = '{"0": "Alpha", "3": "Beta"}'; + const result = parseAiNamesResponse(response, [0, 3]); + expect(result).toEqual({ 0: 'Alpha', 3: 'Beta' }); + }); +}); From b14631ca4318d6d36d62e21df6e03e535a81746e Mon Sep 17 00:00:00 2001 From: Harsh16gupta Date: Mon, 27 Jul 2026 02:02:52 +0530 Subject: [PATCH 2/2] feat: optimize native Joplin AI embedding pipeline and clustering quality --- src/pipeline/EmbeddingWorkerOrchestrator.ts | 9 ++- src/pipeline/nativeEmbeddingPipeline.ts | 36 +++++++++-- src/pipeline/pipelineConfig.ts | 43 ++++++++++++-- src/pipeline/runPipeline.ts | 42 +++++++++---- src/pipeline/vectorAggregator.ts | 66 +++++++++++++++++++++ src/pipeline/vectorCache.ts | 10 +++- src/types/joplinAi.d.ts | 3 +- test/pipeline/pipelineConfig.test.ts | 61 ++++++++++++++++--- test/pipeline/vectorAggregator.test.ts | 63 ++++++++++++++++++++ 9 files changed, 297 insertions(+), 36 deletions(-) diff --git a/src/pipeline/EmbeddingWorkerOrchestrator.ts b/src/pipeline/EmbeddingWorkerOrchestrator.ts index 8274f4a..ca19722 100644 --- a/src/pipeline/EmbeddingWorkerOrchestrator.ts +++ b/src/pipeline/EmbeddingWorkerOrchestrator.ts @@ -5,7 +5,7 @@ import { VectorCache } from './vectorCache'; import { isGenericTitle } from '../utils/titleFilter'; import { log, logErr } from '../utils/logger'; import { averageVectors, blendVectors, computeTitleWeight, cosineSimilarity } from './vectorAggregator'; -import { isValidEmbeddingVector } from './pipelineConfig'; +import { EMBEDDING_DIM, isValidEmbeddingVector } from './pipelineConfig'; const enc = getEncoding('cl100k_base'); const MAX_TOKENS = 200; @@ -171,6 +171,8 @@ export class EmbeddingWorkerOrchestrator { hash, updatedTime: note.updated_time, titleWeight, + modelId: 'local-onnx', + dimension: EMBEDDING_DIM, }); this.reportProgress(); @@ -204,7 +206,8 @@ export class EmbeddingWorkerOrchestrator { const cachedItem = await this.cache.getItem(note.id); if (cachedItem && cachedItem.metadata.hash === this.currentNoteHash) { - if (isValidEmbeddingVector(cachedItem.vector)) { + const isCorrectModel = !cachedItem.metadata.modelId || cachedItem.metadata.modelId === 'local-onnx'; + if (isCorrectModel && isValidEmbeddingVector(cachedItem.vector, EMBEDDING_DIM)) { log( `[${this.currentNoteIndex + 1}/${this.notes.length}] cache hit for "${note.title.slice(0, 30)}"`, ); @@ -220,7 +223,7 @@ export class EmbeddingWorkerOrchestrator { continue; } else { log( - `[${this.currentNoteIndex + 1}/${this.notes.length}] cache invalid (contains null/NaN) for "${note.title.slice(0, 30)}"`, + `[${this.currentNoteIndex + 1}/${this.notes.length}] cache invalid (wrong dimension/model) for "${note.title.slice(0, 30)}"`, ); } } diff --git a/src/pipeline/nativeEmbeddingPipeline.ts b/src/pipeline/nativeEmbeddingPipeline.ts index 02f49b8..5be6c27 100644 --- a/src/pipeline/nativeEmbeddingPipeline.ts +++ b/src/pipeline/nativeEmbeddingPipeline.ts @@ -13,6 +13,12 @@ export interface NativeEmbeddingChunk { vector: number[]; } +export interface NativeEmbeddingResult { + chunks: NativeEmbeddingChunk[]; + modelId: string; + dimension: number; +} + /** * Checks if Joplin's native AI indexing is active and ready. */ @@ -29,15 +35,19 @@ export const isNativeAiReady = async (): Promise => { }; /** - * Pages through Joplin's native index to fetch raw embedding vectors for the requested notes. + * Pages through Joplin's native index to fetch raw embedding vectors for the requested notes, + * returning the chunks along with modelId and vector dimension metadata. */ -export const fetchNativeEmbeddings = async (noteIds: string[]): Promise => { - if (noteIds.length === 0) return []; +export const fetchNativeEmbeddings = async (noteIds: string[]): Promise => { + if (noteIds.length === 0) { + return { chunks: [], modelId: 'unknown', dimension: 384 }; + } log(`Fetching native embeddings for ${noteIds.length} notes...`); const chunks: NativeEmbeddingChunk[] = []; const BATCH_SIZE = 500; let modelId: string | null = null; + let dimension: number | null = null; for (let i = 0; i < noteIds.length; i += BATCH_SIZE) { const batchIds = noteIds.slice(i, i + BATCH_SIZE); @@ -59,11 +69,22 @@ export const fetchNativeEmbeddings = async (noteIds: string[]): Promise 0) { + if (dimension && dimension !== page.dimension) { + throw new Error('Embedding dimension changed mid-fetch. Please restart.'); + } + dimension = page.dimension; + } + for (const chunk of page.chunks) { if (!chunk.noteId || !Array.isArray(chunk.vector)) { log(`Skipping malformed embedding chunk: ${JSON.stringify(chunk).slice(0, 100)}`); continue; } + if (!dimension && chunk.vector.length > 0) { + dimension = chunk.vector.length; + } chunks.push(chunk); } cursor = page.nextCursor; @@ -77,6 +98,11 @@ export const fetchNativeEmbeddings = async (noteIds: string[]): Promise 0 ? chunks[0].vector.length : 384); + const finalModelId = modelId ?? 'native-ai'; + + log( + `Successfully fetched ${chunks.length} native embedding chunks (model: ${finalModelId}, dim: ${finalDimension})`, + ); + return { chunks, modelId: finalModelId, dimension: finalDimension }; }; diff --git a/src/pipeline/pipelineConfig.ts b/src/pipeline/pipelineConfig.ts index a16d3db..16bf758 100644 --- a/src/pipeline/pipelineConfig.ts +++ b/src/pipeline/pipelineConfig.ts @@ -1,14 +1,48 @@ import { CategorizationConfig } from '../types/cluster'; -/** Dimensionality of embedding vectors (all-MiniLM-L6-v2 / multilingual-e5-small). */ +/** Default dimensionality of local ONNX embedding vectors (all-MiniLM-L6-v2 / multilingual-e5-small). */ export const EMBEDDING_DIM = 384; -export function isValidEmbeddingVector(vector: number[] | undefined | null): boolean { - if (!vector) return false; - if (vector.length !== EMBEDDING_DIM) return false; +export function isValidEmbeddingVector(vector: number[] | undefined | null, expectedDim?: number): boolean { + if (!vector || vector.length === 0) return false; + if (expectedDim !== undefined && vector.length !== expectedDim) return false; return vector.every((v) => Number.isFinite(v)); } +/** + * Computes UMAP intermediate dimensionality scaled logarithmically with input embedding dimension. + * Formula: clamp(⌊2·log₂(D)⌋, 5, 50) + */ +export function adaptiveIntermediateDim(inputDim: number): number { + const raw = Math.floor(2 * Math.log2(inputDim)); + if (!Number.isFinite(raw)) return 5; + return Math.max(5, Math.min(50, raw)); +} + +/** + * Computes UMAP neighbor count scaled with square root of note count. + * Formula: clamp(⌊√N⌋, 5, 50) + */ +export function adaptiveNeighbors(noteCount: number): number { + const raw = Math.floor(Math.sqrt(noteCount)); + if (!Number.isFinite(raw)) return 5; + return Math.max(5, Math.min(50, raw)); +} + +export function createAdaptiveConfig(inputDim: number, noteCount: number): CategorizationConfig { + return { + seed: 42, + metric: 'cosine', + intermediateDim: adaptiveIntermediateDim(inputDim), + intermediateNeighbors: adaptiveNeighbors(noteCount), + strategies: [ + { name: 'kmeans-6', algorithm: 'kmeans', K: 6 }, + { name: 'kmedoids-6', algorithm: 'kmedoids', K: 6 }, + { name: 'hdbscan', algorithm: 'hdbscan', minClusterSize: 3, minSamples: 2 }, + ], + }; +} + export const DEFAULT_CONFIG: CategorizationConfig = { seed: 42, metric: 'cosine', @@ -17,7 +51,6 @@ export const DEFAULT_CONFIG: CategorizationConfig = { strategies: [ { name: 'kmeans-6', algorithm: 'kmeans', K: 6 }, { name: 'kmedoids-6', algorithm: 'kmedoids', K: 6 }, - // { name: 'hdbscan-tuned', algorithm: 'hdbscan', minClusterSize: 4, minSamples: 1 }, { name: 'hdbscan', algorithm: 'hdbscan', minClusterSize: 3, minSamples: 2 }, ], }; diff --git a/src/pipeline/runPipeline.ts b/src/pipeline/runPipeline.ts index 1a1a1f1..45339bf 100644 --- a/src/pipeline/runPipeline.ts +++ b/src/pipeline/runPipeline.ts @@ -1,11 +1,11 @@ import { fetchAllNotes } from './noteReader'; import { benchmark } from './clustering/benchmark'; -import { averageVectors } from './vectorAggregator'; +import { weightedAverageVectorsWithNorm } from './vectorAggregator'; import { PanelNote } from '../types/panel'; import { log, logErr } from '../utils/logger'; import { VectorCache } from './vectorCache'; import { isNativeAiReady, fetchNativeEmbeddings } from './nativeEmbeddingPipeline'; -import { DEFAULT_CONFIG, isValidEmbeddingVector } from './pipelineConfig'; +import { DEFAULT_CONFIG, isValidEmbeddingVector, createAdaptiveConfig } from './pipelineConfig'; import { enrichResultsWithTags } from './clustering/postProcess'; import { upgradeClusterNamesWithAi } from './clustering/aiNamingService'; import { EmbeddingWorkerOrchestrator } from './EmbeddingWorkerOrchestrator'; @@ -17,6 +17,11 @@ export interface PipelineCallbacks { onError: (message: string) => void; } +interface IndexedVector { + chunkIndex: number; + vector: number[]; +} + /** * Runs the full embedding + clustering pipeline, reporting progress via callbacks. * @@ -45,13 +50,14 @@ export const runPipeline = async (installDir: string, callbacks: PipelineCallbac try { const noteIds = notes.map((n) => n.id); - const chunks = await fetchNativeEmbeddings(noteIds); + const nativeResult = await fetchNativeEmbeddings(noteIds); + log(`Native AI embeddings retrieved (model: ${nativeResult.modelId}, dim: ${nativeResult.dimension})`); - // Group chunks by noteId - const noteChunksMap = new Map(); - for (const chunk of chunks) { + // Group chunks by noteId preserving chunkIndex for ordering + const noteChunksMap = new Map(); + for (const chunk of nativeResult.chunks) { const list = noteChunksMap.get(chunk.noteId) || []; - list.push(chunk.vector); + list.push({ chunkIndex: chunk.chunkIndex, vector: chunk.vector }); noteChunksMap.set(chunk.noteId, list); } @@ -59,11 +65,22 @@ export const runPipeline = async (installDir: string, callbacks: PipelineCallbac const vectors: number[][] = []; for (const note of notes) { - const chunkVectors = noteChunksMap.get(note.id); - if (chunkVectors && chunkVectors.length > 0) { - const avgVector = averageVectors(chunkVectors); + const chunkEntries = noteChunksMap.get(note.id); + if (chunkEntries && chunkEntries.length > 0) { + // Sort by chunkIndex ascending to ensure lead paragraph/header (chunk 0) gets highest weight + const sortedEntries = chunkEntries.sort((a, b) => a.chunkIndex - b.chunkIndex); + const chunkVectors = sortedEntries.map((e) => e.vector); + + const { vector: avgVector, rawNorm } = weightedAverageVectorsWithNorm(chunkVectors); + + if (rawNorm < 1e-6) { + logErr( + `Native embedding for note "${note.title}" has near-zero L2 norm (${rawNorm}). Skipping as outlier.`, + ); + continue; + } - if (isValidEmbeddingVector(avgVector)) { + if (isValidEmbeddingVector(avgVector, nativeResult.dimension)) { vectors.push(avgVector); validNotes.push(note); } else { @@ -80,7 +97,8 @@ export const runPipeline = async (installDir: string, callbacks: PipelineCallbac log('Too few indexed notes found in native DB. Falling back to local ONNX Web Worker.'); } else { callbacks.onStatus('Clustering...'); - const results = benchmark(vectors, DEFAULT_CONFIG); + const adaptiveConfig = createAdaptiveConfig(nativeResult.dimension, validNotes.length); + const results = benchmark(vectors, adaptiveConfig); // Post-process to extract tags/keywords for each cluster (keep parity with local pipeline) const allPipelineDocuments = validNotes.map((n) => ({ diff --git a/src/pipeline/vectorAggregator.ts b/src/pipeline/vectorAggregator.ts index bc3c37b..6f45b47 100644 --- a/src/pipeline/vectorAggregator.ts +++ b/src/pipeline/vectorAggregator.ts @@ -28,6 +28,72 @@ export const averageVectors = (vectors: number[][]): number[] => { return normalise(avg); }; +export interface WeightedPoolingOptions { + lambda?: number; + leadBoost?: number; +} + +export interface WeightedAverageResult { + vector: number[]; + rawNorm: number; +} + +/** + * Element-wise weighted average of chunk vectors using exponential position decay + * with a lead-chunk boost, returning both the L2-normalised vector and the raw pre-normalisation norm. + * + * w_i = (leadBoost + exp(-lambda * i)) if i == 0 else exp(-lambda * i) + */ +export const weightedAverageVectorsWithNorm = ( + vectors: number[][], + options: WeightedPoolingOptions = {}, +): WeightedAverageResult => { + if (vectors.length === 0) throw new Error('Cannot average zero vectors'); + const dim = vectors[0].length; + for (const vec of vectors) { + if (vec.length !== dim) throw new Error('Cannot average vectors of different dimensions'); + } + if (vectors.length === 1) { + const rawNorm = Math.sqrt(vectors[0].reduce((sum, v) => sum + v * v, 0)); + const vector = rawNorm === 0 ? vectors[0] : vectors[0].map((v) => v / rawNorm); + return { vector, rawNorm }; + } + + const lambda = options.lambda ?? 0.15; + const leadBoost = options.leadBoost ?? 0.5; + + const weights: number[] = new Array(vectors.length); + let weightSum = 0; + + for (let i = 0; i < vectors.length; i++) { + const w = Math.exp(-lambda * i) + (i === 0 ? leadBoost : 0); + weights[i] = w; + weightSum += w; + } + + const weightedSum = new Array(dim).fill(0); + for (let i = 0; i < vectors.length; i++) { + const normW = weights[i] / weightSum; + const vec = vectors[i]; + for (let d = 0; d < dim; d++) { + weightedSum[d] += normW * vec[d]; + } + } + + const rawNorm = Math.sqrt(weightedSum.reduce((sum, v) => sum + v * v, 0)); + const vector = rawNorm === 0 ? weightedSum : weightedSum.map((v) => v / rawNorm); + + return { vector, rawNorm }; +}; + +/** + * Element-wise weighted average of chunk vectors using exponential position decay + * with a lead-chunk boost, then L2-normalised. + */ +export const weightedAverageVectors = (vectors: number[][], options: WeightedPoolingOptions = {}): number[] => { + return weightedAverageVectorsWithNorm(vectors, options).vector; +}; + /** * Cosine similarity between two L2-normalised vectors (= dot product). */ diff --git a/src/pipeline/vectorCache.ts b/src/pipeline/vectorCache.ts index f84614d..1239a03 100644 --- a/src/pipeline/vectorCache.ts +++ b/src/pipeline/vectorCache.ts @@ -9,7 +9,9 @@ export interface CacheMetadata { hash: string; updatedTime: number; titleWeight: number; - [key: string]: MetadataTypes; + modelId?: string; + dimension?: number; + [key: string]: MetadataTypes | undefined; } export class VectorCache { @@ -56,7 +58,9 @@ export class VectorCache { */ public async getItem(id: string) { try { - return await this.index.getItem(id); + const item = await this.index.getItem>(id); + if (!item) return undefined; + return item as unknown as { id: string; vector: number[]; metadata: CacheMetadata }; } catch (err) { log('Error getting cached item:', err); return undefined; @@ -70,7 +74,7 @@ export class VectorCache { return await this.index.upsertItem({ id, vector, - metadata, + metadata: metadata as unknown as Record, }); } diff --git a/src/types/joplinAi.d.ts b/src/types/joplinAi.d.ts index b099518..e93ece7 100644 --- a/src/types/joplinAi.d.ts +++ b/src/types/joplinAi.d.ts @@ -7,7 +7,7 @@ export interface AiIndexStatus { state: string; ready: boolean; - modelId?: string; + modelId: string | null; } export interface AiEmbeddingChunk { @@ -20,6 +20,7 @@ export interface AiEmbeddingChunk { export interface AiEmbeddingsPage { chunks: AiEmbeddingChunk[]; modelId: string; + dimension: number; nextCursor?: string; } diff --git a/test/pipeline/pipelineConfig.test.ts b/test/pipeline/pipelineConfig.test.ts index 6bd1893..16ac510 100644 --- a/test/pipeline/pipelineConfig.test.ts +++ b/test/pipeline/pipelineConfig.test.ts @@ -1,21 +1,36 @@ -import { EMBEDDING_DIM, isValidEmbeddingVector, DEFAULT_CONFIG } from '../../src/pipeline/pipelineConfig'; +import { + EMBEDDING_DIM, + isValidEmbeddingVector, + adaptiveIntermediateDim, + adaptiveNeighbors, + createAdaptiveConfig, + DEFAULT_CONFIG, +} from '../../src/pipeline/pipelineConfig'; describe('isValidEmbeddingVector', () => { - it('accepts a valid 384-dim vector of finite numbers', () => { + it('accepts a valid vector of finite numbers with default dimension check', () => { const vec = new Array(EMBEDDING_DIM).fill(0.1); expect(isValidEmbeddingVector(vec)).toBe(true); }); - it('rejects null', () => { - expect(isValidEmbeddingVector(null)).toBe(false); + it('accepts custom dimension vectors when expectedDim is specified', () => { + const vec768 = new Array(768).fill(0.05); + const vec1536 = new Array(1536).fill(0.02); + expect(isValidEmbeddingVector(vec768, 768)).toBe(true); + expect(isValidEmbeddingVector(vec1536, 1536)).toBe(true); }); - it('rejects undefined', () => { + it('rejects null and undefined', () => { + expect(isValidEmbeddingVector(null)).toBe(false); expect(isValidEmbeddingVector(undefined)).toBe(false); }); - it('rejects wrong dimension', () => { - expect(isValidEmbeddingVector(new Array(383).fill(0.1))).toBe(false); + it('rejects empty array', () => { + expect(isValidEmbeddingVector([])).toBe(false); + }); + + it('rejects vector matching wrong expected dimension', () => { + expect(isValidEmbeddingVector(new Array(768).fill(0.1), 384)).toBe(false); }); it('rejects vector containing NaN', () => { @@ -31,6 +46,38 @@ describe('isValidEmbeddingVector', () => { }); }); +describe('adaptive scaling functions', () => { + it('computes adaptive intermediate dimensions logarithmic with input dimension', () => { + expect(adaptiveIntermediateDim(384)).toBe(17); + expect(adaptiveIntermediateDim(768)).toBe(19); + expect(adaptiveIntermediateDim(1536)).toBe(21); + }); + + it('clamps intermediate dimensions between 5 and 50', () => { + expect(adaptiveIntermediateDim(2)).toBe(5); + expect(adaptiveIntermediateDim(1e12)).toBe(50); + }); + + it('computes adaptive neighbors based on square root of note count', () => { + expect(adaptiveNeighbors(25)).toBe(5); + expect(adaptiveNeighbors(100)).toBe(10); + expect(adaptiveNeighbors(500)).toBe(22); + }); + + it('clamps neighbors between 5 and 50', () => { + expect(adaptiveNeighbors(2)).toBe(5); + expect(adaptiveNeighbors(10000)).toBe(50); + }); + + it('creates adaptive configuration dynamically', () => { + const config = createAdaptiveConfig(768, 100); + expect(config.metric).toBe('cosine'); + expect(config.intermediateDim).toBe(19); + expect(config.intermediateNeighbors).toBe(10); + expect(config.strategies.length).toBe(3); + }); +}); + describe('DEFAULT_CONFIG', () => { it('uses cosine metric and seed 42', () => { expect(DEFAULT_CONFIG.metric).toBe('cosine'); diff --git a/test/pipeline/vectorAggregator.test.ts b/test/pipeline/vectorAggregator.test.ts index 43ea083..b9dc06f 100644 --- a/test/pipeline/vectorAggregator.test.ts +++ b/test/pipeline/vectorAggregator.test.ts @@ -1,5 +1,7 @@ import { averageVectors, + weightedAverageVectors, + weightedAverageVectorsWithNorm, cosineSimilarity, computeTitleWeight, blendVectors, @@ -53,6 +55,67 @@ describe('averageVectors', () => { }); }); +describe('weightedAverageVectors', () => { + it('normalizes a single vector', () => { + const result = weightedAverageVectors([[3, 4]]); + expect(result[0]).toBeCloseTo(0.6, 10); + expect(result[1]).toBeCloseTo(0.8, 10); + }); + + it('gives lead chunk (index 0) higher weight than tail chunks', () => { + const chunk0 = [1, 0]; + const chunk1 = [0, 1]; + const result = weightedAverageVectors([chunk0, chunk1]); + + // Chunk 0 has weight = exp(0) + 0.5 = 1.5 + // Chunk 1 has weight = exp(-0.15) ≈ 0.8607 + // Therefore result[0] > result[1] + expect(result[0]).toBeGreaterThan(result[1]); + }); + + it('supports custom lambda and leadBoost parameters', () => { + const chunk0 = [1, 0]; + const chunk1 = [0, 1]; + const result = weightedAverageVectors([chunk0, chunk1], { lambda: 0.5, leadBoost: 1.0 }); + + // Chunk 0 weight = 2.0, Chunk 1 weight = exp(-0.5) ≈ 0.6065 + expect(result[0]).toBeGreaterThan(result[1]); + }); + + it('throws on empty input', () => { + expect(() => weightedAverageVectors([])).toThrow('Cannot average zero vectors'); + }); + + it('throws on dimension mismatch', () => { + expect(() => + weightedAverageVectors([ + [1, 2], + [1, 2, 3], + ]), + ).toThrow('different dimensions'); + }); +}); + +describe('weightedAverageVectorsWithNorm', () => { + it('returns both normalized vector and pre-normalization raw norm', () => { + const chunk0 = [0.6, 0.8]; + const res = weightedAverageVectorsWithNorm([chunk0]); + expect(res.vector[0]).toBeCloseTo(0.6, 10); + expect(res.vector[1]).toBeCloseTo(0.8, 10); + expect(res.rawNorm).toBeCloseTo(1.0, 10); + }); + + it('detects near-zero raw norm when vectors cancel out', () => { + // Opposing unit vectors with equal weights + const chunk0 = [1, 0]; + const chunk1 = [-1, 0]; + // Weight chunk0 = 1.5, weight chunk1 = exp(-0.15) ≈ 0.8607 + // If lambda=0 and leadBoost=0, weights are equal + const res = weightedAverageVectorsWithNorm([chunk0, chunk1], { lambda: 0, leadBoost: 0 }); + expect(res.rawNorm).toBeCloseTo(0, 10); + }); +}); + describe('cosineSimilarity', () => { it('returns 1 for identical unit vectors', () => { expect(cosineSimilarity([0.6, 0.8], [0.6, 0.8])).toBeCloseTo(1.0, 10);