File size: 5,507 Bytes
55847d9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 | import { ipcMain, dialog } from 'electron'
import fs from 'node:fs'
import {
importDocument, importFolder, importText, searchKnowledge, searchKnowledgeFTS,
listDocuments, removeDocument, getKnowledgeStats,
getVectorlessCount, backfillVectors,
} from '../knowledge-base'
import { readJsonFile, GLOBAL_CONFIG_PATH, DEFAULT_GLOBAL_CONFIG, MODELS_CONFIG_PATH, RECENT_PROJECTS_PATH } from '../utils/config-utils'
import { GlobalConfig, ModelProfile } from '../../src/shared/ipc-channels'
function getEmbeddingConfig(): { protocol: 'openai' | 'gemini'; model: { baseUrl: string; apiKey: string; modelName: string } } | null {
const config = readJsonFile<GlobalConfig>(GLOBAL_CONFIG_PATH, DEFAULT_GLOBAL_CONFIG)
const targetModelId = config.defaultEmbeddingModelId || config.defaultModelId
if (!targetModelId) return null
const models = readJsonFile<ModelProfile[]>(MODELS_CONFIG_PATH, [])
const model = models.find((m) => m.id === targetModelId)
if (!model) return null
return {
protocol: model.protocol as 'openai' | 'gemini',
model: { baseUrl: model.baseUrl, apiKey: model.apiKey, modelName: model.modelName },
}
}
function getCurrentProjectPath(): string | null {
try {
const recent = JSON.parse(fs.readFileSync(RECENT_PROJECTS_PATH, 'utf-8')) as Array<{ path: string }>
return recent[0]?.path ?? null
} catch { return null }
}
export function registerKBController() {
ipcMain.handle('kb:import-document', async (_event, filePath: string) => {
const embConfig = getEmbeddingConfig()
const projectPath = getCurrentProjectPath()
if (!projectPath) return { success: false, error: '未打开项目' }
const protocol = embConfig?.protocol ?? 'openai'
const model = embConfig?.model ?? { baseUrl: '', apiKey: '' }
return importDocument(filePath, projectPath, protocol, model)
})
ipcMain.handle('kb:import-folder', async (_event, folderPath: string) => {
const embConfig = getEmbeddingConfig()
const projectPath = getCurrentProjectPath()
if (!projectPath) return { success: false, error: '未打开项目' }
const protocol = embConfig?.protocol ?? 'openai'
const model = embConfig?.model ?? { baseUrl: '', apiKey: '' }
return importFolder(folderPath, projectPath, protocol, model)
})
ipcMain.handle('kb:import-text', async (_event, text: string, fileName: string, projectPath: string) => {
const embConfig = getEmbeddingConfig()
const protocol = embConfig?.protocol ?? 'openai'
const model = embConfig?.model ?? { baseUrl: '', apiKey: '' }
return importText(text, fileName, projectPath, protocol, model)
})
ipcMain.handle('kb:search', async (_event, query: string, topK?: number) => {
const embConfig = getEmbeddingConfig()
const projectPath = getCurrentProjectPath()
if (!projectPath) return []
if (embConfig) {
return searchKnowledge(query, projectPath, embConfig.protocol, embConfig.model, topK ?? 5)
}
return searchKnowledgeFTS(query, projectPath, topK ?? 5)
})
ipcMain.handle('kb:search-with-scope', async (_event, query: string, fromChapter: number, toChapter: number, topK?: number) => {
const embConfig = getEmbeddingConfig()
const projectPath = getCurrentProjectPath()
if (!projectPath) return []
const scope: [number, number] = [fromChapter, toChapter]
if (embConfig) {
return searchKnowledge(query, projectPath, embConfig.protocol, embConfig.model, topK ?? 5, scope)
}
return searchKnowledgeFTS(query, projectPath, topK ?? 5, scope)
})
ipcMain.handle('kb:list-documents', async () => {
const projectPath = getCurrentProjectPath()
if (!projectPath) return []
return listDocuments(projectPath)
})
ipcMain.handle('kb:remove-document', async (_event, docId: string) => {
const projectPath = getCurrentProjectPath()
if (!projectPath) return { success: false }
return { success: removeDocument(docId, projectPath) }
})
ipcMain.handle('kb:stats', async () => {
const projectPath = getCurrentProjectPath()
if (!projectPath) return { documentCount: 0, totalChunks: 0, vectorDimension: 0 }
return getKnowledgeStats(projectPath)
})
ipcMain.handle('kb:get-vectorless-count', async () => {
const projectPath = getCurrentProjectPath()
if (!projectPath) return { count: 0 }
return getVectorlessCount(projectPath)
})
ipcMain.handle('kb:backfill-vectors', async () => {
const embConfig = getEmbeddingConfig()
if (!embConfig) return { success: false, processed: 0, failed: 0, error: '未配置 Embedding 模型' }
const projectPath = getCurrentProjectPath()
if (!projectPath) return { success: false, processed: 0, failed: 0, error: '未打开项目' }
return backfillVectors(projectPath, embConfig.protocol, embConfig.model)
})
ipcMain.handle('dialog:select-files', async () => {
const result = await dialog.showOpenDialog({
properties: ['openFile', 'multiSelections'],
title: '选择要导入的文档',
filters: [{ name: '文本文件', extensions: ['txt', 'md', 'markdown'] }],
})
if (result.canceled || result.filePaths.length === 0) return null
return result.filePaths
})
ipcMain.handle('dialog:select-import-folder', async () => {
const result = await dialog.showOpenDialog({
properties: ['openDirectory'],
title: '选择要批量导入的文件夹',
})
if (result.canceled || result.filePaths.length === 0) return null
return result.filePaths[0]
})
}
|