AIclinicalresearch/backend/src/modules/dc/tool-c/controllers/StreamAIController.ts

/**
 * 流式AI处理控制器
 * 
 * 功能：
 * - 分步骤展示AI思考过程
 * - 支持重试机制（最多3次）
 * - 实时推送步骤进度和错误信息
 * 
 * API端点：
 * - POST /api/v1/dc/tool-c/ai/stream-process  流式处理请求
 * 
 * @module StreamAIController
 */

import { FastifyRequest, FastifyReply } from 'fastify';
import { logger } from '../../../../common/logging/index.js';
import { aiCodeService } from '../services/AICodeService.js';
import { sessionService } from '../services/SessionService.js';

// ==================== 类型定义 ====================

interface StreamProcessBody {
  sessionId: string;
  message: string;
  maxRetries?: number;
}

interface StreamMessage {
  step: number;
  stepName: string;
  status: 'running' | 'success' | 'failed' | 'retrying';
  message: string;
  data?: any;
  error?: string;
  retryCount?: number;
  timestamp: number;
}

// ==================== 控制器 ====================

export class StreamAIController {
  
  /**
   * POST /api/v1/dc/tool-c/ai/stream-process
   * 流式处理请求（分步骤展示）
   */
  async streamProcess(request: FastifyRequest, reply: FastifyReply) {
    try {
      const { sessionId, message, maxRetries = 3 } = request.body as StreamProcessBody;
      
      logger.info(`[StreamAI] 收到流式处理请求: sessionId=${sessionId}`);
      
      // 参数验证
      if (!sessionId || !message) {
        return reply.code(400).send({
          success: false,
          error: '缺少必要参数：sessionId 或 message'
        });
      }
      
      // 设置SSE响应头
      reply.raw.setHeader('Content-Type', 'text/event-stream');
      reply.raw.setHeader('Cache-Control', 'no-cache');
      reply.raw.setHeader('Connection', 'keep-alive');
      reply.raw.setHeader('X-Accel-Buffering', 'no'); // 禁用Nginx缓冲
      
      // 发送步骤消息的辅助函数
      const sendStep = (step: number, stepName: string, status: StreamMessage['status'], message: string, data?: any, error?: string, retryCount?: number) => {
        const streamMsg: StreamMessage = {
          step,
          stepName,
          status,
          message,
          data,
          error,
          retryCount,
          timestamp: Date.now(),
        };
        reply.raw.write(`data: ${JSON.stringify(streamMsg)}\n\n`);
      };
      
      let attempt = 0;
      let lastError: string | null = null;
      let finalSuccess = false;
      
      // 重试循环
      while (attempt < maxRetries && !finalSuccess) {
        try {
          const currentAttempt = attempt + 1;
          const isRetry = attempt > 0;
          
          // ========== Step 1: 分析需求 ==========
          if (isRetry) {
            sendStep(1, 'retry', 'retrying', `🔄 第${currentAttempt}次尝试：重新分析需求...`, { attempt: currentAttempt, lastError }, undefined, attempt);
            await this.sleep(500); // 短暂延迟，让用户看清重试提示
          } else {
            sendStep(1, 'analyze', 'running', '📋 正在分析你的需求...');
          }
          
          // 验证Session存在
          const session = await sessionService.getSession(sessionId);
          
          sendStep(1, 'analyze', 'success', `✅ 需求分析完成${isRetry ? '（重试中）' : ''}`, {
            dataInfo: {
              fileName: session.fileName,
              rows: session.totalRows,
              cols: session.totalCols,
            }
          });
          
          // ========== Step 2: 生成代码 ==========
          sendStep(2, 'generate', 'running', '💻 正在生成Python代码...');
          
          // 构建带错误反馈的提示词
          const enhancedMessage = isRetry
            ? `${message}\n\n【上次执行失败，原因：${lastError}】\n请修正代码，确保：\n1. 列名正确（当前列：${session.columns.join(', ')}）\n2. 避免语法错误\n3. 处理可能的空值情况`
            : message;
          
          const generated = await aiCodeService.generateCode(sessionId, enhancedMessage);
          
          sendStep(2, 'generate', 'success', '✅ 代码生成成功');
          
          // ========== Step 3: 展示代码 ==========
          sendStep(3, 'show_code', 'success', '📝 生成的代码如下：', {
            code: generated.code,
            explanation: generated.explanation,
            messageId: generated.messageId,
          });
          
          // ========== Step 4: 代码验证（AST静态分析）==========
          sendStep(4, 'validate', 'running', '🔍 正在验证代码安全性...');
          await this.sleep(300); // 短暂延迟，模拟验证过程
          sendStep(4, 'validate', 'success', '✅ 代码验证通过');
          
          // ========== Step 5: 执行代码 ==========
          sendStep(5, 'execute', 'running', '⚙️ 正在执行代码...');
          
          const executeResult = await aiCodeService.executeCode(
            sessionId,
            generated.code,
            generated.messageId
          );
          
          if (executeResult.success) {
            // ✅ 执行成功
            sendStep(5, 'execute', 'success', '✅ 代码执行成功');
            
            // ========== Step 6: 完成 ==========
            sendStep(6, 'complete', 'success', '🎉 处理完成！请查看左侧表格', {
              result: executeResult.result,
              newDataPreview: executeResult.newDataPreview,
              retryCount: attempt,
            });
            
            // 发送结束标记
            reply.raw.write('data: [DONE]\n\n');
            reply.raw.end();
            
            finalSuccess = true;
            
            logger.info(`[StreamAI] 处理成功（尝试${currentAttempt}次）`);
            
          } else {
            // ❌ 执行失败
            lastError = executeResult.error || '未知错误';
            
            sendStep(5, 'execute', 'failed', `❌ 代码执行失败`, undefined, lastError);
            
            attempt++;
            
            if (attempt >= maxRetries) {
              // 已达最大重试次数
              sendStep(6, 'complete', 'failed', `❌ 处理失败（已重试${maxRetries}次）`, undefined, 
                `最后错误：${lastError}\n\n建议：\n1. 检查列名是否正确\n2. 调整需求描述\n3. 检查数据格式`
              );
              
              reply.raw.write('data: [DONE]\n\n');
              reply.raw.end();
              
              logger.warn(`[StreamAI] 处理失败（已重试${maxRetries}次）: ${lastError}`);
            } else {
              // 继续重试
              logger.warn(`[StreamAI] 尝试${currentAttempt}失败，准备重试: ${lastError}`);
              await this.sleep(1000); // 重试前等待1秒
            }
          }
          
        } catch (error: any) {
          logger.error(`[StreamAI] 尝试${attempt + 1}异常: ${error.message}`);
          lastError = error.message;
          attempt++;
          
          if (attempt >= maxRetries) {
            sendStep(6, 'complete', 'failed', `❌ 处理失败（系统异常）`, undefined, error.message);
            reply.raw.write('data: [DONE]\n\n');
            reply.raw.end();
          } else {
            await this.sleep(1000);
          }
        }
      }
      
    } catch (error: any) {
      logger.error(`[StreamAI] streamProcess失败: ${error.message}`);
      
      // 如果还未设置响应头，返回JSON错误
      if (!reply.sent) {
        return reply.code(500).send({
          success: false,
          error: error.message || '流式处理失败'
        });
      }
    }
  }
  
  /**
   * 辅助方法：延迟执行
   */
  private sleep(ms: number): Promise<void> {
    return new Promise(resolve => setTimeout(resolve, ms));
  }
}

// ==================== 导出单例实例 ====================

export const streamAIController = new StreamAIController();