377 lines
14 KiB
Markdown
377 lines
14 KiB
Markdown
# OpenCode RAG 插件设计文档
|
||
|
||
> 本文档详细说明了 ocrag 项目的设计目的、技术选型考量、架构构思和实现细节。
|
||
|
||
---
|
||
|
||
## 1. 设计背景与目标
|
||
|
||
### 1.1 问题背景
|
||
|
||
在大型代码库中进行 AI 辅助开发时，AI 需要理解项目中的大量代码才能给出准确的建议和答案。然而：
|
||
- **上下文窗口有限**：无法将整个代码库放入 AI 提示词
|
||
- **实时性需求**：开发者需要 AI 能立即理解刚编写的代码
|
||
- **本地化优先**：代码不应该上传到外部服务器
|
||
|
||
### 1.2 设计目标
|
||
|
||
为 OpenCode 构建一个**本地代码知识库 RAG 系统**，实现：
|
||
|
||
| 功能 | 描述 |
|
||
|------|------|
|
||
| **实时添加** | 将代码文件或目录添加到本地知识库 |
|
||
| **语义搜索** | 通过自然语言查询获取相关代码片段 |
|
||
| **智能管理** | 支持删除和列出知识库中的条目 |
|
||
|
||
### 1.3 设计原则
|
||
|
||
1. **本地化优先**：所有数据和计算都在本地完成，不依赖外部服务
|
||
2. **轻量高效**：避免引入复杂的服务端组件，保持极低的延迟
|
||
3. **零运维**：嵌入式数据库，无需安装配置，即装即用
|
||
4. **AI 友好**：生成可被 AI 直接理解和使用的上下文
|
||
|
||
---
|
||
|
||
## 2. 技术选型详解
|
||
|
||
### 2.1 为什么选择 Python
|
||
|
||
| 考量因素 | Python | Rust |
|
||
|----------|--------|------|
|
||
| 开发效率 | ✅ 高 | ⚠️ 中 |
|
||
| LLM 生成质量 | ✅ 高（AI 更熟悉 Python） | ⚠️ 中 |
|
||
| 生态丰富度 | ✅ 成熟 | ⚠️ 一般 |
|
||
| 运行时性能 | ⚠️ 中（可通过 C 扩展优化） | ✅ 高 |
|
||
| 社区支持 | ✅ 丰富 | ⚠️ 有限 |
|
||
|
||
**结论**：Python 的高开发效率和 AI 生成质量优势明显，即使运行时性能略低，但对于 RAG 这种 I/O 密集型应用影响有限。
|
||
|
||
### 2.2 为什么选择 LanceDB
|
||
|
||
传统向量数据库对比：
|
||
|
||
| 数据库 | 特点 | 缺点 |
|
||
|--------|------|------|
|
||
| Chroma | 简单易用 | 不支持持久化、不适合生产 |
|
||
| Milvus | 功能强大 | 需要 Docker 部署 |
|
||
| Qdrant | Rust 实现，高性能 | 需要单独部署 |
|
||
| **LanceDB** | **嵌入式、零运维、Python 原生** | **相对较新** |
|
||
|
||
**LanceDB 优势**：
|
||
- **嵌入式**：数据库就是一个文件夹，无需单独服务
|
||
- **零运维**：安装即用，自动管理
|
||
- **Python 优先**：原生 Python SDK，类型提示完善
|
||
- **性能优秀**：基于 Apache Arrow，查询速度快
|
||
|
||
### 2.3 为什么选择 tree-sitter
|
||
|
||
使用 `tree-sitter` 实现 AST 级别的语义感知分块：
|
||
|
||
| 方案 | 优势 | 劣势 |
|
||
|------|------|------|
|
||
| tree-sitter | AST 感知、语义完整、不跨类/函数切割 | 依赖语言 parser |
|
||
| langchain-text-splitters | 轻量、规则简单 | chunk_size 优先，会切割语义单元 |
|
||
|
||
**最终选择 tree-sitter**：
|
||
- 代码以**类、函数**等完整语义单元为最小块
|
||
- 超过 token 限制时，递归在子节点级别拆分
|
||
- 支持 40+ 编程语言（通过 tree-sitter-languages）
|
||
- 底层通过 ctypes 加载 parser，绕过语言绑定兼容性问题
|
||
|
||
### 2.4 为什么使用 Qwen3-Embedding-0.6B
|
||
|
||
| 模型 | 维度 | 优势 | 劣势 |
|
||
|------|------|------|------|
|
||
| all-MiniLM-L6-v2 | 384 | 快速、小巧 | 英文为主 |
|
||
| **Qwen3-Embedding-0.6B** | **1024** | **中文优化、中英双语** | 较大、首次加载慢 |
|
||
|
||
**选择理由**：
|
||
- 开源可本地部署，代码安全
|
||
- 中文支持优秀
|
||
- 适合代码+文档混合场景
|
||
|
||
### 2.5 为什么不用 MCP 服务器
|
||
|
||
**MCP 方案的劣势**：
|
||
- 需要额外部署 MCP 服务器
|
||
- 增加系统复杂度
|
||
- 调试困难
|
||
|
||
**CLI 方案的优势**：
|
||
- 零额外组件
|
||
- 通过 `bash` 工具直接调用
|
||
- 易于调试和扩展
|
||
|
||
---
|
||
|
||
## 3. 架构设计
|
||
|
||
### 3.1 整体架构
|
||
|
||
```
|
||
┌─────────────────────────────────────────────────────────────────┐
|
||
│ OpenCode AI │
|
||
└────────────────────────────┬──────────────────────────────────┘
|
||
│ 1. rag_add / rag_search
|
||
▼
|
||
┌─────────────────────────────────────────────────────────────────┐
|
||
│ TypeScript Plugin (ocrag-plugin.ts) │
|
||
│ │
|
||
│ ┌──────────┐ ┌──────────┐ ┌──────────────┐ │
|
||
│ │ rag_add │ │rag_search│ │ Skill 指令 │ │
|
||
│ └────┬─────┘ └────┬─────┘ └──────────────┘ │
|
||
└───────┼─────────────────┼───────────────────────────────────────┘
|
||
│ │
|
||
▼ ▼
|
||
┌─────────────────────────────────────────────────────────────────┐
|
||
│ Python CLI (ocrag) │
|
||
│ │
|
||
│ ┌─────────┐ ┌─────────┐ ┌─────────┐ ┌─────────────┐ │
|
||
│ │ add │ │ search │ │ remove │ │ list │ │
|
||
│ └───┬─────┘ └────┬─────┘ └───┬─────┘ └──────┬──────┘ │
|
||
└──────┼───────────────┼─────────────┼──────────────────┼──────────┘
|
||
│ │ │ │
|
||
▼ ▼ ▼ │
|
||
┌─────────────────────────────────────────┐ │
|
||
│ Processing Pipeline │ │
|
||
│ │ │
|
||
│ ┌────────────┐ ┌────────────────┐ │ │
|
||
│ │ Chunker │───▶│ Embedder │──┼────────────┤
|
||
│ │ (分块) │ │ (向量化) │ │ │
|
||
│ └────────────┘ └────────────────┘ │ │
|
||
└─────────────────────────────────────────┼────────────┘
|
||
│
|
||
▼
|
||
┌─────────────────────────────────────────────────────────────────┐
|
||
│ LanceDB (向量数据库) │
|
||
│ │
|
||
│ ┌──────────────────────────────────────────────────────────┐ │
|
||
│ │ documents 表 │ │
|
||
│ │ ┌──────────┬──────────────────────┬──────────────────┐ │ │
|
||
│ │ │ text │ vector │ metadata │ │ │
|
||
│ │ │ (文本) │ (1024维向量) │ (JSON元数据) │ │ │
|
||
│ │ └──────────┴──────────────────────┴──────────────────┘ │ │
|
||
│ └──────────────────────────────────────────────────────────┘ │
|
||
└─────────────────────────────────────────────────────────────────┘
|
||
```
|
||
|
||
### 3.2 数据流设计
|
||
|
||
#### 添加文件流程
|
||
|
||
```
|
||
用户请求 → Plugin → CLI:add → Chunker → Embedder → LanceDB → 返回结果
|
||
```
|
||
|
||
**详细步骤**：
|
||
1. **文件收集**：递归遍历目录或单个文件
|
||
2. **内容读取**：以 UTF-8 编码读取文件内容
|
||
3. **代码分块**：按语言和语法结构切分代码
|
||
4. **向量化**：使用 Embedder 将文本转为 1024 维向量
|
||
5. **存储入库**：将文本、向量、元数据存入 LanceDB
|
||
|
||
#### 搜索流程
|
||
|
||
```
|
||
用户查询 → Plugin → CLI:search → Embedder → LanceDB → 返回结果
|
||
```
|
||
|
||
**详细步骤**：
|
||
1. **查询向量化**：将自然语言查询转为向量
|
||
2. **相似度搜索**：LanceDB 执行向量相似度检索
|
||
3. **结果排序**：按相似度距离排序
|
||
4. **元数据提取**：解析 JSON 元数据
|
||
5. **结果返回**：格式化的代码片段列表
|
||
|
||
### 3.3 模块职责划分
|
||
|
||
| 模块 | 职责 | 依赖 |
|
||
|------|------|------|
|
||
| `cli.py` | 命令行入口，参数解析 | click |
|
||
| `chunker.py` | 代码分块处理 | tree-sitter, tree-sitter-languages |
|
||
| `embedder.py` | 文本向量化 | sentence-transformers |
|
||
| `db.py` | 向量数据库操作 | lancedb, pyarrow |
|
||
| `commands/*.py` | 各命令实现 | cli, chunker, embedder, db |
|
||
| `utils.py` | 工具函数 | - |
|
||
|
||
---
|
||
|
||
## 4. 核心算法设计
|
||
|
||
### 4.1 代码分块算法
|
||
|
||
**实现**：使用 tree-sitter 进行 AST 级别的语义分块，支持 40+ 编程语言。
|
||
|
||
**分块策略**：
|
||
- **源码文件**：使用 tree-sitter 解析为 AST，按 `class_definition`、`function_definition` 等语义节点提取完整块
|
||
- **非语义单元文件**（如 Markdown）：使用 tiktoken 按行分块
|
||
- **超大语义单元**（>4000 tokens）：递归在子节点级别拆分，保证每个块不超过大小限制
|
||
- **Parser 加载**：通过 ctypes 从 `languages.so` 直接加载，绕过语言绑定兼容性问题
|
||
|
||
### 4.2 向量化策略
|
||
|
||
**模型选择**：Qwen3-Embedding-0.6B
|
||
- 输出维度：1024
|
||
- 向量归一化：L2 归一化，便于余弦相似度计算
|
||
|
||
**批量处理**：
|
||
- 一次性对多个文本块进行编码
|
||
- 利用 GPU/CPU 的批处理能力提升吞吐量
|
||
|
||
### 4.3 搜索策略
|
||
|
||
**相似度度量**：LanceDB 默认使用余弦相似度
|
||
|
||
**top_k 参数**：控制返回结果数量，默认 5 条
|
||
|
||
---
|
||
|
||
## 5. 安全与性能考量
|
||
|
||
### 5.1 安全性设计
|
||
|
||
| 风险点 | 防护措施 |
|
||
|--------|----------|
|
||
| SQL 注入 | 使用 Pandas 过滤而非 SQL 字符串拼接 |
|
||
| 路径遍历 | 仅处理指定路径，不执行动态导入 |
|
||
| 数据泄露 | 所有数据本地存储，不涉及网络传输 |
|
||
|
||
### 5.2 性能优化
|
||
|
||
**已实现**：
|
||
- ✅ 批量 embedding 减少 I/O 开销
|
||
- ✅ 单例模式避免重复加载模型
|
||
- ✅ 向量归一化加速相似度计算
|
||
|
||
**可优化项**：
|
||
- 文件哈希缓存避免重复添加
|
||
- GPU 加速 embedding 计算
|
||
- 增量索引更新而非全量重建
|
||
|
||
### 5.3 性能基准
|
||
|
||
| 指标 | 实测结果 | 说明 |
|
||
|------|----------|------|
|
||
| 搜索延迟 | **63-70 ms** | 包含 embedding + 向量检索 |
|
||
| 数据库写入 | 2-3 ms/块 | LanceDB 性能优秀 |
|
||
| 分块速度 | <1 ms | 纯规则，无模型加载 |
|
||
| Embedding | ~2.5秒/块 | Qwen3 模型较大 |
|
||
|
||
---
|
||
|
||
## 6. 扩展性设计
|
||
|
||
### 6.1 多语言支持
|
||
|
||
只需在 `LANGUAGE_MAP` 中添加新扩展名即可：
|
||
|
||
```python
|
||
LANGUAGE_MAP = {
|
||
".py": "python",
|
||
".js": "javascript",
|
||
# 添加新语言...
|
||
".kt": "kotlin",
|
||
".swift": "swift",
|
||
}
|
||
```
|
||
|
||
### 6.2 自定义 Embedding 模型
|
||
|
||
修改 `embedder.py` 中的模型路径：
|
||
|
||
```python
|
||
model_path = "path/to/your/model"
|
||
```
|
||
|
||
### 6.3 增量同步（Watch 模式）
|
||
|
||
使用 `watchdog` 库监听文件变化，自动更新知识库：
|
||
|
||
```python
|
||
observer = Observer()
|
||
observer.schedule(RAGSyncHandler(), path, recursive=True)
|
||
observer.start()
|
||
```
|
||
|
||
---
|
||
|
||
## 7. 与 OpenCode 的集成
|
||
|
||
### 7.1 插件架构
|
||
|
||
```
|
||
OpenCode
|
||
├── TypeScript Plugin
|
||
│ ├── rag_add 工具
|
||
│ └── rag_search 工具
|
||
│
|
||
└── Skill 指令
|
||
└── SKILL.md
|
||
```
|
||
|
||
### 7.2 工具定义
|
||
|
||
| 工具名 | 参数 | 功能 |
|
||
|--------|------|------|
|
||
| `rag_add` | `paths: string[]`, `recursive?: boolean` | 添加文件到知识库 |
|
||
| `rag_search` | `query: string`, `top_k?: number` | 搜索知识库 |
|
||
|
||
### 7.3 错误处理
|
||
|
||
所有工具调用都包裹在 try-catch 中：
|
||
|
||
```typescript
|
||
execute: async (args) => {
|
||
try {
|
||
const result = await Bun.$`${cmd}`.text();
|
||
return result;
|
||
} catch (error) {
|
||
return `Error: ${error.message}`;
|
||
}
|
||
}
|
||
```
|
||
|
||
---
|
||
|
||
## 8. 设计总结
|
||
|
||
### 8.1 核心创新
|
||
|
||
1. **零 MCP 架构**：通过 CLI 直接集成，简化系统复杂度
|
||
2. **本地化优先**：数据不出本地，保证代码安全
|
||
3. **轻量高效**：嵌入式数据库，秒级启动
|
||
4. **AI 原生**：输出格式专为 AI 消费设计
|
||
|
||
### 8.2 适用场景
|
||
|
||
| 场景 | 适用性 | 说明 |
|
||
|------|--------|------|
|
||
| 个人项目知识管理 | ✅ 非常适合 | 本地存储，隐私安全 |
|
||
| 小团队代码库问答 | ✅ 适合 | 轻量易部署 |
|
||
| 大型企业代码库 | ⚠️ 需优化 | 可能需要分布式扩展 |
|
||
| 跨语言代码库 | ✅ 支持 | 多语言分块支持 |
|
||
|
||
### 8.3 未来展望
|
||
|
||
1. **语义分块升级**：集成更智能的分块算法
|
||
2. **多模态支持**：支持图片、图表等非文本内容
|
||
3. **增量索引**：支持大型代码库的实时更新
|
||
4. **分布式部署**：支持多节点协同检索
|
||
|
||
---
|
||
|
||
## 附录：术语表
|
||
|
||
| 术语 | 解释 |
|
||
|------|------|
|
||
| RAG | Retrieval-Augmented Generation，检索增强生成 |
|
||
| Embedding | 将文本转为向量的过程 |
|
||
| 向量数据库 | 专门存储和检索向量的数据库 |
|
||
| Chunk | 分块后的文本片段 |
|
||
| LanceDB | 嵌入式向量数据库 |
|
||
|
||
---
|
||
|
||
**文档版本**：1.0
|
||
**最后更新**：2026年4月15日
|