ai-knowledge/src/loaders/queueManager.js

ai-knowledge

class LoaderQueue {
    constructor() {
        this.queue = [];
        this.isProcessing = false;
        this.taskStatus = new Map();
    }
    async process(job) {
        const { ragApplication, loader, source, type, taskId } = job;
        try {
            this.updateTaskStatus(taskId, 'processing', 0);
            const result = await ragApplication.addLoader(loader, job.forceReload);
            this.updateTaskStatus(taskId, 'completed', 100, result);
            console.log(`Successfully processed ${type}: ${source} ${new Date().toISOString()}`);
            return result;
        } catch (error) {
            console.error(`Error processing ${type}: ${source}`, error);
            this.updateTaskStatus(taskId, 'failed', 0, null, error.message);
            throw error;
        }
    }
    async processQueue() {
        if (this.isProcessing || this.queue.length === 0) return;

        this.isProcessing = true;
        while (this.queue.length > 0) {
            const job = this.queue.shift();
            await this.process(job);
        }
        this.isProcessing = false;
    }
    add(data) {
        this.queue.push(data);
        this.updateTaskStatus(data.taskId, 'queued', 0);
        this.processQueue();
    }
    updateTaskStatus(taskId, state, progress, result = null, error = null) {
        this.taskStatus.set(taskId, {
            state,
            progress,
            result,
            error,
            updatedAt: new Date()
        });
    }
    getTaskStatus(taskId) {
        const status = this.taskStatus.get(taskId);
        if (!status) {
            return {
                success: false,
                error: 'Task not found'
            };
        }
        return {
            success: true,
            taskId,
            ...status
        };
    }
}

// 创建单例实例
const loaderQueue = new LoaderQueue();
module.exports = {
    add: (data) => loaderQueue.add(data),
    getTaskStatus: (taskId) => loaderQueue.getTaskStatus(taskId),
    queue: loaderQueue
};