Update index.js

index.js

@@ -1,24 +1,25 @@
 import express from 'express';
 import fetch from 'node-fetch';
+import FormData from 'form-data';
 import dotenv from 'dotenv';
 import cors from 'cors';
-import { launch } from 'puppeteer';
+import puppeteer from 'puppeteer';
 import { v4 as uuidv4 } from 'uuid';
-import path from 'path';
-import fs from 'fs';
 
 dotenv.config();
 // 配置常量
 const CONFIG = {
     MODELS: {
         'grok-latest': 'grok-latest',
-        'grok-latest-image': 'grok-latest'
+        'grok-latest-image': 'grok-latest',
+        'grok-latest-search': 'grok-latest'
     },
     API: {
         BASE_URL: "https://grok.com",
         API_KEY: process.env.API_KEY || "sk-123456",
         SSO_TOKEN: null,//登录时才有的认证cookie,这里暂时用不到，之后可能需要
-        SIGNATURE_COOKIE: null
+        SIGNATURE_COOKIE: null,
+        PICGO_KEY: process.env.PICGO_KEY || null //想要生图的话需要填入这个PICGO图床的key
     },
     SERVER: {
         PORT: process.env.PORT || 3000,
@@ -28,7 +29,7 @@ const CONFIG = {
         MAX_ATTEMPTS: 3,//重试次数
         DELAY_BASE: 1000 // 基础延迟时间（毫秒）
     },
-    CHROME_PATH: process.env.CHROME_PATH || 'C:/Program Files/Google/Chrome/Application/chrome.exe'// 替换为你的 Chrome 实际路径
+    ISSHOW_SEARCH_RESULTS: process.env.ISSHOW_SEARCH_RESULTS || true//是否显示搜索结果,默认开启
 };
 
 
@@ -51,20 +52,24 @@ const DEFAULT_HEADERS = {
     'baggage': 'sentry-public_key=b311e0f2690c81f25e2c4cf6d4f7ce1c'
 };
 
-// 定义签名文件路径为 /tmp 目录
-const SIGNATURE_FILE_PATH = '/tmp/signature.json';
-
 class Utils {
     static async extractGrokHeaders() {
+        console.log("开始提取头信息");
         try {
             // 启动浏览器
-            const browser = await launch({
-                executablePath: CONFIG.CHROME_PATH,
-                headless: true
+            const browser = await puppeteer.launch({
+                headless: true,
+                args: [
+                    '--no-sandbox',
+                    '--disable-setuid-sandbox',
+                    '--disable-dev-shm-usage',
+                    '--disable-gpu'
+                ],
+                executablePath: CONFIG.CHROME_PATH
             });
 
             const page = await browser.newPage();
-            await page.goto(`${CONFIG.API.BASE_URL}`, { waitUntil: 'networkidle0' });
+            await page.goto('https://grok.com/', { waitUntil: 'networkidle0' });
 
             // 获取所有 Cookies
             const cookies = await page.cookies();
@@ -89,6 +94,9 @@ class Utils {
         }
     }
     static async get_signature() {
+        if (CONFIG.API.SIGNATURE_COOKIE) {
+            return CONFIG.API.SIGNATURE_COOKIE;
+        }
         console.log("刷新认证信息");
         let retryCount = 0;
         while (retryCount < CONFIG.RETRY.MAX_ATTEMPTS) {
@@ -96,13 +104,7 @@ class Utils {
             if (headers) {
                 console.log("获取认证信息成功");
                 CONFIG.API.SIGNATURE_COOKIE = { cookie: `x-anonuserid=${headers["x-anonuserid"]}; x-challenge=${headers["x-challenge"]}; x-signature=${headers["x-signature"]}` };
-                try {
-                    fs.writeFileSync(SIGNATURE_FILE_PATH, JSON.stringify(CONFIG.API.SIGNATURE_COOKIE));//保存认证信息
-                    return CONFIG.API.SIGNATURE_COOKIE;
-        } catch (error) {
-            console.error('写入签名文件失败:', error);
-                    return CONFIG.API.SIGNATURE_COOKIE;
-                }
+                return CONFIG.API.SIGNATURE_COOKIE;
             }
             retryCount++;
             if (retryCount >= CONFIG.RETRY.MAX_ATTEMPTS) {
@@ -125,8 +127,6 @@ class Utils {
         }
     }
     static async handleError(error, res, originalRequest = null) {
-        console.error('Error:', error);
-
         // 如果是500错误且提供了原始请求函数，尝试重新获取签名并重试
         if (error.status === 500 && originalRequest) {
             try {
@@ -155,71 +155,22 @@ class Utils {
             }
         });
     }
+    static async organizeSearchResults(searchResults) {
+        // 确保传入的是有效的搜索结果对象
+        if (!searchResults || !searchResults.results) {
+            return '';
+        }
 
-    static async makeRequest(req, res) {
-        try {
-            if (!CONFIG.API.SIGNATURE_COOKIE) {
-                await Utils.get_signature();
-                try {
-                    CONFIG.API.SIGNATURE_COOKIE = JSON.parse(fs.readFileSync(SIGNATURE_FILE_PATH));
-                } catch (error) {
-                    console.error('读取签名文件失败:', error);
-                    await Utils.get_signature(); // 如果读取失败，重新获取签名
-                }
-            }
-            const grokClient = new GrokApiClient(req.body.model);
-            const requestPayload = await grokClient.prepareChatRequest(req.body);
-            //创建新对话        
-            const newMessageReq = await fetch(`${CONFIG.API.BASE_URL}/api/rpc`, {
-                method: 'POST',
-                headers: {
-                    ...DEFAULT_HEADERS,
-                    ...CONFIG.API.SIGNATURE_COOKIE
-                },
-                body: JSON.stringify({
-                    rpc: "createConversation",
-                    req: {
-                        temporary: false
-                    }
-                })
-            });
-
-            if (!newMessageReq.ok) {
-                throw new Error(`上游服务请求失败! status: ${newMessageReq.status}`);
-            }
-
-            // 获取响应文本
-            const responseText = await newMessageReq.json();
-            const conversationId = responseText.conversationId;
-            console.log("会话ID:conversationId", conversationId);
-            if (!conversationId) {
-                throw new Error(`创建会话失败! status: ${newMessageReq.status}`);
-            }
-            //发送对话
-            const response = await fetch(`${CONFIG.API.BASE_URL}/api/conversations/${conversationId}/responses`, {
-                method: 'POST',
-                headers: {
-                    "accept": "text/event-stream",
-                    "baggage": "sentry-public_key=b311e0f2690c81f25e2c4cf6d4f7ce1c",
-                    "content-type": "text/plain;charset=UTF-8",
-                    "Connection": "keep-alive",
-                    ...CONFIG.API.SIGNATURE_COOKIE
-                },
-                body: JSON.stringify(requestPayload)
-            });
-
-            if (!response.ok) {
-                throw new Error(`上游服务请求失败! status: ${response.status}`);
-            }
+        const results = searchResults.results;
+        const formattedResults = results.map((result, index) => {
+            // 处理可能为空的字段
+            const title = result.title || '未知标题';
+            const url = result.url || '#';
+            const preview = result.preview || '无预览内容';
 
-            if (req.body.stream) {
-                await handleStreamResponse(response, req.body.model, res);
-            } else {
-                await handleNormalResponse(response, req.body.model, res);
-            }
-        } catch (error) {
-            throw error;
-        }
+            return `\r\n<details><summary>资料[${index}]: ${title}</summary>\r\n${preview}\r\n\n[Link](${url})\r\n</details>`;
+        });
+        return formattedResults.join('\n\n');
     }
 }
 
@@ -299,6 +250,20 @@ class GrokApiClient {
     }
 
     async prepareChatRequest(request) {
+
+        if (request.model === 'grok-latest-image' && !CONFIG.API.PICGO_KEY) {
+            throw new Error(`该模型需要配置PICGO图床密钥!`);
+        }
+        var todoMessages = request.messages;
+        if (request.model === 'grok-latest-image' || request.model === 'grok-latest-search') {
+            todoMessages = Array.isArray(todoMessages) ? [todoMessages[todoMessages.length - 1]] : [todoMessages];;
+        }
+        let fileAttachments = [];
+        let messages = '';
+        let lastRole = null;
+        let lastContent = '';
+        let search = false;
+
         const processImageUrl = async (content) => {
             if (content.type === 'image_url' && content.image_url.url.includes('data:image')) {
                 const imageResponse = await this.uploadBase64Image(
@@ -309,12 +274,8 @@ class GrokApiClient {
             }
             return null;
         };
-        let fileAttachments = [];
-        let messages = '';
-        let lastRole = null;
-        let lastContent = '';
 
-        for (const current of request.messages) {
+        for (const current of todoMessages) {
             const role = current.role === 'assistant' ? 'assistant' : 'user';
             let textContent = '';
             // 处理消息内容
@@ -323,7 +284,7 @@ class GrokApiClient {
                 for (const item of current.content) {
                     if (item.type === 'image_url') {
                         // 如果是图片且是最后一条消息，则处理图片
-                        if (current === request.messages[request.messages.length - 1]) {
+                        if (current === todoMessages[todoMessages.length - 1]) {
                             const processedImage = await processImageUrl(item);
                             if (processedImage) fileAttachments.push(processedImage);
                         }
@@ -336,7 +297,7 @@ class GrokApiClient {
                 // 处理单个对象内容
                 if (current.content.type === 'image_url') {
                     // 如果是图片且是最后一条消息，则处理图片
-                    if (current === request.messages[request.messages.length - 1]) {
+                    if (current === todoMessages[todoMessages.length - 1]) {
                         const processedImage = await processImageUrl(current.content);
                         if (processedImage) fileAttachments.push(processedImage);
                     }
@@ -361,7 +322,7 @@ class GrokApiClient {
                     lastContent = textContent;
                     lastRole = role;
                 }
-            } else if (current === request.messages[request.messages.length - 1] && fileAttachments.length > 0) {
+            } else if (current === todoMessages[todoMessages.length - 1] && fileAttachments.length > 0) {
                 // 如果是最后一条消息且有图片附件，添加空消息占位
                 messages += `${role.toUpperCase()}: [图片]\n`;
             }
@@ -373,6 +334,9 @@ class GrokApiClient {
 
         messages = messages.trim();
 
+        if (request.model === 'grok-latest-search') {
+            search = true;
+        }
         return {
             message: messages,
             modelName: this.modelId,
@@ -385,11 +349,11 @@ class GrokApiClient {
             imageGenerationCount: 1,
             toolOverrides: {
                 imageGen: request.model === 'grok-latest-image',
-                webSearch: false,
-                xSearch: false,
-                xMediaSearch: false,
-                trendsSearch: false,
-                xPostAnalyze: false
+                webSearch: search,
+                xSearch: search,
+                xMediaSearch: search,
+                trendsSearch: search,
+                xPostAnalyze: search
             }
         };
     }
@@ -416,9 +380,6 @@ class MessageProcessor {
             };
         }
 
-        // 如果是数组（图片响应），直接使用整个数组作为content
-        const messageContent = Array.isArray(message) ? message : message;
-
         return {
             ...baseResponse,
             object: 'chat.completion',
@@ -426,7 +387,7 @@ class MessageProcessor {
                 index: 0,
                 message: {
                     role: 'assistant',
-                    content: messageContent
+                    content: message
                 },
                 finish_reason: 'stop'
             }],
@@ -437,8 +398,8 @@ class MessageProcessor {
 
 // 中间件配置
 const app = express();
-app.use(express.json({ limit: '5mb' }));
-app.use(express.urlencoded({ extended: true, limit: '5mb' }));
+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'],
@@ -462,11 +423,63 @@ app.post('/hf/v1/chat/completions', async (req, res) => {
     if (authToken !== CONFIG.API.API_KEY) {
         return res.status(401).json({ error: 'Unauthorized' });
     }
+    const makeRequest = async () => {
+        if (!CONFIG.API.SIGNATURE_COOKIE) {
+            await Utils.get_signature();
+        }
+        const grokClient = new GrokApiClient(req.body.model);
+        const requestPayload = await grokClient.prepareChatRequest(req.body);
+        //创建新对话        
+        const newMessageReq = await fetch(`${CONFIG.API.BASE_URL}/api/rpc`, {
+            method: 'POST',
+            headers: {
+                ...DEFAULT_HEADERS,
+                ...CONFIG.API.SIGNATURE_COOKIE
+            },
+            body: JSON.stringify({
+                rpc: "createConversation",
+                req: {
+                    temporary: false
+                }
+            })
+        });
+        if (!newMessageReq.ok) {
+            throw new Error(`上游服务请求失败! status: ${newMessageReq.status}`);
+        }
 
+        // 获取响应文本
+        const responseText = await newMessageReq.json();
+        const conversationId = responseText.conversationId;
+        console.log("会话ID:conversationId", conversationId);
+        if (!conversationId) {
+            throw new Error(`创建会话失败! status: ${newMessageReq.status}`);
+        }
+        //发送对话
+        const response = await fetch(`${CONFIG.API.BASE_URL}/api/conversations/${conversationId}/responses`, {
+            method: 'POST',
+            headers: {
+                "accept": "text/event-stream",
+                "baggage": "sentry-public_key=b311e0f2690c81f25e2c4cf6d4f7ce1c",
+                "content-type": "text/plain;charset=UTF-8",
+                "Connection": "keep-alive",
+                ...CONFIG.API.SIGNATURE_COOKIE
+            },
+            body: JSON.stringify(requestPayload)
+        });
+
+        if (!response.ok) {
+            throw new Error(`上游服务请求失败! status: ${response.status}`);
+        }
+        if (req.body.stream) {
+            await handleStreamResponse(response, req.body.model, res);
+        } else {
+            await handleNormalResponse(response, req.body.model, res);
+        }
+    }
     try {
-        await Utils.makeRequest(req, res);
+        await makeRequest();
     } catch (error) {
-        await Utils.handleError(error, res);
+        await Utils.handleError(error, res, makeRequest);
     }
 });
 
@@ -489,8 +502,22 @@ async function handleStreamResponse(response, model, res) {
                 continue;
             }
             let linejosn = JSON.parse(data);
+            if (linejosn?.error?.name === "RateLimitError") {
+                var responseData = MessageProcessor.createChatResponse(`${linejosn.error.name},请重新对话`, model, true);
+                CONFIG.API.SIGNATURE_COOKIE = null;
+                console.log("认证信息已删除");
+                res.write(`data: ${JSON.stringify(responseData)}\n\n`);
+                return res.end();
+            }
             switch (model) {
-                case "grok-latest":
+                case 'grok-latest-image':
+                    if (linejosn.response === "modelResponse" && linejosn?.modelResponse?.generatedImageUrls) {
+                        const dataImage = await handleImageResponse(linejosn.modelResponse.generatedImageUrls);
+                        const responseData = MessageProcessor.createChatResponse(dataImage, model, true);
+                        res.write(`data: ${JSON.stringify(responseData)}\n\n`);
+                    }
+                    break;
+                case 'grok-latest':
                     if (linejosn.response === "token") {
                         const token = linejosn.token;
                         if (token && token.length > 0) {
@@ -499,15 +526,19 @@ async function handleStreamResponse(response, model, res) {
                         }
                     }
                     break;
-                case "grok-latest-image":
-                    if (linejosn.response === "modelResponse") {
-                        if (linejosn?.modelResponse?.generatedImageUrls?.length > 0) {
-                            let imageUrl = linejosn.modelResponse.generatedImageUrls[0]; // 获取第一个URL
-                            const dataImage = await handleImageResponse(imageUrl);
-                            const responseData = MessageProcessor.createChatResponse(dataImage, model, true);
+                case 'grok-latest-search':
+                    if (linejosn.response === "token") {
+                        const token = linejosn.token;
+                        if (token && token.length > 0) {
+                            const responseData = MessageProcessor.createChatResponse(token, model, true);
                             res.write(`data: ${JSON.stringify(responseData)}\n\n`);
                         }
                     }
+                    if (linejosn.response === "webSearchResults" && CONFIG.ISSHOW_SEARCH_RESULTS) {
+                        const searchResults = await Utils.organizeSearchResults(linejosn.webSearchResults);
+                        const responseData = MessageProcessor.createChatResponse(`<thinking>\r\n${searchResults}\r\n</thinking>\r\n`, model, true);
+                        res.write(`data: ${JSON.stringify(responseData)}\n\n`);
+                    }
                     break;
             }
         }
@@ -524,7 +555,6 @@ async function handleNormalResponse(response, model, res) {
 
     try {
         const responseText = await response.text();
-        console.log('原始响应文本:', responseText);
         const lines = responseText.split('\n');
 
         for (const line of lines) {
@@ -533,101 +563,114 @@ async function handleNormalResponse(response, model, res) {
             const data = line.slice(6);
             if (data === '[DONE]') continue;
             let linejosn = JSON.parse(data);
-            console.log('解析的JSON数据:', linejosn);
-            
-            if (linejosn.response === "token") {
-                const token = linejosn.token;
-                if (token && token.length > 0) {
-                    fullResponse += token;
-                }
-            } else if (linejosn.response === "modelResponse" && model === 'grok-latest-image') {
-                console.log('检测到图片响应:', linejosn.modelResponse);
-                if (linejosn?.modelResponse?.generatedImageUrls?.length > 0) {
-                    imageUrl = linejosn.modelResponse.generatedImageUrls[0]; // 获取第一个URL
-                    console.log('获取到图片URL:', imageUrl);
-                }
+            if (linejosn?.error?.name === "RateLimitError") {
+                fullResponse = `${linejosn.error.name},请重新对话`;
+                CONFIG.API.SIGNATURE_COOKIE = null;
+                console.log("认证信息已删除");
+                return res.json(MessageProcessor.createChatResponse(fullResponse, model));
+            }
+            switch (model) {
+                case 'grok-latest-image':
+                    if (linejosn.response === "modelResponse" && linejosn?.modelResponse?.generatedImageUrls) {
+                        imageUrl = linejosn.modelResponse.generatedImageUrls;
+                    }
+                    break;
+                case 'grok-latest':
+                    if (linejosn.response === "token") {
+                        const token = linejosn.token;
+                        if (token && token.length > 0) {
+                            fullResponse += token;
+                        }
+                    }
+                    break;
+                case 'grok-latest-search':
+                    if (linejosn.response === "token") {
+                        const token = linejosn.token;
+                        if (token && token.length > 0) {
+                            fullResponse += token;
+                        }
+                    }
+                    if (linejosn.response === "webSearchResults" && CONFIG.ISSHOW_SEARCH_RESULTS) {
+                        fullResponse += `\r\n<thinking>${await Utils.organizeSearchResults(linejosn.webSearchResults)}</thinking>\r\n`;
+                    }
+                    break;
             }
         }
-
         if (imageUrl) {
-            console.log('开始处理图片URL:', imageUrl);
             const dataImage = await handleImageResponse(imageUrl);
-            
-            const responseData = MessageProcessor.createChatResponse([dataImage], model);
-            console.log('最终的响应对象:', JSON.stringify(responseData, null, 2));
+            const responseData = MessageProcessor.createChatResponse(dataImage, model);
             res.json(responseData);
         } else {
-            console.log('没有图片URL，返回文本响应:', fullResponse);
             const responseData = MessageProcessor.createChatResponse(fullResponse, model);
             res.json(responseData);
         }
     } catch (error) {
-        console.error('处理响应时发生错误:', error);
         Utils.handleError(error, res);
     }
 }
-
 async function handleImageResponse(imageUrl) {
-    try {
-        //对服务器发送图片请求
-        const MAX_RETRIES = 3;
-        let retryCount = 0;
-        let imageBase64Response;
-
-        while (retryCount < MAX_RETRIES) {
-            try {
-                //发送图片请求获取图片
-                imageBase64Response = await fetch(`https://assets.grok.com/${imageUrl}`, {
-                    method: 'GET',
-                    headers: {
-                        ...DEFAULT_HEADERS,
-                        ...CONFIG.API.SIGNATURE_COOKIE
-                    }
-                });
+    //对服务器发送图片请求
+    const MAX_RETRIES = 3;
+    let retryCount = 0;
+    let imageBase64Response;
 
-                if (imageBase64Response.ok) {
-                    break; // 如果请求成功，跳出重试循环
-                }
-
-                retryCount++;
-                if (retryCount === MAX_RETRIES) {
-                    throw new Error(`上游服务请求失败! status: ${imageBase64Response.status}`);
+    while (retryCount < MAX_RETRIES) {
+        try {
+            //发送图片请求获取图片
+            imageBase64Response = await fetch(`https://assets.grok.com/${imageUrl}`, {
+                method: 'GET',
+                headers: {
+                    ...DEFAULT_HEADERS,
+                    ...CONFIG.API.SIGNATURE_COOKIE
                 }
+            });
 
-                // 等待一段时间后重试（可以使用指数退避）
-                await new Promise(resolve => setTimeout(resolve, 500 * retryCount));
-
-            } catch (error) {
-                retryCount++;
-                if (retryCount === MAX_RETRIES) {
-                    throw error;
-                }
-                // 等待一段时间后重试
-                await new Promise(resolve => setTimeout(resolve, 500 * retryCount));
+            if (imageBase64Response.ok) {
+                break; // 如果请求成功，跳出重试循环
             }
-        }
 
-        const arrayBuffer = await imageBase64Response.arrayBuffer();
-        const imagebuffer = Buffer.from(arrayBuffer);
-
-        const base64String = imagebuffer.toString('base64');
+            retryCount++;
+            if (retryCount === MAX_RETRIES) {
+                throw new Error(`上游服务请求失败! status: ${imageBase64Response.status}`);
+            }
 
-        // 获取图片类型（例如：image/jpeg, image/png）
-        const imageContentType = imageBase64Response.headers.get('content-type');
+            // 等待一段时间后重试（可以使用指数退避）
+            await new Promise(resolve => setTimeout(resolve, CONFIG.API.RETRY_TIME * retryCount));
 
-        // 修改返回结构，确保返回Markdown格式的图片标签
-        return {
-            type: "image_url",
-            image_url: {
-                url: `data:image/jpg;base64,${base64String}`
+        } catch (error) {
+            retryCount++;
+            if (retryCount === MAX_RETRIES) {
+                throw error;
             }
-        };
+            // 等待一段时间后重试
+            await new Promise(resolve => setTimeout(resolve, CONFIG.API.RETRY_TIME * retryCount));
+        }
+    }
 
-        // 或者如果需要直接返回可显示的字符串
-        // return `![generated image](data:image/jpg;base64,${base64String})`;
-    } catch (error) {
-        console.error('图片处理失败:', error);
-        return '图片生成失败';
+    const arrayBuffer = await imageBase64Response.arrayBuffer();
+    const imageBuffer = Buffer.from(arrayBuffer);
+    const formData = new FormData();
+
+    formData.append('source', imageBuffer, {
+        filename: '2223.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) {
+        console.error("上传图床失败");
+        return "生图失败，请查看图床密钥是否设置正确"
+    } else {
+        const result = await responseURL.json();
+        return `![image](${result.image.url})`
     }
 }