Update index.js

index.js

@@ -7,8 +7,6 @@ import puppeteer from 'puppeteer-extra'
 import StealthPlugin from 'puppeteer-extra-plugin-stealth'
 import { v4 as uuidv4 } from 'uuid';
 import Logger from './logger.js';
-import fs from 'fs';
-import path from 'path';
 
 dotenv.config();
 
@@ -19,6 +17,7 @@ const CONFIG = {
         'grok-2-imageGen': 'grok-latest',
         'grok-2-search': 'grok-latest',
         "grok-3": "grok-3",
+        "grok-3-search": "grok-3",
         "grok-3-imageGen": "grok-3",
         "grok-3-deepsearch": "grok-3",
         "grok-3-reasoning": "grok-3"
@@ -28,21 +27,21 @@ const CONFIG = {
         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
     },
     SERVER: {
         PORT: process.env.PORT || 3000,
         BODY_LIMIT: '5mb'
     },
     RETRY: {
-        MAX_ATTEMPTS: 2,//重试次数
-        DELAY_BASE: 1000 // 基础延迟时间（毫秒）
+        MAX_ATTEMPTS: 2//重试次数
     },
+    SHOW_THINKING:process.env.SHOW_THINKING === 'true',
     IS_THINKING: false,
     IS_IMG_GEN: false,
     IS_IMG_GEN2: false,
     SSO_INDEX: 0,//sso的索引
-    ISSHOW_SEARCH_RESULTS: process.env.ISSHOW_SEARCH_RESULTS === 'true',//是否显示搜索结果,默认关闭
+    ISSHOW_SEARCH_RESULTS: process.env.ISSHOW_SEARCH_RESULTS === 'true',//是否显示搜索结果
     CHROME_PATH: process.env.CHROME_PATH || "/usr/bin/chromium"//chrome路径
 };
 puppeteer.use(StealthPlugin())
