zadam/trilium
1
1
Fork 0
mirror of https://github.com/zadam/trilium.git synced 2025-10-31 11:16:05 +08:00
Code Issues Projects Releases Packages Wiki Activity

feat(llm): also add functions to clear/unregister embedding providers

Browse source
This commit is contained in:
perf3ct 2025-06-05 18:59:32 +00:00
parent 49e123f399
commit c1b10d70b8
No known key found for this signature in database
GPG key ID: 569C4EEC436F5232
3 changed files with 43 additions and 12 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
apps/server/src/services/llm/embeddings/init.ts
View file

@ -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 {

17
apps/server/src/services/llm/index_service.ts
View file

@ -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;

33
apps/server/src/services/llm/providers/providers.ts
View file

@ -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,