From b6df3a721c184408a7692b28632e24c700f15dfe Mon Sep 17 00:00:00 2001
From: perf3ct <jonfuller2012@gmail.com>
Date: Wed, 12 Mar 2025 18:02:51 +0000
Subject: [PATCH] allow user to select *where* they want to generate embeddings

---
 db/migrations/0230__vector_embeddings.sql     |   4 +-
 .../type_widgets/options/ai_settings.ts       |  15 +++
 src/public/translations/en/translation.json   |   4 +
 src/services/llm/embeddings/queue.ts          |  10 ++
 src/services/llm/index_service.ts             | 105 ++++++++++++++++--
 src/services/options_interface.ts             |   1 +
 6 files changed, 127 insertions(+), 12 deletions(-)

diff --git a/db/migrations/0230__vector_embeddings.sql b/db/migrations/0230__vector_embeddings.sql
index a614948a1..07c425fd2 100644
--- a/db/migrations/0230__vector_embeddings.sql
+++ b/db/migrations/0230__vector_embeddings.sql
@@ -52,4 +52,6 @@ VALUES ('embeddingUpdateInterval', '5000', 1, strftime('%Y-%m-%d %H:%M:%f', 'now
 INSERT INTO options (name, value, isSynced, utcDateModified) 
 VALUES ('embeddingBatchSize', '10', 1, strftime('%Y-%m-%d %H:%M:%f', 'now'));
 INSERT INTO options (name, value, isSynced, utcDateModified) 
-VALUES ('embeddingDefaultDimension', '1536', 1, strftime('%Y-%m-%d %H:%M:%f', 'now')); 
\ No newline at end of file
+VALUES ('embeddingDefaultDimension', '1536', 1, strftime('%Y-%m-%d %H:%M:%f', 'now')); 
+INSERT INTO options (name, value, isSynced, utcDateModified) 
+VALUES ('embeddingGenerationLocation', 'client', 1, strftime('%Y-%m-%dT%H:%M:%fZ', 'now')); 
\ No newline at end of file
diff --git a/src/public/app/widgets/type_widgets/options/ai_settings.ts b/src/public/app/widgets/type_widgets/options/ai_settings.ts
index 6f43b4319..92fdfa24c 100644
--- a/src/public/app/widgets/type_widgets/options/ai_settings.ts
+++ b/src/public/app/widgets/type_widgets/options/ai_settings.ts
@@ -186,6 +186,15 @@ export default class AiSettingsWidget extends OptionsWidget {
                     <div class="help-text">${t("ai_llm.embedding_default_provider_description")}</div>
                 </div>
 
+                <div class="form-group">
+                    <label>${t("ai_llm.embedding_generation_location")}</label>
+                    <select class="embedding-generation-location form-control">
+                        <option value="client">${t("ai_llm.embedding_generation_location_client")}</option>
+                        <option value="sync_server">${t("ai_llm.embedding_generation_location_sync_server")}</option>
+                    </select>
+                    <div class="help-text">${t("ai_llm.embedding_generation_location_description")}</div>
+                </div>
+
                 <div class="form-group">
                     <label>
                         <input class="embedding-auto-update-enabled" type="checkbox">
@@ -446,6 +455,11 @@ export default class AiSettingsWidget extends OptionsWidget {
             await this.updateOption('embeddingsDefaultProvider', $embeddingDefaultProvider.val() as string);
         });
 
+        const $embeddingGenerationLocation = this.$widget.find('.embedding-generation-location');
+        $embeddingGenerationLocation.on('change', async () => {
+            await this.updateOption('embeddingGenerationLocation', $embeddingGenerationLocation.val() as string);
+        });
+
         const $embeddingBatchSize = this.$widget.find('.embedding-batch-size');
         $embeddingBatchSize.on('change', async () => {
             await this.updateOption('embeddingBatchSize', $embeddingBatchSize.val() as string);
@@ -541,6 +555,7 @@ export default class AiSettingsWidget extends OptionsWidget {
 
         // Load embedding options
         this.$widget.find('.embedding-default-provider').val(options.embeddingsDefaultProvider || 'openai');
+        this.$widget.find('.embedding-generation-location').val(options.embeddingGenerationLocation || 'client');
         this.setCheckboxState(this.$widget.find('.embedding-auto-update-enabled'), options.embeddingAutoUpdateEnabled || 'true');
         this.setCheckboxState(this.$widget.find('.enable-automatic-indexing'), options.enableAutomaticIndexing || 'true');
         this.$widget.find('.embedding-similarity-threshold').val(options.embeddingSimilarityThreshold || '0.65');
diff --git a/src/public/translations/en/translation.json b/src/public/translations/en/translation.json
index bcb42939f..3ba3ad2d4 100644
--- a/src/public/translations/en/translation.json
+++ b/src/public/translations/en/translation.json
@@ -1157,6 +1157,10 @@
     "embedding_configuration": "Embeddings Configuration",
     "embedding_default_provider": "Default Provider",
     "embedding_default_provider_description": "Select the default provider used for generating note embeddings",
+    "embedding_generation_location": "Generation Location",
+    "embedding_generation_location_description": "Select where embedding generation should happen",
+    "embedding_generation_location_client": "Client/Server",
+    "embedding_generation_location_sync_server": "Sync Server",
     "enable_auto_update_embeddings": "Auto-update Embeddings",
     "enable_auto_update_embeddings_description": "Automatically update embeddings when notes are modified",
     "auto_update_embeddings": "Auto-update Embeddings",
diff --git a/src/services/llm/embeddings/queue.ts b/src/services/llm/embeddings/queue.ts
index 76ae6ad22..9ea9fc3a7 100644
--- a/src/services/llm/embeddings/queue.ts
+++ b/src/services/llm/embeddings/queue.ts
@@ -157,6 +157,16 @@ export async function processEmbeddingQueue() {
         return;
     }
 
+    // Check if this instance should process embeddings
+    const embeddingLocation = await options.getOption('embeddingGenerationLocation') || 'client';
+    const isSyncServer = await indexService.isSyncServerForEmbeddings();
+    const shouldProcessEmbeddings = embeddingLocation === 'client' || isSyncServer;
+
+    if (!shouldProcessEmbeddings) {
+        // This instance is not configured to process embeddings
+        return;
+    }
+
     const batchSize = parseInt(await options.getOption('embeddingBatchSize') || '10', 10);
     const enabledProviders = await getEnabledEmbeddingProviders();
 
diff --git a/src/services/llm/index_service.ts b/src/services/llm/index_service.ts
index 131ec8de9..64e42c453 100644
--- a/src/services/llm/index_service.ts
+++ b/src/services/llm/index_service.ts
@@ -57,9 +57,17 @@ class IndexService {
                 throw new Error("No embedding providers available");
             }
 
-            // Setup automatic indexing if enabled
-            if (await options.getOptionBool('embeddingAutoUpdate')) {
+            // Check if this instance should process embeddings
+            const embeddingLocation = await options.getOption('embeddingGenerationLocation') || 'client';
+            const isSyncServer = await this.isSyncServerForEmbeddings();
+            const shouldProcessEmbeddings = embeddingLocation === 'client' || isSyncServer;
+
+            // Setup automatic indexing if enabled and this instance should process embeddings
+            if (await options.getOptionBool('embeddingAutoUpdate') && shouldProcessEmbeddings) {
                 this.setupAutomaticIndexing();
+                log.info(`Index service: Automatic indexing enabled, processing embeddings ${isSyncServer ? 'as sync server' : 'as client'}`);
+            } else if (await options.getOptionBool('embeddingAutoUpdate')) {
+                log.info("Index service: Automatic indexing enabled, but this instance is not configured to process embeddings");
             }
 
             // Listen for note changes to update index
@@ -78,32 +86,33 @@ class IndexService {
      */
     private setupEventListeners() {
         // Listen for note content changes
-        eventService.subscribe(eventService.NOTE_CONTENT_CHANGE, ({ entity }) => {
+        eventService.subscribe(eventService.NOTE_CONTENT_CHANGE, async ({ entity }) => {
             if (entity && entity.noteId) {
-                this.queueNoteForIndexing(entity.noteId);
+                // Always queue notes for indexing, but the actual processing will depend on configuration
+                await this.queueNoteForIndexing(entity.noteId);
             }
         });
 
         // Listen for new notes
-        eventService.subscribe(eventService.ENTITY_CREATED, ({ entityName, entity }) => {
+        eventService.subscribe(eventService.ENTITY_CREATED, async ({ entityName, entity }) => {
             if (entityName === "notes" && entity && entity.noteId) {
-                this.queueNoteForIndexing(entity.noteId);
+                await this.queueNoteForIndexing(entity.noteId);
             }
         });
 
         // Listen for note title changes
-        eventService.subscribe(eventService.NOTE_TITLE_CHANGED, ({ noteId }) => {
+        eventService.subscribe(eventService.NOTE_TITLE_CHANGED, async ({ noteId }) => {
             if (noteId) {
-                this.queueNoteForIndexing(noteId);
+                await this.queueNoteForIndexing(noteId);
             }
         });
 
         // Listen for changes in AI settings
-        eventService.subscribe(eventService.ENTITY_CHANGED, ({ entityName, entity }) => {
+        eventService.subscribe(eventService.ENTITY_CHANGED, async ({ entityName, entity }) => {
             if (entityName === "options" && entity && entity.name) {
                 if (entity.name.startsWith('ai') || entity.name.startsWith('embedding')) {
                     log.info("AI settings changed, updating index service configuration");
-                    this.updateConfiguration();
+                    await this.updateConfiguration();
                 }
             }
         });
@@ -122,6 +131,16 @@ class IndexService {
         this.automaticIndexingInterval = setInterval(async () => {
             try {
                 if (!this.indexingInProgress) {
+                    // Check if this instance should process embeddings
+                    const embeddingLocation = await options.getOption('embeddingGenerationLocation') || 'client';
+                    const isSyncServer = await this.isSyncServerForEmbeddings();
+                    const shouldProcessEmbeddings = embeddingLocation === 'client' || isSyncServer;
+
+                    if (!shouldProcessEmbeddings) {
+                        // This instance is not configured to process embeddings
+                        return;
+                    }
+
                     const stats = await vectorStore.getEmbeddingStats();
 
                     // Only run automatic indexing if we're below 95% completion
@@ -147,10 +166,20 @@ class IndexService {
             const intervalMs = parseInt(await options.getOption('embeddingUpdateInterval') || '3600000', 10);
             this.indexUpdateInterval = intervalMs;
 
+            // Check if this instance should process embeddings
+            const embeddingLocation = await options.getOption('embeddingGenerationLocation') || 'client';
+            const isSyncServer = await this.isSyncServerForEmbeddings();
+            const shouldProcessEmbeddings = embeddingLocation === 'client' || isSyncServer;
+
             // Update automatic indexing setting
             const autoIndexing = await options.getOptionBool('embeddingAutoUpdate');
-            if (autoIndexing && !this.automaticIndexingInterval) {
+            if (autoIndexing && shouldProcessEmbeddings && !this.automaticIndexingInterval) {
                 this.setupAutomaticIndexing();
+                log.info(`Index service: Automatic indexing enabled, processing embeddings ${isSyncServer ? 'as sync server' : 'as client'}`);
+            } else if (autoIndexing && !shouldProcessEmbeddings && this.automaticIndexingInterval) {
+                clearInterval(this.automaticIndexingInterval);
+                this.automaticIndexingInterval = undefined;
+                log.info("Index service: Automatic indexing disabled for this instance based on configuration");
             } else if (!autoIndexing && this.automaticIndexingInterval) {
                 clearInterval(this.automaticIndexingInterval);
                 this.automaticIndexingInterval = undefined;
@@ -179,6 +208,8 @@ class IndexService {
         }
 
         try {
+            // Always queue notes for indexing, regardless of where embedding generation happens
+            // The actual processing will be determined when the queue is processed
             await vectorStore.queueNoteForEmbedding(noteId, 'UPDATE');
             return true;
         } catch (error: any) {
@@ -201,6 +232,17 @@ class IndexService {
         }
 
         try {
+            // Check if this instance should process embeddings
+            const embeddingLocation = await options.getOption('embeddingGenerationLocation') || 'client';
+            const isSyncServer = await this.isSyncServerForEmbeddings();
+            const shouldProcessEmbeddings = embeddingLocation === 'client' || isSyncServer;
+
+            if (!shouldProcessEmbeddings) {
+                // This instance is not configured to process embeddings
+                log.info("Skipping full indexing as this instance is not configured to process embeddings");
+                return false;
+            }
+
             this.indexingInProgress = true;
             this.indexRebuildInProgress = true;
             this.indexRebuildProgress = 0;
@@ -321,6 +363,17 @@ class IndexService {
         try {
             this.indexingInProgress = true;
 
+            // Check if this instance should process embeddings
+            const embeddingLocation = await options.getOption('embeddingGenerationLocation') || 'client';
+            const isSyncServer = await this.isSyncServerForEmbeddings();
+            const shouldProcessEmbeddings = embeddingLocation === 'client' || isSyncServer;
+
+            if (!shouldProcessEmbeddings) {
+                // This instance is not configured to process embeddings
+                log.info("Skipping batch indexing as this instance is not configured to process embeddings");
+                return false;
+            }
+
             // Process the embedding queue
             await vectorStore.processEmbeddingQueue();
 
@@ -618,6 +671,21 @@ class IndexService {
         }
     }
 
+    /**
+     * Check if this instance is a sync server and should generate embeddings
+     */
+    async isSyncServerForEmbeddings() {
+        // Check if this is a sync server (no syncServerHost means this is a sync server)
+        const syncServerHost = await options.getOption('syncServerHost');
+        const isSyncServer = !syncServerHost;
+
+        // Check if embedding generation should happen on the sync server
+        const embeddingLocation = await options.getOption('embeddingGenerationLocation') || 'client';
+        const shouldGenerateOnSyncServer = embeddingLocation === 'sync_server';
+
+        return isSyncServer && shouldGenerateOnSyncServer;
+    }
+
     /**
      * Generate a comprehensive index entry for a note
      * This prepares all metadata and contexts for optimal LLM retrieval
@@ -633,6 +701,21 @@ class IndexService {
                 throw new Error(`Note ${noteId} not found`);
             }
 
+            // Check where embedding generation should happen
+            const embeddingLocation = await options.getOption('embeddingGenerationLocation') || 'client';
+
+            // If embedding generation should happen on the sync server and we're not the sync server,
+            // just queue the note for embedding but don't generate it
+            const isSyncServer = await this.isSyncServerForEmbeddings();
+            const shouldSkipGeneration = embeddingLocation === 'sync_server' && !isSyncServer;
+
+            if (shouldSkipGeneration) {
+                // We're not the sync server, so just queue the note for embedding
+                // The sync server will handle the actual embedding generation
+                log.info(`Note ${noteId} queued for embedding generation on sync server`);
+                return true;
+            }
+
             // Get complete note context for indexing
             const context = await vectorStore.getNoteEmbeddingContext(noteId);
 
diff --git a/src/services/options_interface.ts b/src/services/options_interface.ts
index 2d7b5a66b..8a704b0ca 100644
--- a/src/services/options_interface.ts
+++ b/src/services/options_interface.ts
@@ -72,6 +72,7 @@ export interface OptionDefinitions extends KeyboardShortcutsOptions<KeyboardActi
     enableAutomaticIndexing: boolean;
     embeddingSimilarityThreshold: string;
     maxNotesPerLlmQuery: string;
+    embeddingGenerationLocation: string;
 
     lastSyncedPull: number;
     lastSyncedPush: number;