mirror of
				https://github.com/zadam/trilium.git
				synced 2025-10-31 02:16:05 +01:00 
			
		
		
		
	feat(llm): also add functions to clear/unregister embedding providers
This commit is contained in:
		| @@ -43,11 +43,10 @@ export async function initializeEmbeddings() { | |||||||
|         // Reset any stuck embedding queue items from previous server shutdown |         // Reset any stuck embedding queue items from previous server shutdown | ||||||
|         await resetStuckEmbeddingQueue(); |         await resetStuckEmbeddingQueue(); | ||||||
|  |  | ||||||
|         // Initialize default embedding providers |  | ||||||
|         await providerManager.initializeDefaultProviders(); |  | ||||||
|  |  | ||||||
|         // Start the embedding system if AI is enabled |         // Start the embedding system if AI is enabled | ||||||
|         if (await options.getOptionBool('aiEnabled')) { |         if (await options.getOptionBool('aiEnabled')) { | ||||||
|  |             // Initialize default embedding providers when AI is enabled | ||||||
|  |             await providerManager.initializeDefaultProviders(); | ||||||
|             await initEmbeddings(); |             await initEmbeddings(); | ||||||
|             log.info("Embedding system initialized successfully."); |             log.info("Embedding system initialized successfully."); | ||||||
|         } else { |         } else { | ||||||
|   | |||||||
| @@ -845,17 +845,21 @@ export class IndexService { | |||||||
|         try { |         try { | ||||||
|             log.info("Starting embedding generation system"); |             log.info("Starting embedding generation system"); | ||||||
|              |              | ||||||
|             // Re-initialize if needed |  | ||||||
|             if (!this.initialized) { |  | ||||||
|                 await this.initialize(); |  | ||||||
|             } |  | ||||||
|              |  | ||||||
|             const aiEnabled = options.getOptionOrNull('aiEnabled') === "true"; |             const aiEnabled = options.getOptionOrNull('aiEnabled') === "true"; | ||||||
|             if (!aiEnabled) { |             if (!aiEnabled) { | ||||||
|                 log.error("Cannot start embedding generation - AI features are disabled"); |                 log.error("Cannot start embedding generation - AI features are disabled"); | ||||||
|                 throw new Error("AI features must be enabled first"); |                 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 |             // Check if this instance should process embeddings | ||||||
|             const embeddingLocation = await options.getOption('embeddingGenerationLocation') || 'client'; |             const embeddingLocation = await options.getOption('embeddingGenerationLocation') || 'client'; | ||||||
|             const isSyncServer = await this.isSyncServerForEmbeddings(); |             const isSyncServer = await this.isSyncServerForEmbeddings(); | ||||||
| @@ -902,6 +906,9 @@ export class IndexService { | |||||||
|             // Stop the background processing from embeddings/events.ts |             // Stop the background processing from embeddings/events.ts | ||||||
|             vectorStore.stopEmbeddingBackgroundProcessing(); |             vectorStore.stopEmbeddingBackgroundProcessing(); | ||||||
|  |  | ||||||
|  |             // Clear all embedding providers to clean up resources | ||||||
|  |             providerManager.clearAllEmbeddingProviders(); | ||||||
|  |  | ||||||
|             // Mark as not indexing |             // Mark as not indexing | ||||||
|             this.indexingInProgress = false; |             this.indexingInProgress = false; | ||||||
|             this.indexRebuildInProgress = false; |             this.indexRebuildInProgress = false; | ||||||
|   | |||||||
| @@ -86,6 +86,29 @@ export function registerEmbeddingProvider(provider: EmbeddingProvider) { | |||||||
|     log.info(`Registered embedding provider: ${provider.name}`); |     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 |  * Get all registered embedding providers | ||||||
|  */ |  */ | ||||||
| @@ -296,9 +319,9 @@ export async function initializeDefaultProviders() { | |||||||
|             } |             } | ||||||
|         } |         } | ||||||
|  |  | ||||||
|         // Register Ollama provider if base URL is configured |         // Register Ollama embedding provider if embedding base URL is configured | ||||||
|         const ollamaBaseUrl = await options.getOption('ollamaBaseUrl'); |         const ollamaEmbeddingBaseUrl = await options.getOption('ollamaEmbeddingBaseUrl'); | ||||||
|         if (ollamaBaseUrl) { |         if (ollamaEmbeddingBaseUrl) { | ||||||
|             // Use specific embedding models if available |             // Use specific embedding models if available | ||||||
|             const embeddingModel = await options.getOption('ollamaEmbeddingModel'); |             const embeddingModel = await options.getOption('ollamaEmbeddingModel'); | ||||||
|  |  | ||||||
| @@ -308,7 +331,7 @@ export async function initializeDefaultProviders() { | |||||||
|                     model: embeddingModel, |                     model: embeddingModel, | ||||||
|                     dimension: 768, // Initial value, will be updated during initialization |                     dimension: 768, // Initial value, will be updated during initialization | ||||||
|                     type: 'float32', |                     type: 'float32', | ||||||
|                     baseUrl: ollamaBaseUrl |                     baseUrl: ollamaEmbeddingBaseUrl | ||||||
|                 }); |                 }); | ||||||
|  |  | ||||||
|                 // Register the provider |                 // Register the provider | ||||||
| @@ -362,6 +385,8 @@ export async function initializeDefaultProviders() { | |||||||
|  |  | ||||||
| export default { | export default { | ||||||
|     registerEmbeddingProvider, |     registerEmbeddingProvider, | ||||||
|  |     unregisterEmbeddingProvider, | ||||||
|  |     clearAllEmbeddingProviders, | ||||||
|     getEmbeddingProviders, |     getEmbeddingProviders, | ||||||
|     getEmbeddingProvider, |     getEmbeddingProvider, | ||||||
|     getEnabledEmbeddingProviders, |     getEnabledEmbeddingProviders, | ||||||
|   | |||||||
		Reference in New Issue
	
	Block a user