## 背景介绍
在日常开发中，为代码添加注释是一件费时费力的工作。很多开发者，要么是因为赶工期，要么是觉得写注释太枯燥，最终导致代码缺乏必要的文档。当团队协作或几个月后回顾代码时，理解逻辑变成了一项艰巨的任务。
有没有一种方法，可以利用大语言模型自动为代码生成注释？本文将介绍如何使用 Go 语言结合 LLM API，构建一个实用的代码注释生成工具。
## 问题描述
传统的代码注释方式存在几个痛点：一是手动编写注释耗时，尤其是在大型项目中；二是注释质量参差不齐，有些甚至是过时的；三是不同开发者的注释风格不统一，影响代码可读性。
利用 LLM 的理解能力，我们可以让 AI 分析代码逻辑，自动生成准确、易读的注释。这比开发者手动编写的注释更快，而且通常更加规范。
## 解决方案
我们将使用 Go 语言开发一个命令行工具，通过调用 LLM API 来分析 Go 代码并生成注释。整体架构如下：
“`
用户输入 → Go 程序 → LLM API → 解析响应 → 输出带注释的代码
“`
## 准备工作
首先需要安装 Go 环境（1.18 以上版本）。然后获取一个 LLM API Key，市面上有多种选择，如 OpenAI API、Claude API 或国产的硅基流动 API 等。
接下来创建项目目录并初始化：
“`bash
mkdir -p code-comment-generator/{cmd,internal}
cd code-comment-generator
go mod init code-comment-generator
“`
## 核心代码实现
### 1. 配置管理
首先定义配置结构，用于存储 API 相关配置：
“`go
// internal/config.go
package config
import (
“os”
)
type Config struct {
APIKey string
BaseURL string
Model string
MaxTokens int
}
func Load() *Config {
return &Config{
APIKey: os.Getenv(“LLM_API_KEY”),
BaseURL: getenv(“LLM_BASE_URL”, “https://api.siliconflow.cn/v1”),
Model: getenv(“LLM_MODEL”, “Qwen/Qwen2.5-Coder-32B-Instruct”),
MaxTokens: 4096,
}
}
func getenv(key, defaultValue string) string {
if value := os.Getenv(key); value != “” {
return value
}
return defaultValue
}
“`
### 2. LLM 客户端
接下来实现 LLM API 调用逻辑：
“`go
// internal/client.go
package internal
import (
“bytes”
“encoding/json”
“fmt”
“net/http”
“time”
“code-comment-generator/internal/config”
)
type ChatClient struct {
cfg *config.Config
client *http.Client
}
type Message struct {
Role string `json:”role”`
Content string `json:”content”`
}
type Request struct {
Model string `json:”model”`
Messages []Message `json:”messages”`
MaxTokens int `json:”max_tokens,omitempty”`
Temperature float64 `json:”temperature,omitempty”`
}
type Response struct {
Choices []Choice `json:”choices”`
}
type Choice struct {
Message Message `json:”message”`
}
func NewChatClient(cfg *config.Config) *ChatClient {
return &ChatClient{
cfg: cfg,
client: &http.Client{Timeout: 120 * time.Second},
}
}
func (c *ChatClient) Chat(system, user string) (string, error) {
url := fmt.Sprintf(“%s/chat/completions”, c.cfg.BaseURL)
reqBody := Request{
Model: c.cfg.Model,
MaxTokens: c.cfg.MaxTokens,
Messages: []Message{
{Role: “system”, Content: system},
{Role: “user”, Content: user},
},
}
body, err := json.Marshal(reqBody)
if err != nil {
return “”, err
}
req, err := http.NewRequest(“POST”, url, bytes.NewReader(body))
if err != nil {
return “”, err
}
req.Header.Set(“Content-Type”, “application/json”)
req.Header.Set(“Authorization”, “Bearer “+c.cfg.APIKey)
resp, err := c.client.Do(req)
if err != nil {
return “”, err
}
defer resp.Body.Close()
var result Response
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
return “”, err
}
if len(result.Choices) == 0 {
return “”, fmt.Errorf(“empty response”)
}
return result.Choices[0].Message.Content, nil
}
// 生成代码注释的提示词
func BuildCommentPrompt(code string) (system, user string) {
system = `你是一个专业的代码文档助手。你的任务是为输入的代码生成准确、简洁的注释。
要求：
1. 为每个重要的函数、方法、结构体生成注释
2. 注释应该解释”做什么”而不是”怎么做”
3. 使用英文注释，遵循 Go 文档惯例
4. 注释格式应该符合 go/doc 规范
5. 只输出带注释的代码，不要输出其他内容`
user = fmt.Sprintf(“请为以下 Go 代码生成注释：
%s
只返回带注释的代码，不要解释。”, code)
return system, user
}
“`
### 3. 代码解析与处理
现在实现代码解析逻辑，将输入代码按函数或方法分割：
“`go
// internal/parser.go
package internal
import (
“go/ast”
“go/parser”
“go/token”
“strings”
)
type CodeParser struct {
fset *token.FileSet
}
func NewCodeParser() *CodeParser {
return &CodeParser{
fset: token.NewFileSet(),
}
}
// 解析代码，获取函数和方法的列表
func (p *CodeParser) Parse(code string) ([]FunctionInfo, error) {
file, err := parser.ParseFile(p.fset, “”, code, parser.ParseComments)
if err != nil {
return nil, err
}
var functions []FunctionInfo
ast.Inspect(file, func(n ast.Node) bool {
switch node := n.(type) {
case *ast.FuncDecl:
if node.Name != nil {
info := FunctionInfo{
Name: node.Name.Name,
IsMethod: node.Recv != nil,
}
functions = append(functions, info)
}
case *ast.TypeSpec:
if node.Name != nil {
info := FunctionInfo{
Name: node.Name.Name,
IsType: true,
}
functions = append(functions, info)
}
}
return true
})
return functions, nil
}
type FunctionInfo struct {
Name string
IsMethod bool
IsType bool
}
// 按函数级别分割代码，便于逐个生成注释
func (p *CodeParser) SplitByFunction(code string) ([]string, error) {
lines := strings.Split(code, ”
“)
var blocks []string
var current []string
indentLevel := 0
inFunction := false
for _, line := range lines {
trimmed := strings.TrimSpace(line)
// 检测函数开始
if strings.HasPrefix(trimmed, “func “) {
if len(current) > 0 && inFunction {
blocks = append(blocks, strings.Join(current, ”
“))
}
current = []string{line}
inFunction = true
indentLevel = countIndent(line)
continue
}
// 检测函数结束
if inFunction && trimmed != “” {
currentIndent := countIndent(line)
if currentIndent <= indentLevel && !strings.HasPrefix(trimmed, "//") {
blocks = append(blocks, strings.Join(current, "
"))
current = []string{line}
indentLevel = currentIndent
if indentLevel == 0 {
inFunction = false
}
continue
}
}
current = append(current, line)
}
if len(current) > 0 {
blocks = append(blocks, strings.Join(current, ”
“))
}
if len(blocks) == 0 {
return []string{code}, nil
}
return blocks, nil
}
func countIndent(s string) int {
count := 0
for _, c := range s {
if c == ‘ ‘ {
count++
} else if c == ‘ ‘ {
count++
} else {
break
}
}
return count
}
“`
### 4. 主程序入口
最后是主程序逻辑：
“`go
// cmd/main.go
package main
import (
“bufio”
“flag”
“fmt”
“os”
“strings”
“code-comment-generator/internal”
“code-comment-generator/internal/config”
)
func main() {
var inputFile string
var outputFile string
var apiKey string
flag.StringVar(&inputFile, “i”, “”, “Input Go file”)
flag.StringVar(&outputFile, “o”, “”, “Output file”)
flag.StringVar(&apiKey, “key”, “”, “LLM API Key”)
flag.Parse()
// 从环境变量或命令行获取 API Key
if apiKey == “” {
apiKey = os.Getenv(“LLM_API_KEY”)
}
if apiKey == “” {
fmt.Fprintln(os.Stderr, “Error: API key required, use -key or LLM_API_KEY”)
os.Exit(1)
}
// 加载配置
cfg := config.Load()
cfg.APIKey = apiKey
// 读取输入
var code string
if inputFile != “” {
data, err := os.ReadFile(inputFile)
if err != nil {
fmt.Fprintf(os.Stderr, “Error reading file: %v
“, err)
os.Exit(1)
}
code = string(data)
} else {
// 从 stdin 读取
var lines []string
scanner := bufio.NewScanner(os.Stdin)
for scanner.Scan() {
lines = append(lines, scanner.Text())
}
code = strings.Join(lines, ”
“)
}
if code == “” {
fmt.Fprintln(os.Stderr, “Error: no code provided”)
os.Exit(1)
}
// 创建客户端和解析器
client := internal.NewChatClient(cfg)
parser := internal.NewCodeParser()
// 分割代码块
blocks, err := parser.SplitByFunction(code)
if err != nil {
fmt.Fprintf(os.Stderr, “Error parsing code: %v
“, err)
os.Exit(1)
}
fmt.Printf(“Processing %d code blocks…
“, len(blocks))
var results []string
for i, block := range blocks {
fmt.Printf(“Processing block %d/%d…
“, i+1, len(blocks))
systemPrompt, userPrompt := internal.BuildCommentPrompt(block)
commented, err := client.Chat(systemPrompt, userPrompt)
if err != nil {
fmt.Fprintf(os.Stderr, “Error processing block %d: %v
“, i+1, err)
results = append(results, block)
continue
}
results = append(results, commented)
}
output := strings.Join(results, ”
“)
// 输出结果
if outputFile != “” {
err := os.WriteFile(outputFile, []byte(output), 0644)
if err != nil {
fmt.Fprintf(os.Stderr, “Error writing output: %v
“, err)
os.Exit(1)
}
fmt.Printf(“Output written to %s
“, outputFile)
} else {
fmt.Println(output)
}
}
“`
## 运行结果
假设我们有一个待注释的代码文件 input.go：
“`go
package main
func add(a, b int) int {
return a + b
}
func multiply(a, b int) int {
return a * b
}
type Calculator struct {
result int
}
func (c *Calculator) Add(a, b int) int {
c.result = a + b
return c.result
}
“`
运行命令：
“`bash
LLM_API_KEY=your-key go run cmd/main.go -i input.go -o output.go
“`
生成的带注释代码：
“`go
package main
// add returns the sum of two integers.
func add(a, b int) int {
return a + b
}
// multiply returns the product of two integers.
func multiply(a, b int) int {
return a * b
}
// Calculator performs basic arithmetic operations.
type Calculator struct {
result int
}
// Add adds two numbers and stores the result.
func (c *Calculator) Add(a, b int) int {
c.result = a + b
return c.result
}
“`
## 总结
本文介绍了如何使用 Go 语言结合 LLM API 构建代码注释生成工具。主要完成了以下工作：
1. 实现了配置管理模块，支持通过环境变量或命令行参数配置 API
2. 实现了 LLM 客户端，支持多种 LLM API 服务
3. 实现了代码解析器，能够按函数级别分割代码
4. 实现了主程序，支持从文件或标准输入读取代码，输出带注释的代码
这个工具可以显著提高代码注释的效率和质量。当然，目前的实现还有一些改进空间，比如支持更多编程语言、优化注释风格、支持批量处理等。有兴趣的读者可以在此基础上继续扩展。