Update index.js

index.js

@@ -10,165 +10,155 @@ dotenv.config();
 
 // 配置常量
 const CONFIG = {
-  API: {
-      BASE_URL: "wss://api.inkeep.com/graphql",
-      API_KEY: process.env.API_KEY || "sk-123456",
-      SYSTEM_MESSAGE: process.env.SYSTEM_MESSAGE || null
-  },
-  MODELS: {
-      'claude-3-5-sonnet-20241022': 'claude-3-5-sonnet-20241022',
-  },
-  SERVER: {
-      PORT: process.env.PORT || 3000,
-      BODY_LIMIT: '5mb'
-  },
-  DEFAULT_HEADERS: {
-      'Host': 'api.inkeep.com',
-      'Connection': 'Upgrade',
-      'Pragma': 'no-cache',
-      'Cache-Control': 'no-cache',
-      'User-Agent': 'Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/127.0.6533.72 Safari/537.36',
-      'Upgrade': 'websocket',
-      'Origin': 'https://docs.anthropic.com',
-      'Sec-WebSocket-Version': '13',
-      'Accept-Encoding': 'gzip, deflate, br, zstd',
-      'Accept-Language': 'zh-CN,zh;q=0.9',
-      'Sec-WebSocket-Extensions': 'permessage-deflate; client_max_window_bits',
-      'Sec-WebSocket-Protocol': 'graphql-transport-ws'
-  }
+    API: {
+        BASE_URL: "wss://api.inkeep.com/graphql",
+        API_KEY: process.env.API_KEY || "sk-123456",
+        SYSTEM_MESSAGE: process.env.SYSTEM_MESSAGE || null
+    },
+    MODELS: {
+        'claude-3-5-sonnet-20241022': 'claude-3-5-sonnet-20241022',
+    },
+    SERVER: {
+        PORT: process.env.PORT || 3000,
+        BODY_LIMIT: '5mb'
+    },
+    DEFAULT_HEADERS: {
+        'Host': 'api.inkeep.com',
+        'Connection': 'Upgrade',
+        'Pragma': 'no-cache',
+        'Cache-Control': 'no-cache',
+        'User-Agent': 'Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/127.0.6533.72 Safari/537.36',
+        'Upgrade': 'websocket',
+        'Origin': 'https://docs.anthropic.com',
+        'Sec-WebSocket-Version': '13',
+        'Accept-Encoding': 'gzip, deflate, br, zstd',
+        'Accept-Language': 'zh-CN,zh;q=0.9',
+        'Sec-WebSocket-Extensions': 'permessage-deflate; client_max_window_bits',
+        'Sec-WebSocket-Protocol': 'graphql-transport-ws'
+    }
 };
 
 // AI API 客户端类
 class AiApiClient {
-  constructor(modelId) {
-      this.modelId = CONFIG.MODELS[modelId];
-      if (!this.modelId) {
-          throw new Error(`不支持的模型: ${modelId}`);
-      }
-  }
-
-  // 处理消息内容
-  processMessageContent(content) {
-      if (typeof content === 'string') return content;
-      if (Array.isArray(content)) {
-          return content
-              .filter(item => item.type === 'text')
-              .map(item => item.text)
-              .join('\n');
-      }
-      return typeof content === 'object' ? content.text || null : null;
-  }
-
-  // 转换消息格式
-  async transformMessages(request) {
-      let systemMessageList = [];
-      let systemMergeMode = false;
-      let closedSystemMergeMode = false;
-
-      const contextMessages = await request.messages.reduce(async (accPromise, current) => {
-          const acc = await accPromise;
-          const currentContent = this.processMessageContent(current.content);
-          
-          if (currentContent === null) return acc;
-
-          const currentMessageRole = current.role === "system" || current.role === "user" ? "Human" : "Assistant";
-
-          // 系统消息处理逻辑
-          if (current.role === "system") {
-              if (!closedSystemMergeMode) {
-                  systemMergeMode = true;
-                  const lastSystemMessage = systemMessageList[systemMessageList.length - 1];
-
-                  if (!lastSystemMessage) {
-                      systemMessageList.push(currentContent);
-                  } else {
-                      systemMessageList[systemMessageList.length - 1] = `${lastSystemMessage}\n${currentContent}`;
-                  }
-                  return acc;
-              }
-          }
-
-          // 关闭系统消息合并模式
-          if (current.role !== "system" && systemMergeMode) {
-              systemMergeMode = false;
-              closedSystemMergeMode = true;
-          }
-
-          // 消息合并逻辑
-          const previousMessage = acc[acc.length - 1];
-          const newMessage = `${currentMessageRole}: ${currentContent}`;
-
-          if (!previousMessage || previousMessage.startsWith(currentMessageRole)) {
-              return previousMessage
-                  ? [...acc.slice(0, -1), `${previousMessage}\n${currentContent}`]
-                  : [...acc, newMessage];
-          }
-
-          return [...acc, newMessage];
-      }, Promise.resolve([]));
-
-      return {
-          contextMessages: contextMessages.join('\n'),
-          systemMessage: systemMessageList.join('\n')
-      };
-  }
+    constructor(modelId) {
+        this.modelId = CONFIG.MODELS[modelId];
+        if (!this.modelId) {
+            throw new Error(`不支持的模型: ${modelId}`);
+        }
+    }
+
+    // 处理消息内容
+    processMessageContent(content) {
+        if (typeof content === 'string') return content;
+        if (Array.isArray(content)) {
+            return content
+                .filter(item => item.type === 'text')
+                .map(item => item.text)
+                .join('\n');
+        }
+        return typeof content === 'object' ? content.text || null : null;
+    }
+
+    // 转换消息格式
+    async transformMessages(request) {
+        let systemMessageList = [];
+        let systemMergeMode = false;
+        let closedSystemMergeMode = false;
+
+        const contextMessages = await request.messages.reduce(async (accPromise, current) => {
+            const acc = await accPromise;
+            const currentContent = this.processMessageContent(current.content);
+
+            if (currentContent === null) return acc;
+
+            const currentMessageRole = current.role === "system" || current.role === "user" ? "Human" : "Assistant";
+
+            // 系统消息处理逻辑
+            if (current.role === "system") {
+                if (!closedSystemMergeMode) {
+                    systemMergeMode = true;
+                    const lastSystemMessage = systemMessageList[systemMessageList.length - 1];
+
+                    if (!lastSystemMessage) {
+                        systemMessageList.push(currentContent);
+                    } else {
+                        systemMessageList[systemMessageList.length - 1] = `${lastSystemMessage}\n${currentContent}`;
+                    }
+                    return acc;
+                }
+            }
+
+            // 关闭系统消息合并模式
+            if (current.role !== "system" && systemMergeMode) {
+                systemMergeMode = false;
+                closedSystemMergeMode = true;
+            }
+
+            // 消息合并逻辑
+            const previousMessage = acc[acc.length - 1];
+            const newMessage = `${currentMessageRole}: ${currentContent}`;
+
+            if (!previousMessage || previousMessage.startsWith(currentMessageRole)) {
+                return previousMessage
+                    ? [...acc.slice(0, -1), `${previousMessage}\n${currentContent}`]
+                    : [...acc, newMessage];
+            }
+
+            return [...acc, newMessage];
+        }, Promise.resolve([]));
+
+        return {
+            contextMessages: contextMessages.join('\n'),
+            systemMessage: systemMessageList.join('\n')
+        };
+    }
 }
 
 // 响应处理类
 class ResponseHandler {
-  // 流式响应处理
-  static async handleStreamResponse(responseContent, model, res) {
-      res.setHeader('Content-Type', 'text/event-stream');
-      res.setHeader('Cache-Control', 'no-cache');
-      res.setHeader('Connection', 'keep-alive');
-      
-      let index = 0;
-      while (index < responseContent.length) {
-          const chunkSize = Math.floor(Math.random() * (30 - 16)) + 15;
-          const chunk = responseContent.slice(index, index + chunkSize);
-          
-          res.write(`data: ${JSON.stringify({
-              id: uuidv4(),
-              object: 'chat.completion.chunk',
-              created: Math.floor(Date.now() / 1000),
-              model: model,
-              choices: [{
-                  index: 0,
-                  delta: { content: chunk },
-                  finish_reason: null
-              }]
-          })}\n\n`);
-          
-          index += chunkSize;
-          await new Promise(resolve => setTimeout(resolve, 50));
-      }
-      
-      res.write('data: [DONE]\n\n');
-      res.end();
-  }
-
-  // 普通响应处理
-  static async handleNormalResponse(userMessage, responseContent, model, res) {
-      res.json({
-          id: uuidv4(),
-          object: "chat.completion",
-          created: Math.floor(Date.now() / 1000),
-          model: model,
-          choices: [{
-              index: 0,
-              message: {
-                  role: "assistant",
-                  content: responseContent
-              },
-              finish_reason: "stop"
-          }],
-          usage: {
-              prompt_tokens: userMessage.length,
-              completion_tokens: responseContent.length,
-              total_tokens: userMessage.length + responseContent.length
-          }
-      });
-  }
+    // 流式响应处理
+    static async handleStreamResponse(responseContent, model, res) {
+        if (!res.headersSent) {
+            res.setHeader('Content-Type', 'text/event-stream');
+            res.setHeader('Cache-Control', 'no-cache');
+            res.setHeader('Connection', 'keep-alive');
+        }
+
+        res.write(`data: ${JSON.stringify({
+            id: uuidv4(),
+            object: 'chat.completion.chunk',
+            created: Math.floor(Date.now() / 1000),
+            model: model,
+            choices: [{
+                index: 0,
+                delta: { content: responseContent },
+                finish_reason: null
+            }]
+        })}\n\n`);
+    }
+
+    // 普通响应处理
+    static async handleNormalResponse(userMessage, responseContent, model, res) {
+        res.json({
+            id: uuidv4(),
+            object: "chat.completion",
+            created: Math.floor(Date.now() / 1000),
+            model: model,
+            choices: [{
+                index: 0,
+                message: {
+                    role: "assistant",
+                    content: responseContent
+                },
+                finish_reason: "stop"
+            }],
+            usage: {
+                prompt_tokens: userMessage.length,
+                completion_tokens: responseContent.length,
+                total_tokens: userMessage.length + responseContent.length
+            }
+        });
+    }
 }
 
 // WebSocket工具类
