api-server/src/modules/knowledge-source/knowledge-source.service.ts

import { Injectable, NotFoundException, BadRequestException } from '@nestjs/common';
import { KnowledgeSourceRepository } from './knowledge-source.repository';
import { DocumentImportRepository } from '../document-import/document-import.repository';
import { QueueService } from '../../infrastructure/queue/queue.service';
import { RedisService } from '../../infrastructure/redis/redis.service';
import { FilesService } from '../files/files.service';

@Injectable()
export class KnowledgeSourceService {
  constructor(
    private readonly repository: KnowledgeSourceRepository,
    private readonly importRepo: DocumentImportRepository,
    private readonly queue: QueueService,
    private readonly redis: RedisService,
    private readonly filesService: FilesService,
  ) {}

  async addSource(userId: string, knowledgeBaseId: string, dto: {
    fileId?: string;
    type?: string;
    title?: string;
    originalFilename?: string;
    mimeType?: string;
    sizeBytes?: number;
    originalObjectKey?: string;
  }) {
    const source = await this.repository.create(userId, knowledgeBaseId, dto);

    // 自动创建 import 任务并入队
    const importJob = await this.importRepo.create({
      userId,
      knowledgeBaseId,
      sourceId: source.id,
      fileId: dto.fileId,
      sourceType: dto.type ?? 'file',
      sourceName: dto.originalFilename ?? dto.title ?? '',
      status: 'QUEUED',
    });

    // 设置 Redis 状态
    await this.redis.set(`job:document-import:${importJob.id}:status`, 'pending', 86400);
    await this.redis.set(`job:document-import:${importJob.id}:progress`, '0', 86400);
    await this.redis.set(`job:document-import:${importJob.id}:message`, '任务已加入队列', 86400);

    // 入队处理
    await this.queue.add('document-import', {
      importId: importJob.id,
      userId,
      knowledgeBaseId,
      sourceId: source.id,
      fileName: dto.originalFilename ?? dto.title,
    });

    return source;
  }

  async findByKnowledgeBase(knowledgeBaseId: string, pagination: { page?: number; limit?: number }) {
    return this.repository.findByKnowledgeBase(knowledgeBaseId, pagination);
  }

  async findOne(id: string) {
    const source = await this.repository.findById(id);
    if (!source) throw new NotFoundException('资料来源不存在');
    return source;
  }

  async remove(userId: string, id: string) {
    const source = await this.repository.findById(id);
    if (!source) throw new NotFoundException('资料来源不存在');
    if (source.fileId) {
      await this.filesService.deleteFile(userId, source.fileId);
    }
    return this.repository.softDelete(id);
  }

  async triggerParse(sourceId: string, userId: string) {
    const source = await this.repository.findById(sourceId);
    if (!source) throw new NotFoundException('资料来源不存在');

    // 创建新的 import 任务
    const importJob = await this.importRepo.create({
      userId,
      knowledgeBaseId: source.knowledgeBaseId,
      sourceId: source.id,
      sourceType: source.type ?? 'file',
      sourceName: source.originalFilename ?? source.title ?? '',
      status: 'QUEUED',
    });

    // 设置 Redis 状态
    await this.redis.set(`job:document-import:${importJob.id}:status`, 'pending', 86400);
    await this.redis.set(`job:document-import:${importJob.id}:progress`, '0', 86400);
    await this.redis.set(`job:document-import:${importJob.id}:message`, '手动触发重新解析', 86400);

    // 入队
    await this.queue.add('document-import', {
      importId: importJob.id,
      userId,
      knowledgeBaseId: source.knowledgeBaseId,
      sourceId: source.id,
      fileName: source.originalFilename ?? source.title,
    });

    return { id: importJob.id, status: 'queued' };
  }

  async updateParseStatus(id: string, parseStatus: string, data?: any) {
    return this.repository.updateParseStatus(id, parseStatus, data);
  }

  async updateIndexStatus(id: string, indexStatus: string, errorCode?: string, errorMessage?: string) {
    return this.repository.updateIndexStatus(id, indexStatus, errorCode, errorMessage);
  }
}