feat: remove faiss database (#10178)
This commit is contained in:
Chen Tao
+331
-93
@@ -1,3 +1,18 @@
|
||||
/**
|
||||
* Knowledge Service - Manages knowledge bases using RAG (Retrieval-Augmented Generation)
|
||||
*
|
||||
* This service handles creation, management, and querying of knowledge bases from various sources
|
||||
* including files, directories, URLs, sitemaps, and notes.
|
||||
*
|
||||
* Features:
|
||||
* - Concurrent task processing with workload management
|
||||
* - Multiple data source support
|
||||
* - Vector database integration
|
||||
*
|
||||
* For detailed documentation, see:
|
||||
* @see {@link ../../../docs/technical/KnowledgeService.md}
|
||||
*/
|
||||
|
||||
import * as fs from 'node:fs'
|
||||
import path from 'node:path'
|
||||
|
||||
@@ -9,32 +24,87 @@ import { loggerService } from '@logger'
|
||||
import Embeddings from '@main/knowledge/embedjs/embeddings/Embeddings'
|
||||
import { addFileLoader } from '@main/knowledge/embedjs/loader'
|
||||
import { NoteLoader } from '@main/knowledge/embedjs/loader/noteLoader'
|
||||
import { preprocessingService } from '@main/knowledge/preprocess/PreprocessingService'
|
||||
import PreprocessProvider from '@main/knowledge/preprocess/PreprocessProvider'
|
||||
import Reranker from '@main/knowledge/reranker/Reranker'
|
||||
import { fileStorage } from '@main/services/FileStorage'
|
||||
import { windowService } from '@main/services/WindowService'
|
||||
import { getDataPath } from '@main/utils'
|
||||
import { getAllFiles } from '@main/utils/file'
|
||||
import { TraceMethod } from '@mcp-trace/trace-core'
|
||||
import { MB } from '@shared/config/constant'
|
||||
import { LoaderReturn } from '@shared/config/types'
|
||||
import type { LoaderReturn } from '@shared/config/types'
|
||||
import { IpcChannel } from '@shared/IpcChannel'
|
||||
import { FileMetadata, KnowledgeBaseParams, KnowledgeSearchResult } from '@types'
|
||||
import { FileMetadata, KnowledgeBaseParams, KnowledgeItem, KnowledgeSearchResult } from '@types'
|
||||
import { v4 as uuidv4 } from 'uuid'
|
||||
|
||||
import { windowService } from '../WindowService'
|
||||
import {
|
||||
IKnowledgeFramework,
|
||||
KnowledgeBaseAddItemOptionsNonNullableAttribute,
|
||||
LoaderDoneReturn,
|
||||
LoaderTask,
|
||||
LoaderTaskItem,
|
||||
LoaderTaskItemState
|
||||
} from './IKnowledgeFramework'
|
||||
|
||||
const logger = loggerService.withContext('MainKnowledgeService')
|
||||
|
||||
export class EmbedJsFramework implements IKnowledgeFramework {
|
||||
private storageDir: string
|
||||
private ragApplications: Map<string, RAGApplication> = new Map()
|
||||
private pendingDeleteFile: string
|
||||
private dbInstances: Map<string, LibSqlDb> = new Map()
|
||||
export interface KnowledgeBaseAddItemOptions {
|
||||
base: KnowledgeBaseParams
|
||||
item: KnowledgeItem
|
||||
forceReload?: boolean
|
||||
userId?: string
|
||||
}
|
||||
|
||||
interface KnowledgeBaseAddItemOptionsNonNullableAttribute {
|
||||
base: KnowledgeBaseParams
|
||||
item: KnowledgeItem
|
||||
forceReload: boolean
|
||||
userId: string
|
||||
}
|
||||
|
||||
interface EvaluateTaskWorkload {
|
||||
workload: number
|
||||
}
|
||||
|
||||
type LoaderDoneReturn = LoaderReturn | null
|
||||
|
||||
enum LoaderTaskItemState {
|
||||
PENDING,
|
||||
PROCESSING,
|
||||
DONE
|
||||
}
|
||||
|
||||
interface LoaderTaskItem {
|
||||
state: LoaderTaskItemState
|
||||
task: () => Promise<unknown>
|
||||
evaluateTaskWorkload: EvaluateTaskWorkload
|
||||
}
|
||||
|
||||
interface LoaderTask {
|
||||
loaderTasks: LoaderTaskItem[]
|
||||
loaderDoneReturn: LoaderDoneReturn
|
||||
}
|
||||
|
||||
interface LoaderTaskOfSet {
|
||||
loaderTasks: Set<LoaderTaskItem>
|
||||
loaderDoneReturn: LoaderDoneReturn
|
||||
}
|
||||
|
||||
interface QueueTaskItem {
|
||||
taskPromise: () => Promise<unknown>
|
||||
resolve: () => void
|
||||
evaluateTaskWorkload: EvaluateTaskWorkload
|
||||
}
|
||||
|
||||
const loaderTaskIntoOfSet = (loaderTask: LoaderTask): LoaderTaskOfSet => {
|
||||
return {
|
||||
loaderTasks: new Set(loaderTask.loaderTasks),
|
||||
loaderDoneReturn: loaderTask.loaderDoneReturn
|
||||
}
|
||||
}
|
||||
|
||||
class KnowledgeService {
|
||||
private storageDir = path.join(getDataPath(), 'KnowledgeBase')
|
||||
private pendingDeleteFile = path.join(this.storageDir, 'knowledge_pending_delete.json')
|
||||
// Byte based
|
||||
private workload = 0
|
||||
private processingItemCount = 0
|
||||
private knowledgeItemProcessingQueueMappingPromise: Map<LoaderTaskOfSet, () => void> = new Map()
|
||||
private ragApplications: Map<string, RAGApplication> = new Map()
|
||||
private dbInstances: Map<string, LibSqlDb> = new Map()
|
||||
private static MAXIMUM_WORKLOAD = 80 * MB
|
||||
private static MAXIMUM_PROCESSING_ITEM_COUNT = 30
|
||||
private static ERROR_LOADER_RETURN: LoaderReturn = {
|
||||
entriesAdded: 0,
|
||||
uniqueId: '',
|
||||
@@ -43,9 +113,7 @@ export class EmbedJsFramework implements IKnowledgeFramework {
|
||||
status: 'failed'
|
||||
}
|
||||
|
||||
constructor(storageDir: string) {
|
||||
this.storageDir = storageDir
|
||||
this.pendingDeleteFile = path.join(this.storageDir, 'knowledge_pending_delete.json')
|
||||
constructor() {
|
||||
this.initStorageDir()
|
||||
this.cleanupOnStartup()
|
||||
}
|
||||
@@ -160,28 +228,33 @@ export class EmbedJsFramework implements IKnowledgeFramework {
|
||||
logger.info(`Startup cleanup completed: ${deletedCount}/${pendingDeleteIds.length} knowledge bases deleted`)
|
||||
}
|
||||
|
||||
private async getRagApplication(base: KnowledgeBaseParams): Promise<RAGApplication> {
|
||||
if (this.ragApplications.has(base.id)) {
|
||||
return this.ragApplications.get(base.id)!
|
||||
private getRagApplication = async ({
|
||||
id,
|
||||
embedApiClient,
|
||||
dimensions,
|
||||
documentCount
|
||||
}: KnowledgeBaseParams): Promise<RAGApplication> => {
|
||||
if (this.ragApplications.has(id)) {
|
||||
return this.ragApplications.get(id)!
|
||||
}
|
||||
|
||||
let ragApplication: RAGApplication
|
||||
const embeddings = new Embeddings({
|
||||
embedApiClient: base.embedApiClient,
|
||||
dimensions: base.dimensions
|
||||
embedApiClient,
|
||||
dimensions
|
||||
})
|
||||
try {
|
||||
const libSqlDb = new LibSqlDb({ path: path.join(this.storageDir, base.id) })
|
||||
const libSqlDb = new LibSqlDb({ path: path.join(this.storageDir, id) })
|
||||
// Save database instance for later closing
|
||||
this.dbInstances.set(base.id, libSqlDb)
|
||||
this.dbInstances.set(id, libSqlDb)
|
||||
|
||||
ragApplication = await new RAGApplicationBuilder()
|
||||
.setModel('NO_MODEL')
|
||||
.setEmbeddingModel(embeddings)
|
||||
.setVectorDatabase(libSqlDb)
|
||||
.setSearchResultCount(base.documentCount || 30)
|
||||
.setSearchResultCount(documentCount || 30)
|
||||
.build()
|
||||
this.ragApplications.set(base.id, ragApplication)
|
||||
this.ragApplications.set(id, ragApplication)
|
||||
} catch (e) {
|
||||
logger.error('Failed to create RAGApplication:', e as Error)
|
||||
throw new Error(`Failed to create RAGApplication: ${e}`)
|
||||
@@ -189,14 +262,17 @@ export class EmbedJsFramework implements IKnowledgeFramework {
|
||||
|
||||
return ragApplication
|
||||
}
|
||||
async initialize(base: KnowledgeBaseParams): Promise<void> {
|
||||
|
||||
public create = async (_: Electron.IpcMainInvokeEvent, base: KnowledgeBaseParams): Promise<void> => {
|
||||
await this.getRagApplication(base)
|
||||
}
|
||||
async reset(base: KnowledgeBaseParams): Promise<void> {
|
||||
const ragApp = await this.getRagApplication(base)
|
||||
await ragApp.reset()
|
||||
|
||||
public reset = async (_: Electron.IpcMainInvokeEvent, base: KnowledgeBaseParams): Promise<void> => {
|
||||
const ragApplication = await this.getRagApplication(base)
|
||||
await ragApplication.reset()
|
||||
}
|
||||
async delete(id: string): Promise<void> {
|
||||
|
||||
public async delete(_: Electron.IpcMainInvokeEvent, id: string): Promise<void> {
|
||||
logger.debug(`delete id: ${id}`)
|
||||
|
||||
await this.cleanupKnowledgeResources(id)
|
||||
@@ -209,41 +285,15 @@ export class EmbedJsFramework implements IKnowledgeFramework {
|
||||
this.pendingDeleteManager.add(id)
|
||||
}
|
||||
}
|
||||
getLoaderTask(options: KnowledgeBaseAddItemOptionsNonNullableAttribute): LoaderTask {
|
||||
const { item } = options
|
||||
const getRagApplication = () => this.getRagApplication(options.base)
|
||||
switch (item.type) {
|
||||
case 'file':
|
||||
return this.fileTask(getRagApplication, options)
|
||||
case 'directory':
|
||||
return this.directoryTask(getRagApplication, options)
|
||||
case 'url':
|
||||
return this.urlTask(getRagApplication, options)
|
||||
case 'sitemap':
|
||||
return this.sitemapTask(getRagApplication, options)
|
||||
case 'note':
|
||||
return this.noteTask(getRagApplication, options)
|
||||
default:
|
||||
return {
|
||||
loaderTasks: [],
|
||||
loaderDoneReturn: null
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async remove(options: { uniqueIds: string[]; base: KnowledgeBaseParams }): Promise<void> {
|
||||
const ragApp = await this.getRagApplication(options.base)
|
||||
for (const id of options.uniqueIds) {
|
||||
await ragApp.deleteLoader(id)
|
||||
}
|
||||
private maximumLoad() {
|
||||
return (
|
||||
this.processingItemCount >= KnowledgeService.MAXIMUM_PROCESSING_ITEM_COUNT ||
|
||||
this.workload >= KnowledgeService.MAXIMUM_WORKLOAD
|
||||
)
|
||||
}
|
||||
async search(options: { search: string; base: KnowledgeBaseParams }): Promise<KnowledgeSearchResult[]> {
|
||||
const ragApp = await this.getRagApplication(options.base)
|
||||
return await ragApp.search(options.search)
|
||||
}
|
||||
|
||||
private fileTask(
|
||||
getRagApplication: () => Promise<RAGApplication>,
|
||||
ragApplication: RAGApplication,
|
||||
options: KnowledgeBaseAddItemOptionsNonNullableAttribute
|
||||
): LoaderTask {
|
||||
const { base, item, forceReload, userId } = options
|
||||
@@ -256,8 +306,7 @@ export class EmbedJsFramework implements IKnowledgeFramework {
|
||||
task: async () => {
|
||||
try {
|
||||
// Add preprocessing logic
|
||||
const ragApplication = await getRagApplication()
|
||||
const fileToProcess: FileMetadata = await preprocessingService.preprocessFile(file, base, item, userId)
|
||||
const fileToProcess: FileMetadata = await this.preprocessing(file, base, item, userId)
|
||||
|
||||
// Use processed file for loading
|
||||
return addFileLoader(ragApplication, fileToProcess, base, forceReload)
|
||||
@@ -268,7 +317,7 @@ export class EmbedJsFramework implements IKnowledgeFramework {
|
||||
.catch((e) => {
|
||||
logger.error(`Error in addFileLoader for ${file.name}: ${e}`)
|
||||
const errorResult: LoaderReturn = {
|
||||
...EmbedJsFramework.ERROR_LOADER_RETURN,
|
||||
...KnowledgeService.ERROR_LOADER_RETURN,
|
||||
message: e.message,
|
||||
messageSource: 'embedding'
|
||||
}
|
||||
@@ -278,7 +327,7 @@ export class EmbedJsFramework implements IKnowledgeFramework {
|
||||
} catch (e: any) {
|
||||
logger.error(`Preprocessing failed for ${file.name}: ${e}`)
|
||||
const errorResult: LoaderReturn = {
|
||||
...EmbedJsFramework.ERROR_LOADER_RETURN,
|
||||
...KnowledgeService.ERROR_LOADER_RETURN,
|
||||
message: e.message,
|
||||
messageSource: 'preprocess'
|
||||
}
|
||||
@@ -295,7 +344,7 @@ export class EmbedJsFramework implements IKnowledgeFramework {
|
||||
return loaderTask
|
||||
}
|
||||
private directoryTask(
|
||||
getRagApplication: () => Promise<RAGApplication>,
|
||||
ragApplication: RAGApplication,
|
||||
options: KnowledgeBaseAddItemOptionsNonNullableAttribute
|
||||
): LoaderTask {
|
||||
const { base, item, forceReload } = options
|
||||
@@ -322,9 +371,8 @@ export class EmbedJsFramework implements IKnowledgeFramework {
|
||||
for (const file of files) {
|
||||
loaderTasks.push({
|
||||
state: LoaderTaskItemState.PENDING,
|
||||
task: async () => {
|
||||
const ragApplication = await getRagApplication()
|
||||
return addFileLoader(ragApplication, file, base, forceReload)
|
||||
task: () =>
|
||||
addFileLoader(ragApplication, file, base, forceReload)
|
||||
.then((result) => {
|
||||
loaderDoneReturn.entriesAdded += 1
|
||||
processedFiles += 1
|
||||
@@ -335,12 +383,11 @@ export class EmbedJsFramework implements IKnowledgeFramework {
|
||||
.catch((err) => {
|
||||
logger.error('Failed to add dir loader:', err)
|
||||
return {
|
||||
...EmbedJsFramework.ERROR_LOADER_RETURN,
|
||||
...KnowledgeService.ERROR_LOADER_RETURN,
|
||||
message: `Failed to add dir loader: ${err.message}`,
|
||||
messageSource: 'embedding'
|
||||
}
|
||||
})
|
||||
},
|
||||
}),
|
||||
evaluateTaskWorkload: { workload: file.size }
|
||||
})
|
||||
}
|
||||
@@ -352,7 +399,7 @@ export class EmbedJsFramework implements IKnowledgeFramework {
|
||||
}
|
||||
|
||||
private urlTask(
|
||||
getRagApplication: () => Promise<RAGApplication>,
|
||||
ragApplication: RAGApplication,
|
||||
options: KnowledgeBaseAddItemOptionsNonNullableAttribute
|
||||
): LoaderTask {
|
||||
const { base, item, forceReload } = options
|
||||
@@ -362,8 +409,7 @@ export class EmbedJsFramework implements IKnowledgeFramework {
|
||||
loaderTasks: [
|
||||
{
|
||||
state: LoaderTaskItemState.PENDING,
|
||||
task: async () => {
|
||||
const ragApplication = await getRagApplication()
|
||||
task: () => {
|
||||
const loaderReturn = ragApplication.addLoader(
|
||||
new WebLoader({
|
||||
urlOrContent: content,
|
||||
@@ -387,7 +433,7 @@ export class EmbedJsFramework implements IKnowledgeFramework {
|
||||
.catch((err) => {
|
||||
logger.error('Failed to add url loader:', err)
|
||||
return {
|
||||
...EmbedJsFramework.ERROR_LOADER_RETURN,
|
||||
...KnowledgeService.ERROR_LOADER_RETURN,
|
||||
message: `Failed to add url loader: ${err.message}`,
|
||||
messageSource: 'embedding'
|
||||
}
|
||||
@@ -402,7 +448,7 @@ export class EmbedJsFramework implements IKnowledgeFramework {
|
||||
}
|
||||
|
||||
private sitemapTask(
|
||||
getRagApplication: () => Promise<RAGApplication>,
|
||||
ragApplication: RAGApplication,
|
||||
options: KnowledgeBaseAddItemOptionsNonNullableAttribute
|
||||
): LoaderTask {
|
||||
const { base, item, forceReload } = options
|
||||
@@ -412,9 +458,8 @@ export class EmbedJsFramework implements IKnowledgeFramework {
|
||||
loaderTasks: [
|
||||
{
|
||||
state: LoaderTaskItemState.PENDING,
|
||||
task: async () => {
|
||||
const ragApplication = await getRagApplication()
|
||||
return ragApplication
|
||||
task: () =>
|
||||
ragApplication
|
||||
.addLoader(
|
||||
new SitemapLoader({ url: content, chunkSize: base.chunkSize, chunkOverlap: base.chunkOverlap }) as any,
|
||||
forceReload
|
||||
@@ -432,12 +477,11 @@ export class EmbedJsFramework implements IKnowledgeFramework {
|
||||
.catch((err) => {
|
||||
logger.error('Failed to add sitemap loader:', err)
|
||||
return {
|
||||
...EmbedJsFramework.ERROR_LOADER_RETURN,
|
||||
...KnowledgeService.ERROR_LOADER_RETURN,
|
||||
message: `Failed to add sitemap loader: ${err.message}`,
|
||||
messageSource: 'embedding'
|
||||
}
|
||||
})
|
||||
},
|
||||
}),
|
||||
evaluateTaskWorkload: { workload: 20 * MB }
|
||||
}
|
||||
],
|
||||
@@ -447,7 +491,7 @@ export class EmbedJsFramework implements IKnowledgeFramework {
|
||||
}
|
||||
|
||||
private noteTask(
|
||||
getRagApplication: () => Promise<RAGApplication>,
|
||||
ragApplication: RAGApplication,
|
||||
options: KnowledgeBaseAddItemOptionsNonNullableAttribute
|
||||
): LoaderTask {
|
||||
const { base, item, forceReload } = options
|
||||
@@ -460,8 +504,7 @@ export class EmbedJsFramework implements IKnowledgeFramework {
|
||||
loaderTasks: [
|
||||
{
|
||||
state: LoaderTaskItemState.PENDING,
|
||||
task: async () => {
|
||||
const ragApplication = await getRagApplication()
|
||||
task: () => {
|
||||
const loaderReturn = ragApplication.addLoader(
|
||||
new NoteLoader({
|
||||
text: content,
|
||||
@@ -484,7 +527,7 @@ export class EmbedJsFramework implements IKnowledgeFramework {
|
||||
.catch((err) => {
|
||||
logger.error('Failed to add note loader:', err)
|
||||
return {
|
||||
...EmbedJsFramework.ERROR_LOADER_RETURN,
|
||||
...KnowledgeService.ERROR_LOADER_RETURN,
|
||||
message: `Failed to add note loader: ${err.message}`,
|
||||
messageSource: 'embedding'
|
||||
}
|
||||
@@ -497,4 +540,199 @@ export class EmbedJsFramework implements IKnowledgeFramework {
|
||||
}
|
||||
return loaderTask
|
||||
}
|
||||
|
||||
private processingQueueHandle() {
|
||||
const getSubtasksUntilMaximumLoad = (): QueueTaskItem[] => {
|
||||
const queueTaskList: QueueTaskItem[] = []
|
||||
that: for (const [task, resolve] of this.knowledgeItemProcessingQueueMappingPromise) {
|
||||
for (const item of task.loaderTasks) {
|
||||
if (this.maximumLoad()) {
|
||||
break that
|
||||
}
|
||||
|
||||
const { state, task: taskPromise, evaluateTaskWorkload } = item
|
||||
|
||||
if (state !== LoaderTaskItemState.PENDING) {
|
||||
continue
|
||||
}
|
||||
|
||||
const { workload } = evaluateTaskWorkload
|
||||
this.workload += workload
|
||||
this.processingItemCount += 1
|
||||
item.state = LoaderTaskItemState.PROCESSING
|
||||
queueTaskList.push({
|
||||
taskPromise: () =>
|
||||
taskPromise().then(() => {
|
||||
this.workload -= workload
|
||||
this.processingItemCount -= 1
|
||||
task.loaderTasks.delete(item)
|
||||
if (task.loaderTasks.size === 0) {
|
||||
this.knowledgeItemProcessingQueueMappingPromise.delete(task)
|
||||
resolve()
|
||||
}
|
||||
this.processingQueueHandle()
|
||||
}),
|
||||
resolve: () => {},
|
||||
evaluateTaskWorkload
|
||||
})
|
||||
}
|
||||
}
|
||||
return queueTaskList
|
||||
}
|
||||
const subTasks = getSubtasksUntilMaximumLoad()
|
||||
if (subTasks.length > 0) {
|
||||
const subTaskPromises = subTasks.map(({ taskPromise }) => taskPromise())
|
||||
Promise.all(subTaskPromises).then(() => {
|
||||
subTasks.forEach(({ resolve }) => resolve())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
private appendProcessingQueue(task: LoaderTask): Promise<LoaderReturn> {
|
||||
return new Promise((resolve) => {
|
||||
this.knowledgeItemProcessingQueueMappingPromise.set(loaderTaskIntoOfSet(task), () => {
|
||||
resolve(task.loaderDoneReturn!)
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
public add = (_: Electron.IpcMainInvokeEvent, options: KnowledgeBaseAddItemOptions): Promise<LoaderReturn> => {
|
||||
return new Promise((resolve) => {
|
||||
const { base, item, forceReload = false, userId = '' } = options
|
||||
const optionsNonNullableAttribute = { base, item, forceReload, userId }
|
||||
this.getRagApplication(base)
|
||||
.then((ragApplication) => {
|
||||
const task = (() => {
|
||||
switch (item.type) {
|
||||
case 'file':
|
||||
return this.fileTask(ragApplication, optionsNonNullableAttribute)
|
||||
case 'directory':
|
||||
return this.directoryTask(ragApplication, optionsNonNullableAttribute)
|
||||
case 'url':
|
||||
return this.urlTask(ragApplication, optionsNonNullableAttribute)
|
||||
case 'sitemap':
|
||||
return this.sitemapTask(ragApplication, optionsNonNullableAttribute)
|
||||
case 'note':
|
||||
return this.noteTask(ragApplication, optionsNonNullableAttribute)
|
||||
default:
|
||||
return null
|
||||
}
|
||||
})()
|
||||
|
||||
if (task) {
|
||||
this.appendProcessingQueue(task).then(() => {
|
||||
resolve(task.loaderDoneReturn!)
|
||||
})
|
||||
this.processingQueueHandle()
|
||||
} else {
|
||||
resolve({
|
||||
...KnowledgeService.ERROR_LOADER_RETURN,
|
||||
message: 'Unsupported item type',
|
||||
messageSource: 'embedding'
|
||||
})
|
||||
}
|
||||
})
|
||||
.catch((err) => {
|
||||
logger.error('Failed to add item:', err)
|
||||
resolve({
|
||||
...KnowledgeService.ERROR_LOADER_RETURN,
|
||||
message: `Failed to add item: ${err.message}`,
|
||||
messageSource: 'embedding'
|
||||
})
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
@TraceMethod({ spanName: 'remove', tag: 'Knowledge' })
|
||||
public async remove(
|
||||
_: Electron.IpcMainInvokeEvent,
|
||||
{ uniqueId, uniqueIds, base }: { uniqueId: string; uniqueIds: string[]; base: KnowledgeBaseParams }
|
||||
): Promise<void> {
|
||||
const ragApplication = await this.getRagApplication(base)
|
||||
logger.debug(`Remove Item UniqueId: ${uniqueId}`)
|
||||
for (const id of uniqueIds) {
|
||||
await ragApplication.deleteLoader(id)
|
||||
}
|
||||
}
|
||||
|
||||
@TraceMethod({ spanName: 'RagSearch', tag: 'Knowledge' })
|
||||
public async search(
|
||||
_: Electron.IpcMainInvokeEvent,
|
||||
{ search, base }: { search: string; base: KnowledgeBaseParams }
|
||||
): Promise<KnowledgeSearchResult[]> {
|
||||
const ragApplication = await this.getRagApplication(base)
|
||||
return await ragApplication.search(search)
|
||||
}
|
||||
|
||||
@TraceMethod({ spanName: 'rerank', tag: 'Knowledge' })
|
||||
public async rerank(
|
||||
_: Electron.IpcMainInvokeEvent,
|
||||
{ search, base, results }: { search: string; base: KnowledgeBaseParams; results: KnowledgeSearchResult[] }
|
||||
): Promise<KnowledgeSearchResult[]> {
|
||||
if (results.length === 0) {
|
||||
return results
|
||||
}
|
||||
return await new Reranker(base).rerank(search, results)
|
||||
}
|
||||
|
||||
public getStorageDir = (): string => {
|
||||
return this.storageDir
|
||||
}
|
||||
|
||||
private preprocessing = async (
|
||||
file: FileMetadata,
|
||||
base: KnowledgeBaseParams,
|
||||
item: KnowledgeItem,
|
||||
userId: string
|
||||
): Promise<FileMetadata> => {
|
||||
let fileToProcess: FileMetadata = file
|
||||
if (base.preprocessProvider && file.ext.toLowerCase() === '.pdf') {
|
||||
try {
|
||||
const provider = new PreprocessProvider(base.preprocessProvider.provider, userId)
|
||||
const filePath = fileStorage.getFilePathById(file)
|
||||
// Check if file has already been preprocessed
|
||||
const alreadyProcessed = await provider.checkIfAlreadyProcessed(file)
|
||||
if (alreadyProcessed) {
|
||||
logger.debug(`File already preprocess processed, using cached result: ${filePath}`)
|
||||
return alreadyProcessed
|
||||
}
|
||||
|
||||
// Execute preprocessing
|
||||
logger.debug(`Starting preprocess processing for scanned PDF: ${filePath}`)
|
||||
const { processedFile, quota } = await provider.parseFile(item.id, file)
|
||||
fileToProcess = processedFile
|
||||
const mainWindow = windowService.getMainWindow()
|
||||
mainWindow?.webContents.send('file-preprocess-finished', {
|
||||
itemId: item.id,
|
||||
quota: quota
|
||||
})
|
||||
} catch (err) {
|
||||
logger.error(`Preprocess processing failed: ${err}`)
|
||||
// If preprocessing fails, use original file
|
||||
// fileToProcess = file
|
||||
throw new Error(`Preprocess processing failed: ${err}`)
|
||||
}
|
||||
}
|
||||
|
||||
return fileToProcess
|
||||
}
|
||||
|
||||
public checkQuota = async (
|
||||
_: Electron.IpcMainInvokeEvent,
|
||||
base: KnowledgeBaseParams,
|
||||
userId: string
|
||||
): Promise<number> => {
|
||||
try {
|
||||
if (base.preprocessProvider && base.preprocessProvider.type === 'preprocess') {
|
||||
const provider = new PreprocessProvider(base.preprocessProvider.provider, userId)
|
||||
return await provider.checkQuota()
|
||||
}
|
||||
throw new Error('No preprocess provider configured')
|
||||
} catch (err) {
|
||||
logger.error(`Failed to check quota: ${err}`)
|
||||
throw new Error(`Failed to check quota: ${err}`)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
export default new KnowledgeService()
|
||||
@@ -1,72 +0,0 @@
|
||||
import { LoaderReturn } from '@shared/config/types'
|
||||
import { KnowledgeBaseParams, KnowledgeItem, KnowledgeSearchResult } from '@types'
|
||||
|
||||
export interface KnowledgeBaseAddItemOptions {
|
||||
base: KnowledgeBaseParams
|
||||
item: KnowledgeItem
|
||||
forceReload?: boolean
|
||||
userId?: string
|
||||
}
|
||||
|
||||
export interface KnowledgeBaseAddItemOptionsNonNullableAttribute {
|
||||
base: KnowledgeBaseParams
|
||||
item: KnowledgeItem
|
||||
forceReload: boolean
|
||||
userId: string
|
||||
}
|
||||
|
||||
export interface EvaluateTaskWorkload {
|
||||
workload: number
|
||||
}
|
||||
|
||||
export type LoaderDoneReturn = LoaderReturn | null
|
||||
|
||||
export enum LoaderTaskItemState {
|
||||
PENDING,
|
||||
PROCESSING,
|
||||
DONE
|
||||
}
|
||||
|
||||
export interface LoaderTaskItem {
|
||||
state: LoaderTaskItemState
|
||||
task: () => Promise<unknown>
|
||||
evaluateTaskWorkload: EvaluateTaskWorkload
|
||||
}
|
||||
|
||||
export interface LoaderTask {
|
||||
loaderTasks: LoaderTaskItem[]
|
||||
loaderDoneReturn: LoaderDoneReturn
|
||||
}
|
||||
|
||||
export interface LoaderTaskOfSet {
|
||||
loaderTasks: Set<LoaderTaskItem>
|
||||
loaderDoneReturn: LoaderDoneReturn
|
||||
}
|
||||
|
||||
export interface QueueTaskItem {
|
||||
taskPromise: () => Promise<unknown>
|
||||
resolve: () => void
|
||||
evaluateTaskWorkload: EvaluateTaskWorkload
|
||||
}
|
||||
|
||||
export const loaderTaskIntoOfSet = (loaderTask: LoaderTask): LoaderTaskOfSet => {
|
||||
return {
|
||||
loaderTasks: new Set(loaderTask.loaderTasks),
|
||||
loaderDoneReturn: loaderTask.loaderDoneReturn
|
||||
}
|
||||
}
|
||||
|
||||
export interface IKnowledgeFramework {
|
||||
/** 为给定知识库初始化框架资源 */
|
||||
initialize(base: KnowledgeBaseParams): Promise<void>
|
||||
/** 重置知识库，删除其所有内容 */
|
||||
reset(base: KnowledgeBaseParams): Promise<void>
|
||||
/** 删除与知识库关联的资源，包括文件 */
|
||||
delete(id: string): Promise<void>
|
||||
/** 生成用于添加条目的任务对象，由队列处理 */
|
||||
getLoaderTask(options: KnowledgeBaseAddItemOptionsNonNullableAttribute): LoaderTask
|
||||
/** 从知识库中删除特定条目 */
|
||||
remove(options: { uniqueIds: string[]; base: KnowledgeBaseParams }): Promise<void>
|
||||
/** 搜索知识库 */
|
||||
search(options: { search: string; base: KnowledgeBaseParams }): Promise<KnowledgeSearchResult[]>
|
||||
}
|
||||
@@ -1,48 +0,0 @@
|
||||
import path from 'node:path'
|
||||
|
||||
import { KnowledgeBaseParams } from '@types'
|
||||
import { app } from 'electron'
|
||||
|
||||
import { EmbedJsFramework } from './EmbedJsFramework'
|
||||
import { IKnowledgeFramework } from './IKnowledgeFramework'
|
||||
import { LangChainFramework } from './LangChainFramework'
|
||||
class KnowledgeFrameworkFactory {
|
||||
private static instance: KnowledgeFrameworkFactory
|
||||
private frameworks: Map<string, IKnowledgeFramework> = new Map()
|
||||
private storageDir: string
|
||||
|
||||
private constructor(storageDir: string) {
|
||||
this.storageDir = storageDir
|
||||
}
|
||||
|
||||
public static getInstance(storageDir: string): KnowledgeFrameworkFactory {
|
||||
if (!KnowledgeFrameworkFactory.instance) {
|
||||
KnowledgeFrameworkFactory.instance = new KnowledgeFrameworkFactory(storageDir)
|
||||
}
|
||||
return KnowledgeFrameworkFactory.instance
|
||||
}
|
||||
|
||||
public getFramework(base: KnowledgeBaseParams): IKnowledgeFramework {
|
||||
const frameworkType = base.framework || 'embedjs' // 如果未指定，默认为 embedjs
|
||||
if (this.frameworks.has(frameworkType)) {
|
||||
return this.frameworks.get(frameworkType)!
|
||||
}
|
||||
let framework: IKnowledgeFramework
|
||||
switch (frameworkType) {
|
||||
case 'langchain':
|
||||
framework = new LangChainFramework(this.storageDir)
|
||||
break
|
||||
case 'embedjs':
|
||||
default:
|
||||
framework = new EmbedJsFramework(this.storageDir)
|
||||
break
|
||||
}
|
||||
|
||||
this.frameworks.set(frameworkType, framework)
|
||||
return framework
|
||||
}
|
||||
}
|
||||
|
||||
export const knowledgeFrameworkFactory = KnowledgeFrameworkFactory.getInstance(
|
||||
path.join(app.getPath('userData'), 'Data', 'KnowledgeBase')
|
||||
)
|
||||
@@ -1,190 +0,0 @@
|
||||
import * as fs from 'node:fs'
|
||||
import path from 'node:path'
|
||||
|
||||
import { loggerService } from '@logger'
|
||||
import { preprocessingService } from '@main/knowledge/preprocess/PreprocessingService'
|
||||
import Reranker from '@main/knowledge/reranker/Reranker'
|
||||
import { TraceMethod } from '@mcp-trace/trace-core'
|
||||
import { MB } from '@shared/config/constant'
|
||||
import { LoaderReturn } from '@shared/config/types'
|
||||
import { KnowledgeBaseParams, KnowledgeSearchResult } from '@types'
|
||||
import { app } from 'electron'
|
||||
|
||||
import {
|
||||
KnowledgeBaseAddItemOptions,
|
||||
LoaderTask,
|
||||
loaderTaskIntoOfSet,
|
||||
LoaderTaskItemState,
|
||||
LoaderTaskOfSet,
|
||||
QueueTaskItem
|
||||
} from './IKnowledgeFramework'
|
||||
import { knowledgeFrameworkFactory } from './KnowledgeFrameworkFactory'
|
||||
|
||||
const logger = loggerService.withContext('MainKnowledgeService')
|
||||
|
||||
class KnowledgeService {
|
||||
private storageDir = path.join(app.getPath('userData'), 'Data', 'KnowledgeBase')
|
||||
|
||||
private workload = 0
|
||||
private processingItemCount = 0
|
||||
private knowledgeItemProcessingQueueMappingPromise: Map<LoaderTaskOfSet, () => void> = new Map()
|
||||
private static MAXIMUM_WORKLOAD = 80 * MB
|
||||
private static MAXIMUM_PROCESSING_ITEM_COUNT = 30
|
||||
private static ERROR_LOADER_RETURN: LoaderReturn = {
|
||||
entriesAdded: 0,
|
||||
uniqueId: '',
|
||||
uniqueIds: [''],
|
||||
loaderType: '',
|
||||
status: 'failed'
|
||||
}
|
||||
|
||||
constructor() {
|
||||
this.initStorageDir()
|
||||
}
|
||||
|
||||
private initStorageDir = (): void => {
|
||||
if (!fs.existsSync(this.storageDir)) {
|
||||
fs.mkdirSync(this.storageDir, { recursive: true })
|
||||
}
|
||||
}
|
||||
|
||||
private maximumLoad() {
|
||||
return (
|
||||
this.processingItemCount >= KnowledgeService.MAXIMUM_PROCESSING_ITEM_COUNT ||
|
||||
this.workload >= KnowledgeService.MAXIMUM_WORKLOAD
|
||||
)
|
||||
}
|
||||
|
||||
private processingQueueHandle() {
|
||||
const getSubtasksUntilMaximumLoad = (): QueueTaskItem[] => {
|
||||
const queueTaskList: QueueTaskItem[] = []
|
||||
that: for (const [task, resolve] of this.knowledgeItemProcessingQueueMappingPromise) {
|
||||
for (const item of task.loaderTasks) {
|
||||
if (this.maximumLoad()) {
|
||||
break that
|
||||
}
|
||||
|
||||
const { state, task: taskPromise, evaluateTaskWorkload } = item
|
||||
|
||||
if (state !== LoaderTaskItemState.PENDING) {
|
||||
continue
|
||||
}
|
||||
|
||||
const { workload } = evaluateTaskWorkload
|
||||
this.workload += workload
|
||||
this.processingItemCount += 1
|
||||
item.state = LoaderTaskItemState.PROCESSING
|
||||
queueTaskList.push({
|
||||
taskPromise: () =>
|
||||
taskPromise().then(() => {
|
||||
this.workload -= workload
|
||||
this.processingItemCount -= 1
|
||||
task.loaderTasks.delete(item)
|
||||
if (task.loaderTasks.size === 0) {
|
||||
this.knowledgeItemProcessingQueueMappingPromise.delete(task)
|
||||
resolve()
|
||||
}
|
||||
this.processingQueueHandle()
|
||||
}),
|
||||
resolve: () => {},
|
||||
evaluateTaskWorkload
|
||||
})
|
||||
}
|
||||
}
|
||||
return queueTaskList
|
||||
}
|
||||
const subTasks = getSubtasksUntilMaximumLoad()
|
||||
if (subTasks.length > 0) {
|
||||
const subTaskPromises = subTasks.map(({ taskPromise }) => taskPromise())
|
||||
Promise.all(subTaskPromises).then(() => {
|
||||
subTasks.forEach(({ resolve }) => resolve())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
private appendProcessingQueue(task: LoaderTask): Promise<LoaderReturn> {
|
||||
return new Promise((resolve) => {
|
||||
this.knowledgeItemProcessingQueueMappingPromise.set(loaderTaskIntoOfSet(task), () => {
|
||||
resolve(task.loaderDoneReturn!)
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
public async create(_: Electron.IpcMainInvokeEvent, base: KnowledgeBaseParams): Promise<void> {
|
||||
logger.info(`Creating knowledge base: ${JSON.stringify(base)}`)
|
||||
const framework = knowledgeFrameworkFactory.getFramework(base)
|
||||
await framework.initialize(base)
|
||||
}
|
||||
public async reset(_: Electron.IpcMainInvokeEvent, base: KnowledgeBaseParams): Promise<void> {
|
||||
const framework = knowledgeFrameworkFactory.getFramework(base)
|
||||
await framework.reset(base)
|
||||
}
|
||||
|
||||
public async delete(_: Electron.IpcMainInvokeEvent, base: KnowledgeBaseParams, id: string): Promise<void> {
|
||||
logger.info(`Deleting knowledge base: ${JSON.stringify(base)}`)
|
||||
const framework = knowledgeFrameworkFactory.getFramework(base)
|
||||
await framework.delete(id)
|
||||
}
|
||||
|
||||
public add = async (_: Electron.IpcMainInvokeEvent, options: KnowledgeBaseAddItemOptions): Promise<LoaderReturn> => {
|
||||
logger.info(`Adding item to knowledge base: ${JSON.stringify(options)}`)
|
||||
return new Promise((resolve) => {
|
||||
const { base, item, forceReload = false, userId = '' } = options
|
||||
const framework = knowledgeFrameworkFactory.getFramework(base)
|
||||
|
||||
const task = framework.getLoaderTask({ base, item, forceReload, userId })
|
||||
|
||||
if (task) {
|
||||
this.appendProcessingQueue(task).then(() => {
|
||||
resolve(task.loaderDoneReturn!)
|
||||
})
|
||||
this.processingQueueHandle()
|
||||
} else {
|
||||
resolve({
|
||||
...KnowledgeService.ERROR_LOADER_RETURN,
|
||||
message: 'Unsupported item type',
|
||||
messageSource: 'embedding'
|
||||
})
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
public async remove(
|
||||
_: Electron.IpcMainInvokeEvent,
|
||||
{ uniqueIds, base }: { uniqueIds: string[]; base: KnowledgeBaseParams }
|
||||
): Promise<void> {
|
||||
logger.info(`Removing items from knowledge base: ${JSON.stringify({ uniqueIds, base })}`)
|
||||
const framework = knowledgeFrameworkFactory.getFramework(base)
|
||||
await framework.remove({ uniqueIds, base })
|
||||
}
|
||||
public async search(
|
||||
_: Electron.IpcMainInvokeEvent,
|
||||
{ search, base }: { search: string; base: KnowledgeBaseParams }
|
||||
): Promise<KnowledgeSearchResult[]> {
|
||||
logger.info(`Searching knowledge base: ${JSON.stringify({ search, base })}`)
|
||||
const framework = knowledgeFrameworkFactory.getFramework(base)
|
||||
return framework.search({ search, base })
|
||||
}
|
||||
|
||||
@TraceMethod({ spanName: 'rerank', tag: 'Knowledge' })
|
||||
public async rerank(
|
||||
_: Electron.IpcMainInvokeEvent,
|
||||
{ search, base, results }: { search: string; base: KnowledgeBaseParams; results: KnowledgeSearchResult[] }
|
||||
): Promise<KnowledgeSearchResult[]> {
|
||||
logger.info(`Reranking knowledge base: ${JSON.stringify({ search, base, results })}`)
|
||||
if (results.length === 0) {
|
||||
return results
|
||||
}
|
||||
return await new Reranker(base).rerank(search, results)
|
||||
}
|
||||
|
||||
public getStorageDir = (): string => {
|
||||
return this.storageDir
|
||||
}
|
||||
|
||||
public async checkQuota(_: Electron.IpcMainInvokeEvent, base: KnowledgeBaseParams, userId: string): Promise<number> {
|
||||
return preprocessingService.checkQuota(base, userId)
|
||||
}
|
||||
}
|
||||
|
||||
export default new KnowledgeService()
|
||||
@@ -1,557 +0,0 @@
|
||||
import * as fs from 'node:fs'
|
||||
import path from 'node:path'
|
||||
|
||||
import { FaissStore } from '@langchain/community/vectorstores/faiss'
|
||||
import type { Document } from '@langchain/core/documents'
|
||||
import { loggerService } from '@logger'
|
||||
import TextEmbeddings from '@main/knowledge/langchain/embeddings/TextEmbeddings'
|
||||
import {
|
||||
addFileLoader,
|
||||
addNoteLoader,
|
||||
addSitemapLoader,
|
||||
addVideoLoader,
|
||||
addWebLoader
|
||||
} from '@main/knowledge/langchain/loader'
|
||||
import { RetrieverFactory } from '@main/knowledge/langchain/retriever'
|
||||
import { preprocessingService } from '@main/knowledge/preprocess/PreprocessingService'
|
||||
import { getAllFiles } from '@main/utils/file'
|
||||
import { getUrlSource } from '@main/utils/knowledge'
|
||||
import { MB } from '@shared/config/constant'
|
||||
import { LoaderReturn } from '@shared/config/types'
|
||||
import { IpcChannel } from '@shared/IpcChannel'
|
||||
import {
|
||||
FileMetadata,
|
||||
isKnowledgeDirectoryItem,
|
||||
isKnowledgeFileItem,
|
||||
isKnowledgeNoteItem,
|
||||
isKnowledgeSitemapItem,
|
||||
isKnowledgeUrlItem,
|
||||
isKnowledgeVideoItem,
|
||||
KnowledgeBaseParams,
|
||||
KnowledgeSearchResult
|
||||
} from '@types'
|
||||
import { uuidv4 } from 'zod'
|
||||
|
||||
import { windowService } from '../WindowService'
|
||||
import {
|
||||
IKnowledgeFramework,
|
||||
KnowledgeBaseAddItemOptionsNonNullableAttribute,
|
||||
LoaderDoneReturn,
|
||||
LoaderTask,
|
||||
LoaderTaskItem,
|
||||
LoaderTaskItemState
|
||||
} from './IKnowledgeFramework'
|
||||
|
||||
const logger = loggerService.withContext('LangChainFramework')
|
||||
|
||||
export class LangChainFramework implements IKnowledgeFramework {
|
||||
private storageDir: string
|
||||
|
||||
private static ERROR_LOADER_RETURN: LoaderReturn = {
|
||||
entriesAdded: 0,
|
||||
uniqueId: '',
|
||||
uniqueIds: [''],
|
||||
loaderType: '',
|
||||
status: 'failed'
|
||||
}
|
||||
|
||||
constructor(storageDir: string) {
|
||||
this.storageDir = storageDir
|
||||
this.initStorageDir()
|
||||
}
|
||||
private initStorageDir = (): void => {
|
||||
if (!fs.existsSync(this.storageDir)) {
|
||||
fs.mkdirSync(this.storageDir, { recursive: true })
|
||||
}
|
||||
}
|
||||
|
||||
private async createDatabase(base: KnowledgeBaseParams): Promise<void> {
|
||||
const dbPath = path.join(this.storageDir, base.id)
|
||||
const embeddings = this.getEmbeddings(base)
|
||||
const vectorStore = new FaissStore(embeddings, {})
|
||||
|
||||
const mockDocument: Document = {
|
||||
pageContent: 'Create Database Document',
|
||||
metadata: {}
|
||||
}
|
||||
|
||||
await vectorStore.addDocuments([mockDocument], { ids: ['1'] })
|
||||
await vectorStore.save(dbPath)
|
||||
await vectorStore.delete({ ids: ['1'] })
|
||||
await vectorStore.save(dbPath)
|
||||
}
|
||||
|
||||
private getEmbeddings(base: KnowledgeBaseParams): TextEmbeddings {
|
||||
return new TextEmbeddings({
|
||||
embedApiClient: base.embedApiClient,
|
||||
dimensions: base.dimensions
|
||||
})
|
||||
}
|
||||
|
||||
private async getVectorStore(base: KnowledgeBaseParams): Promise<FaissStore> {
|
||||
const embeddings = this.getEmbeddings(base)
|
||||
const vectorStore = await FaissStore.load(path.join(this.storageDir, base.id), embeddings)
|
||||
|
||||
return vectorStore
|
||||
}
|
||||
|
||||
async initialize(base: KnowledgeBaseParams): Promise<void> {
|
||||
await this.createDatabase(base)
|
||||
}
|
||||
async reset(base: KnowledgeBaseParams): Promise<void> {
|
||||
const dbPath = path.join(this.storageDir, base.id)
|
||||
if (fs.existsSync(dbPath)) {
|
||||
fs.rmSync(dbPath, { recursive: true })
|
||||
}
|
||||
// 立即重建空索引，避免随后加载时报错
|
||||
await this.createDatabase(base)
|
||||
}
|
||||
|
||||
async delete(id: string): Promise<void> {
|
||||
const dbPath = path.join(this.storageDir, id)
|
||||
if (fs.existsSync(dbPath)) {
|
||||
fs.rmSync(dbPath, { recursive: true })
|
||||
}
|
||||
}
|
||||
getLoaderTask(options: KnowledgeBaseAddItemOptionsNonNullableAttribute): LoaderTask {
|
||||
const { item } = options
|
||||
const getStore = () => this.getVectorStore(options.base)
|
||||
switch (item.type) {
|
||||
case 'file':
|
||||
return this.fileTask(getStore, options)
|
||||
case 'directory':
|
||||
return this.directoryTask(getStore, options)
|
||||
case 'url':
|
||||
return this.urlTask(getStore, options)
|
||||
case 'sitemap':
|
||||
return this.sitemapTask(getStore, options)
|
||||
case 'note':
|
||||
return this.noteTask(getStore, options)
|
||||
case 'video':
|
||||
return this.videoTask(getStore, options)
|
||||
default:
|
||||
return {
|
||||
loaderTasks: [],
|
||||
loaderDoneReturn: null
|
||||
}
|
||||
}
|
||||
}
|
||||
async remove(options: { uniqueIds: string[]; base: KnowledgeBaseParams }): Promise<void> {
|
||||
const { uniqueIds, base } = options
|
||||
const vectorStore = await this.getVectorStore(base)
|
||||
logger.info(`[ KnowledgeService Remove Item UniqueIds: ${uniqueIds}]`)
|
||||
|
||||
await vectorStore.delete({ ids: uniqueIds })
|
||||
await vectorStore.save(path.join(this.storageDir, base.id))
|
||||
}
|
||||
async search(options: { search: string; base: KnowledgeBaseParams }): Promise<KnowledgeSearchResult[]> {
|
||||
const { search, base } = options
|
||||
logger.info(`search base: ${JSON.stringify(base)}`)
|
||||
|
||||
try {
|
||||
const vectorStore = await this.getVectorStore(base)
|
||||
|
||||
// 如果是 bm25 或 hybrid 模式，则从数据库获取所有文档
|
||||
const documents: Document[] = await this.getAllDocuments(base)
|
||||
if (documents.length === 0) return []
|
||||
|
||||
const retrieverFactory = new RetrieverFactory()
|
||||
const retriever = retrieverFactory.createRetriever(base, vectorStore, documents)
|
||||
|
||||
const results = await retriever.invoke(search)
|
||||
logger.info(`Search Results: ${JSON.stringify(results)}`)
|
||||
|
||||
// VectorStoreRetriever 和 EnsembleRetriever 会将分数附加到 metadata.score
|
||||
// BM25Retriever 默认不返回分数，所以我们需要处理这种情况
|
||||
return results.map((item) => {
|
||||
return {
|
||||
pageContent: item.pageContent,
|
||||
metadata: item.metadata,
|
||||
// 如果 metadata 中没有 score，提供一个默认值
|
||||
score: typeof item.metadata.score === 'number' ? item.metadata.score : 0
|
||||
}
|
||||
})
|
||||
} catch (error: any) {
|
||||
logger.error(`Error during search in knowledge base ${base.id}: ${error.message}`)
|
||||
return []
|
||||
}
|
||||
}
|
||||
|
||||
private fileTask(
|
||||
getVectorStore: () => Promise<FaissStore>,
|
||||
options: KnowledgeBaseAddItemOptionsNonNullableAttribute
|
||||
): LoaderTask {
|
||||
const { base, item, userId } = options
|
||||
|
||||
if (!isKnowledgeFileItem(item)) {
|
||||
logger.error(`Invalid item type for fileTask: expected 'file', got '${item.type}'`)
|
||||
return {
|
||||
loaderTasks: [],
|
||||
loaderDoneReturn: {
|
||||
...LangChainFramework.ERROR_LOADER_RETURN,
|
||||
message: `Invalid item type: expected 'file', got '${item.type}'`,
|
||||
messageSource: 'validation'
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const file = item.content
|
||||
|
||||
const loaderTask: LoaderTask = {
|
||||
loaderTasks: [
|
||||
{
|
||||
state: LoaderTaskItemState.PENDING,
|
||||
task: async () => {
|
||||
try {
|
||||
const vectorStore = await getVectorStore()
|
||||
|
||||
// 添加预处理逻辑
|
||||
const fileToProcess: FileMetadata = await preprocessingService.preprocessFile(file, base, item, userId)
|
||||
|
||||
// 使用处理后的文件进行加载
|
||||
return addFileLoader(base, vectorStore, fileToProcess)
|
||||
.then((result) => {
|
||||
loaderTask.loaderDoneReturn = result
|
||||
return result
|
||||
})
|
||||
.then(async () => {
|
||||
await vectorStore.save(path.join(this.storageDir, base.id))
|
||||
})
|
||||
.catch((e) => {
|
||||
logger.error(`Error in addFileLoader for ${file.name}: ${e}`)
|
||||
const errorResult: LoaderReturn = {
|
||||
...LangChainFramework.ERROR_LOADER_RETURN,
|
||||
message: e.message,
|
||||
messageSource: 'embedding'
|
||||
}
|
||||
loaderTask.loaderDoneReturn = errorResult
|
||||
return errorResult
|
||||
})
|
||||
} catch (e: any) {
|
||||
logger.error(`Preprocessing failed for ${file.name}: ${e}`)
|
||||
const errorResult: LoaderReturn = {
|
||||
...LangChainFramework.ERROR_LOADER_RETURN,
|
||||
message: e.message,
|
||||
messageSource: 'preprocess'
|
||||
}
|
||||
loaderTask.loaderDoneReturn = errorResult
|
||||
return errorResult
|
||||
}
|
||||
},
|
||||
evaluateTaskWorkload: { workload: file.size }
|
||||
}
|
||||
],
|
||||
loaderDoneReturn: null
|
||||
}
|
||||
|
||||
return loaderTask
|
||||
}
|
||||
private directoryTask(
|
||||
getVectorStore: () => Promise<FaissStore>,
|
||||
options: KnowledgeBaseAddItemOptionsNonNullableAttribute
|
||||
): LoaderTask {
|
||||
const { base, item } = options
|
||||
|
||||
if (!isKnowledgeDirectoryItem(item)) {
|
||||
logger.error(`Invalid item type for directoryTask: expected 'directory', got '${item.type}'`)
|
||||
return {
|
||||
loaderTasks: [],
|
||||
loaderDoneReturn: {
|
||||
...LangChainFramework.ERROR_LOADER_RETURN,
|
||||
message: `Invalid item type: expected 'directory', got '${item.type}'`,
|
||||
messageSource: 'validation'
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const directory = item.content
|
||||
const files = getAllFiles(directory)
|
||||
const totalFiles = files.length
|
||||
let processedFiles = 0
|
||||
|
||||
const sendDirectoryProcessingPercent = (totalFiles: number, processedFiles: number) => {
|
||||
const mainWindow = windowService.getMainWindow()
|
||||
mainWindow?.webContents.send(IpcChannel.DirectoryProcessingPercent, {
|
||||
itemId: item.id,
|
||||
percent: (processedFiles / totalFiles) * 100
|
||||
})
|
||||
}
|
||||
|
||||
const loaderDoneReturn: LoaderDoneReturn = {
|
||||
entriesAdded: 0,
|
||||
uniqueId: `DirectoryLoader_${uuidv4()}`,
|
||||
uniqueIds: [],
|
||||
loaderType: 'DirectoryLoader'
|
||||
}
|
||||
const loaderTasks: LoaderTaskItem[] = []
|
||||
for (const file of files) {
|
||||
loaderTasks.push({
|
||||
state: LoaderTaskItemState.PENDING,
|
||||
task: async () => {
|
||||
const vectorStore = await getVectorStore()
|
||||
return addFileLoader(base, vectorStore, file)
|
||||
.then((result) => {
|
||||
loaderDoneReturn.entriesAdded += 1
|
||||
processedFiles += 1
|
||||
sendDirectoryProcessingPercent(totalFiles, processedFiles)
|
||||
loaderDoneReturn.uniqueIds.push(result.uniqueId)
|
||||
return result
|
||||
})
|
||||
.then(async () => {
|
||||
await vectorStore.save(path.join(this.storageDir, base.id))
|
||||
})
|
||||
.catch((err) => {
|
||||
logger.error(err)
|
||||
return {
|
||||
...LangChainFramework.ERROR_LOADER_RETURN,
|
||||
message: `Failed to add dir loader: ${err.message}`,
|
||||
messageSource: 'embedding'
|
||||
}
|
||||
})
|
||||
},
|
||||
evaluateTaskWorkload: { workload: file.size }
|
||||
})
|
||||
}
|
||||
|
||||
return {
|
||||
loaderTasks,
|
||||
loaderDoneReturn
|
||||
}
|
||||
}
|
||||
|
||||
private urlTask(
|
||||
getVectorStore: () => Promise<FaissStore>,
|
||||
options: KnowledgeBaseAddItemOptionsNonNullableAttribute
|
||||
): LoaderTask {
|
||||
const { base, item } = options
|
||||
|
||||
if (!isKnowledgeUrlItem(item)) {
|
||||
logger.error(`Invalid item type for urlTask: expected 'url', got '${item.type}'`)
|
||||
return {
|
||||
loaderTasks: [],
|
||||
loaderDoneReturn: {
|
||||
...LangChainFramework.ERROR_LOADER_RETURN,
|
||||
message: `Invalid item type: expected 'url', got '${item.type}'`,
|
||||
messageSource: 'validation'
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const url = item.content
|
||||
|
||||
const loaderTask: LoaderTask = {
|
||||
loaderTasks: [
|
||||
{
|
||||
state: LoaderTaskItemState.PENDING,
|
||||
task: async () => {
|
||||
// 使用处理后的网页进行加载
|
||||
const vectorStore = await getVectorStore()
|
||||
return addWebLoader(base, vectorStore, url, getUrlSource(url))
|
||||
.then((result) => {
|
||||
loaderTask.loaderDoneReturn = result
|
||||
return result
|
||||
})
|
||||
.then(async () => {
|
||||
await vectorStore.save(path.join(this.storageDir, base.id))
|
||||
})
|
||||
.catch((e) => {
|
||||
logger.error(`Error in addWebLoader for ${url}: ${e}`)
|
||||
const errorResult: LoaderReturn = {
|
||||
...LangChainFramework.ERROR_LOADER_RETURN,
|
||||
message: e.message,
|
||||
messageSource: 'embedding'
|
||||
}
|
||||
loaderTask.loaderDoneReturn = errorResult
|
||||
return errorResult
|
||||
})
|
||||
},
|
||||
evaluateTaskWorkload: { workload: 2 * MB }
|
||||
}
|
||||
],
|
||||
loaderDoneReturn: null
|
||||
}
|
||||
return loaderTask
|
||||
}
|
||||
|
||||
private sitemapTask(
|
||||
getVectorStore: () => Promise<FaissStore>,
|
||||
options: KnowledgeBaseAddItemOptionsNonNullableAttribute
|
||||
): LoaderTask {
|
||||
const { base, item } = options
|
||||
|
||||
if (!isKnowledgeSitemapItem(item)) {
|
||||
logger.error(`Invalid item type for sitemapTask: expected 'sitemap', got '${item.type}'`)
|
||||
return {
|
||||
loaderTasks: [],
|
||||
loaderDoneReturn: {
|
||||
...LangChainFramework.ERROR_LOADER_RETURN,
|
||||
message: `Invalid item type: expected 'sitemap', got '${item.type}'`,
|
||||
messageSource: 'validation'
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const url = item.content
|
||||
|
||||
const loaderTask: LoaderTask = {
|
||||
loaderTasks: [
|
||||
{
|
||||
state: LoaderTaskItemState.PENDING,
|
||||
task: async () => {
|
||||
// 使用处理后的网页进行加载
|
||||
const vectorStore = await getVectorStore()
|
||||
return addSitemapLoader(base, vectorStore, url)
|
||||
.then((result) => {
|
||||
loaderTask.loaderDoneReturn = result
|
||||
return result
|
||||
})
|
||||
.then(async () => {
|
||||
await vectorStore.save(path.join(this.storageDir, base.id))
|
||||
})
|
||||
.catch((e) => {
|
||||
logger.error(`Error in addWebLoader for ${url}: ${e}`)
|
||||
const errorResult: LoaderReturn = {
|
||||
...LangChainFramework.ERROR_LOADER_RETURN,
|
||||
message: e.message,
|
||||
messageSource: 'embedding'
|
||||
}
|
||||
loaderTask.loaderDoneReturn = errorResult
|
||||
return errorResult
|
||||
})
|
||||
},
|
||||
evaluateTaskWorkload: { workload: 2 * MB }
|
||||
}
|
||||
],
|
||||
loaderDoneReturn: null
|
||||
}
|
||||
return loaderTask
|
||||
}
|
||||
|
||||
private noteTask(
|
||||
getVectorStore: () => Promise<FaissStore>,
|
||||
options: KnowledgeBaseAddItemOptionsNonNullableAttribute
|
||||
): LoaderTask {
|
||||
const { base, item } = options
|
||||
|
||||
if (!isKnowledgeNoteItem(item)) {
|
||||
logger.error(`Invalid item type for noteTask: expected 'note', got '${item.type}'`)
|
||||
return {
|
||||
loaderTasks: [],
|
||||
loaderDoneReturn: {
|
||||
...LangChainFramework.ERROR_LOADER_RETURN,
|
||||
message: `Invalid item type: expected 'note', got '${item.type}'`,
|
||||
messageSource: 'validation'
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const content = item.content
|
||||
const sourceUrl = item.sourceUrl ?? ''
|
||||
|
||||
logger.info(`noteTask ${content}, ${sourceUrl}`)
|
||||
|
||||
const encoder = new TextEncoder()
|
||||
const contentBytes = encoder.encode(content)
|
||||
const loaderTask: LoaderTask = {
|
||||
loaderTasks: [
|
||||
{
|
||||
state: LoaderTaskItemState.PENDING,
|
||||
task: async () => {
|
||||
// 使用处理后的笔记进行加载
|
||||
const vectorStore = await getVectorStore()
|
||||
return addNoteLoader(base, vectorStore, content, sourceUrl)
|
||||
.then((result) => {
|
||||
loaderTask.loaderDoneReturn = result
|
||||
return result
|
||||
})
|
||||
.then(async () => {
|
||||
await vectorStore.save(path.join(this.storageDir, base.id))
|
||||
})
|
||||
.catch((e) => {
|
||||
logger.error(`Error in addNoteLoader for ${sourceUrl}: ${e}`)
|
||||
const errorResult: LoaderReturn = {
|
||||
...LangChainFramework.ERROR_LOADER_RETURN,
|
||||
message: e.message,
|
||||
messageSource: 'embedding'
|
||||
}
|
||||
loaderTask.loaderDoneReturn = errorResult
|
||||
return errorResult
|
||||
})
|
||||
},
|
||||
evaluateTaskWorkload: { workload: contentBytes.length }
|
||||
}
|
||||
],
|
||||
loaderDoneReturn: null
|
||||
}
|
||||
return loaderTask
|
||||
}
|
||||
|
||||
private videoTask(
|
||||
getVectorStore: () => Promise<FaissStore>,
|
||||
options: KnowledgeBaseAddItemOptionsNonNullableAttribute
|
||||
): LoaderTask {
|
||||
const { base, item } = options
|
||||
|
||||
if (!isKnowledgeVideoItem(item)) {
|
||||
logger.error(`Invalid item type for videoTask: expected 'video', got '${item.type}'`)
|
||||
return {
|
||||
loaderTasks: [],
|
||||
loaderDoneReturn: {
|
||||
...LangChainFramework.ERROR_LOADER_RETURN,
|
||||
message: `Invalid item type: expected 'video', got '${item.type}'`,
|
||||
messageSource: 'validation'
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const files = item.content
|
||||
|
||||
const loaderTask: LoaderTask = {
|
||||
loaderTasks: [
|
||||
{
|
||||
state: LoaderTaskItemState.PENDING,
|
||||
task: async () => {
|
||||
const vectorStore = await getVectorStore()
|
||||
return addVideoLoader(base, vectorStore, files)
|
||||
.then((result) => {
|
||||
loaderTask.loaderDoneReturn = result
|
||||
return result
|
||||
})
|
||||
.then(async () => {
|
||||
await vectorStore.save(path.join(this.storageDir, base.id))
|
||||
})
|
||||
.catch((e) => {
|
||||
logger.error(`Preprocessing failed for ${files[0].name}: ${e}`)
|
||||
const errorResult: LoaderReturn = {
|
||||
...LangChainFramework.ERROR_LOADER_RETURN,
|
||||
message: e.message,
|
||||
messageSource: 'preprocess'
|
||||
}
|
||||
loaderTask.loaderDoneReturn = errorResult
|
||||
return errorResult
|
||||
})
|
||||
},
|
||||
evaluateTaskWorkload: { workload: files[0].size }
|
||||
}
|
||||
],
|
||||
loaderDoneReturn: null
|
||||
}
|
||||
return loaderTask
|
||||
}
|
||||
|
||||
private async getAllDocuments(base: KnowledgeBaseParams): Promise<Document[]> {
|
||||
logger.info(`Fetching all documents from database for knowledge base: ${base.id}`)
|
||||
|
||||
try {
|
||||
const results = (await this.getVectorStore(base)).docstore._docs
|
||||
|
||||
const documents: Document[] = Array.from(results.values())
|
||||
logger.info(`Fetched ${documents.length} documents for BM25/Hybrid retriever.`)
|
||||
return documents
|
||||
} catch (e) {
|
||||
logger.error(`Could not fetch documents from database for base ${base.id}: ${e}`)
|
||||
// 如果表不存在或查询失败，返回空数组
|
||||
return []
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user