@@ -181,9 +171,11 @@ class WebSocketUtils {
     static generateWebSocketKey() {
         return randomBytes(16).toString('base64');
     }
-
+    static getMessageDiff(prevContent, newContent) {
+        return newContent.slice(prevContent.length);
+    }
     // 创建WebSocket客户端
-    static async createWebSocketClient(requestPayload) {
+    static async createWebSocketClient(requestPayload, stream = false, res = null) {
         // 检查当前连接数是否达到上限
         if (this.activeConnections.size >= this.MAX_CONNECTIONS) {
             throw new Error(`当前连接数已达到上限 (${this.MAX_CONNECTIONS})，请稍后重试！`);
@@ -217,6 +209,7 @@ class WebSocketUtils {
                 }, this.TIMEOUT);
 
                 let responseContent = '';
+                let prevContent = '';
                 let isComplete = false;
 
                 ws.on('open', () => {
@@ -235,7 +228,7 @@ class WebSocketUtils {
                 ws.on('message', async (data) => {
                     const message = data.toString();
                     const parsedMessage = JSON.parse(message);
-                    
+
                     switch (parsedMessage.type) {
                         case 'connection_ack':
                             console.log('WebSocket连接请求中');
@@ -245,14 +238,22 @@ class WebSocketUtils {
                             const chatResponse = await this.handleChatResponse(parsedMessage);
                             if (chatResponse) {
                                 responseContent = chatResponse;
-                                // 获取到响应后立即关闭连接
-                                isComplete = true;
-                                ws.close();
-                                resolve(responseContent);
+
+                                if (stream && res) {
+                                    const diff = this.getMessageDiff(prevContent, responseContent);
+                                    if (diff) {
+                                        await ResponseHandler.handleStreamResponse(diff, "claude-3-5-sonnet-20241022", res);
+                                        prevContent = responseContent;
+                                    }
+                                }
                             }
                             break;
                         case 'complete':
                             isComplete = true;
+                            if (stream && res) {
+                                res.write('data: [DONE]\n\n');
+                                res.end();
+                            }
                             ws.close();
                             resolve(responseContent);
                             break;
@@ -332,8 +333,6 @@ class WebSocketUtils {
                 }`
             }
         };
-
-        //console.log(JSON.stringify(subscribeMessage,null,2));
         if (ws.readyState === WebSocket.OPEN) {
             ws.send(JSON.stringify(subscribeMessage));
         }
@@ -343,7 +342,7 @@ class WebSocketUtils {
     static async handleChatResponse(message) {
         if (message.payload && message.payload.data) {
             const chatResult = message.payload.data.newSessionChatResult;
-            if (chatResult && chatResult.isEnd == true && chatResult.message) {
+            if (chatResult && chatResult.message) {
                 return chatResult.message.content;
             }
         }
@@ -363,62 +362,60 @@ const app = express();
 app.use(express.json({ limit: CONFIG.SERVER.BODY_LIMIT }));
 app.use(express.urlencoded({ extended: true, limit: CONFIG.SERVER.BODY_LIMIT }));
 app.use(cors({
-  origin: '*',
-  methods: ['GET', 'POST', 'OPTIONS'],
-  allowedHeaders: ['Content-Type', 'Authorization']
+    origin: '*',
+    methods: ['GET', 'POST', 'OPTIONS'],
+    allowedHeaders: ['Content-Type', 'Authorization']
 }));
 
 
 // 获取模型列表路由
 app.get('/hf/v1/models', (req, res) => {
-  res.json({
-      object: "list",
-      data: [{
-          id: "claude-3-5-sonnet-20241022",
-          object: "model",
-          created: Math.floor(Date.now() / 1000),
-          owned_by: "claude",
-      }]
-  });
+    res.json({
+        object: "list",
+        data: [{
+            id: "claude-3-5-sonnet-20241022",
+            object: "model",
+            created: Math.floor(Date.now() / 1000),
+            owned_by: "claude",
+        }]
+    });
 });
 
 // 聊天完成路由
 app.post('/hf/v1/chat/completions', async (req, res) => {
-  try {
-      const { messages, model, stream } = req.body;
-      const authToken = req.headers.authorization?.replace('Bearer ', '');
-      
-      if (authToken !== CONFIG.API.API_KEY) {
-          return res.status(401).json({ error: "Unauthorized" });
-      }
-
-      const apiClient = new AiApiClient(req.body.model);
-      const requestPayload = await apiClient.transformMessages(req.body);
-
-      const userMessage = messages.reverse().find(message => message.role === 'user')?.content;
-      if (!userMessage) {
-          return res.status(400).json({ error: "缺失用户消息" });
-      }
-
-      const responseContent = await WebSocketUtils.createWebSocketClient(requestPayload);
-
-      if (stream) {
-          await ResponseHandler.handleStreamResponse(responseContent, model, res);
-      } else {
-          await ResponseHandler.handleNormalResponse(userMessage, responseContent, model, res);
-      }
-  } catch (error) {
-      console.error('处理请求时发生错误:', error);
-      res.status(500).json({ error: "内部服务器错误，请查询日志记录！", details: error.message });
-  }
+    try {
+        const { messages, model, stream } = req.body;
+        const authToken = req.headers.authorization?.replace('Bearer ', '');
+
+        if (authToken !== CONFIG.API.API_KEY) {
+            return res.status(401).json({ error: "Unauthorized" });
+        }
+
+        const apiClient = new AiApiClient(req.body.model);
+        const requestPayload = await apiClient.transformMessages(req.body);
+
+        const userMessage = messages.reverse().find(message => message.role === 'user')?.content;
+        if (!userMessage) {
+            return res.status(400).json({ error: "缺失用户消息" });
+        }
+
+        const responseContent = await WebSocketUtils.createWebSocketClient(requestPayload, stream, stream ? res : null);
+
+        if (!stream) {
+            await ResponseHandler.handleNormalResponse(userMessage, responseContent, model, res);
+        }
+    } catch (error) {
+        console.error('处理请求时发生错误:', error);
+        res.status(500).json({ error: "内部服务器错误，请查询日志记录！", details: error.message });
+    }
 });
 
 // 404处理
 app.use((req, res) => {
-  res.status(404).json({ message: "服务创建成功运行中，请根据规则使用正确请求路径" });
+    res.status(404).json({ message: "服务创建成功运行中，请根据规则使用正确请求路径" });
 });
 
 // 启动服务器
 app.listen(CONFIG.SERVER.PORT, () => {
-  console.log(`服务器运行在 http://localhost:${CONFIG.SERVER.PORT}`);
+    console.log(`服务器运行在 http://localhost:${CONFIG.SERVER.PORT}`);
 });