Update index.js

index.js

@@ -23,11 +23,13 @@ const CONFIG = {
         "grok-3-reasoning": "grok-3"
     },
     API: {
+        IS_CUSTOM_SSO: process.env.IS_CUSTOM_SSO === 'true',
         BASE_URL: "https://grok.com",
         API_KEY: process.env.API_KEY || "sk-123456",
         SIGNATURE_COOKIE: null,
         TEMP_COOKIE: null,
-        PICGO_KEY: process.env.PICGO_KEY || null //想要流式生图的话需要填入这个PICGO图床的key
+        PICGO_KEY: process.env.PICGO_KEY || null, //想要流式生图的话需要填入这个PICGO图床的key
+        TUMY_KEY: process.env.TUMY_KEY || null //想要流式生图的话需要填入这个TUMY图床的key 两个图床二选一，默认使用PICGO
     },
     SERVER: {
         PORT: process.env.PORT || 3000,
@@ -36,7 +38,7 @@ const CONFIG = {
     RETRY: {
         MAX_ATTEMPTS: 2//重试次数
     },
-    SHOW_THINKING:process.env.SHOW_THINKING === 'true',
+    SHOW_THINKING: process.env.SHOW_THINKING === 'true',
     IS_THINKING: false,
     IS_IMG_GEN: false,
     IS_IMG_GEN2: false,
@@ -66,10 +68,13 @@ const DEFAULT_HEADERS = {
 
 
 async function initialization() {
+    if (CONFIG.API.IS_CUSTOM_SSO) {
+        await Utils.get_signature()
+        return;
+    }
     const ssoArray = process.env.SSO.split(',');
-    const ssorwArray = process.env.SSO_RW.split(',');
-    ssoArray.forEach((sso, index) => {
-        tokenManager.addToken(`sso-rw=${ssorwArray[index]};sso=${sso}`);
+    ssoArray.forEach((sso) => {
+        tokenManager.addToken(`sso-rw=${sso};sso=${sso}`);
     });
     console.log(JSON.stringify(tokenManager.getActiveTokens(), null, 2));
     await Utils.get_signature()
@@ -99,9 +104,17 @@ class AuthTokenManager {
             });
         }
     }
+    setToken(token) {
+        this.activeTokens = [token];
+        this.tokenModelFrequency.set(token, {
+            "grok-3": 0,
+            "grok-3-deepsearch": 0,
+            "grok-3-reasoning": 0
+        });
+    }
 
     getTokenByIndex(index, model) {
-        if(this.activeTokens.length === 0){
+        if (this.activeTokens.length === 0) {
             return null;
         }
         const token = this.activeTokens[index];
@@ -110,7 +123,7 @@ class AuthTokenManager {
     }
 
     recordModelRequest(token, model) {
-        if(model === 'grok-3-search' || model === 'grok-3-imageGen'){
+        if (model === 'grok-3-search' || model === 'grok-3-imageGen') {
             model = 'grok-3';
         }
 
@@ -122,7 +135,7 @@ class AuthTokenManager {
         this.checkAndRemoveTokenIfLimitReached(token);
     }
     setModelLimit(index, model) {
-        if(model === 'grok-3-search' || model === 'grok-3-imageGen'){
+        if (model === 'grok-3-search' || model === 'grok-3-imageGen') {
             model = 'grok-3';
         }
         if (!this.modelRateLimit[model]) return;
@@ -130,7 +143,7 @@ class AuthTokenManager {
         tokenFrequency[model] = 9999;
     }
     isTokenModelLimitReached(index, model) {
-        if(model === 'grok-3-search' || model === 'grok-3-imageGen'){
+        if (model === 'grok-3-search' || model === 'grok-3-imageGen') {
             model = 'grok-3';
         }
         if (!this.modelRateLimit[model]) return;
@@ -170,6 +183,7 @@ class AuthTokenManager {
     }
 
     startTokenRecoveryProcess() {
+        if (CONFIG.API.IS_CUSTOM_SSO) return;
         setInterval(() => {
             const now = Date.now();
             for (const [token, expiredTime] of this.expiredTokens.entries()) {
@@ -200,7 +214,6 @@ class Utils {
     static async extractGrokHeaders() {
         Logger.info("开始提取头信息", 'Server');
         try {
-            // 启动浏览器
             const browser = await puppeteer.launch({
                 headless: true,
                 args: [
@@ -221,16 +234,12 @@ class Utils {
             const cookies = await page.cookies();
             const targetHeaders = ['x-anonuserid', 'x-challenge', 'x-signature'];
             const extractedHeaders = {};
-            // 遍历 Cookies
             for (const cookie of cookies) {
-                // 检查是否为目标头信息
                 if (targetHeaders.includes(cookie.name.toLowerCase())) {
                     extractedHeaders[cookie.name.toLowerCase()] = cookie.value;
                 }
             }
-            // 关闭浏览器
             await browser.close();
-            // 打印并返回提取的头信息
             Logger.info('提取的头信息:', JSON.stringify(extractedHeaders, null, 2), 'Server');
             return extractedHeaders;
 
@@ -358,10 +367,10 @@ class GrokApiClient {
     }
 
     async prepareChatRequest(request) {
-        if ((request.model === 'grok-2-imageGen' || request.model === 'grok-3-imageGen') && !CONFIG.API.PICGO_KEY && request.stream) {
-            throw new Error(`该模型流式输出需要配置PICGO图床密钥!`);
+        if ((request.model === 'grok-2-imageGen' || request.model === 'grok-3-imageGen') && !CONFIG.API.PICGO_KEY && !CONFIG.API.TUMY_KEY && request.stream) {
+            throw new Error(`该模型流式输出需要配置PICGO或者PICUI图床密钥!`);
         }
-        
+
         // 处理画图模型的消息限制
         let todoMessages = request.messages;
         if (request.model === 'grok-2-imageGen' || request.model === 'grok-3-imageGen') {
@@ -369,22 +378,22 @@ class GrokApiClient {
             if (lastMessage.role !== 'user') {
                 throw new Error('画图模型的最后一条消息必须是用户消息!');
             }
-            todoMessages = [lastMessage]; 
+            todoMessages = [lastMessage];
         }
-        
+
         const fileAttachments = [];
         let messages = '';
         let lastRole = null;
         let lastContent = '';
         const search = request.model === 'grok-2-search' || request.model === 'grok-3-search';
-    
+
         // 移除<think>标签及其内容和base64图片
         const removeThinkTags = (text) => {
             text = text.replace(/<think>[\s\S]*?<\/think>/g, '').trim();
             text = text.replace(/!\[image\]\(data:.*?base64,.*?\)/g, '[图片]');
             return text;
         };
-    
+
         const processImageUrl = async (content) => {
             if (content.type === 'image_url' && content.image_url.url.includes('data:image')) {
                 const imageResponse = await this.uploadBase64Image(
@@ -395,7 +404,7 @@ class GrokApiClient {
             }
             return null;
         };
-    
+
         const processContent = async (content) => {
             if (Array.isArray(content)) {
                 let textContent = '';
@@ -416,11 +425,11 @@ class GrokApiClient {
             }
             return removeThinkTags(this.processMessageContent(content));
         };
-    
+
         for (const current of todoMessages) {
             const role = current.role === 'assistant' ? 'assistant' : 'user';
             const isLastMessage = current === todoMessages[todoMessages.length - 1];
-            
+
             // 处理图片附件
             if (isLastMessage && current.content) {
                 if (Array.isArray(current.content)) {
@@ -435,10 +444,10 @@ class GrokApiClient {
                     if (processedImage) fileAttachments.push(processedImage);
                 }
             }
-    
+
             // 处理文本内容
             const textContent = await processContent(current.content);
-            
+
             if (textContent || (isLastMessage && fileAttachments.length > 0)) {
                 if (role === lastRole && textContent) {
                     lastContent += '\n' + textContent;
@@ -451,7 +460,7 @@ class GrokApiClient {
                 }
             }
         }
-    
+
         return {
             modelName: this.modelId,
             message: messages.trim(),
@@ -519,7 +528,7 @@ class MessageProcessor {
     }
 }
 async function processModelResponse(linejosn, model) {
-    let result = { token: '', imageUrl: null }
+    let result = { token: null, imageUrl: null }
     if (CONFIG.IS_IMG_GEN) {
         if (linejosn?.cachedImageGenerationResponse && !CONFIG.IS_IMG_GEN2) {
             result.imageUrl = linejosn.cachedImageGenerationResponse.imageUrl;
@@ -549,12 +558,12 @@ async function processModelResponse(linejosn, model) {
             }
             return result;
         case 'grok-3-reasoning':
-            if(linejosn?.isThinking && !CONFIG.SHOW_THINKING)return result;
+            if (linejosn?.isThinking && !CONFIG.SHOW_THINKING) return result;
 
             if (linejosn?.isThinking && !CONFIG.IS_THINKING) {
                 result.token = "<think>" + linejosn?.token;
                 CONFIG.IS_THINKING = true;
-            } else if (CONFIG.IS_THINKING && !linejosn.isThinking) {
+            } else if (!linejosn.isThinking && CONFIG.IS_THINKING) {
                 result.token = "</think>" + linejosn?.token;
                 CONFIG.IS_THINKING = false;
             } else {
@@ -562,6 +571,7 @@ async function processModelResponse(linejosn, model) {
             }
             return result;
     }
+    return result;
 }
 
 async function handleResponse(response, model, res, isStream) {
@@ -570,8 +580,14 @@ async function handleResponse(response, model, res, isStream) {
         let buffer = '';
         let fullResponse = '';
         const dataPromises = [];
-        
-        return new Promise((resolve, reject) => { 
+        if (isStream) {
+            res.setHeader('Content-Type', 'text/event-stream');
+            res.setHeader('Cache-Control', 'no-cache');
+            res.setHeader('Connection', 'keep-alive');
+        }
+        Logger.info("开始处理流式响应", 'Server');
+
+        return new Promise((resolve, reject) => {
             stream.on('data', async (chunk) => {
                 buffer += chunk.toString();
                 const lines = buffer.split('\n');
@@ -584,12 +600,15 @@ async function handleResponse(response, model, res, isStream) {
                         const data = trimmedLine.substring(6);
                         try {
                             if (!data.trim()) continue;
-                            if(data === "[DONE]") continue;
+                            if (data === "[DONE]") continue;                     
                             const linejosn = JSON.parse(data);
-                            if (linejosn?.error) {
-                                Logger.error(JSON.stringify(linejosn,null,2), 'Server');
+                            if (linejosn?.error) {                          
+                                Logger.error(JSON.stringify(linejosn, null, 2), 'Server');
+                                if(linejosn.error?.name === "RateLimitError"){
+                                    CONFIG.API.TEMP_COOKIE = null;
+                                }
                                 stream.destroy();
-                                reject(new Error("RateLimitError")); 
+                                reject(new Error("RateLimitError"));
                                 return;
                             }
                             if (linejosn?.doImgGen || linejosn?.imageAttachmentInfo) {
@@ -597,6 +616,7 @@ async function handleResponse(response, model, res, isStream) {
                             }
                             const processPromise = (async () => {
                                 const result = await processModelResponse(linejosn, model);
+
                                 if (result.token) {
                                     if (isStream) {
                                         res.write(`data: ${JSON.stringify(MessageProcessor.createChatResponse(result.token, model, true))}\n\n`);
@@ -616,6 +636,7 @@ async function handleResponse(response, model, res, isStream) {
                             })();
                             dataPromises.push(processPromise);
                         } catch (error) {
+                            Logger.error(error, 'Server');
                             continue;
                         }
                     }
@@ -631,25 +652,27 @@ async function handleResponse(response, model, res, isStream) {
                     } else {
                         if (!CONFIG.IS_IMG_GEN2) {
                             res.json(MessageProcessor.createChatResponse(fullResponse, model));
-                        }                        
+                        }
                     }
                     CONFIG.IS_IMG_GEN = false;
                     CONFIG.IS_IMG_GEN2 = false;
-                    resolve(); 
+                    resolve();
                 } catch (error) {
-                    reject(error); 
+                    Logger.error(error, 'Server');
+                    reject(error);
                 }
             });
 
             stream.on('error', (error) => {
-                reject(error); 
+                Logger.error(error, 'Server');
+                reject(error);
             });
         });
     } catch (error) {
         Logger.error(error, 'Server');
         CONFIG.IS_IMG_GEN = false;
         CONFIG.IS_IMG_GEN2 = false;
-        throw error;
+        throw new Error(error);
     }
 }
 
@@ -668,7 +691,7 @@ async function handleImageResponse(imageUrl) {
                 }
             });
 
-            if (imageBase64Response.ok) break; 
+            if (imageBase64Response.ok) break;
             retryCount++;
             if (retryCount === MAX_RETRIES) {
                 throw new Error(`上游服务请求失败! status: ${imageBase64Response.status}`);
@@ -676,6 +699,7 @@ async function handleImageResponse(imageUrl) {
             await new Promise(resolve => setTimeout(resolve, CONFIG.API.RETRY_TIME * retryCount));
 
         } catch (error) {
+            Logger.error(error, 'Server');
             retryCount++;
             if (retryCount === MAX_RETRIES) {
                 throw error;
@@ -688,34 +712,64 @@ async function handleImageResponse(imageUrl) {
     const arrayBuffer = await imageBase64Response.arrayBuffer();
     const imageBuffer = Buffer.from(arrayBuffer);
 
-    if(!CONFIG.API.PICGO_KEY){
+    if (!CONFIG.API.PICGO_KEY && !CONFIG.API.TUMY_KEY) {
         const base64Image = imageBuffer.toString('base64');
         const imageContentType = imageBase64Response.headers.get('content-type');
         return `![image](data:${imageContentType};base64,${base64Image})`
     }
 
+    Logger.info("开始请求图床", 'Server');
     const formData = new FormData();
-
-    formData.append('source', imageBuffer, {
-        filename: 'new.jpg',
-        contentType: 'image/jpeg'
-    });
-    const formDataHeaders = formData.getHeaders();
-    const responseURL = await fetch("https://www.picgo.net/api/1/upload", {
-        method: "POST",
-        headers: {
+    if(CONFIG.API.PICGO_KEY){
+        formData.append('source', imageBuffer, {
+            filename: `image-${Date.now()}.jpg`,
+            contentType: 'image/jpeg'
+        });
+        const formDataHeaders = formData.getHeaders();
+        const responseURL = await fetch("https://www.picgo.net/api/1/upload", {
+            method: "POST",
+            headers: {
             ...formDataHeaders,
             "Content-Type": "multipart/form-data",
             "X-API-Key": CONFIG.API.PICGO_KEY
         },
         body: formData
-    });
-    if (!responseURL.ok) {
-        return "生图失败，请查看图床密钥是否设置正确"
-    } else {
-        Logger.info("生图成功", 'Server');
-        const result = await responseURL.json();
-        return `![image](${result.image.url})`
+        });
+        if (!responseURL.ok) {
+            return "生图失败，请查看PICGO图床密钥是否设置正确"
+        } else {
+            Logger.info("生图成功", 'Server');
+            const result = await responseURL.json();
+            return `![image](${result.image.url})`
+        }
+    }else if(CONFIG.API.TUMY_KEY){
+        const formData = new FormData();
+        formData.append('file', imageBuffer, {
+            filename: `image-${Date.now()}.jpg`,
+            contentType: 'image/jpeg'
+        });
+        const formDataHeaders = formData.getHeaders();
+        const responseURL = await fetch("https://tu.my/api/v1/upload", {
+            method: "POST",
+            headers: {
+                ...formDataHeaders,
+                "Accept": "application/json",
+                'Authorization': `Bearer ${CONFIG.API.TUMY_KEY}`
+            },
+            body: formData
+        });
+        if (!responseURL.ok) {
+            return "生图失败，请查看TUMY图床密钥是否设置正确"
+        } else {
+            Logger.info("生图成功", 'Server');
+            const result = await responseURL.json();
+            try {
+              var url = `![image](${result.data.links.url})`
+              return url;
+            } catch (error) {
+              return "生图失败，请查看TUMY图床密钥是否设置正确"
+            }
+        }
     }
 }
 
@@ -749,14 +803,22 @@ app.get('/hf/v1/models', (req, res) => {
 app.post('/hf/v1/chat/completions', async (req, res) => {
     try {
         const authToken = req.headers.authorization?.replace('Bearer ', '');
-        if (authToken !== CONFIG.API.API_KEY) {
+        if (CONFIG.API.IS_CUSTOM_SSO) {
+            if (authToken) {
+                const result = `sso=${authToken};ssp_rw=${authToken}`;
+                tokenManager.setToken(result);
+            } else {
+                return res.status(401).json({ error: '自定义的SSO令牌缺失' });
+            }
+        } else if (authToken !== CONFIG.API.API_KEY) {
             return res.status(401).json({ error: 'Unauthorized' });
         }
         let isTempCookie = req.body.model.includes("grok-2");
         let retryCount = 0;
         const grokClient = new GrokApiClient(req.body.model);
         const requestPayload = await grokClient.prepareChatRequest(req.body);
-        
+        Logger.info(`请求体: ${JSON.stringify(requestPayload,null,2)}`, 'Server');
+
         while (retryCount < CONFIG.RETRY.MAX_ATTEMPTS) {
             retryCount++;
             if (!CONFIG.API.TEMP_COOKIE) {
@@ -769,7 +831,11 @@ app.post('/hf/v1/chat/completions', async (req, res) => {
             } else {
                 CONFIG.API.SIGNATURE_COOKIE = await Utils.createAuthHeaders(req.body.model);
             }
+            if (!CONFIG.API.SIGNATURE_COOKIE) {
+                throw new Error('该模型无可用令牌');
+            }
             Logger.info(`当前令牌索引: ${CONFIG.SSO_INDEX}`, 'Server');
+            Logger.info(`当前令牌: ${JSON.stringify(CONFIG.API.SIGNATURE_COOKIE,null,2)}`, 'Server');
             const newMessageReq = await fetch(`${CONFIG.API.BASE_URL}/api/rpc`, {
                 method: 'POST',
                 headers: {
@@ -783,9 +849,15 @@ app.post('/hf/v1/chat/completions', async (req, res) => {
                     }
                 })
             });
-
-            const responseText = await newMessageReq.json();
-            const conversationId = responseText.conversationId;
+            let conversationId;
+            var responseText2 = await newMessageReq.clone().text();
+            if (newMessageReq.status === 200) {
+                const responseText = await newMessageReq.json();
+                conversationId = responseText.conversationId;
+            } else {
+                Logger.error(`创建会话响应错误: ${responseText2}`, 'Server');
+                throw new Error(`创建会话响应错误: ${responseText2}`);
+            }
 
             const response = await fetch(`${CONFIG.API.BASE_URL}/api/conversations/${conversationId}/responses`, {
                 method: 'POST',
@@ -798,19 +870,20 @@ app.post('/hf/v1/chat/completions', async (req, res) => {
                 },
                 body: JSON.stringify(requestPayload)
             });
-          
-            
+
             if (response.ok) {
-                CONFIG.SSO_INDEX = (CONFIG.SSO_INDEX + 1) % tokenManager.getTokenCount();
                 Logger.info(`请求成功`, 'Server');
+                CONFIG.SSO_INDEX = (CONFIG.SSO_INDEX + 1) % tokenManager.getTokenCount();
                 Logger.info(`当前剩余可用令牌数: ${tokenManager.getTokenCount()}`, 'Server');
                 try {
                     await handleResponse(response, req.body.model, res, req.body.stream);
-                    return; 
+                    Logger.info(`请求结束`, 'Server');
+                    return;
                 } catch (error) {
-                    if(isTempCookie){
+                    Logger.error(error, 'Server');
+                    if (isTempCookie) {
                         await Utils.get_signature();
-                    }else{
+                    } else {
                         tokenManager.setModelLimit(CONFIG.SSO_INDEX, req.body.model);
                         for (let i = 1; i <= tokenManager.getTokenCount(); i++) {
                             CONFIG.SSO_INDEX = (CONFIG.SSO_INDEX + 1) % tokenManager.getTokenCount();
@@ -826,7 +899,7 @@ app.post('/hf/v1/chat/completions', async (req, res) => {
                 if (response.status === 429) {
                     if (isTempCookie) {
                         await Utils.get_signature();
-                    } else {
+                    } else {                      
                         tokenManager.setModelLimit(CONFIG.SSO_INDEX, req.body.model);
                         for (let i = 1; i <= tokenManager.getTokenCount(); i++) {
                             CONFIG.SSO_INDEX = (CONFIG.SSO_INDEX + 1) % tokenManager.getTokenCount();
@@ -837,7 +910,7 @@ app.post('/hf/v1/chat/completions', async (req, res) => {
                             }
                         }
                     }
-                } else {
+                } else {                
                     // 非429错误直接抛出
                     if (isTempCookie) {
                         await Utils.get_signature();
@@ -852,14 +925,11 @@ app.post('/hf/v1/chat/completions', async (req, res) => {
         }
         throw new Error('当前模型所有令牌都已耗尽');
     } catch (error) {
-        CONFIG.SSO_INDEX = (CONFIG.SSO_INDEX + 1) % tokenManager.getTokenCount();
         Logger.error(error, 'ChatAPI');
         res.status(500).json({
             error: {
-                message: error.message,
-                type: 'server_error',
-                param: null,
-                code: error.code || null
+                message: error.message || error,
+                type: 'server_error'
             }
         });
     }
@@ -873,4 +943,4 @@ app.use((req, res) => {
 
 app.listen(CONFIG.SERVER.PORT, () => {
     Logger.info(`服务器已启动，监听端口: ${CONFIG.SERVER.PORT}`, 'Server');
-});
+});