Update index.mjs

index.mjs

@@ -5,17 +5,17 @@ import localtunnel from "localtunnel";
 import ngrok from 'ngrok';
 import {v4 as uuidv4} from "uuid";
 import './proxyAgent.mjs';
-import SessionManager from './sessionManager.mjs';
 import {storeImage} from './imageStorage.mjs';
 import fetch from 'node-fetch';
 import path from 'path';
+import geoip from 'geoip-lite';
+import RequestLogger from './requestLogger.mjs';
 
 const app = express();
 const port = process.env.PORT || 8080;
 const validApiKey = process.env.PASSWORD;
 const availableModels = [
     "openai_o1",
-    "openai_o1_mini",
     "openai_o1_preview",
     "gpt_4o",
     "gpt_4_turbo",
@@ -34,7 +34,8 @@ const availableModels = [
     "command_r_plus",
     "zephyr",
     "qwen2p5_72b",
-    "llama3_1_405b"
+    "llama3_1_405b",
+    "grok_2"
 ];
 const modelMappping = {
     "claude-3-5-sonnet-latest": "claude_3_5_sonnet",
@@ -49,9 +50,8 @@ const modelMappping = {
     "gpt-4": "gpt_4",
     "gpt-4o": "gpt_4o",
     "gpt-4-turbo": "gpt_4_turbo",
-    "o1": "openai_o1",
+    "openai-o1": "openai_o1",
-    "o1-preview": "openai_o1_preview",
+    "o1-preview": "openai_o1",
-    "o1-mini": "openai_o1_mini"
 };
 
 // import config.mjs
@@ -69,7 +69,10 @@ const provider = new YouProvider(config);
 await provider.init(config);
 
 // 初始化 SessionManager
-const sessionManager = new SessionManager(provider.provider);
+const sessionManager = provider.getSessionManager();
+
+// 初始化 RequestLogger
+const requestLogger = new RequestLogger();
 
 // handle preflight request
 app.use((req, res, next) => {
@@ -144,11 +147,12 @@ app.post("/v1/chat/completions", OpenAIApiKeyAuth, (req, res) => {
         let releaseSessionCalled = false;
         let completion;
         let cancel;
+        let selectedBrowserId;
         // 定义释放会话
         const releaseSession = () => {
-            if (selectedSession && !releaseSessionCalled) {
+            if (selectedSession && selectedBrowserId && !releaseSessionCalled) {
-                sessionManager.releaseSession(selectedSession);
+                sessionManager.releaseSession(selectedSession, selectedBrowserId);
-                console.log(`释放会话 ${selectedSession}`);
+                console.log(`释放会话 ${selectedSession} 和浏览器实例 ${selectedBrowserId}`);
                 releaseSessionCalled = true;
             }
         };
@@ -167,14 +171,35 @@ app.post("/v1/chat/completions", OpenAIApiKeyAuth, (req, res) => {
         });
 
         try {
-            // 获取并锁定可用会话
+            // 获取客户端 IP
-            const { selectedUsername, modeSwitched } = await sessionManager.getSessionByStrategy('round_robin');
+            const clientIpAddress = req.headers["x-forwarded-for"] || req.socket.remoteAddress;
+            const geo = geoip.lookup(clientIpAddress) || {};
+            const locationInfo = `${geo.country || 'Unknown'}-${geo.region || 'Unknown'}-${geo.city || 'Unknown'}`;
+            const requestTime = new Date();
+
+            // 获取并锁定可用会话和浏览器实例
+            const {
+                selectedUsername,
+                modeSwitched,
+                browserInstance
+            } = await sessionManager.getSessionByStrategy('round_robin');
             selectedSession = selectedUsername;
+            selectedBrowserId = browserInstance.id;
             console.log("Using session " + selectedSession);
 
+            // 记录请求信息
+            await requestLogger.logRequest({
+                time: requestTime,
+                ip: clientIpAddress,
+                location: locationInfo,
+                model: jsonBody.model,
+                session: selectedSession
+            });
+
             ({completion, cancel} = await provider.getCompletion({
                 username: selectedSession,
                 messages: jsonBody.messages,
+                browserInstance: browserInstance,
                 stream: !!jsonBody.stream,
                 proxyModel: jsonBody.model,
                 useCustomMode: process.env.USE_CUSTOM_MODE === "true",
@@ -268,6 +293,7 @@ app.post("/v1/chat/completions", OpenAIApiKeyAuth, (req, res) => {
                     res.write(createEvent("data", "[DONE]"));
                     res.end();
                 }
+
                 releaseSession();
             });
 
@@ -534,21 +560,28 @@ app.post("/v1/messages", AnthropicApiKeyAuth, (req, res) => {
             proxyModel = process.env.AI_MODEL;
         } else if (jsonBody.model && modelMappping[jsonBody.model]) {
             proxyModel = modelMappping[jsonBody.model];
+        } else if (jsonBody.model) {
+            proxyModel = jsonBody.model;
         } else {
             proxyModel = "claude_3_opus";
         }
         console.log(`Using model ${proxyModel}`);
 
+        if (proxyModel && !availableModels.includes(proxyModel)) {
+            res.json({error: {code: 404, message: "Invalid Model"}});
+            return;
+        }
+
         let selectedSession;
         let releaseSessionCalled = false;
         let completion;
         let cancel;
-
+        let selectedBrowserId;
         // 定义释放会话
         const releaseSession = () => {
-            if (selectedSession && !releaseSessionCalled) {
+            if (selectedSession && selectedBrowserId && !releaseSessionCalled) {
-                sessionManager.releaseSession(selectedSession);
+                sessionManager.releaseSession(selectedSession, selectedBrowserId);
-                console.log(`释放会话 ${selectedSession}`);
+                console.log(`释放会话 ${selectedSession} 和浏览器实例 ${selectedBrowserId}`);
                 releaseSessionCalled = true;
             }
         };
@@ -567,14 +600,35 @@ app.post("/v1/messages", AnthropicApiKeyAuth, (req, res) => {
         });
 
         try {
-            // 获取并锁定会话
+            // 获取客户端 IP
-            const { selectedUsername, modeSwitched } = await sessionManager.getSessionByStrategy('round_robin');
+            const clientIpAddress = req.headers["x-forwarded-for"] || req.socket.remoteAddress;
+            const geo = geoip.lookup(clientIpAddress) || {};
+            const locationInfo = `${geo.country || 'Unknown'}-${geo.region || 'Unknown'}-${geo.city || 'Unknown'}`;
+            const requestTime = new Date();
+
+            // 获取并锁定可用会话和浏览器实例
+            const {
+                selectedUsername,
+                modeSwitched,
+                browserInstance
+            } = await sessionManager.getSessionByStrategy('round_robin');
             selectedSession = selectedUsername;
+            selectedBrowserId = browserInstance.id;
             console.log("Using session " + selectedSession);
 
+            // 记录请求信息
+            await requestLogger.logRequest({
+                time: requestTime,
+                ip: clientIpAddress,
+                location: locationInfo,
+                model: jsonBody.model,
+                session: selectedSession
+            });
+
             ({completion, cancel} = await provider.getCompletion({
                 username: selectedSession,
                 messages: jsonBody.messages,
+                browserInstance: browserInstance,
                 stream: !!jsonBody.stream,
                 proxyModel: proxyModel,
                 useCustomMode: process.env.USE_CUSTOM_MODE === "true",
@@ -758,6 +812,13 @@ const createLocaltunnel = async (port, subdomain) => {
     }
 };
 
+/*
+    * 创建ngrok隧道
+    * @param {number} port - 本地端口
+    * @param {string} authToken - ngrok的认证token
+    * @param {string} customDomain - 自定义域名
+    * @param {string} subdomain - 子域名
+ */
 const createNgrok = async (port, authToken, customDomain, subdomain) => {
     const ngrokOptions = {addr: port, authtoken: authToken};
 
@@ -793,7 +854,7 @@ const createTunnel = async (tunnelType, port) => {
     if (tunnelType === "localtunnel") {
         return createLocaltunnel(port, process.env.SUBDOMAIN);
     } else if (tunnelType === "ngrok") {
-        return createNgrok(port, process.env.NGROK_AUTH_TOKEN, process.env.NGROK_CUSTOM_DOMAIN, process.env.SUBDOMAIN);
+        return createNgrok(port, process.env.NGROK_AUTH_TOKEN, process.env.NGROK_CUSTOM_DOMAIN, process.env.NGROK_SUBDOMAIN);
     }
 };