基于 Ollama 与 TypeScript 实现本地私有化代码助手

概述

基于大模型的能力可以理解文本或代码的上下文，然后根据代码上下文实现代码注释、代码补全等功能，从而构建一个代码助手。

本文实现了一个简单的 VS Code 代码助手插件，通过该插件可以对代码进行补全。后台大模型服务是在本地 CPU 机器上搭建的，通过 Ollama 实现通过接口来访问大模型。

插件功能设计

插件实现的功能很简单：在 VS Code 中选择一段代码，然后发起一个命令（后续可优化为快捷键），根据代码上下文对代码进行补全。

设计考虑

设计分为两个层面：前端和后端服务。这个插件是一个简单的可用 Demo，包含两个端：一个前端，一个基于大模型的后端服务。 流程：前端 -> Ollama -> 大模型

环境准备

1. 安装 Ollama：访问官网下载并安装 Ollama，启动服务。
2. 拉取模型：使用命令行拉取模型，例如 ollama pull llama3.2。
3. 开发环境：确保已安装 Node.js 和 TypeScript 环境。
4. VS Code 扩展开发：安装 Yeoman Generator for VS Code 等工具。

前端开发

通过 TypeScript 开发 VS Code 插件。目前前端实现的功能非常简单：选择一段代码，然后发起 ask Ollama 的命令，此时就会把代码片段放到提示词模板中发送给后台大模型服务。并等待后台大模型的返回，然后使用返回的代码覆盖现有代码。

前端要实现的步骤和逻辑如下：

1. 获取当前选择的代码片段；
2. 根据获取到的代码片段，构建提示词请求；
3. 发送补全代码的提示词请求；
4. 获取大模型的返回代码；
5. 覆盖现有代码。

根据目前实现的功能，还可以对功能进行继续扩展，比如实现：可以撤销返回的代码，获取光标所在的代码片段补全光标所在代码，基于代码实现智能聊天功能，等等。

代码实现

插件入口文件

import * as vscode from 'vscode';
import axios from 'axios';

export function activate(context: vscode.ExtensionContext) {
   let disposable = vscode.commands.registerCommand('extension.askOllama', async () => {
       const editor = vscode.window.activeTextEditor;
       if (!editor) {
           vscode.window.showErrorMessage('请打开一个编辑器并选中代码');
           return;
      }

       const selection = editor.selection;
       const selectedText = editor.document.getText(selection);

       if (!selectedText) {
           vscode.window.showWarningMessage('未选中任何代码');
           return;
      }

       try {
           const result = await getCompletionFromOllama(selectedText);
           console.log("debug: complete result: ", result);
           // 在 VS Code 中显示补全结果
           vscode.window.showInformationMessage('代码补全结果已生成');
           editor.edit(editBuilder => {
               editBuilder.replace(selection, result);
          });
       } catch (error: any) {
           vscode.window.showErrorMessage(`请求 Ollama 出错：${error.message}`);
       }
   });

   context.subscriptions.push(disposable);
}

export function deactivate() {}

// 调用 Ollama API 获取补全
async function getCompletionFromOllama(code: string): Promise<string> {
   const url = 'http://localhost:11434/api/chat'; // 替换为你的 Ollama 服务地址
   const payload = {
       model: 'llama3.2', // 替换为实际模型名称
       messages: [
          {
               role: 'user',
               content: `补全以下代码：\n ${code}` // 将用户输入的代码放入 content 字段
          }
      ],
       stream: false
  };

   try {
       // 发送 POST 请求到 Ollama API
       const response = await axios.post(url, payload);

       // 获取返回的 content 字段
       const completion = response.data.message.content;
         
       // 如果返回的 content 为空，抛出错误
       if (!completion) {
           throw new Error('Ollama API 返回为空的补全结果');
      }

       return completion;
   } catch (error) {
       console.error('Error calling Ollama API:', error);
       throw new Error('请求 Ollama API 出错');
   }
}

说明：后台的 URL 可以换成自己的 Ollama 本地访问的 URL 地址。

后端服务

后端使用 Ollama 提供的框架来实现通过接口请求大模型，该接口在 Ollama 的官方网站上有说明。

API 示例

curl http://localhost:11434/api/chat -d '{
 "model":"llama3.2",
 "messages": [
  {
     "role":"user",
     "content":"why is the sky blue?"
  }
 ],
 "stream":false
}'

返回值示例

{
 "model":"llama3.2",
 "created_at":"2023-12-12T14:13:43.416799Z",
 "message": {
   "role":"assistant",
   "content":"Hello! How are you today?"
 },
 "done":true,
 "total_duration":5191566416,
 "load_duration":2154458,
 "prompt_eval_count":26,
 "prompt_eval_duration":383809000,
 "eval_count":298,
 "eval_duration":4799921000
}

效果演示

（1）编写一个函数，选中代码后触发命令。 （2）按下 Ctrl+Shift+P，输入 Ask Ollama for Code Complete，得到补全结果。

后续改进

该插件只是一个 Demo，通过本地 Ollama 来搭建大模型服务，可以实现私有化部署的代码助手。 但这里要注意的是：这里是完全基于大模型的能力，而在实际的使用过程中，还需要结合外部知识库和本地代码库才能更好的完成代码补全，也就是要借助 RAG 框架的能力。通用大模型的能力是有限的，通过本地代码库来构建外部知识库，这也是代码助手的难点。

参考资料

• https://github.com/ollama/ollama/blob/main/docs/api.md