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 { 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,
  ) {}

  @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;

    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,
      } : 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 });
    }
    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 || [],
    );
    return { success: true, count: body.candidates?.length || 0 };
  }
}