KnowledgeBase/server/src/api/api-v1.controller.ts

import {
    Body,
    Controller,
    Delete,
    Get,
    Param,
    Post,
    Request,
    Res,
    UploadedFile,
    UseGuards,
    UseInterceptors,
} from '@nestjs/common';
import { FileInterceptor } from '@nestjs/platform-express';
import { Response } from 'express';
import { ApiKeyGuard } from '../auth/api-key.guard';
import { RagService } from '../rag/rag.service';
import { ChatService } from '../chat/chat.service';
import { KnowledgeBaseService } from '../knowledge-base/knowledge-base.service';
import { ModelConfigService } from '../model-config/model-config.service';
import { UserSettingService } from '../user-setting/user-setting.service';

@Controller('v1')
@UseGuards(ApiKeyGuard)
export class ApiV1Controller {
    constructor(
        private readonly ragService: RagService,
        private readonly chatService: ChatService,
        private readonly knowledgeBaseService: KnowledgeBaseService,
        private readonly modelConfigService: ModelConfigService,
        private readonly userSettingService: UserSettingService,
    ) { }

    // ========== Chat / RAG ==========
    /**
     * POST /api/v1/chat
     * Tenant-scoped RAG chat. Supports both streaming (SSE) and standard JSON responses.
     * Body: { message, stream?, selectedGroups?, selectedFiles? }
     */
    @Post('chat')
    async chat(
        @Request() req,
        @Body()
        body: {
            message: string;
            stream?: boolean;
            selectedGroups?: string[];
            selectedFiles?: string[];
        },
        @Res() res: Response,
    ) {
        const { message, stream = false, selectedGroups, selectedFiles } = body;
        const user = req.user;

        if (!message) {
            return res.status(400).json({ error: 'message is required' });
        }

        // Get user settings and model configuration
        const userSetting = await this.userSettingService.findOrCreate(user.id);
        const models = await this.modelConfigService.findAll(user.id, user.tenantId);
        const llmModel = models.find((m) => m.id === userSetting?.selectedLLMId) ?? models.find((m) => m.type === 'llm' && m.isDefault);

        if (!llmModel) {
            return res.status(400).json({ error: 'No LLM model configured for this user' });
        }

        const modelConfig = llmModel as any;

        if (stream) {
            res.setHeader('Content-Type', 'text/event-stream');
            res.setHeader('Cache-Control', 'no-cache');
            res.setHeader('Connection', 'keep-alive');

            try {
                const stream = this.chatService.streamChat(
                    message,
                    [],                                       // history
                    user.id,
                    modelConfig,
                    userSetting?.language ?? 'zh',            // userLanguage
                    userSetting?.selectedEmbeddingId,          // selectedEmbeddingId
                    selectedGroups,                           // selectedGroups
                    selectedFiles,                            // selectedFiles
                    undefined,                                // historyId
                    userSetting?.enableRerank ?? false,        // enableRerank
                    userSetting?.selectedRerankId,             // selectedRerankId
                    userSetting?.temperature,                  // temperature
                    userSetting?.maxTokens,                    // maxTokens
                    userSetting?.topK ?? 5,                    // topK
                    userSetting?.similarityThreshold ?? 0.3,   // similarityThreshold
                    userSetting?.rerankSimilarityThreshold ?? 0.5, // rerankSimilarityThreshold
                    userSetting?.enableQueryExpansion ?? false, // enableQueryExpansion
                    userSetting?.enableHyDE ?? false,           // enableHyDE
                    user.tenantId,                             // Passing tenantId correctly
                );

                for await (const chunk of stream) {
                    res.write(`data: ${JSON.stringify(chunk)}\n\n`);
                }
                res.write('data: [DONE]\n\n');
                res.end();
            } catch (error) {
                res.write(`data: ${JSON.stringify({ type: 'error', data: error.message })}\n\n`);
                res.end();
            }
        } else {
            // Non-streaming: collect all chunks and return as JSON
            try {
                let fullContent = '';
                const sources: any[] = [];
                let historyId: string | undefined;

                const chatStream = this.chatService.streamChat(
                    message,
                    [],
                    user.id,
                    modelConfig,
                    userSetting?.language ?? 'zh',
                    userSetting?.selectedEmbeddingId,
                    selectedGroups,
                    selectedFiles,
                    undefined,                                // historyId
                    userSetting?.enableRerank ?? false,
                    userSetting?.selectedRerankId,
                    userSetting?.temperature,
                    userSetting?.maxTokens,
                    userSetting?.topK ?? 5,
                    userSetting?.similarityThreshold ?? 0.3,
                    userSetting?.rerankSimilarityThreshold ?? 0.5,
                    userSetting?.enableQueryExpansion ?? false,
                    userSetting?.enableHyDE ?? false,
                    user.tenantId,                            // Passing tenantId correctly
                );

                for await (const chunk of chatStream) {
                    if (chunk.type === 'content') fullContent += chunk.data;
                    else if (chunk.type === 'sources') sources.push(...chunk.data);
                    else if (chunk.type === 'historyId') historyId = chunk.data;
                }

                return res.json({ content: fullContent, sources, historyId });
            } catch (error) {
                return res.status(500).json({ error: error.message });
            }
        }
    }

