代码编辑器添加AI助手，支持AI代码编写

2025-10-10 07:23:23 +08:00 · 2025-02-20 01:04:11 +08:00
parent ac19d06845
commit d82c735cfe
14 changed files with 1033 additions and 325 deletions
--- a/plugin/lib/ai.js
+++ b/plugin/lib/ai.js
@@ -1,155 +1,303 @@
-const axios = require("axios");
-
 // 支持的模型类型
-const MODEL_TYPES = {
+const API_TYPES = {
  OPENAI: "openai",
  OLLAMA: "ollama",
 };

-// 预设提示词
-const PRESET_PROMPTS = {
+// 角色提示词
+const ROLE_PROMPTS = {
  // 翻译
-  translate: `请将以下内容翻译成地道的中文，要求：
-1. 保持原文的专业性和准确性
-2. 符合中文的表达习惯
-3. 对于专业术语保留英文原文，并在括号中给出中文翻译
-4. 保持原文的段落格式
-
-原文：`,
+  translate: `你是一名翻译专家，请将我给你的内容进行翻译，要求：
+1. 无论给的内容长短，请直接翻译，不要进行任何解释
+2. 提供中文时，翻译成地道的英文，符合英文的表达习惯
+3. 提供英文时，翻译成地道的中文，符合中文的表达习惯
+4. 保持原文的专业性和准确性
+5. 对于专业术语保留原文，并在括号中给出对应的中文翻译
+6. 保持原文的段落格式
+`,

  // 生成SHELL命令
-  shell: `请根据以下描述生成一个 shell 命令，要求：
+  shell: `你是一名shell命令专家，请根据我的描述生成 shell 命令，要求：
 1. 命令应当简洁高效
 2. 优先使用常见的命令行工具
 3. 确保命令的安全性和可靠性
 4. 对于复杂操作，添加注释说明
 5. 如果需要多个命令，使用 && 连接或使用脚本格式
 6. 直接输出命令，不要输出任何解释，不要使用markdown格式
-
-需求描述：`,
+`,

  // 总结
-  summarize: `请总结以下内容的要点，要求：
+  summarize: `你是一名总结专家，请总结我给你的内容的要点，要求：
 1. 提取最重要和最有价值的信息
 2. 使用简洁的语言
 3. 按重要性排序
 4. 保持逻辑性和连贯性
 5. 如果有专业术语，保留并解释
-
-原文：`,
+`,
 };

+// API URL 处理
+const API_ENDPOINTS = {
+  [API_TYPES.OPENAI]: {
+    chat: "/v1/chat/completions",
+    models: "/v1/models",
+  },
+  [API_TYPES.OLLAMA]: {
+    chat: "/api/chat",
+    models: "/api/tags",
+  },
+};
+
+// 构建API URL
+function buildApiUrl(baseUrl, endpoint) {
+  if (!baseUrl.endsWith(endpoint)) {
+    return baseUrl.replace(/\/?$/, endpoint);
+  }
+  return baseUrl;
+}
+
+// 构建请求配置
+function buildRequestConfig(apiConfig) {
+  const config = {
+    headers: {
+      "Content-Type": "application/json",
+    },
+  };
+
+  if (apiConfig.apiType === API_TYPES.OPENAI && apiConfig.apiToken) {
+    config.headers["Authorization"] = `Bearer ${apiConfig.apiToken}`;
+  }
+
+  return config;
+}
+
+// 构建请求数据
+function buildRequestData(content, apiConfig, stream = false) {
+  const { model } = apiConfig;
+  const { prompt, role, context = [] } = content;
+  const rolePrompt = ROLE_PROMPTS[role] || role;
+
+  const roleMessage = rolePrompt
+    ? [
+        {
+          role: "user",
+          content: rolePrompt,
+        },
+      ]
+    : [];
+
+  // 统一的消息格式处理
+  const messages = [
+    // 添加系统角色消息（如果有）
+    ...roleMessage,
+    // 添加上下文消息
+    ...context.map((msg) => ({
+      role: msg.role || "user",
+      content: msg.content,
+    })),
+    // 添加当前用户消息
+    {
+      role: "user",
+      content: prompt,
+    },
+  ];
+
+  return {
+    model,
+    messages,
+    stream,
+  };
+}
+
+// 处理普通响应
+function parseResponse(response, apiType) {
+  if (apiType === API_TYPES.OPENAI) {
+    if (!response.data.choices || !response.data.choices[0]) {
+      throw new Error("OpenAI 响应格式错误");
+    }
+    return response.data.choices[0].message.content;
+  } else {
+    if (!response.data.message) {
+      throw new Error("Ollama 响应格式错误");
+    }
+    return response.data.message.content;
+  }
+}
+
+// 处理模型列表响应
+function parseModelsResponse(response, apiType) {
+  if (apiType === API_TYPES.OPENAI) {
+    if (!response.data.data) {
+      throw new Error("OpenAI 响应格式错误");
+    }
+    return response.data.data.map((model) => model.id);
+  } else {
+    if (!response.data.models) {
+      throw new Error("Ollama 响应格式错误");
+    }
+    return response.data.models.map((model) => model.name);
+  }
+}
+
+// 处理 OpenAI 流式响应
+async function handleOpenAIStreamResponse(line, controller, onStream) {
+  if (line.startsWith("data: ")) {
+    const jsonStr = line.replace(/^data: /, "");
+    if (jsonStr === "[DONE]") {
+      onStream("", controller, true);
+      return;
+    }
+    const json = JSON.parse(jsonStr);
+    const content = json.choices[0]?.delta?.content;
+    if (content) {
+      onStream(content, controller, false);
+    }
+  }
+}
+
+// 处理 Ollama 流式响应
+async function handleOllamaStreamResponse(line, controller, onStream) {
+  const json = JSON.parse(line);
+  if (json.done) {
+    onStream("", controller, true);
+    return;
+  }
+  if (json.message?.content) {
+    onStream(json.message.content, controller, false);
+  }
+}
+
+// 处理流式响应
+async function handleStreamResponse(response, apiConfig, controller, onStream) {
+  const reader = response.body.getReader();
+  const decoder = new TextDecoder();
+  let buffer = "";
+
+  try {
+    while (true) {
+      const { value, done } = await reader.read();
+      if (done) break;
+
+      buffer += decoder.decode(value, { stream: true });
+      const lines = buffer.split("\n");
+      buffer = lines.pop() || "";
+
+      for (const line of lines) {
+        if (line.trim()) {
+          try {
+            if (apiConfig.apiType === API_TYPES.OPENAI) {
+              await handleOpenAIStreamResponse(line, controller, onStream);
+            } else {
+              await handleOllamaStreamResponse(line, controller, onStream);
+            }
+          } catch (e) {
+            console.error("解析响应失败:", e);
+          }
+        }
+      }
+    }
+
+    // 处理剩余的缓冲区
+    if (buffer.trim()) {
+      try {
+        if (apiConfig.apiType === API_TYPES.OPENAI) {
+          await handleOpenAIStreamResponse(buffer, controller, onStream);
+        } else {
+          await handleOllamaStreamResponse(buffer, controller, onStream);
+        }
+      } catch (e) {
+        console.error("解析剩余响应失败:", e);
+      }
+    }
+  } catch (error) {
+    if (error.name === "AbortError") {
+      return {
+        success: false,
+        error: "请求已取消",
+        cancelled: true,
+      };
+    }
+    throw error;
+  } finally {
+    reader.releaseLock();
+  }
+
+  return { success: true, result: "流式请求完成" };
+}
+
 /**
 * 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 - 预设提示词类型
+ * @param {Object} apiConfig - API配置参数
+ * @param {Object} options - 其他选项
 * @returns {Promise<Object>} 对话响应
 */
-async function chat(content, apiConfig) {
+async function chat(content, apiConfig, options = {}) {
  try {
-    const { modelType, apiUrl, apiToken, model } = apiConfig;
-    const { prompt, presetPrompt } = content;
+    const { showLoadingBar = true, stream = false, onStream } = options;

    // 验证必要参数
-    if (!apiUrl || !prompt || !model) {
+    if (!apiConfig.apiUrl || !content.prompt || !apiConfig.model) {
      throw new Error("API地址、模型名称和提示词不能为空");
    }

-    // 构建完整提示词
-    const fullPrompt = presetPrompt
-      ? `${PRESET_PROMPTS[presetPrompt]}\n${prompt}`
-      : prompt;
+    if (stream && !onStream) {
+      throw new Error("使用流式请求时必须提供onStream回调函数");
+    }

-    // 准备请求配置
-    const config = {
-      headers: {
-        "Content-Type": "application/json",
-      },
-    };
+    // 构建请求URL和配置
+    const url = buildApiUrl(
+      apiConfig.apiUrl,
+      API_ENDPOINTS[apiConfig.apiType].chat
+    );
+    const config = buildRequestConfig(apiConfig, stream);
+    const requestData = buildRequestData(content, apiConfig, stream);

-    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,
+    // 显示加载条
+    const loadingBar = showLoadingBar
+      ? await quickcommand.showLoadingBar({
+          text: "AI思考中...",
+          onClose: () => {
+            if (controller) {
+              controller.abort();
+            }
          },
-        ],
-      };
-    } else if (modelType === MODEL_TYPES.OLLAMA) {
-      // Ollama API
-      // 如果用户没有指定完整的 API 路径，添加 /api/generate
-      if (!url.endsWith("/api/generate")) {
-        url = url.replace(/\/?$/, "/api/generate");
-      }
+        })
+      : null;

-      requestData = {
-        model: model,
-        prompt: fullPrompt,
-        stream: false,
-      };
-    } else {
-      throw new Error("不支持的模型类型");
+    // 统一使用 fetch 处理请求
+    const controller = new AbortController();
+    const response = await fetch(url, {
+      method: "POST",
+      headers: config.headers,
+      body: JSON.stringify(requestData),
+      signal: controller.signal,
+    });
+
+    if (!response.ok) {
+      throw new Error(`HTTP error! status: ${response.status}`);
    }

-    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;
+    if (stream) {
+      result = await handleStreamResponse(
+        response,
+        apiConfig,
+        controller,
+        onStream
+      );
    } else {
-      // Ollama 响应格式
-      if (!response.data.response) {
-        throw new Error("Ollama 响应格式错误");
-      }
-      result = response.data.response;
+      const responseData = await response.json();
+      result = {
+        success: true,
+        result: parseResponse({ data: responseData }, apiConfig.apiType),
+      };
    }

-    return {
-      success: true,
-      result,
-    };
+    loadingBar?.close();
+    return result;
  } catch (error) {
-    // 如果是用户取消的请求，返回特定的错误信息
-    if (axios.isCancel(error)) {
+    if (error.name === "AbortError") {
      return {
        success: false,
        error: "请求已取消",
@@ -166,69 +314,33 @@ async function chat(content, apiConfig) {
 /**
 * 获取API支持的模型列表
 * @param {Object} apiConfig - API配置参数
- * @param {string} apiConfig.modelType - 模型类型(openai/ollama)
- * @param {string} apiConfig.apiUrl - API地址
- * @param {string} apiConfig.apiToken - API令牌
 * @returns {Promise<Object>} 模型列表响应
 */
 async function getModels(apiConfig) {
  try {
-    const { modelType, apiUrl, apiToken } = apiConfig;
-
-    // 验证必要参数
-    if (!apiUrl) {
+    if (!apiConfig.apiUrl) {
      throw new Error("API地址不能为空");
    }

-    // 准备请求配置
-    const config = {
-      headers: {
-        "Content-Type": "application/json",
-      },
-    };
+    const url = buildApiUrl(
+      apiConfig.apiUrl,
+      API_ENDPOINTS[apiConfig.apiType].models
+    );
+    const config = buildRequestConfig(apiConfig);

-    let url = apiUrl;
+    const response = await fetch(url, {
+      method: "GET",
+      headers: config.headers,
+    });

-    // 根据不同的模型类型构建请求
-    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);
+    if (!response.ok) {
+      throw new Error(`HTTP error! status: ${response.status}`);
    }

+    const responseData = await response.json();
    return {
      success: true,
-      result: models,
+      result: parseModelsResponse({ data: responseData }, apiConfig.apiType),
    };
  } catch (error) {
    return {
--- a/plugin/lib/quickcommand.js
+++ b/plugin/lib/quickcommand.js
@@ -188,8 +188,8 @@ const quickcommand = {
    return null;
  },

-  askAI: async function (content, apiConfig) {
-    return await chat(content, apiConfig);
+  askAI: async function (content, apiConfig, options) {
+    return await chat(content, apiConfig, options);
  },

  ...systemDialog,