const axios = require("axios"); // 支持的模型类型 const MODEL_TYPES = { OPENAI: "openai", OLLAMA: "ollama", }; // 预设提示词 const PRESET_PROMPTS = { // 翻译 translate: `请将以下内容翻译成地道的中文,要求: 1. 保持原文的专业性和准确性 2. 符合中文的表达习惯 3. 对于专业术语保留英文原文,并在括号中给出中文翻译 4. 保持原文的段落格式 原文:`, // 生成SHELL命令 shell: `请根据以下描述生成一个 shell 命令,要求: 1. 命令应当简洁高效 2. 优先使用常见的命令行工具 3. 确保命令的安全性和可靠性 4. 对于复杂操作,添加注释说明 5. 如果需要多个命令,使用 && 连接或使用脚本格式 6. 直接输出命令,不要输出任何解释,不要使用markdown格式 需求描述:`, // 总结 summarize: `请总结以下内容的要点,要求: 1. 提取最重要和最有价值的信息 2. 使用简洁的语言 3. 按重要性排序 4. 保持逻辑性和连贯性 5. 如果有专业术语,保留并解释 原文:`, }; /** * AI对话功能 * @param {Object} apiConfig - API配置参数 * @param {string} apiConfig.modelType - 模型类型(openai/ollama) * @param {string} apiConfig.apiUrl - API地址 * @param {string} apiConfig.apiToken - API令牌 * @param {string} apiConfig.model - 模型名称 * @param {Object} content - 对话内容参数 * @param {string} content.prompt - 用户输入的提示词 * @param {string} content.presetPrompt - 预设提示词类型 * @returns {Promise} 对话响应 */ async function chat(content, apiConfig) { try { const { modelType, apiUrl, apiToken, model } = apiConfig; const { prompt, presetPrompt } = content; // 验证必要参数 if (!apiUrl || !prompt || !model) { throw new Error("API地址、模型名称和提示词不能为空"); } // 构建完整提示词 const fullPrompt = presetPrompt ? `${PRESET_PROMPTS[presetPrompt]}\n${prompt}` : prompt; // 准备请求配置 const config = { headers: { "Content-Type": "application/json", }, }; let requestData; let url = apiUrl; // 根据不同的模型类型构建请求数据 if (modelType === MODEL_TYPES.OPENAI) { // OpenAI API config.headers["Authorization"] = `Bearer ${apiToken}`; requestData = { model: model, messages: [ { role: "user", content: fullPrompt, }, ], }; } else if (modelType === MODEL_TYPES.OLLAMA) { // Ollama API // 如果用户没有指定完整的 API 路径,添加 /api/generate if (!url.endsWith("/api/generate")) { url = url.replace(/\/?$/, "/api/generate"); } requestData = { model: model, prompt: fullPrompt, stream: false, }; } else { throw new Error("不支持的模型类型"); } const loadingBar = await quickcommand.showLoadingBar({ text: "AI思考中...", onClose: () => { // 取消请求 if (source) { source.cancel("操作已取消"); } }, }); // 创建取消令牌 const CancelToken = axios.CancelToken; const source = CancelToken.source(); // 发送请求 const response = await axios.post(url, requestData, { ...config, cancelToken: source.token, }); loadingBar.close(); // 解析不同模型的响应 let result; if (modelType === MODEL_TYPES.OPENAI) { // OpenAI 响应格式 if (!response.data.choices || !response.data.choices[0]) { throw new Error("OpenAI 响应格式错误"); } result = response.data.choices[0].message.content; } else { // Ollama 响应格式 if (!response.data.response) { throw new Error("Ollama 响应格式错误"); } result = response.data.response; } return { success: true, result, }; } catch (error) { // 如果是用户取消的请求,返回特定的错误信息 if (axios.isCancel(error)) { return { success: false, error: "请求已取消", cancelled: true, }; } return { success: false, error: error.response?.data?.error?.message || error.message, }; } } /** * 获取API支持的模型列表 * @param {Object} apiConfig - API配置参数 * @param {string} apiConfig.modelType - 模型类型(openai/ollama) * @param {string} apiConfig.apiUrl - API地址 * @param {string} apiConfig.apiToken - API令牌 * @returns {Promise} 模型列表响应 */ async function getModels(apiConfig) { try { const { modelType, apiUrl, apiToken } = apiConfig; // 验证必要参数 if (!apiUrl) { throw new Error("API地址不能为空"); } // 准备请求配置 const config = { headers: { "Content-Type": "application/json", }, }; let url = apiUrl; // 根据不同的模型类型构建请求 if (modelType === MODEL_TYPES.OPENAI) { // OpenAI API config.headers["Authorization"] = `Bearer ${apiToken}`; // OpenAI的模型列表接口是 /v1/models if (!url.endsWith("/models")) { url = "https://api.openai.com/v1/models"; } } else if (modelType === MODEL_TYPES.OLLAMA) { // Ollama API // Ollama的模型列表接口是 /api/tags if (!url.endsWith("/api/tags")) { url = url.replace(/\/?$/, "/api/tags"); } } else { throw new Error("不支持的模型类型"); } // 发送请求 const response = await axios.get(url, config); // 解析不同模型的响应 let models; if (modelType === MODEL_TYPES.OPENAI) { // OpenAI 响应格式 if (!response.data.data) { throw new Error("OpenAI 响应格式错误"); } models = response.data.data.map((model) => model.id); } else { // Ollama 响应格式 if (!response.data.models) { throw new Error("Ollama 响应格式错误"); } models = response.data.models.map((model) => model.name); } return { success: true, result: models, }; } catch (error) { return { success: false, error: error.response?.data?.error?.message || error.message, }; } } module.exports = { chat, getModels };