KnowledgeBase/server/src/chat/chat.controller.ts

import {
  Body,
  Controller,
  Post,
  Request,
  Res,
  UseGuards,
} from '@nestjs/common';
import { Response } from 'express';
import { ChatMessage, ChatService } from './chat.service';
import { CombinedAuthGuard } from '../auth/combined-auth.guard';
import { ModelConfigService } from '../model-config/model-config.service';
import { TenantService } from '../tenant/tenant.service';
import { ModelType } from '../types';

class StreamChatDto {
  message: string;
  history: ChatMessage[];
  userLanguage?: string;
  selectedEmbeddingId?: string;
  selectedLLMId?: string;
  selectedGroups?: string[];
  selectedFiles?: string[];
  historyId?: string;
  enableRerank?: boolean;
  selectedRerankId?: string;
  temperature?: number;
  maxTokens?: number;
  topK?: number;
  similarityThreshold?: number;
  rerankSimilarityThreshold?: number;
  enableQueryExpansion?: boolean;
  enableHyDE?: boolean;
}

@Controller('chat')
@UseGuards(CombinedAuthGuard)
export class ChatController {
  constructor(
    private chatService: ChatService,
    private modelConfigService: ModelConfigService,
    private tenantService: TenantService,
  ) { }

  @Post('stream')
  async streamChat(
    @Request() req,
    @Body() body: StreamChatDto,
    @Res() res: Response,
  ) {
    try {
      console.log('Full Request Body:', JSON.stringify(body, null, 2));
      const { message, history = [], userLanguage = 'zh', selectedEmbeddingId, selectedLLMId, selectedGroups, selectedFiles, historyId, enableRerank, selectedRerankId, temperature, maxTokens, topK, similarityThreshold, rerankSimilarityThreshold, enableQueryExpansion, enableHyDE } = body;
      const userId = req.user.id;

      console.log('=== Chat Debug Info ===');
      console.log('User ID:', userId);
      console.log('Message:', message);
      console.log('User Language:', userLanguage);
      console.log('Selected Embedding ID:', selectedEmbeddingId);
      console.log('Selected LLM ID:', selectedLLMId);
      console.log('Selected Groups:', selectedGroups);
      console.log('Selected Files:', selectedFiles);
      console.log('History ID:', historyId);
      console.log('Temperature:', temperature);
      console.log('Max Tokens:', maxTokens);
      console.log('Top K:', topK);
      console.log('Similarity Threshold:', similarityThreshold);
      console.log('Rerank Similarity Threshold:', rerankSimilarityThreshold);
      console.log('Query Expansion:', enableQueryExpansion);
      console.log('HyDE:', enableHyDE);

      const role = req.user.role;
      const tenantId = req.user.tenantId;


      let models = await this.modelConfigService.findAll(userId, tenantId);

      if (role !== 'SUPER_ADMIN') {
        const tenantSettings = await this.tenantService.getSettings(tenantId);
        const enabledIds = tenantSettings.enabledModelIds || [];
        // Only allow models that are enabled by the tenant admin
        models = models.filter(m => enabledIds.includes(m.id));
      }

      let llmModel;
      if (selectedLLMId) {
        // Find specifically selected model
        llmModel = await this.modelConfigService.findOne(selectedLLMId, userId, tenantId);
        console.log('Using selected LLM model:', llmModel.name);
      } else {
        // Use organization's default LLM from Index Chat Config (strict)
        llmModel = await this.modelConfigService.findDefaultByType(tenantId, ModelType.LLM);
        console.log('Final LLM model used (default):', llmModel ? llmModel.name : 'None');
      }


      res.setHeader('Content-Type', 'text/event-stream');
      res.setHeader('Cache-Control', 'no-cache');
      res.setHeader('Connection', 'keep-alive');
      res.setHeader('Access-Control-Allow-Origin', '*');

      if (!llmModel) {
        res.write(
          `data: ${JSON.stringify({ type: 'error', data: 'Please add LLM model and configure API key in model management' })}\n\n`,
        );
        res.write('data: [DONE]\n\n');
        res.end();
        return;
      }

      const stream = this.chatService.streamChat(
        message,
        history,
        userId,
        llmModel as any,
        userLanguage,
        selectedEmbeddingId,
        selectedGroups,
        selectedFiles,
        historyId,
        enableRerank,
        selectedRerankId,
        temperature,
        maxTokens,
        topK,
        similarityThreshold,
        rerankSimilarityThreshold,
        enableQueryExpansion,
        enableHyDE,
        req.user.tenantId // Pass tenant ID
      );

      for await (const chunk of stream) {
        res.write(`data: ${JSON.stringify(chunk)}\n\n`);
      }

      res.write('data: [DONE]\n\n');
      res.end();
    } catch (error) {
      console.error('Stream chat error:', error);
      try {
        res.write(
          `data: ${JSON.stringify({ type: 'error', data: error.message || 'Server Error' })}\n\n`,
        );
        res.write('data: [DONE]\n\n');
        res.end();
      } catch (writeError) {
        console.error('Failed to write error response:', writeError);
      }
    }
  }

  @Post('assist')
  async streamAssist(
    @Request() req,
    @Body() body: { instruction: string; context: string },
    @Res() res: Response,
  ) {
    try {
      const { instruction, context } = body;
      const userId = req.user.id; // Corrected to use req.user.id
      const tenantId = req.user.tenantId;
      const role = req.user.role;

      // Use organization's default LLM from Index Chat Config (strict)
      const llmModel = await this.modelConfigService.findDefaultByType(tenantId, ModelType.LLM);

      res.setHeader('Content-Type', 'text/event-stream');
      res.setHeader('Cache-Control', 'no-cache');
      res.setHeader('Connection', 'keep-alive');
      res.setHeader('Access-Control-Allow-Origin', '*');

      if (!llmModel) {
        res.write(
          `data: ${JSON.stringify({ type: 'error', data: 'LLM model configuration not found' })}\n\n`,
        );
        res.write('data: [DONE]\n\n');
        res.end();
        return;
      }

      const stream = this.chatService.streamAssist(
        instruction,
        context,
        llmModel as any,
      );

      for await (const chunk of stream) {
        res.write(`data: ${JSON.stringify(chunk)}\n\n`);
      }

      res.write('data: [DONE]\n\n');
      res.end();
    } catch (error) {
      console.error('Stream assist error:', error);
      res.write(
        `data: ${JSON.stringify({ type: 'error', data: error.message || 'Server Error' })}\n\n`,
      );
      res.write('data: [DONE]\n\n');
      res.end();
    }
  }
}