api-server/src/modules/rag/internal-rag.controller.ts

import { Controller, Get, Post, Body, Param, UseGuards } from '@nestjs/common';
import { ApiTags } from '@nestjs/swagger';
import { DocumentImportRepository } from '../document-import/document-import.repository';
import { KnowledgeSourceRepository } from '../knowledge-source/knowledge-source.repository';
import { ImportCandidateRepository } from '../import-candidate/import-candidate.repository';
import { PrismaService } from '../../infrastructure/database/prisma.service';
import { StorageService } from '../../infrastructure/storage/storage.service';
import { InternalAuthGuard } from '../../common/guards/internal-auth.guard';

@ApiTags('internal-rag')
@Controller('internal/rag')
@UseGuards(InternalAuthGuard)
export class InternalRagController {
  constructor(
    private readonly importRepo: DocumentImportRepository,
    private readonly sourceRepo: KnowledgeSourceRepository,
    private readonly candidateRepo: ImportCandidateRepository,
    private readonly prisma: PrismaService,
    private readonly storage: StorageService,
  ) {}

  @Get('jobs/next')
  async getNextJob() {
    const job = await this.importRepo.claimNext('');  // 先查询，不认领
    if (!job) return { job: null };
    return {
      job: {
        id: job.id,
        userId: job.userId,
        knowledgeBaseId: job.knowledgeBaseId,
        sourceId: job.sourceId,
        fileId: job.fileId,
        sourceType: job.sourceType,
        sourceName: job.sourceName,
        rawText: job.rawText,
        status: job.status,
      },
    };
  }

  @Get('jobs/:id')
  async getJobDetail(@Param('id') id: string) {
    const job = await this.importRepo.findById(id);
    if (!job) return { job: null, source: null };

    const source = job.sourceId
      ? await this.sourceRepo.findById(job.sourceId)
      : null;

    let downloadUrl = '';
    if (source?.originalObjectKey) {
      try { downloadUrl = await this.storage.getDownloadUrl(source.originalObjectKey); } catch {}
    }

    return {
      job: {
        id: job.id,
        userId: job.userId,
        knowledgeBaseId: job.knowledgeBaseId,
        sourceId: job.sourceId,
        fileId: job.fileId,
        rawText: job.rawText,
        status: job.status,
      },
      source: source ? {
        id: source.id,
        type: source.type,
        originalFilename: source.originalFilename,
        mimeType: source.mimeType,
        sizeBytes: Number(source.sizeBytes),
        originalObjectKey: source.originalObjectKey,
        downloadUrl,
      } : null,
    };
  }

  @Post('jobs/:id/claim')
  async claimJob(@Param('id') id: string, @Body() body: { workerId?: string }) {
    const workerId = body.workerId || 'unknown';
    const result = await this.importRepo.claim(id, workerId);
    return { success: result.count > 0 };
  }

  @Post('jobs/:id/heartbeat')
  async heartbeat(@Param('id') id: string) {
    await this.importRepo.heartbeat(id);
    return { success: true };
  }

  @Post('jobs/:id/status')
  async updateStatus(
    @Param('id') id: string,
    @Body() body: { status: string; progress?: number; errorCode?: string; errorMessage?: string },
  ) {
    await this.importRepo.updateStatus(id, body.status, {
      step: body.status,
      progress: body.progress,
      errorCode: body.errorCode,
      errorMessage: body.errorMessage,
    });
    return { success: true };
  }

  @Post('chunks')
  async saveChunks(@Body() body: { chunks: any[] }) {
    const chunks = body.chunks || [];
    if (chunks.length > 0) {
      await this.prisma.knowledgeChunk.createMany({ data: chunks });

      // 更新 source 的 textLength 和 parseStatus
      const sourceIds = [...new Set(chunks.map(c => c.sourceId).filter(Boolean))] as string[];
      if (sourceIds.length > 0) {
        for (const sourceId of sourceIds) {
          const sourceChunks = chunks.filter(c => c.sourceId === sourceId);
          const totalChars = sourceChunks.reduce((sum, c) => sum + (c.tokenCount || c.content?.length || 0), 0);
          await this.prisma.knowledgeSource.updateMany({
            where: { id: sourceId },
            data: { textLength: totalChars, parseStatus: 'completed', indexStatus: 'completed', updatedAt: new Date() },
          });
        }
      }
    }
    return { success: true, count: chunks.length };
  }

  @Post('candidates')
  async saveCandidates(
    @Body() body: {
      userId: string;
      knowledgeBaseId: string;
      sourceId: string;
      importId: string;
      candidates: any[];
    },
  ) {
    await this.candidateRepo.createMany(
      body.userId,
      body.knowledgeBaseId,
      body.sourceId,
      body.importId,
      body.candidates || [],
    );

    // 更新 source 的学习状态
    if (body.sourceId && body.candidates?.length > 0) {
      await this.prisma.knowledgeSource.updateMany({
        where: { id: body.sourceId },
        data: { learningStatus: 'completed', updatedAt: new Date() },
      });
    }

    return { success: true, count: body.candidates?.length || 0 };
  }
}