From c1b10d70b8b65a6d9fa8853972bf66acf09a0d1f Mon Sep 17 00:00:00 2001
From: perf3ct <jonfuller2012@gmail.com>
Date: Thu, 5 Jun 2025 18:59:32 +0000
Subject: [PATCH] feat(llm): also add functions to clear/unregister embedding
 providers

---
 .../src/services/llm/embeddings/init.ts       |  5 ++-
 apps/server/src/services/llm/index_service.ts | 17 +++++++---
 .../src/services/llm/providers/providers.ts   | 33 ++++++++++++++++---
 3 files changed, 43 insertions(+), 12 deletions(-)

diff --git a/apps/server/src/services/llm/embeddings/init.ts b/apps/server/src/services/llm/embeddings/init.ts
index 50724c05e..94188fa69 100644
--- a/apps/server/src/services/llm/embeddings/init.ts
+++ b/apps/server/src/services/llm/embeddings/init.ts
@@ -43,11 +43,10 @@ export async function initializeEmbeddings() {
         // Reset any stuck embedding queue items from previous server shutdown
         await resetStuckEmbeddingQueue();
 
-        // Initialize default embedding providers
-        await providerManager.initializeDefaultProviders();
-
         // Start the embedding system if AI is enabled
         if (await options.getOptionBool('aiEnabled')) {
+            // Initialize default embedding providers when AI is enabled
+            await providerManager.initializeDefaultProviders();
             await initEmbeddings();
             log.info("Embedding system initialized successfully.");
         } else {
diff --git a/apps/server/src/services/llm/index_service.ts b/apps/server/src/services/llm/index_service.ts
index 992739947..a179431ee 100644
--- a/apps/server/src/services/llm/index_service.ts
+++ b/apps/server/src/services/llm/index_service.ts
@@ -845,17 +845,21 @@ export class IndexService {
         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");
             }
 
+            // Re-initialize providers first in case they weren't available when server started
+            log.info("Re-initializing embedding providers");
+            await providerManager.initializeDefaultProviders();
+
+            // Re-initialize if needed
+            if (!this.initialized) {
+                await this.initialize();
+            }
+
             // Check if this instance should process embeddings
             const embeddingLocation = await options.getOption('embeddingGenerationLocation') || 'client';
             const isSyncServer = await this.isSyncServerForEmbeddings();
@@ -902,6 +906,9 @@ export class IndexService {
             // Stop the background processing from embeddings/events.ts
             vectorStore.stopEmbeddingBackgroundProcessing();
 
+            // Clear all embedding providers to clean up resources
+            providerManager.clearAllEmbeddingProviders();
+
             // Mark as not indexing
             this.indexingInProgress = false;
             this.indexRebuildInProgress = false;
diff --git a/apps/server/src/services/llm/providers/providers.ts b/apps/server/src/services/llm/providers/providers.ts
index f4d69801c..5b4955ce6 100644
--- a/apps/server/src/services/llm/providers/providers.ts
+++ b/apps/server/src/services/llm/providers/providers.ts
@@ -86,6 +86,29 @@ export function registerEmbeddingProvider(provider: EmbeddingProvider) {
     log.info(`Registered embedding provider: ${provider.name}`);
 }
 
+/**
+ * Unregister an embedding provider
+ */
+export function unregisterEmbeddingProvider(name: string): boolean {
+    const existed = providers.has(name);
+    if (existed) {
+        providers.delete(name);
+        log.info(`Unregistered embedding provider: ${name}`);
+    }
+    return existed;
+}
+
+/**
+ * Clear all embedding providers
+ */
+export function clearAllEmbeddingProviders(): void {
+    const providerNames = Array.from(providers.keys());
+    providers.clear();
+    if (providerNames.length > 0) {
+        log.info(`Cleared all embedding providers: ${providerNames.join(', ')}`);
+    }
+}
+
 /**
  * Get all registered embedding providers
  */
@@ -296,9 +319,9 @@ export async function initializeDefaultProviders() {
             }
         }
 
-        // Register Ollama provider if base URL is configured
-        const ollamaBaseUrl = await options.getOption('ollamaBaseUrl');
-        if (ollamaBaseUrl) {
+        // Register Ollama embedding provider if embedding base URL is configured
+        const ollamaEmbeddingBaseUrl = await options.getOption('ollamaEmbeddingBaseUrl');
+        if (ollamaEmbeddingBaseUrl) {
             // Use specific embedding models if available
             const embeddingModel = await options.getOption('ollamaEmbeddingModel');
 
@@ -308,7 +331,7 @@ export async function initializeDefaultProviders() {
                     model: embeddingModel,
                     dimension: 768, // Initial value, will be updated during initialization
                     type: 'float32',
-                    baseUrl: ollamaBaseUrl
+                    baseUrl: ollamaEmbeddingBaseUrl
                 });
 
                 // Register the provider
@@ -362,6 +385,8 @@ export async function initializeDefaultProviders() {
 
 export default {
     registerEmbeddingProvider,
+    unregisterEmbeddingProvider,
+    clearAllEmbeddingProviders,
     getEmbeddingProviders,
     getEmbeddingProvider,
     getEnabledEmbeddingProviders,