import { Controller, Get, Post, Body, Param } 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'; @ApiTags('internal-rag') @Controller('internal/rag') 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 }; } }