如何开发一个 IDEA 插件通过 Ollama 调用大模型为方法生成仙侠风格的注释

本文将介绍如何创建一个 IntelliJ IDEA 插件，该插件可以通过 Ollama 调用大型语言模型为 Java 方法生成具有仙侠风格的注释。

项目结构与配置

首先我们需要建立一个基本的 IntelliJ Platform Plugin 项目，并在 plugin.xml中定义插件的基本信息：

xml 复制代码

<idea-plugin>
    <id>cn.wubo.xianxiacomment</id>
    <name>Xianxiacomment</name>
    <vendor url="https://gitee.com/wb04307201/XianxiaComment">wubo</vendor>
    <description><![CDATA[
        通过 Ollama 调用大模型为方法生成仙侠风格注释并自动插入<br>
        <em>右键点击方法，选择生成修仙注释即可使用</em>
    ]]></description>
    
    <depends>com.intellij.modules.platform</depends>
    <depends>com.intellij.modules.java</depends>
</idea-plugin>

实现核心功能

配置服务

为了使用户能够自定义 Ollama 地址和使用的模型名称，我们实现了 OllamaConfigService和 OllamaConfigurable类。这些类允许用户在设置中输入他们的 Ollama URL 和希望使用的大模型名称。

动作处理器

接下来是实现主要逻辑的部分 ------ 创建一个动作处理器 GenerateXianXiaCommentAction来响应用户的操作请求。这个类继承自 AnAction 并重写了几个关键的方法：

actionPerformed: 处理实际的动作事件。
update: 控制菜单项何时可见以及是否启用。
getActionUpdateThread: 定义更新 UI 的线程。

当用户右键单击某个 Java 方法并选择"生成修仙注释"时，会触发 actionPerformed方法。此方法检查当前选中的元素是否为 PsiMethod 类型，如果是，则开始执行后台任务来获取注释。

kotlin 复制代码

if (element is PsiMethod) {
    val methodText = element.text
    val settings = OllamaConfigService.instance
    
    object : Task.Backgroundable(project, "正在生成仙侠注释...", true) {
        override fun run(indicator: ProgressIndicator) {
            indicator.text = "正在调用 Ollama 服务..."
            val comment = callOllama(settings.ollamaUrl, settings.modelName, methodText)
            
            WriteCommandAction.runWriteCommandAction(project) {
                insertComment(element, comment)
            }
        }
        
        override fun onThrowable(error: Throwable) {
            // Handle error
        }
    }.queue()
}

调用 Ollama API

核心部分在于 callOllama 函数，它负责构建请求体并向 Ollama 发送 POST 请求以获得生成的注释内容。这里使用了 OkHttp 库来进行网络通信。

kotlin 复制代码

private fun callOllama(url: String, model: String, methodText: String): String {
    val prompt = "请为以下方法生成仙侠风格的注释，只返回注释内容，不要包含任何其他内容：\n\n$methodText"
    
    val json = JSONObject()
    json.put("model", model)
    json.put("prompt", prompt)
    json.put("stream", false)
    
    val body = json.toString().toRequestBody("application/json; charset=utf-8".toMediaType())
    val request = Request.Builder()
        .url("$url/api/generate")
        .post(body)
        .build()
        
    client.newBuilder()
        .connectTimeout(600, java.util.concurrent.TimeUnit.SECONDS)
        .build()
        .newCall(request).execute().use { response ->
        val responseBody = response.body?.string()
        if (response.isSuccessful && responseBody != null) {
            val jsonResponse = JSONObject(responseBody)
            return jsonResponse.getString("response")
        }
        return "// 未能生成注释"
    }
}

最后，在收到返回的结果后，利用 insertComment将生成的注释放置到对应的方法上方。

kotlin 复制代码

private fun insertComment(method: PsiMethod, comment: String) {
    method.addCommentBefore(comment)
}

以上就是整个插件的核心流程。通过这种方式，我们可以轻松地扩展更多有趣的功能，例如支持不同的编程语言或者提供更多样式的注释模板等。

示例

我们用一段两数之和算法进行测试

java 复制代码

    public int[] twoSum(int[] nums, int target) {
        if (nums == null || nums.length < 2) {
            return new int[0];
        }

        Map<Integer, Integer> map = new HashMap<>();
        for (int i = 0; i < nums.length; i++) {
            int temp = target - nums[i];
            if (map.containsKey(temp)) {
                return new int[]{map.get(temp), i};
            }
            map.put(nums[i], i);
        }
        return new int[0];
    }

$示例代码$ （https://gitee.com/wb04307201/XianxiaComment）