From 49e123f39990f6fd33a2a7b7b2ef9ff9e2ad4b3d Mon Sep 17 00:00:00 2001 From: perf3ct Date: Thu, 5 Jun 2025 18:47:25 +0000 Subject: [PATCH] feat(llm): create endpoints for starting/stopping embeddings --- .../options/ai_settings/ai_settings_widget.ts | 29 +++++++ apps/server/src/routes/api/embeddings.ts | 47 +++++++++++- apps/server/src/routes/routes.ts | 2 + .../src/services/llm/ai_service_manager.ts | 24 +++++- .../src/services/llm/embeddings/events.ts | 22 +++++- .../src/services/llm/embeddings/index.ts | 2 + apps/server/src/services/llm/index_service.ts | 75 +++++++++++++++++++ 7 files changed, 197 insertions(+), 4 deletions(-) diff --git a/apps/client/src/widgets/type_widgets/options/ai_settings/ai_settings_widget.ts b/apps/client/src/widgets/type_widgets/options/ai_settings/ai_settings_widget.ts index 67cc80cbb..4f63bfcb1 100644 --- a/apps/client/src/widgets/type_widgets/options/ai_settings/ai_settings_widget.ts +++ b/apps/client/src/widgets/type_widgets/options/ai_settings/ai_settings_widget.ts @@ -51,6 +51,35 @@ export default class AiSettingsWidget extends OptionsWidget { await this.updateOption(optionName, value); + // Special handling for aiEnabled option + if (optionName === 'aiEnabled') { + try { + const isEnabled = value === 'true'; + + if (isEnabled) { + // Start embedding generation + await server.post('llm/embeddings/start'); + toastService.showMessage(t("ai_llm.embeddings_started") || "Embedding generation started"); + + // Start polling for stats updates + this.refreshEmbeddingStats(); + } else { + // Stop embedding generation + await server.post('llm/embeddings/stop'); + toastService.showMessage(t("ai_llm.embeddings_stopped") || "Embedding generation stopped"); + + // Clear any active polling intervals + if (this.indexRebuildRefreshInterval) { + clearInterval(this.indexRebuildRefreshInterval); + this.indexRebuildRefreshInterval = null; + } + } + } catch (error) { + console.error('Error toggling embeddings:', error); + toastService.showError(t("ai_llm.embeddings_toggle_error") || "Error toggling embeddings"); + } + } + if (validateAfter) { await this.displayValidationWarnings(); } diff --git a/apps/server/src/routes/api/embeddings.ts b/apps/server/src/routes/api/embeddings.ts index 012a9c82f..8fd9b475a 100644 --- a/apps/server/src/routes/api/embeddings.ts +++ b/apps/server/src/routes/api/embeddings.ts @@ -782,6 +782,49 @@ async function getIndexRebuildStatus(req: Request, res: Response) { }; } +/** + * Start embedding generation when AI is enabled + */ +async function startEmbeddings(req: Request, res: Response) { + try { + log.info("Starting embedding generation system"); + + // Initialize the index service if not already initialized + await indexService.initialize(); + + // Start automatic indexing + await indexService.startEmbeddingGeneration(); + + return { + success: true, + message: "Embedding generation started" + }; + } catch (error: any) { + log.error(`Error starting embeddings: ${error.message || 'Unknown error'}`); + throw new Error(`Failed to start embeddings: ${error.message || 'Unknown error'}`); + } +} + +/** + * Stop embedding generation when AI is disabled + */ +async function stopEmbeddings(req: Request, res: Response) { + try { + log.info("Stopping embedding generation system"); + + // Stop automatic indexing + await indexService.stopEmbeddingGeneration(); + + return { + success: true, + message: "Embedding generation stopped" + }; + } catch (error: any) { + log.error(`Error stopping embeddings: ${error.message || 'Unknown error'}`); + throw new Error(`Failed to stop embeddings: ${error.message || 'Unknown error'}`); + } +} + export default { findSimilarNotes, searchByText, @@ -794,5 +837,7 @@ export default { retryFailedNote, retryAllFailedNotes, rebuildIndex, - getIndexRebuildStatus + getIndexRebuildStatus, + startEmbeddings, + stopEmbeddings }; diff --git a/apps/server/src/routes/routes.ts b/apps/server/src/routes/routes.ts index d9ac2c2f8..73beb28e2 100644 --- a/apps/server/src/routes/routes.ts +++ b/apps/server/src/routes/routes.ts @@ -400,6 +400,8 @@ function register(app: express.Application) { asyncApiRoute(PST, "/api/llm/embeddings/retry-all-failed", embeddingsRoute.retryAllFailedNotes); asyncApiRoute(PST, "/api/llm/embeddings/rebuild-index", embeddingsRoute.rebuildIndex); asyncApiRoute(GET, "/api/llm/embeddings/index-rebuild-status", embeddingsRoute.getIndexRebuildStatus); + asyncApiRoute(PST, "/api/llm/embeddings/start", embeddingsRoute.startEmbeddings); + asyncApiRoute(PST, "/api/llm/embeddings/stop", embeddingsRoute.stopEmbeddings); // LLM provider endpoints - moved under /api/llm/providers hierarchy asyncApiRoute(GET, "/api/llm/providers/ollama/models", ollamaRoute.listModels); diff --git a/apps/server/src/services/llm/ai_service_manager.ts b/apps/server/src/services/llm/ai_service_manager.ts index aeb6e9763..9e84d1e28 100644 --- a/apps/server/src/services/llm/ai_service_manager.ts +++ b/apps/server/src/services/llm/ai_service_manager.ts @@ -605,6 +605,7 @@ export class AIServiceManager implements IAIServiceManager { private setupProviderChangeListener(): void { // List of AI-related options that should trigger service recreation const aiRelatedOptions = [ + 'aiEnabled', 'aiSelectedProvider', 'embeddingSelectedProvider', 'openaiApiKey', @@ -618,10 +619,29 @@ export class AIServiceManager implements IAIServiceManager { 'voyageApiKey' ]; - eventService.subscribe(['entityChanged'], ({ entityName, entity }) => { + eventService.subscribe(['entityChanged'], async ({ entityName, entity }) => { if (entityName === 'options' && entity && aiRelatedOptions.includes(entity.name)) { log.info(`AI-related option '${entity.name}' changed, recreating LLM services`); - this.recreateServices(); + + // Special handling for aiEnabled toggle + if (entity.name === 'aiEnabled') { + const isEnabled = entity.value === 'true'; + + if (isEnabled) { + log.info('AI features enabled, initializing AI service and embeddings'); + // Initialize the AI service + await this.initialize(); + // Initialize embeddings through index service + await indexService.startEmbeddingGeneration(); + } else { + log.info('AI features disabled, stopping embeddings'); + // Stop embeddings through index service + await indexService.stopEmbeddingGeneration(); + } + } else { + // For other AI-related options, just recreate services + this.recreateServices(); + } } }); } diff --git a/apps/server/src/services/llm/embeddings/events.ts b/apps/server/src/services/llm/embeddings/events.ts index a078b2c32..2b8eac7d9 100644 --- a/apps/server/src/services/llm/embeddings/events.ts +++ b/apps/server/src/services/llm/embeddings/events.ts @@ -9,6 +9,9 @@ import becca from "../../../becca/becca.js"; // Add mutex to prevent concurrent processing let isProcessingEmbeddings = false; +// Store interval reference for cleanup +let backgroundProcessingInterval: NodeJS.Timeout | null = null; + /** * Setup event listeners for embedding-related events */ @@ -53,9 +56,15 @@ export function setupEmbeddingEventListeners() { * Setup background processing of the embedding queue */ export async function setupEmbeddingBackgroundProcessing() { + // Clear any existing interval + if (backgroundProcessingInterval) { + clearInterval(backgroundProcessingInterval); + backgroundProcessingInterval = null; + } + const interval = parseInt(await options.getOption('embeddingUpdateInterval') || '200', 10); - setInterval(async () => { + backgroundProcessingInterval = setInterval(async () => { try { // Skip if already processing if (isProcessingEmbeddings) { @@ -78,6 +87,17 @@ export async function setupEmbeddingBackgroundProcessing() { }, interval); } +/** + * Stop background processing of the embedding queue + */ +export function stopEmbeddingBackgroundProcessing() { + if (backgroundProcessingInterval) { + clearInterval(backgroundProcessingInterval); + backgroundProcessingInterval = null; + log.info("Embedding background processing stopped"); + } +} + /** * Initialize embeddings system */ diff --git a/apps/server/src/services/llm/embeddings/index.ts b/apps/server/src/services/llm/embeddings/index.ts index 89d0f711e..2757f8808 100644 --- a/apps/server/src/services/llm/embeddings/index.ts +++ b/apps/server/src/services/llm/embeddings/index.ts @@ -58,6 +58,7 @@ export const processNoteWithChunking = async ( export const { setupEmbeddingEventListeners, setupEmbeddingBackgroundProcessing, + stopEmbeddingBackgroundProcessing, initEmbeddings } = events; @@ -100,6 +101,7 @@ export default { // Event handling setupEmbeddingEventListeners: events.setupEmbeddingEventListeners, setupEmbeddingBackgroundProcessing: events.setupEmbeddingBackgroundProcessing, + stopEmbeddingBackgroundProcessing: events.stopEmbeddingBackgroundProcessing, initEmbeddings: events.initEmbeddings, // Stats and maintenance diff --git a/apps/server/src/services/llm/index_service.ts b/apps/server/src/services/llm/index_service.ts index bbf372860..992739947 100644 --- a/apps/server/src/services/llm/index_service.ts +++ b/apps/server/src/services/llm/index_service.ts @@ -837,6 +837,81 @@ export class IndexService { return false; } } + + /** + * Start embedding generation (called when AI is enabled) + */ + async startEmbeddingGeneration() { + try { + log.info("Starting embedding generation system"); + + // Re-initialize if needed + if (!this.initialized) { + await this.initialize(); + } + + const aiEnabled = options.getOptionOrNull('aiEnabled') === "true"; + if (!aiEnabled) { + log.error("Cannot start embedding generation - AI features are disabled"); + throw new Error("AI features must be enabled first"); + } + + // 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) { + log.info("This instance is not configured to process embeddings"); + return; + } + + // Setup automatic indexing if enabled + if (await options.getOptionBool('embeddingAutoUpdateEnabled')) { + this.setupAutomaticIndexing(); + log.info(`Automatic embedding indexing started ${isSyncServer ? 'as sync server' : 'as client'}`); + } + + // Re-initialize event listeners + this.setupEventListeners(); + + // Start processing the queue immediately + await this.runBatchIndexing(20); + + log.info("Embedding generation started successfully"); + } catch (error: any) { + log.error(`Error starting embedding generation: ${error.message || "Unknown error"}`); + throw error; + } + } + + /** + * Stop embedding generation (called when AI is disabled) + */ + async stopEmbeddingGeneration() { + try { + log.info("Stopping embedding generation system"); + + // Clear automatic indexing interval + if (this.automaticIndexingInterval) { + clearInterval(this.automaticIndexingInterval); + this.automaticIndexingInterval = undefined; + log.info("Automatic indexing stopped"); + } + + // Stop the background processing from embeddings/events.ts + vectorStore.stopEmbeddingBackgroundProcessing(); + + // Mark as not indexing + this.indexingInProgress = false; + this.indexRebuildInProgress = false; + + log.info("Embedding generation stopped successfully"); + } catch (error: any) { + log.error(`Error stopping embedding generation: ${error.message || "Unknown error"}`); + throw error; + } + } } // Create singleton instance