Update index.js

index.js

@@ -172,127 +172,164 @@ class ResponseHandler {
 
 // WebSocket工具类
 class WebSocketUtils {
-  // 生成WebSocket密钥
-  static generateWebSocketKey() {
-      return randomBytes(16).toString('base64');
-  }
+    static isProcessing = false; // 添加静态属性用于并行控制
+    static TIMEOUT = 5 * 60 * 1000; // 5分钟超时时间
+
+    // 生成WebSocket密钥
+    static generateWebSocketKey() {
+        return randomBytes(16).toString('base64');
+    }
+
+    // 创建WebSocket客户端
+    static async createWebSocketClient(requestPayload) {
+        if (this.isProcessing) {
+            throw new Error('上游服务繁忙，只允许单次请求，请稍后重试！');
+        }
+
+        this.isProcessing = true;
+        let timeoutId;
+        let ws;
+
+        try {
+            return await new Promise((resolve, reject) => {
+                const websocketKey = this.generateWebSocketKey();
+                ws = new WebSocket(CONFIG.API.BASE_URL, 'graphql-transport-ws', {
+                    headers: {
+                        ...CONFIG.DEFAULT_HEADERS,
+                        'Sec-WebSocket-Key': websocketKey,
+                    }
+                });
 
-  // 创建WebSocket客户端
-  static createWebSocketClient(requestPayload) {
-      return new Promise((resolve, reject) => {
-          const websocketKey = this.generateWebSocketKey();
-          const ws = new WebSocket(CONFIG.API.BASE_URL, 'graphql-transport-ws', {
-              headers: {
-                  ...CONFIG.DEFAULT_HEADERS,
-                  'Sec-WebSocket-Key': websocketKey,
-              }
-          });
-
-          let responseContent = '';
-          let isComplete = false;
-
-          ws.on('open', () => {
-              console.log('WebSocket连接已建立');
-              const connectionInitMessage = {
-                  type: 'connection_init',
-                  payload: {
-                      headers: {
-                          Authorization: 'Bearer ee5b7c15ed3553cd6abc407340aad09ac7cb3b9f76d8613a'
-                      }
-                  }
-              };
-              ws.send(JSON.stringify(connectionInitMessage));
-          });
-
-          ws.on('message', async (data) => {
-              const message = data.toString();
-              const parsedMessage = JSON.parse(message);
-              
-              switch (parsedMessage.type) {
-                  case 'connection_ack':
-                      console.log('WebSocket连接请求中');
-                      this.sendChatSubscription(ws, requestPayload);
-                      break;
-                  case 'next':
-                      const chatResponse = await this.handleChatResponse(parsedMessage);
-                      if (chatResponse) {
-                        responseContent = chatResponse;
-                        // 获取到响应后立即关闭连接
-                        isComplete = true;
+                // 设置超时处理
+                timeoutId = setTimeout(() => {
+                    if (ws.readyState === WebSocket.OPEN) {
                         ws.close();
-                        resolve(responseContent);
                     }
-                      break;
-                  case 'complete':
-                      isComplete = true;
-                      ws.close();
-                      resolve(responseContent);
-                      break;
-              }
-          });
-
-          ws.on('error', (err) => {
-              console.error('WebSocket错误:', err);
-              reject(err);
-          });
-
-          ws.on('close', (code, reason) => {
-              console.log('请求完毕，关闭连接');
-              if (!isComplete) {
-                  reject(new Error('WebSocket closed unexpectedly'));
-              }
-          });
-      });
-  }
-
-  // 发送聊天订阅
-  static sendChatSubscription(ws, requestPayload) {
-      const subscribeMessage = {
-          id: uuidv4(),
-          type: 'subscribe',
-          payload: {
-              variables: {
-                  messageInput: requestPayload.contextMessages,
-                  messageContext: null,
-                  organizationId: 'org_JfjtEvzbwOikUEUn',
-                  integrationId: 'clwtqz9sq001izszu8ms5g4om',
-                  chatMode: 'AUTO',
-                  context: requestPayload.systemMessage,
-                  messageAttributes: {},
-                  includeAIAnnotations: false,
-                  environment: 'production'
-              },
-              extensions: {},
-              operationName: 'OnNewSessionChatResult',
-              query: `subscription OnNewSessionChatResult($messageInput: String!, $messageContext: String, $organizationId: ID!, $integrationId: ID, $chatMode: ChatMode, $filters: ChatFiltersInput, $messageAttributes: JSON, $tags: [String!], $workflowId: String, $context: String, $guidance: String, $includeAIAnnotations: Boolean!, $environment: String) {
-                  newSessionChatResult(
-                      input: {messageInput: $messageInput, messageContext: $messageContext, organizationId: $organizationId, integrationId: $integrationId, chatMode: $chatMode, filters: $filters, messageAttributes: $messageAttributes, tags: $tags, workflowId: $workflowId, context: $context, guidance: $guidance, environment: $environment}
-                  ) {
-                      isEnd
-                      sessionId
-                      message {
-                          id
-                          content
-                      }
-                      __typename
-                  }
-              }`
-          }
-      };
-
-      ws.send(JSON.stringify(subscribeMessage));
-  }
+                    reject(new Error('WebSocket连接超时（5分钟）'));
+                }, this.TIMEOUT);
+
+                let responseContent = '';
+                let isComplete = false;
+
+                ws.on('open', () => {
+                    console.log('WebSocket连接已建立');
+                    const connectionInitMessage = {
+                        type: 'connection_init',
+                        payload: {
+                            headers: {
+                                Authorization: 'Bearer ee5b7c15ed3553cd6abc407340aad09ac7cb3b9f76d8613a'
+                            }
+                        }
+                    };
+                    ws.send(JSON.stringify(connectionInitMessage));
+                });
+
+                ws.on('message', async (data) => {
+                    const message = data.toString();
+                    const parsedMessage = JSON.parse(message);
+                    
+                    switch (parsedMessage.type) {
+                        case 'connection_ack':
+                            console.log('WebSocket连接请求中');
+                            this.sendChatSubscription(ws, requestPayload);
+                            break;
+                        case 'next':
+                            const chatResponse = await this.handleChatResponse(parsedMessage);
+                            if (chatResponse) {
+                                responseContent = chatResponse;
+                                // 获取到响应后立即关闭连接
+                                isComplete = true;
+                                ws.close();
+                                resolve(responseContent);
+                            }
+                            break;
+                        case 'complete':
+                            isComplete = true;
+                            ws.close();
+                            resolve(responseContent);
+                            break;
+                    }
+                });
 
-  // 处理聊天响应
-  static async handleChatResponse(message) {
-      if (message.payload && message.payload.data) {
-          const chatResult = message.payload.data.newSessionChatResult;
-          if (chatResult && chatResult.isEnd == true && chatResult.message) {
-              return chatResult.message.content;
-          }
-      }
-      return null;
-  }
+                ws.on('error', (err) => {
+                    console.error('WebSocket错误:', err);
+                    clearTimeout(timeoutId);
+                    this.isProcessing = false;
+                    if (ws.readyState === WebSocket.OPEN) {
+                        ws.close();
+                    }
+                    reject(err);
+                });
+
+                ws.on('close', (code, reason) => {
+                    console.log('请求完毕，关闭连接');
+                    clearTimeout(timeoutId);
+                    this.isProcessing = false;
+                    if (!isComplete) {
+                        reject(new Error('WebSocket closed unexpectedly'));
+                    }
+                });
+            });
+        } catch (error) {
+            clearTimeout(timeoutId);
+            this.isProcessing = false;
+            if (ws && ws.readyState === WebSocket.OPEN) {
+                ws.close();
+            }
+            throw error;
+        }
+    }
+
+    // 发送聊天订阅
+    static sendChatSubscription(ws, requestPayload) {
+        const subscribeMessage = {
+            id: uuidv4(),
+            type: 'subscribe',
+            payload: {
+                variables: {
+                    messageInput: requestPayload.contextMessages,
+                    messageContext: null,
+                    organizationId: 'org_JfjtEvzbwOikUEUn',
+                    integrationId: 'clwtqz9sq001izszu8ms5g4om',
+                    chatMode: 'AUTO',
+                    context: requestPayload.systemMessage,
+                    messageAttributes: {},
+                    includeAIAnnotations: false,
+                    environment: 'production'
+                },
+                extensions: {},
+                operationName: 'OnNewSessionChatResult',
+                query: `subscription OnNewSessionChatResult($messageInput: String!, $messageContext: String, $organizationId: ID!, $integrationId: ID, $chatMode: ChatMode, $filters: ChatFiltersInput, $messageAttributes: JSON, $tags: [String!], $workflowId: String, $context: String, $guidance: String, $includeAIAnnotations: Boolean!, $environment: String) {
+                    newSessionChatResult(
+                        input: {messageInput: $messageInput, messageContext: $messageContext, organizationId: $organizationId, integrationId: $integrationId, chatMode: $chatMode, filters: $filters, messageAttributes: $messageAttributes, tags: $tags, workflowId: $workflowId, context: $context, guidance: $guidance, environment: $environment}
+                    ) {
+                        isEnd
+                        sessionId
+                        message {
+                            id
+                            content
+                        }
+                        __typename
+                    }
+                }`
+            }
+        };
+
+        if (ws.readyState === WebSocket.OPEN) {
+            ws.send(JSON.stringify(subscribeMessage));
+        }
+    }
+
+    // 处理聊天响应
+    static async handleChatResponse(message) {
+        if (message.payload && message.payload.data) {
+            const chatResult = message.payload.data.newSessionChatResult;
+            if (chatResult && chatResult.isEnd == true && chatResult.message) {
+                return chatResult.message.content;
+            }
+        }
+        return null;
+    }
 }
 
 // 创建Express应用
@@ -308,7 +345,7 @@ app.use(cors({
 }));
 
 // 获取模型列表路由
-app.get('/hf/v1/models', (req, res) => {
+app.get('/v1/models', (req, res) => {
   res.json({
       object: "list",
       data: [{
@@ -321,7 +358,7 @@ app.get('/hf/v1/models', (req, res) => {
 });
 
 // 聊天完成路由
-app.post('/hf/v1/chat/completions', async (req, res) => {
+app.post('/v1/chat/completions', async (req, res) => {
   try {
       const { messages, model, stream } = req.body;
       const authToken = req.headers.authorization?.replace('Bearer ', '');