用 Ollama 搭一个本地 AI 代码解释器：复杂代码秒变大白话

pip install ollama tree-sitter tree-sitter-python graphviz
brew install graphviz  # macOS
# 或 sudo apt install graphviz  # Ubuntu

ollama pull qwen2.5-coder:7b
ollama serve

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import sys
import json
import ollama
import yaml
from pathlib import Path
from parser import parse_code_structure
from visualizer import generate_flowchart
from prompts import (
    EXPLAIN_PROMPT,
    COMMENT_PROMPT,
    FLOW_PROMPT
)


def load_config(config_path: str = "config.yaml") -> dict:
    """加载配置文件"""
    with open(config_path, "r", encoding="utf-8") as f:
        return yaml.safe_load(f)


def explain_code(code: str, model: str = "qwen2.5-coder:7b", config: dict = None) -> dict:
    """调用 Ollama 解释代码"""
    cfg = config or load_config()
    
    # 1. 大白话解释
    explain_resp = ollama.chat(
        model=model,
        messages=[{
            "role": "user",
            "content": EXPLAIN_PROMPT.format(code=code, language=cfg.get("language", "python"))
        }]
    )
    plain_explain = explain_resp["message"]["content"].strip()
    
    # 2. 逐行注释
    comment_resp = ollama.chat(
        model=model,
        messages=[{
            "role": "user",
            "content": COMMENT_PROMPT.format(code=code, language=cfg.get("language", "python"))
        }]
    )
    commented_code = comment_resp["message"]["content"].strip()
    
    # 3. 流程图描述（给 visualizer 用）
    flow_resp = ollama.chat(
        model=model,
        messages=[{
            "role": "user",
            "content": FLOW_PROMPT.format(code=code)
        }]
    )
    flow_desc = flow_resp["message"]["content"].strip()
    
    return {
        "plain_explain": plain_explain,
        "commented_code": commented_code,
        "flow_desc": flow_desc
    }


def process_file(file_path: str, output_dir: str = "output"):
    """处理单个代码文件"""
    code = Path(file_path).read_text(encoding="utf-8")
    model = load_config().get("model", "qwen2.5-coder:7b")
    
    print(f"📄 正在分析: {file_path}")
    print(f"🤖 使用模型: {model}")
    print("-" * 50)
    
    result = explain_code(code, model=model)
    
    # 生成输出目录
    out = Path(output_dir)
    out.mkdir(exist_ok=True)
    
    base_name = Path(file_path).stem
    
    # 保存大白话解释
    explain_file = out / f"{base_name}_explain.md"
    explain_file.write_text(f"# {base_name} 代码解释\n\n{result['plain_explain']}\n", encoding="utf-8")
    print(f"✅ 解释已保存: {explain_file}")
    
    # 保存带注释的代码
    comment_file = out / f"{base_name}_commented.py"
    comment_file.write_text(result["commented_code"], encoding="utf-8")
    print(f"✅ 注释代码已保存: {comment_file}")
    
    # 生成流程图
    flow_file = out / f"{base_name}_flowchart.dot"
    flow_file.write_text(result["flow_desc"], encoding="utf-8")
    print(f"✅ 流程图描述已保存: {flow_file}")
    
    # 尝试渲染流程图
    try:
        import graphviz as gv
        dot = gv.Source(result["flow_desc"])
        dot.render(out / f"{base_name}_flowchart", format="png", cleanup=True)
        print(f"✅ 流程图已渲染: {out / f'{base_name}_flowchart.png'}")
    except Exception as e:
        print(f"⚠️ 流程图渲染失败: {e}")
    
    # 打印到终端
    print("\n" + "=" * 50)
    print("📖 大白话解释:")
    print("=" * 50)
    print(result["plain_explain"])
    print("\n" + "=" * 50)
    print("💬 终端输出完成，结果已保存到 output/ 目录")


def main():
    if len(sys.argv) < 2:
        print("用法: python explainer.py <代码文件路径> [输出目录]")
        print("示例: python explainer.py samples/example.py ./output")
        sys.exit(1)
    
    code_file = sys.argv[1]
    output_dir = sys.argv[2] if len(sys.argv) > 2 else "output"
    
    if not Path(code_file).exists():
        print(f"❌ 文件不存在: {code_file}")
        sys.exit(1)
    
    process_file(code_file, output_dir)


if __name__ == "__main__":
    main()

# 大白话解释提示词
EXPLAIN_PROMPT = """你是一位经验丰富的程序员，擅长把复杂代码翻译成普通人能听懂的话。

请用中文解释以下 {language} 代码的功能和逻辑：

```{language}
{code}
```

要求：
1. 用一段话概括这段代码的核心功能（不超过100字）
2. 列出代码的输入、输出和关键步骤
3. 指出代码中可能存在的性能问题或潜在 bug
4. 给出改进建议（不超过3条）

只输出解释内容，不要重复代码。"""

# 逐行注释提示词
COMMENT_PROMPT = """你是代码注释专家。请为以下 {language} 代码添加详细的中文注释。

```{language}
{code}
```

要求：
1. 文件头注释说明文件用途
2. 每个函数/方法添加 docstring
3. 关键逻辑行添加行内注释
4. 复杂算法添加步骤说明

直接输出带注释的完整代码，不要添加额外解释。"""

# 流程图提示词
FLOW_PROMPT = """请分析以下代码的逻辑流程，输出 DOT 格式的有向图描述。

```python
{code}
```

要求：
1. 用节点表示关键步骤（函数调用、条件判断、循环等）
2. 用箭头表示执行顺序
3. 条件分支标注 true/false
4. 使用标准的 DOT 语法

只输出 DOT 代码，格式如下：
digraph {
  rankdir=TB;
  node [shape=box];
  ...
}"""

# 先准备一段示例代码
python explainer.py samples/example.py ./output

def quick_sort(arr):
    if len(arr) <= 1:
        return arr
    pivot = arr[len(arr) // 2]
    left = [x for x in arr if x < pivot]
    middle = [x for x in arr if x == pivot]
    right = [x for x in arr if x > pivot]
    return quick_sort(left) + middle + quick_sort(right)