aurak/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;

      // 获取用户的LLM模型配置
      let models = await this.modelConfigService.findAll();

      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);
        console.log('使用选中的LLM模型:', llmModel.name);
      } else {
        // Use organization's default LLM from Index Chat Config (strict)
        llmModel = await this.modelConfigService.findDefaultByType(
          tenantId,
          ModelType.LLM,
        );
        console.log(
          '最终使用的LLM模型 (默认):',
          llmModel ? llmModel.name : '无',
        );
      }

      // 设置 SSE 响应头
      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,
        userLanguage,
        selectedEmbeddingId,
        selectedGroups,
        selectedFiles,
        historyId,
        enableRerank,
        selectedRerankId,
        temperature, // 传递 temperature 参数
        maxTokens, // 传递 maxTokens 参数
        topK, // 传递 topK 参数
        similarityThreshold, // 传递 similarityThreshold 参数
        rerankSimilarityThreshold, // 传递 rerankSimilarityThreshold 参数
        enableQueryExpansion, // 传递 enableQueryExpansion
        enableHyDE, // 传递 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();
    }
  }
}