feat: 初步支持 Gemini Business 并重构和整理

2026-06-16 21:03:59 +08:00 · 2025-11-28 17:53:48 +08:00
parent 811e175054
commit 3a08ccadf4
18 changed files with 2590 additions and 1031 deletions
@@ -1,83 +1,357 @@
-import readline from 'readline';
-import config from './config.js';
-import { initBrowser, generateImage } from './lmarena.js';
-import { MODEL_MAPPING } from './models.js';
+import { getBackend } from './backend/index.js';
+import { getModelsForBackend, resolveModelId } from './backend/models.js';
+import { select, input } from '@inquirer/prompts';
+import fs from 'fs';
+import path from 'path';
+import http from 'http';
+import { logger } from './logger.js';
+
+// 使用统一后端获取配置和函数
+const { config, name, initBrowser, generateImage, TEMP_DIR } = getBackend();
+
+logger.info('CLI/Test', `测试工具启动 (后端适配器: ${name})`);

 /**
- * 创建命令行交互接口
+ * 选择测试模式
 */
-const rl = readline.createInterface({
-    input: process.stdin,
-    output: process.stdout
-});
-
-/**
- * 封装 readline 为 Promise
- * @param {string} query 提示问题
- * @returns {Promise<string>} 用户输入
- */
-const ask = (query) => new Promise((resolve) => rl.question(query, resolve));
-
-async function main() {
-    console.log('>>> [CLI] LMArena CLI 测试工具');
-    console.log('>>> [CLI] 正在启动浏览器...');
-
-    let browserContext;
-    try {
-        // 传入配置对象
-        browserContext = await initBrowser(config);
-        console.log('>>> [CLI] 浏览器已就绪。');
-    } catch (err) {
-        console.error('>>> [Error] 浏览器启动失败:', err);
-        process.exit(1);
-    }
-
-    while (true) {
-        console.log('-----------------------------');
-
-        // 1. 获取图片路径
-        const imgInput = await ask('>>> [CLI] 请输入图片路径 (多张用逗号隔开，回车跳过): ');
-        const imagePaths = imgInput.trim()
-            ? imgInput.split(',').map(p => p.trim()).filter(p => p)
-            : [];
-
-        // 2. 获取提示词
-        const prompt = await ask('>>> [CLI] 请输入提示词: ');
-        if (!prompt.trim()) {
-            console.log('>>> [Error] 提示词不能为空，请重试。');
-            continue;
-        }
-
-        // 3. 获取模型 ID
-        const modelInput = await ask('>>> [CLI] 请输入模型 ID (回车跳过使用默认): ');
-        const modelName = modelInput.trim();
-        let modelId = null;
-
-        if (modelName) {
-            if (MODEL_MAPPING[modelName]) {
-                modelId = MODEL_MAPPING[modelName];
-                console.log(`>>> [CLI] 使用模型: ${modelName} (${modelId})`);
-            } else {
-                console.log(`>>> [Warn] 未找到模型 "${modelName}"，将尝试直接使用默认模型。`);
-            }
-        } else {
-            console.log('>>> [CLI] 未指定模型，使用默认值。');
-        }
-
-        console.log(`>>> [CLI] 开始任务: Prompt="${prompt}", Images=${imagePaths.length}`);
-
-        // 4. 调用生图逻辑
-        const result = await generateImage(browserContext, prompt, imagePaths, modelId);
-
-        // 5. 显示结果
-        if (result.error) {
-            console.error('>>> [Error]', result.error);
-        } else if (result.image) {
-            console.log('>>> [Success] 图片 URL:', result.image);
-        } else {
-            console.log('>>> [CLI] AI 使用文本回复:', result.text);
-        }
-    }
+async function selectTestMode() {
+    const mode = await select({
+        message: '选择测试模式',
+        choices: [
+            { name: 'HTTP 服务器测试（需先启动服务器）', value: 'http' },
+            { name: '直接调用适配器', value: 'direct' }
+        ]
+    });
+    return mode;
 }

-main();
+/**
+ * 选择模型
+ */
+async function selectModel() {
+    const models = getModelsForBackend(name);
+    const choices = [
+        { name: 'Skip（使用默认模型）', value: null },
+        ...models.data.map(m => ({ name: m.id, value: m.id }))
+    ];
+
+    const modelId = await select({
+        message: '选择模型',
+        choices,
+        pageSize: 15
+    });
+
+    return modelId;
+}
+
+/**
+ * 输入提示词
+ */
+async function promptForInput() {
+    const prompt = await input({
+        message: '输入提示词（回车使用默认）',
+        default: 'A cute cat'
+    });
+    return prompt.trim();
+}
+
+/**
+ * 输入图片路径
+ */
+async function promptForImages() {
+    const imagesInput = await input({
+        message: '输入图片路径（逗号分隔，回车跳过）',
+        default: ''
+    });
+
+    if (!imagesInput.trim()) {
+        return [];
+    }
+
+    return imagesInput.split(',').map(p => p.trim()).filter(p => p);
+}
+
+/**
+ * HTTP 测试模式 - OpenAI 格式
+ */
+async function testViaHttpOpenAI(prompt, modelId, imagePaths) {
+    const PORT = config.server.port || 3000;
+    const AUTH_TOKEN = config.server.auth;
+
+    logger.info('CLI/Test', 'HTTP 测试 - OpenAI 模式');
+
+    return new Promise((resolve, reject) => {
+        // 构造请求体
+        const messages = [];
+        const lastMessage = { role: 'user', content: [] };
+
+        // 添加文本
+        if (prompt) {
+            lastMessage.content.push({ type: 'text', text: prompt });
+        }
+
+        // 添加图片
+        for (const imgPath of imagePaths) {
+            if (fs.existsSync(imgPath)) {
+                const buffer = fs.readFileSync(imgPath);
+                const base64 = buffer.toString('base64');
+                const ext = path.extname(imgPath).slice(1).toLowerCase();
+                const mimeType = ext === 'jpg' ? 'jpeg' : ext;
+                lastMessage.content.push({
+                    type: 'image_url',
+                    image_url: { url: `data:image/${mimeType};base64,${base64}` }
+                });
+            } else {
+                logger.warn('CLI/Test', `图片不存在，已跳过: ${imgPath}`);
+            }
+        }
+
+        messages.push(lastMessage);
+
+        const body = {
+            messages,
+            ...(modelId && { model: modelId })
+        };
+
+        const bodyStr = JSON.stringify(body);
+
+        const options = {
+            hostname: '127.0.0.1',
+            port: PORT,
+            path: '/v1/chat/completions',
+            method: 'POST',
+            headers: {
+                'Content-Type': 'application/json',
+                'Content-Length': Buffer.byteLength(bodyStr),
+                'Authorization': `Bearer ${AUTH_TOKEN}`
+            }
+        };
+
+        const req = http.request(options, (res) => {
+            let data = '';
+            res.on('data', chunk => data += chunk);
+            res.on('end', () => {
+                if (res.statusCode === 200) {
+                    const response = JSON.parse(data);
+                    resolve(response);
+                } else {
+                    reject(new Error(`HTTP ${res.statusCode}: ${data}`));
+                }
+            });
+        });
+
+        req.on('error', reject);
+        req.write(bodyStr);
+        req.end();
+    });
+}
+
+/**
+ * HTTP 测试模式 - Queue 格式
+ */
+async function testViaHttpQueue(prompt, modelId, imagePaths) {
+    const PORT = config.server.port || 3000;
+    const AUTH_TOKEN = config.server.auth;
+
+    logger.info('CLI/Test', 'HTTP 测试 - Queue 模式');
+
+    return new Promise((resolve, reject) => {
+        // 构造请求体
+        const messages = [];
+        const lastMessage = { role: 'user', content: [] };
+
+        if (prompt) {
+            lastMessage.content.push({ type: 'text', text: prompt });
+        }
+
+        for (const imgPath of imagePaths) {
+            if (fs.existsSync(imgPath)) {
+                const buffer = fs.readFileSync(imgPath);
+                const base64 = buffer.toString('base64');
+                const ext = path.extname(imgPath).slice(1).toLowerCase();
+                const mimeType = ext === 'jpg' ? 'jpeg' : ext;
+                lastMessage.content.push({
+                    type: 'image_url',
+                    image_url: { url: `data:image/${mimeType};base64,${base64}` }
+                });
+            } else {
+                logger.warn('CLI/Test', `图片不存在，已跳过: ${imgPath}`);
+            }
+        }
+
+        messages.push(lastMessage);
+
+        const body = {
+            messages,
+            ...(modelId && { model: modelId })
+        };
+
+        const bodyStr = JSON.stringify(body);
+
+        const options = {
+            hostname: '127.0.0.1',
+            port: PORT,
+            path: '/v1/queue/join',
+            method: 'POST',
+            headers: {
+                'Content-Type': 'application/json',
+                'Content-Length': Buffer.byteLength(bodyStr),
+                'Authorization': `Bearer ${AUTH_TOKEN}`
+            }
+        };
+
+        const req = http.request(options, (res) => {
+            let buffer = '';
+            res.on('data', chunk => {
+                buffer += chunk.toString();
+                const lines = buffer.split('\n');
+                buffer = lines.pop(); // 保留未完成的行
+
+                for (const line of lines) {
+                    if (!line.trim() || !line.startsWith('data:')) continue;
+
+                    const data = line.slice(5).trim();
+                    if (data === '[DONE]') continue;
+
+                    try {
+                        const event = JSON.parse(data);
+                        if (event.status === 'error') {
+                            reject(new Error(event.msg));
+                        } else if (event.status === 'completed') {
+                            resolve(event);
+                        }
+                    } catch (e) {
+                        // 忽略解析错误
+                    }
+                }
+            });
+
+            res.on('end', () => {
+                // SSE 结束
+            });
+        });
+
+        req.on('error', reject);
+        req.write(bodyStr);
+        req.end();
+    });
+}
+
+/**
+ * 直接调用适配器测试
+ */
+async function testViaDirect(prompt, modelId, imagePaths) {
+    logger.info('CLI/Test', '直接调用适配器测试');
+
+    // 初始化浏览器
+    const context = await initBrowser(config);
+
+    // 解析模型 ID
+    const resolvedModelId = modelId ? resolveModelId(name, modelId) : null;
+
+    // 执行生图
+    const result = await generateImage(context, prompt, imagePaths, resolvedModelId);
+
+    if (result.error) {
+        throw new Error(result.error);
+    }
+
+    return result;
+}
+
+/**
+ * 保存图片
+ */
+function saveImage(base64Data) {
+    const testSaveDir = path.join(TEMP_DIR, 'testSave');
+    if (!fs.existsSync(testSaveDir)) {
+        fs.mkdirSync(testSaveDir, { recursive: true });
+    }
+
+    const timestamp = Date.now();
+    const savePath = path.join(testSaveDir, `test_${timestamp}.png`);
+
+    // 移除 Data URI 前缀（如果有）
+    const cleanBase64 = base64Data.replace(/^data:image\/\w+;base64,/, '');
+    fs.writeFileSync(savePath, Buffer.from(cleanBase64, 'base64'));
+
+    logger.info('CLI/Test', `图片已保存: ${savePath}`);
+    return savePath;
+}
+
+/**
+ * 主流程
+ */
+(async () => {
+    try {
+        // 1. 选择测试模式
+        const testMode = await selectTestMode();
+        logger.info('CLI/Test', `测试模式: ${testMode === 'http' ? 'HTTP 服务器' : '直接调用'}`);
+
+        // 2. 选择模型
+        const modelId = await selectModel();
+        if (modelId) {
+            logger.info('CLI/Test', `选择模型: ${modelId}`);
+        } else {
+            logger.info('CLI/Test', '跳过模型选择，使用默认');
+        }
+
+        // 3. 输入提示词
+        const prompt = await promptForInput();
+        logger.info('CLI/Test', `提示词: ${prompt}`);
+
+        // 4. 输入图片路径
+        const imagePaths = await promptForImages();
+        if (imagePaths.length > 0) {
+            logger.info('CLI/Test', `参考图片: ${imagePaths.join(', ')}`);
+        }
+
+        // 5. 执行测试
+        let result;
+        if (testMode === 'http') {
+            const serverType = config.server.type || 'openai';
+            if (serverType === 'queue') {
+                result = await testViaHttpQueue(prompt, modelId, imagePaths);
+            } else {
+                result = await testViaHttpOpenAI(prompt, modelId, imagePaths);
+            }
+
+            // 处理 HTTP 响应
+            if (result.choices) {
+                // OpenAI 格式
+                const content = result.choices[0].message.content;
+                logger.info('CLI/Test', `响应内容: ${content.slice(0, 100)}...`);
+
+                // 提取图片（如果有）
+                const match = content.match(/!\[.*?\]\((data:image\/[^)]+)\)/);
+                if (match) {
+                    saveImage(match[1]);
+                } else {
+                    logger.info('CLI/Test', `文本回复: ${content}`);
+                }
+            } else if (result.image) {
+                // Queue 格式
+                saveImage(result.image);
+            } else if (result.msg) {
+                logger.info('CLI/Test', `文本回复: ${result.msg}`);
+            }
+
+        } else {
+            // 直接调用
+            result = await testViaDirect(prompt, modelId, imagePaths);
+
+            if (result.image) {
+                saveImage(result.image);
+            } else if (result.text) {
+                logger.info('CLI/Test', `文本回复: ${result.text}`);
+            }
+        }
+
+        logger.info('CLI/Test', '测试完成');
+        process.exit(0);
+
+    } catch (err) {
+        logger.error('CLI/Test', '测试失败', { error: err.message });
+        process.exit(1);
+    }
+})();