    // ========== Search ==========
    /**
     * POST /api/v1/search
     * Tenant-scoped hybrid search across knowledge base.
     * Body: { query, topK?, threshold?, selectedGroups?, selectedFiles? }
     */
    @Post('search')
    async search(
        @Request() req,
        @Body()
        body: {
            query: string;
            topK?: number;
            threshold?: number;
            selectedGroups?: string[];
            selectedFiles?: string[];
        },
    ) {
        const { query, topK = 5, threshold = 0.3, selectedGroups, selectedFiles } = body;
        const user = req.user;

        if (!query) return { error: 'query is required' };

        const userSetting = await this.userSettingService.findOrCreate(user.id);

        const results = await this.ragService.searchKnowledge(
            query,
            user.id,
            topK,
            threshold,
            userSetting?.selectedEmbeddingId,
            userSetting?.enableFullTextSearch ?? false,
            userSetting?.enableRerank ?? false,
            userSetting?.selectedRerankId,
            selectedGroups,
            selectedFiles,
            userSetting?.rerankSimilarityThreshold ?? 0.5,
            user.tenantId,
            userSetting?.enableQueryExpansion ?? false,
            userSetting?.enableHyDE ?? false,
        );

        return { results, total: results.length };
    }

    // ========== Knowledge Base ==========
    /**
     * GET /api/v1/knowledge-bases
     * List all files belonging to the caller's tenant.
     */
    @Get('knowledge-bases')
    async listFiles(@Request() req) {
        const user = req.user;
        const files = await this.knowledgeBaseService.findAll(user.id, user.tenantId);
        return {
            files: files.map((f) => ({
                id: f.id,
                name: f.originalName,
                title: f.title,
                status: f.status,
                size: f.size,
                mimetype: f.mimetype,
                createdAt: f.createdAt,
            })),
            total: files.length,
        };
    }

    /**
     * POST /api/v1/knowledge-bases/upload
     * Upload and index a file into the caller's tenant knowledge base.
     */
    @Post('knowledge-bases/upload')
    @UseInterceptors(FileInterceptor('file'))
    async uploadFile(
        @Request() req,
        @UploadedFile() file: Express.Multer.File,
        @Body() body: { mode?: 'fast' | 'precise'; chunkSize?: number; chunkOverlap?: number },
    ) {
        if (!file) return { error: 'file is required' };
        const user = req.user;

        const kb = await this.knowledgeBaseService.createAndIndex(
            file,
            user.id,
            user.tenantId,
            {
                mode: body.mode ?? 'fast',
                chunkSize: body.chunkSize ? Number(body.chunkSize) : 1000,
                chunkOverlap: body.chunkOverlap ? Number(body.chunkOverlap) : 200,
            },
        );

        return {
            id: kb.id,
            name: kb.originalName,
            status: kb.status,
            message: 'File uploaded and indexing started',
        };
    }

    /**
     * DELETE /api/v1/knowledge-bases/:id
     * Delete a specific file from the knowledge base.
     */
    @Delete('knowledge-bases/:id')
    async deleteFile(@Request() req, @Param('id') id: string) {
        const user = req.user;
        await this.knowledgeBaseService.deleteFile(id, user.id, user.tenantId);
        return { message: 'File deleted successfully' };
    }

    @Get('knowledge-bases/:id')
    async getFile(@Request() req, @Param('id') id: string) {
        const user = req.user;
        const files = await this.knowledgeBaseService.findAll(user.id, user.tenantId);
        const file = files.find((f) => f.id === id);
        if (!file) return { error: 'File not found' };
        return file;
    }
}