update for custom model, ollama

2025-03-20 12:44:53 +08:00
parent 76ecca0da9
commit 570e8d9564
9 changed files with 155 additions and 115 deletions
--- a/src/core/rag/embedding.ts
+++ b/src/core/rag/embedding.ts
@@ -159,10 +159,9 @@ export const getEmbeddingModel = (
 				dangerouslyAllowBrowser: true,
 				baseURL: `${settings.ollamaProvider.baseUrl}/v1`,
 			})
-			const modelInfo = GetEmbeddingModelInfo(settings.embeddingModelProvider, settings.embeddingModelId)
 			return {
 				id: settings.embeddingModelId,
-				dimension: modelInfo.dimensions,
+				dimension: 0,
 				getEmbedding: async (text: string) => {
 					if (!settings.ollamaProvider.baseUrl) {
 						throw new LLMBaseUrlNotSetException(
--- a/src/core/rag/rag-engine.ts
+++ b/src/core/rag/rag-engine.ts
@@ -5,128 +5,145 @@ import { DBManager } from '../../database/database-manager'
 import { VectorManager } from '../../database/modules/vector/vector-manager'
 import { SelectVector } from '../../database/schema'
 import { EmbeddingModel } from '../../types/embedding'
+import { ApiProvider } from '../../types/llm/model'
 import { InfioSettings } from '../../types/settings'

 import { getEmbeddingModel } from './embedding'

 export class RAGEngine {
-  private app: App
-  private settings: InfioSettings
-  private vectorManager: VectorManager
+	private app: App
+	private settings: InfioSettings
+	private vectorManager: VectorManager
 	private embeddingModel: EmbeddingModel | null = null
 	private initialized = false

-  constructor(
-    app: App,
-    settings: InfioSettings,
-    dbManager: DBManager,
-  ) {
-    this.app = app
-    this.settings = settings
-    this.vectorManager = dbManager.getVectorManager()
+	constructor(
+		app: App,
+		settings: InfioSettings,
+		dbManager: DBManager,
+	) {
+		this.app = app
+		this.settings = settings
+		this.vectorManager = dbManager.getVectorManager()
 		this.embeddingModel = getEmbeddingModel(settings)
-  }
+	}

-  setSettings(settings: InfioSettings) {
-    this.settings = settings
-    this.embeddingModel = getEmbeddingModel(settings)
-  }
+	setSettings(settings: InfioSettings) {
+		this.settings = settings
+		this.embeddingModel = getEmbeddingModel(settings)
+	}

-  // TODO: Implement automatic vault re-indexing when settings are changed.
-  // Currently, users must manually re-index the vault.
-  async updateVaultIndex(
-    options: { reindexAll: boolean },
-    onQueryProgressChange?: (queryProgress: QueryProgressState) => void,
+	async initializeDimension(): Promise<void> {
+		if (this.embeddingModel.dimension === 0 && this.settings.embeddingModelProvider === ApiProvider.Ollama) {
+			this.embeddingModel.dimension = (await this.embeddingModel.getEmbedding("hello world")).length
+		}
+	}
+
+	// TODO: Implement automatic vault re-indexing when settings are changed.
+	// Currently, users must manually re-index the vault.
+	async updateVaultIndex(
+		options: { reindexAll: boolean },
+		onQueryProgressChange?: (queryProgress: QueryProgressState) => void,
 	): Promise<void> {
 		if (!this.embeddingModel) {
 			throw new Error('Embedding model is not set')
-    }
-    await this.vectorManager.updateVaultIndex(
+		}
+		await this.initializeDimension()
+
+		await this.vectorManager.updateVaultIndex(
 			this.embeddingModel,
-      {
-        chunkSize: this.settings.ragOptions.chunkSize,
-        excludePatterns: this.settings.ragOptions.excludePatterns,
-        includePatterns: this.settings.ragOptions.includePatterns,
-        reindexAll: options.reindexAll,
-      },
-      (indexProgress) => {
-        onQueryProgressChange?.({
-          type: 'indexing',
-          indexProgress,
-        })
-      },
-    )
+			{
+				chunkSize: this.settings.ragOptions.chunkSize,
+				excludePatterns: this.settings.ragOptions.excludePatterns,
+				includePatterns: this.settings.ragOptions.includePatterns,
+				reindexAll: options.reindexAll,
+			},
+			(indexProgress) => {
+				onQueryProgressChange?.({
+					type: 'indexing',
+					indexProgress,
+				})
+			},
+		)
 		this.initialized = true
 	}

-  async updateFileIndex(file: TFile) {
-    await this.vectorManager.UpdateFileVectorIndex(
-      this.embeddingModel,
-      this.settings.ragOptions.chunkSize,
-      file,
-    )
+	async updateFileIndex(file: TFile) {
+		if (!this.embeddingModel) {
+			throw new Error('Embedding model is not set')
+		}
+
+		await this.initializeDimension()
+
+		await this.vectorManager.UpdateFileVectorIndex(
+			this.embeddingModel,
+			this.settings.ragOptions.chunkSize,
+			file,
+		)
 	}
-	
+
 	async deleteFileIndex(file: TFile) {
+		if (!this.embeddingModel) {
+			throw new Error('Embedding model is not set')
+		}
+
+		await this.initializeDimension()
+
 		await this.vectorManager.DeleteFileVectorIndex(
 			this.embeddingModel,
 			file,
 		)
 	}

-  async processQuery({
-    query,
-    scope,
-    onQueryProgressChange,
-  }: {
-    query: string
-    scope?: {
-      files: string[]
-      folders: string[]
-    }
-    onQueryProgressChange?: (queryProgress: QueryProgressState) => void
-  }): Promise<
-    (Omit<SelectVector, 'embedding'> & {
-      similarity: number
-    })[]
-  > {
-    if (!this.embeddingModel) {
-      throw new Error('Embedding model is not set')
-    }
+	async processQuery({
+		query,
+		scope,
+		onQueryProgressChange,
+	}: {
+		query: string
+		scope?: {
+			files: string[]
+			folders: string[]
+		}
+		onQueryProgressChange?: (queryProgress: QueryProgressState) => void
+	}): Promise<
+		(Omit<SelectVector, 'embedding'> & {
+			similarity: number
+		})[]
+	> {
+		if (!this.embeddingModel) {
+			throw new Error('Embedding model is not set')
+		}
+
+		await this.initializeDimension()

 		if (!this.initialized) {
-      await this.updateVaultIndex({ reindexAll: false }, onQueryProgressChange)
-    }
-    const queryEmbedding = await this.getEmbedding(query)
-    onQueryProgressChange?.({
-      type: 'querying',
+			await this.updateVaultIndex({ reindexAll: false }, onQueryProgressChange)
+		}
+		const queryEmbedding = await this.getEmbedding(query)
+		onQueryProgressChange?.({
+			type: 'querying',
 		})
-		console.log('query, ', {
-			minSimilarity: this.settings.ragOptions.minSimilarity,
-			limit: this.settings.ragOptions.limit,
-			scope,
-		})
-    const queryResult = await this.vectorManager.performSimilaritySearch(
-      queryEmbedding,
-      this.embeddingModel,
-      {
-        minSimilarity: this.settings.ragOptions.minSimilarity,
-        limit: this.settings.ragOptions.limit,
-        scope,
-      },
+		const queryResult = await this.vectorManager.performSimilaritySearch(
+			queryEmbedding,
+			this.embeddingModel,
+			{
+				minSimilarity: this.settings.ragOptions.minSimilarity,
+				limit: this.settings.ragOptions.limit,
+				scope,
+			},
 		)
-		console.log('queryResult', queryResult)
-    onQueryProgressChange?.({
-      type: 'querying-done',
-      queryResult,
-    })
-    return queryResult
-  }
+		onQueryProgressChange?.({
+			type: 'querying-done',
+			queryResult,
+		})
+		return queryResult
+	}

-  async getEmbedding(query: string): Promise<number[]> {
-    if (!this.embeddingModel) {
-      throw new Error('Embedding model is not set')
-    }
-    return this.embeddingModel.getEmbedding(query)
-  }
+	async getEmbedding(query: string): Promise<number[]> {
+		if (!this.embeddingModel) {
+			throw new Error('Embedding model is not set')
+		}
+		return this.embeddingModel.getEmbedding(query)
+	}
 }