@@ -102,7 +101,6 @@ class AuthTokenManager {
     }
 
     getTokenByIndex(index, model) {
-        if (!this.modelRateLimit[model]) return;
         if(this.activeTokens.length === 0){
             return null;
         }
@@ -112,6 +110,10 @@ class AuthTokenManager {
     }
 
     recordModelRequest(token, model) {
+        if(model === 'grok-3-search' || model === 'grok-3-imageGen'){
+            model = 'grok-3';
+        }
+
         if (!this.modelRateLimit[model]) return;
         const tokenFrequency = this.tokenModelFrequency.get(token);
         if (tokenFrequency && tokenFrequency[model] !== undefined) {
@@ -120,11 +122,17 @@ class AuthTokenManager {
         this.checkAndRemoveTokenIfLimitReached(token);
     }
     setModelLimit(index, model) {
+        if(model === 'grok-3-search' || model === 'grok-3-imageGen'){
+            model = 'grok-3';
+        }
         if (!this.modelRateLimit[model]) return;
         const tokenFrequency = this.tokenModelFrequency.get(this.activeTokens[index]);
         tokenFrequency[model] = 9999;
     }
     isTokenModelLimitReached(index, model) {
+        if(model === 'grok-3-search' || model === 'grok-3-imageGen'){
+            model = 'grok-3';
+        }
         if (!this.modelRateLimit[model]) return;
         const token = this.activeTokens[index];
         const tokenFrequency = this.tokenModelFrequency.get(token);
@@ -168,7 +176,6 @@ class AuthTokenManager {
                 if (now - expiredTime >= 2 * 60 * 60 * 1000) {
                     this.tokenModelUsage.set(token, {
                         "grok-3": 0,
-                        "grok-3-imageGen": 0,
                         "grok-3-deepsearch": 0,
                         "grok-3-reasoning": 0
                     });
@@ -191,7 +198,7 @@ class AuthTokenManager {
 
 class Utils {
     static async extractGrokHeaders() {
-        console.log("开始提取头信息");
+        Logger.info("开始提取头信息", 'Server');
         try {
             // 启动浏览器
             const browser = await puppeteer.launch({
@@ -210,7 +217,6 @@ class Utils {
             await page.evaluate(() => {
                 return new Promise(resolve => setTimeout(resolve, 5000))
             })
-            console.log('等待5秒完成')
             // 获取所有 Cookies
             const cookies = await page.cookies();
             const targetHeaders = ['x-anonuserid', 'x-challenge', 'x-signature'];
@@ -225,11 +231,11 @@ class Utils {
             // 关闭浏览器
             await browser.close();
             // 打印并返回提取的头信息
-            console.log('提取的头信息:', JSON.stringify(extractedHeaders, null, 2));
+            Logger.info('提取的头信息:', JSON.stringify(extractedHeaders, null, 2), 'Server');
             return extractedHeaders;
 
         } catch (error) {
-            console.error('获取头信息出错:', error);
+            Logger.error('获取头信息出错:', error, 'Server');
             return null;
         }
     }
@@ -237,12 +243,12 @@ class Utils {
         if (CONFIG.API.TEMP_COOKIE) {
             return CONFIG.API.TEMP_COOKIE;
         }
-        console.log("刷新认证信息");
+        Logger.info("刷新认证信息", 'Server');
         let retryCount = 0;
         while (!CONFIG.API.TEMP_COOKIE || retryCount < CONFIG.RETRY.MAX_ATTEMPTS) {
             let headers = await Utils.extractGrokHeaders();
             if (headers) {
-                console.log("获取认证信息成功");
+                Logger.info("获取认证信息成功", 'Server');
                 CONFIG.API.TEMP_COOKIE = { cookie: `x-anonuserid=${headers["x-anonuserid"]}; x-challenge=${headers["x-challenge"]}; x-signature=${headers["x-signature"]}` };
                 return;
             }
@@ -252,53 +258,6 @@ class Utils {
             }
         }
     }
-    static async handleError(error, res) {
-        // 如果是500错误且提供了原始请求函数，尝试重新获取签名并重试
-        if (String(error).includes("status: 500")) {
-            try {
-                await Utils.get_signature();
-                if (CONFIG.API.SIGNATURE_COOKIE) {
-                    return res.status(500).json({
-                        error: {
-                            message: `${error.message}已重新刷新认证信息，请重新对话`,
-                            type: 'server_error',
-                            param: null,
-                            code: retryError.code || null
-                        }
-                    });
-                } else {
-                    return res.status(500).json({
-                        error: {
-                            message: "认证信息获取失败",
-                            type: 'server_error',
-                            param: null,
-                            code: null
-                        }
-                    });
-                }
-            } catch (error) {
-                console.error('重试失败:', error);
-                return res.status(500).json({
-                    error: {
-                        message: `${error.message},认证信息获取失败`,
-                        type: 'server_error',
-                        param: null,
-                        code: error.code || null
-                    }
-                });
-            }
-        }
-
-        // 其他错误直接返回
-        res.status(500).json({
-            error: {
-                message: error.message,
-                type: 'server_error',
-                param: null,
-                code: error.code || null
-            }
-        });
-    }
     static async organizeSearchResults(searchResults) {
         // 确保传入的是有效的搜索结果对象
         if (!searchResults || !searchResults.results) {
@@ -316,9 +275,9 @@ class Utils {
         });
         return formattedResults.join('\n\n');
     }
-    static createAuthHeaders(model) {
+    static async createAuthHeaders(model) {
         return {
-            'cookie': `${tokenManager.getTokenByIndex(CONFIG.SSO_INDEX, model)}`
+            'cookie': `${await tokenManager.getTokenByIndex(CONFIG.SSO_INDEX, model)}`
         };
     }
 }
@@ -372,7 +331,7 @@ class GrokApiClient {
                     content: imageBuffer
                 }
             };
-            console.log("发送图片请求");
+            Logger.info("发送图片请求", 'Server');
             // 发送请求
             const response = await fetch(url, {
                 method: 'POST',
@@ -384,35 +343,46 @@ class GrokApiClient {
             });
 
             if (!response.ok) {
-                console.error(`上传图片失败,状态码:${response.status},原因:${response.error}`);
+                Logger.error(`上传图片失败,状态码:${response.status},原因:${response.error}`, 'Server');
                 return '';
             }
 
             const result = await response.json();
-            console.log('上传图片成功:', result);
+            Logger.info('上传图片成功:', result, 'Server');
             return result.fileMetadataId;
 
         } catch (error) {
-            console.error('上传图片失败:', error);
+            Logger.error(error, 'Server');
             return '';
         }
     }
 
     async prepareChatRequest(request) {
-        if ((request.model === 'grok-2-imageGen' || request.model === 'grok-3-imageGen') && !CONFIG.API.PICGO_KEY) {
-            throw new Error(`该模型需要配置PICGO图床密钥!`);
+        if ((request.model === 'grok-2-imageGen' || request.model === 'grok-3-imageGen') && !CONFIG.API.PICGO_KEY && request.stream) {
+            throw new Error(`该模型流式输出需要配置PICGO图床密钥!`);
+        }
+        
+        // 处理画图模型的消息限制
+        let todoMessages = request.messages;
+        if (request.model === 'grok-2-imageGen' || request.model === 'grok-3-imageGen') {
+            const lastMessage = todoMessages[todoMessages.length - 1];
+            if (lastMessage.role !== 'user') {
+                throw new Error('画图模型的最后一条消息必须是用户消息!');
+            }
+            todoMessages = [lastMessage]; 
         }
         
-        const todoMessages = request.messages;
         const fileAttachments = [];
         let messages = '';
         let lastRole = null;
         let lastContent = '';
-        const search = request.model === 'grok-2-search';
+        const search = request.model === 'grok-2-search' || request.model === 'grok-3-search';
     
-        // 移除<think>标签及其内容
+        // 移除<think>标签及其内容和base64图片
         const removeThinkTags = (text) => {
-            return text.replace(/<think>[\s\S]*?<\/think>/g, '').trim();
+            text = text.replace(/<think>[\s\S]*?<\/think>/g, '').trim();
+            text = text.replace(/!\[image\]\(data:.*?base64,.*?\)/g, '[图片]');
+            return text;
         };
     
         const processImageUrl = async (content) => {
@@ -485,7 +455,7 @@ class GrokApiClient {
         return {
             modelName: this.modelId,
             message: messages.trim(),
-            fileAttachments: fileAttachments.slice(0, 4), // 最多保留4张图片
+            fileAttachments: fileAttachments.slice(0, 4),
             imageAttachments: [],
             disableSearch: false,
             enableImageGeneration: true,
@@ -556,37 +526,39 @@ async function processModelResponse(linejosn, model) {
         }
         return result;
     }
-    var token = linejosn?.token
-    if (!token) return result;
+
     //非生图模型的处理
     switch (model) {
         case 'grok-2':
-            result.token = token;
+            result.token = linejosn?.token;
             return result;
         case 'grok-2-search':
+        case 'grok-3-search':
             if (linejosn?.webSearchResults && CONFIG.ISSHOW_SEARCH_RESULTS) {
                 result.token = `\r\n<think>${await Utils.organizeSearchResults(linejosn.webSearchResults)}</think>\r\n`;
             } else {
-                result.token = token;
+                result.token = linejosn?.token;
             }
             return result;
         case 'grok-3':
-            result.token = token;
+            result.token = linejosn?.token;
             return result;
         case 'grok-3-deepsearch':
             if (linejosn.messageTag === "final") {
-                result.token = token;
+                result.token = linejosn?.token;
             }
             return result;
         case 'grok-3-reasoning':
+            if(linejosn?.isThinking && !CONFIG.SHOW_THINKING)return result;
+
             if (linejosn?.isThinking && !CONFIG.IS_THINKING) {
-                result.token = "<think>" + token;
+                result.token = "<think>" + linejosn?.token;
                 CONFIG.IS_THINKING = true;
             } else if (CONFIG.IS_THINKING && !linejosn.isThinking) {
-                result.token = "</think>" + token;
+                result.token = "</think>" + linejosn?.token;
                 CONFIG.IS_THINKING = false;
             } else {
-                result.token = token;
+                result.token = linejosn?.token;
             }
             return result;
     }
@@ -598,8 +570,8 @@ async function handleResponse(response, model, res, isStream) {
         let buffer = '';
         let fullResponse = '';
         const dataPromises = [];
-
-        return new Promise((resolve, reject) => {
+        
+        return new Promise((resolve, reject) => { 
             stream.on('data', async (chunk) => {
                 buffer += chunk.toString();
                 const lines = buffer.split('\n');
@@ -608,15 +580,16 @@ async function handleResponse(response, model, res, isStream) {
                 for (const line of lines) {
                     if (!line.trim()) continue;
                     const trimmedLine = line.trim();
-                    fs.appendFileSync(path.resolve(process.cwd(), 'response4.txt'), trimmedLine + '\n');
                     if (trimmedLine.startsWith('data: ')) {
                         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) {
-                                reject(new Error(linejosn));
+                                Logger.error(JSON.stringify(linejosn,null,2), 'Server');
+                                stream.destroy();
+                                reject(new Error("RateLimitError")); 
                                 return;
                             }
                             if (linejosn?.doImgGen || linejosn?.imageAttachmentInfo) {
@@ -643,7 +616,6 @@ async function handleResponse(response, model, res, isStream) {
                             })();
                             dataPromises.push(processPromise);
                         } catch (error) {
-                            console.log(error);
                             continue;
                         }
                     }
@@ -663,14 +635,14 @@ async function handleResponse(response, model, res, isStream) {
                     }
                     CONFIG.IS_IMG_GEN = false;
                     CONFIG.IS_IMG_GEN2 = false;
-                    resolve();
+                    resolve(); 
                 } catch (error) {
-                    reject(error);
+                    reject(error); 
                 }
             });
+
             stream.on('error', (error) => {
-                Logger.error(error, 'Server');
-                reject(error);
+                reject(error); 
             });
         });
     } catch (error) {
@@ -679,18 +651,15 @@ async function handleResponse(response, model, res, isStream) {
         CONFIG.IS_IMG_GEN2 = false;
         throw error;
     }
-
 }
 
 async function handleImageResponse(imageUrl) {
-    //对服务器发送图片请求
-    const MAX_RETRIES = 3;
+    const MAX_RETRIES = 2;
     let retryCount = 0;
     let imageBase64Response;
 
     while (retryCount < MAX_RETRIES) {
         try {
-            //发送图片请求获取图片
             imageBase64Response = await fetch(`https://assets.grok.com/${imageUrl}`, {
                 method: 'GET',
                 headers: {
@@ -699,16 +668,11 @@ async function handleImageResponse(imageUrl) {
                 }
             });
 
-            if (imageBase64Response.ok) {
-                break; // 如果请求成功，跳出重试循环
-            }
-
+            if (imageBase64Response.ok) break; 
             retryCount++;
             if (retryCount === MAX_RETRIES) {
                 throw new Error(`上游服务请求失败! status: ${imageBase64Response.status}`);
             }
-
-            // 等待一段时间后重试（可以使用指数退避）
             await new Promise(resolve => setTimeout(resolve, CONFIG.API.RETRY_TIME * retryCount));
 
         } catch (error) {
@@ -716,13 +680,20 @@ async function handleImageResponse(imageUrl) {
             if (retryCount === MAX_RETRIES) {
                 throw error;
             }
-            // 等待一段时间后重试
             await new Promise(resolve => setTimeout(resolve, CONFIG.API.RETRY_TIME * retryCount));
         }
     }
 
+
     const arrayBuffer = await imageBase64Response.arrayBuffer();
     const imageBuffer = Buffer.from(arrayBuffer);
+
+    if(!CONFIG.API.PICGO_KEY){
+        const base64Image = imageBuffer.toString('base64');
+        const imageContentType = imageBase64Response.headers.get('content-type');
+        return `![image](data:${imageContentType};base64,${base64Image})`
+    }
+
     const formData = new FormData();
 
     formData.append('source', imageBuffer, {
@@ -742,7 +713,7 @@ async function handleImageResponse(imageUrl) {
     if (!responseURL.ok) {
         return "生图失败，请查看图床密钥是否设置正确"
     } else {
-        console.log("生图成功");
+        Logger.info("生图成功", 'Server');
         const result = await responseURL.json();
         return `![image](${result.image.url})`
     }
@@ -750,6 +721,7 @@ async function handleImageResponse(imageUrl) {
 
 const tokenManager = new AuthTokenManager();
 await initialization();
+
 // 中间件配置
 const app = express();
 app.use(Logger.requestLogger);
@@ -760,8 +732,8 @@ app.use(cors({
     methods: ['GET', 'POST', 'OPTIONS'],
     allowedHeaders: ['Content-Type', 'Authorization']
 }));
-// API路由
-app.get('/hf/v1/models', (req, res) => {
+
+app.get('/v1/models', (req, res) => {
     res.json({
         object: "list",
         data: Object.keys(CONFIG.MODELS).map((model, index) => ({
@@ -774,7 +746,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 authToken = req.headers.authorization?.replace('Bearer ', '');
         if (authToken !== CONFIG.API.API_KEY) {
@@ -782,22 +754,20 @@ app.post('/hf/v1/chat/completions', async (req, res) => {
         }
         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);
         
         while (retryCount < CONFIG.RETRY.MAX_ATTEMPTS) {
             retryCount++;
-            const grokClient = new GrokApiClient(req.body.model);
-            const requestPayload = await grokClient.prepareChatRequest(req.body);
-            console.log(JSON.stringify(requestPayload,null,2));
-
             if (!CONFIG.API.TEMP_COOKIE) {
                 await Utils.get_signature();
             }
-            //获取验证cookie
+
             if (isTempCookie) {
                 CONFIG.API.SIGNATURE_COOKIE = CONFIG.API.TEMP_COOKIE;
                 Logger.info(`已切换为临时令牌`, 'Server');
             } else {
-                CONFIG.API.SIGNATURE_COOKIE = Utils.createAuthHeaders(req.body.model);
+                CONFIG.API.SIGNATURE_COOKIE = await Utils.createAuthHeaders(req.body.model);
             }
             Logger.info(`当前令牌索引: ${CONFIG.SSO_INDEX}`, 'Server');
             const newMessageReq = await fetch(`${CONFIG.API.BASE_URL}/api/rpc`, {
@@ -816,7 +786,7 @@ app.post('/hf/v1/chat/completions', async (req, res) => {
 
             const responseText = await newMessageReq.json();
             const conversationId = responseText.conversationId;
-            //发送对话
+
             const response = await fetch(`${CONFIG.API.BASE_URL}/api/conversations/${conversationId}/responses`, {
                 method: 'POST',
                 headers: {
@@ -835,7 +805,7 @@ app.post('/hf/v1/chat/completions', async (req, res) => {
                 Logger.info(`当前剩余可用令牌数: ${tokenManager.getTokenCount()}`, 'Server');
                 try {
                     await handleResponse(response, req.body.model, res, req.body.stream);
-                    return; // 成功后直接返回
+                    return; 
                 } catch (error) {
                     if(isTempCookie){
                         await Utils.get_signature();
@@ -852,7 +822,6 @@ app.post('/hf/v1/chat/completions', async (req, res) => {
                     }
                 }
             } else {
-                // 处理429错误
                 if (response.status === 429) {
                     if (isTempCookie) {
                         await Utils.get_signature();
@@ -895,12 +864,11 @@ app.post('/hf/v1/chat/completions', async (req, res) => {
 });
 
 
-// 404处理
 app.use((req, res) => {
     res.status(200).send('api运行正常');
 });
 
-// 启动服务器
+
 app.listen(CONFIG.SERVER.PORT, () => {
-    console.log(`服务器已启动，监听端口: ${CONFIG.SERVER.PORT}`);
+    Logger.info(`服务器已启动，监听端口: ${CONFIG.SERVER.PORT}`, 'Server');
 });