gitea/infio-copilot
1
0
Fork 0
mirror of https://github.com/EthanMarti/infio-copilot.git synced 2026-01-16 16:31:56 +00:00
Code Issues Packages Projects Releases Wiki Activity

更新嵌入管理器以支持 GPU 加速，调整批处理大小，优化内容处理逻辑，并添加获取数据库最大修改时间的功能以提高文件索引效率。同时修复了向量管理器中的类型问题，确保模型加载和嵌入过程的稳定性。

Browse Source
This commit is contained in:
duanfuxiang 2025-07-05 07:40:54 +08:00
parent 558e3b3fe4
commit c657a50563
7 changed files with 398 additions and 340 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
src/core/rag/embedding.ts
View File

@ -64,7 +64,7 @@ export const getEmbeddingModel = (
// 确保模型已加载
if (!embeddingManager.modelLoaded || embeddingManager.currentModel !== settings.embeddingModelId) {
console.log(`Loading model: ${settings.embeddingModelId}`)
await embeddingManager.loadModel(settings.embeddingModelId, false)
await embeddingManager.loadModel(settings.embeddingModelId, true)
}
const results = await embeddingManager.embedBatch(texts)

161
src/database/modules/vector/vector-manager.ts
View File

@ -27,7 +27,7 @@ export class VectorManager {
constructor(app: App, dbManager: DBManager) {
this.app = app
this.dbManager = dbManager
this.repository = new VectorRepository(app, dbManager.getPgClient())
this.repository = new VectorRepository(app, dbManager.getPgClient() as any)
}
async performSimilaritySearch(
@ -88,6 +88,7 @@ export class VectorManager {
): Promise<void> {
let filesToIndex: TFile[]
if (options.reindexAll) {
console.log("updateVaultIndex reindexAll")
filesToIndex = await this.getFilesToIndex({
embeddingModel: embeddingModel,
excludePatterns: options.excludePatterns,
@ -96,17 +97,22 @@ export class VectorManager {
})
await this.repository.clearAllVectors(embeddingModel)
} else {
console.log("updateVaultIndex for update files")
await this.cleanVectorsForDeletedFiles(embeddingModel)
console.log("updateVaultIndex cleanVectorsForDeletedFiles")
filesToIndex = await this.getFilesToIndex({
embeddingModel: embeddingModel,
excludePatterns: options.excludePatterns,
includePatterns: options.includePatterns,
})
console.log("get files to index: ", filesToIndex.length)
await this.repository.deleteVectorsForMultipleFiles(
filesToIndex.map((file) => file.path),
embeddingModel,
)
console.log("delete vectors for multiple files: ", filesToIndex.length)
}
console.log("get files to index: ", filesToIndex.length)
if (filesToIndex.length === 0) {
return
@ -131,6 +137,7 @@ export class VectorManager {
"",
],
});
console.log("textSplitter chunkSize: ", options.chunkSize, "overlap: ", overlap)
const skippedFiles: string[] = []
const contentChunks: InsertVector[] = (
@ -145,15 +152,16 @@ export class VectorManager {
])
return fileDocuments
.map((chunk): InsertVector | null => {
const content = removeMarkdown(chunk.pageContent).replace(/\0/g, '')
if (!content || content.trim().length === 0) {
// 保存原始内容，不在此处调用 removeMarkdown
const rawContent = chunk.pageContent.replace(/\0/g, '')
if (!rawContent || rawContent.trim().length === 0) {
console.log("skipped chunk", chunk.pageContent)
return null
}
return {
path: file.path,
mtime: file.stat.mtime,
content,
content: rawContent, // 保存原始内容
embedding: [],
metadata: {
startLine: Number(chunk.metadata.loc.lines.from),
@ -171,6 +179,8 @@ export class VectorManager {
)
).flat()
console.log("contentChunks: ", contentChunks.length)
if (skippedFiles.length > 0) {
console.warn(`跳过了 ${skippedFiles.length} 个有问题的文件:`, skippedFiles)
new Notice(`跳过了 ${skippedFiles.length} 个有问题的文件`)
@ -195,22 +205,33 @@ export class VectorManager {
for (let i = 0; i < contentChunks.length; i += embeddingBatchSize) {
batchCount++
const batchChunks = contentChunks.slice(i, Math.min(i + embeddingBatchSize, contentChunks.length))
const batchTexts = batchChunks.map(chunk => chunk.content)
const embeddedBatch: InsertVector[] = []
await backOff(
async () => {
// 在嵌入之前处理 markdown，只处理一次
const cleanedBatchData = batchChunks.map(chunk => {
const cleanContent = removeMarkdown(chunk.content).replace(/\0/g, '')
return { chunk, cleanContent }
}).filter(({ cleanContent }) => cleanContent && cleanContent.trim().length > 0)
if (cleanedBatchData.length === 0) {
return
}
const batchTexts = cleanedBatchData.map(({ cleanContent }) => cleanContent)
const batchEmbeddings = await embeddingModel.getBatchEmbeddings(batchTexts)
// 合并embedding结果到chunk数据
for (let j = 0; j < batchChunks.length; j++) {
for (let j = 0; j < cleanedBatchData.length; j++) {
const { chunk, cleanContent } = cleanedBatchData[j]
const embeddedChunk: InsertVector = {
path: batchChunks[j].path,
mtime: batchChunks[j].mtime,
content: batchChunks[j].content,
path: chunk.path,
mtime: chunk.mtime,
content: cleanContent, // 使用已经清理过的内容
embedding: batchEmbeddings[j],
metadata: batchChunks[j].metadata,
metadata: chunk.metadata,
}
embeddedBatch.push(embeddedChunk)
}
@ -263,11 +284,18 @@ export class VectorManager {
try {
await backOff(
async () => {
const embedding = await embeddingModel.getEmbedding(chunk.content)
// 在嵌入之前处理 markdown
const cleanContent = removeMarkdown(chunk.content).replace(/\0/g, '')
// 跳过清理后为空的内容
if (!cleanContent || cleanContent.trim().length === 0) {
return
}
const embedding = await embeddingModel.getEmbedding(cleanContent)
const embeddedChunk = {
path: chunk.path,
mtime: chunk.mtime,
content: chunk.content,
content: cleanContent, // 使用清理后的内容
embedding,
metadata: chunk.metadata,
}
@ -339,9 +367,23 @@ export class VectorManager {
)
// Embed the files
const textSplitter = new MarkdownTextSplitter({
const overlap = Math.floor(chunkSize * 0.15)
const textSplitter = new RecursiveCharacterTextSplitter({
chunkSize: chunkSize,
chunkOverlap: Math.floor(chunkSize * 0.15)
chunkOverlap: overlap,
separators: [
"\n\n",
"\n",
".",
",",
" ",
"\u200b", // Zero-width space
"\uff0c", // Fullwidth comma
"\u3001", // Ideographic comma
"\uff0e", // Fullwidth full stop
"\u3002", // Ideographic full stop
"",
],
});
let fileContent = await this.app.vault.cachedRead(file)
// 清理null字节，防止PostgreSQL UTF8编码错误
@ -352,14 +394,15 @@ export class VectorManager {
const contentChunks: InsertVector[] = fileDocuments
.map((chunk): InsertVector | null => {
const content = removeMarkdown(chunk.pageContent).replace(/\0/g, '')
if (!content || content.trim().length === 0) {
// 保存原始内容，不在此处调用 removeMarkdown
const rawContent = String(chunk.pageContent || '').replace(/\0/g, '')
if (!rawContent || rawContent.trim().length === 0) {
return null
}
return {
path: file.path,
mtime: file.stat.mtime,
content,
content: rawContent, // 保存原始内容
embedding: [],
metadata: {
startLine: Number(chunk.metadata.loc.lines.from),
@ -382,22 +425,33 @@ export class VectorManager {
batchCount++
console.log(`Embedding batch ${batchCount} of ${Math.ceil(contentChunks.length / embeddingBatchSize)}`)
const batchChunks = contentChunks.slice(i, Math.min(i + embeddingBatchSize, contentChunks.length))
const batchTexts = batchChunks.map(chunk => chunk.content)
const embeddedBatch: InsertVector[] = []
await backOff(
async () => {
// 在嵌入之前处理 markdown，只处理一次
const cleanedBatchData = batchChunks.map(chunk => {
const cleanContent = removeMarkdown(chunk.content).replace(/\0/g, '')
return { chunk, cleanContent }
}).filter(({ cleanContent }) => cleanContent && cleanContent.trim().length > 0)
if (cleanedBatchData.length === 0) {
return
}
const batchTexts = cleanedBatchData.map(({ cleanContent }) => cleanContent)
const batchEmbeddings = await embeddingModel.getBatchEmbeddings(batchTexts)
// 合并embedding结果到chunk数据
for (let j = 0; j < batchChunks.length; j++) {
for (let j = 0; j < cleanedBatchData.length; j++) {
const { chunk, cleanContent } = cleanedBatchData[j]
const embeddedChunk: InsertVector = {
path: batchChunks[j].path,
mtime: batchChunks[j].mtime,
content: batchChunks[j].content,
path: chunk.path,
mtime: chunk.mtime,
content: cleanContent, // 使用已经清理过的内容
embedding: batchEmbeddings[j],
metadata: batchChunks[j].metadata,
metadata: chunk.metadata,
}
embeddedBatch.push(embeddedChunk)
}
@ -443,11 +497,18 @@ export class VectorManager {
try {
await backOff(
async () => {
const embedding = await embeddingModel.getEmbedding(chunk.content)
// 在嵌入之前处理 markdown
const cleanContent = removeMarkdown(chunk.content).replace(/\0/g, '')
// 跳过清理后为空的内容
if (!cleanContent || cleanContent.trim().length === 0) {
return
}
const embedding = await embeddingModel.getEmbedding(cleanContent)
const embeddedChunk = {
path: chunk.path,
mtime: chunk.mtime,
content: chunk.content,
content: cleanContent, // 使用清理后的内容
embedding,
metadata: chunk.metadata,
}
@ -524,6 +585,7 @@ export class VectorManager {
reindexAll?: boolean
}): Promise<TFile[]> {
let filesToIndex = this.app.vault.getMarkdownFiles()
console.log("get all vault files: ", filesToIndex.length)
filesToIndex = filesToIndex.filter((file) => {
return !excludePatterns.some((pattern) => minimatch(file.path, pattern))
@ -538,39 +600,24 @@ export class VectorManager {
if (reindexAll) {
return filesToIndex
}
// Check for updated or new files
filesToIndex = await Promise.all(
filesToIndex.map(async (file) => {
try {
const fileChunks = await this.repository.getVectorsByFilePath(
file.path,
embeddingModel,
)
if (fileChunks.length === 0) {
// File is not indexed, so we need to index it
let fileContent = await this.app.vault.cachedRead(file)
// 清理null字节，防止PostgreSQL UTF8编码错误
fileContent = fileContent.replace(/\0/g, '')
if (fileContent.length === 0) {
// Ignore empty files
return null
}
return file
}
const outOfDate = file.stat.mtime > fileChunks[0].mtime
if (outOfDate) {
// File has changed, so we need to re-index it
console.log("File has changed, so we need to re-index it", file.path)
return file
}
return null
} catch (error) {
console.warn(`跳过文件 ${file.path}:`, error.message)
return null
}
}),
).then((files) => files.filter(Boolean))
// 优化流程：使用数据库最大mtime来过滤需要更新的文件
try {
const maxMtime = await this.repository.getMaxMtime(embeddingModel)
console.log("Database max mtime:", maxMtime)
if (maxMtime === null) {
// 数据库中没有任何向量，需要索引所有文件
return filesToIndex
}
// 筛选出在数据库最后更新时间之后修改的文件
return filesToIndex.filter((file) => {
return file.stat.mtime > maxMtime
})
} catch (error) {
console.error("Error getting max mtime from database:", error)
return []
}
}
}

11
src/database/modules/vector/vector-repository.ts
View File

@ -33,6 +33,17 @@ export class VectorRepository {
return result.rows.map((row: { path: string }) => row.path)
}
async getMaxMtime(embeddingModel: EmbeddingModel): Promise<number | null> {
if (!this.db) {
throw new DatabaseNotInitializedException()
}
const tableName = this.getTableName(embeddingModel)
const result = await this.db.query<{ max_mtime: number | null }>(
`SELECT MAX(mtime) as max_mtime FROM "${tableName}"`
)
return result.rows[0]?.max_mtime || null
}
async getVectorsByFilePath(
filePath: string,
embeddingModel: EmbeddingModel,

1
src/embedworker/EmbeddingManager.ts
View File

@ -243,6 +243,7 @@ export class EmbeddingManager {
*/
public getSupportedModels(): string[] {
return [
'TaylorAI/bge-micro-v2',
'Xenova/all-MiniLM-L6-v2',
'Xenova/bge-small-en-v1.5',
'Xenova/bge-base-en-v1.5',

4
src/embedworker/embed.worker.ts
View File

@ -48,7 +48,7 @@ async function loadTransformers() {
env.allowRemoteModels = true;
// 配置 WASM 后端 - 修复线程配置
env.backends.onnx.wasm.numThreads = 4; // 在 Worker 中使用单线程，避免竞态条件
env.backends.onnx.wasm.numThreads = 1; // 在 Worker 中使用单线程，避免竞态条件
env.backends.onnx.wasm.simd = true;
// 禁用 Node.js 特定功能
@ -201,7 +201,7 @@ async function embedBatch(inputs: EmbedInput[]): Promise<EmbedResult[]> {
}
// 批处理大小（可以根据需要调整）
const batchSize = 1;
const batchSize = 8;
if (filteredInputs.length > batchSize) {
console.log(`Processing ${filteredInputs.length} inputs in batches of ${batchSize}`);

3
src/utils/api.ts
View File

@ -1641,6 +1641,7 @@ export const localProviderDefaultAutoCompleteModelId = null // this is not suppo
export const localProviderDefaultEmbeddingModelId: keyof typeof localProviderEmbeddingModels = "TaylorAI/bge-micro-v2"
export const localProviderEmbeddingModels = {
'TaylorAI/bge-micro-v2': { dimensions: 384, description: 'BGE-micro-v2 (本地，512令牌，384维)' },
'Xenova/all-MiniLM-L6-v2': { dimensions: 384, description: 'All-MiniLM-L6-v2 (推荐，轻量级)' },
'Xenova/bge-small-en-v1.5': { dimensions: 384, description: 'BGE-small-en-v1.5' },
'Xenova/bge-base-en-v1.5': { dimensions: 768, description: 'BGE-base-en-v1.5 (更高质量)' },
@ -1651,8 +1652,6 @@ export const localProviderEmbeddingModels = {
'Xenova/gte-small': { dimensions: 384, description: 'GTE-small' },
'Xenova/e5-small-v2': { dimensions: 384, description: 'E5-small-v2' },
'Xenova/e5-base-v2': { dimensions: 768, description: 'E5-base-v2 (更高质量)' },
// 新增的模型
'TaylorAI/bge-micro-v2': { dimensions: 384, description: 'BGE-micro-v2 (本地，512令牌，384维)' },
'Snowflake/snowflake-arctic-embed-xs': { dimensions: 384, description: 'Snowflake Arctic Embed XS (本地，512令牌，384维)' },
'Snowflake/snowflake-arctic-embed-s': { dimensions: 384, description: 'Snowflake Arctic Embed Small (本地，512令牌，384维)' },
'Snowflake/snowflake-arctic-embed-m': { dimensions: 768, description: 'Snowflake Arctic Embed Medium (本地，512令牌，768维)' },