RAG 符号感知实现

IfAI 的检索增强生成 (RAG) 系统超越了简单的文本搜索。它通过符号感知解析理解代码结构，使 AI 能够理解函数、类、trait 和其他代码构造之间的关系。

概述

传统的 RAG 系统将代码视为纯文本，错过了使代码有意义的结构关系。IfAI 的符号感知 RAG：

1. 使用 Tree-sitter 将代码解析为 AST
2. 提取语义符号（函数、类、trait）
3. 构建跨文件的关系图
4. 使用嵌入进行索引以实现语义搜索
5. 为 AI 查询检索上下文感知的结果

架构

┌─────────────────────────────────────────────────────────┐
│                    用户查询                             │
└─────────────────────────────────────────────────────────┘
                         │
                         ▼
┌─────────────────────────────────────────────────────────┐
│              查询理解                                   │
│  • 提取提到的符号                                       │
│  • 识别文件上下文                                       │
│  • 解析意图（解释、修复、重构）                         │
└─────────────────────────────────────────────────────────┘
                         │
                         ▼
┌─────────────────────────────────────────────────────────┐
│              符号检索                                   │
│  • 查找符号定义                                         │
│  • 获取实现详情                                         │
│  • 获取相关符号（导入、实现）                           │
└─────────────────────────────────────────────────────────┘
                         │
                         ▼
┌─────────────────────────────────────────────────────────┐
│              上下文构建                                 │
│  • 收集文件内容                                         │
│  • 包含符号 AST                                         │
│  • 添加使用示例                                         │
└─────────────────────────────────────────────────────────┘
                         │
                         ▼
┌─────────────────────────────────────────────────────────┐
│              AI 响应生成                                │
│  • 符号感知的代码生成                                   │
│  • 保持关系的重构                                       │
│  • 结构有效的建议                                       │
└─────────────────────────────────────────────────────────┘

符号引擎

Tree-sitter 集成

IfAI 使用 Tree-sitter 进行解析：

// src-tauri/src/symbol_engine.rs
use tree_sitter::{Parser, Language};

pub struct SymbolEngine {
    parser: Parser,
    symbols: HashMap<String, Symbol>,
}

impl SymbolEngine {
    pub fn parse_file(&mut self, content: &str, lang: &str) -> Vec<Symbol> {
        let language = self.get_language(lang);
        self.parser.set_language(language).unwrap();

        let tree = self.parser.parse(content, None).unwrap();
        self.extract_symbols(&tree)
    }
}

支持的语言

语言	解析器	状态
Rust	tree-sitter-rust	✅ 完整
TypeScript	tree-sitter-typescript	✅ 完整
Python	tree-sitter-python	✅ 完整
Go	tree-sitter-go	✅ 完整
Java	tree-sitter-java	✅ 完整
C/C++	tree-sitter-cpp	✅ 完整
JavaScript	tree-sitter-javascript	✅ 完整
35+ 更多	各种	🚧 部分

符号类型

索引内容

1. 函数：函数定义、方法
2. 类：类定义、结构体
3. Trait：trait 定义、接口
4. 变量：全局变量、常量
5. 类型：类型别名、枚举
6. 模块：模块声明、命名空间

符号元数据

每个符号存储：

pub struct Symbol {
    pub name: String,
    pub kind: SymbolKind,
    pub file_path: PathBuf,
    pub range: Range,
    pub documentation: Option<String>,
    pub signatures: Vec<Signature>,
    pub relationships: Vec<Relationship>,
}

pub enum Relationship {
    Imports(SymbolId),
    Implements(SymbolId),
    Calls(SymbolId),
    References(SymbolId),
}

检索过程

步骤 1：查询分析

// 用户问："auth 中间件是如何工作的？"

// 解析为：
{
  symbols: ["auth_middleware", "authenticate"],
  intent: "explain",
  context: "中间件实现"
}

步骤 2：符号查找

pub fn find_symbol(&self, name: &str) -> Vec<Symbol> {
    let mut results = Vec::new();

    // 精确匹配
    if let Some(symbol) = self.symbols.get(name) {
        results.push(symbol.clone());
    }

    // 模糊匹配
    results.extend(self.fuzzy_search(name));

    // 语义搜索（嵌入）
    results.extend(self.semantic_search(name));

    results
}

步骤 3：上下文扩展

// 对于每个找到的符号，收集：
const context = {
  symbol: symbol,
  definition: getDefinition(symbol),
  implementations: getImplementations(symbol),
  usages: getUsages(symbol),
  relatedSymbols: getRelated(symbol),
  examples: getExamples(symbol)
}

步骤 4：响应生成

将丰富的上下文发送到 AI 模型：

您正在使用符号感知的上下文分析代码。

符号：auth_middleware
类型：函数
文件：src/middleware/auth.ts:45-78

定义：
async function auth_middleware(req: Request, res: Response, next: Next) {
  // 实现...
}

相关符号：
- authenticate() (被调用)
- verify_token() (调用)
- AuthError (引用)

用户问题：auth 中间件是如何工作的？

请考虑上面显示的符号关系提供解释。

嵌入与向量搜索

FastEmbed 集成

use fastembed::{EmbeddingModel, InitOptions};

let model = EmbeddingModel::try_new(InitOptions {
    model_name: "BAAI/bge-small-en-v1.5",
    ..Default::default()
})?;

// 为符号生成嵌入
let embeddings = model.embed(symbol_docs, None)?;

语义搜索

pub fn semantic_search(&self, query: &str, limit: usize) -> Vec<Symbol> {
    let query_embedding = self.embeddings.embed(query, None)?;

    let similarities: Vec<_> = self.symbol_embeddings
        .iter()
        .map(|(id, emb)| {
            let similarity = cosine_similarity(&query_embedding, emb);
            (*id, similarity)
        })
        .filter(|(_, sim)| *sim > 0.7)
        .sorted_by(|a, b| b.1.partial_cmp(&a.1).unwrap())
        .take(limit)
        .collect();

    similarities.into_iter()
        .map(|(id, _)| self.symbols.get(&id).unwrap().clone())
        .collect()
}

性能优化

增量索引

pub fn index_file(&mut self, path: &Path, content: &str) {
    // 仅在文件更改时重新索引
    if self.file_hash(path) != self.last_hash(path) {
        let symbols = self.parse_file(content, self.detect_lang(path));
        self.update_index(path, symbols);
    }
}

缓存策略

• 符号缓存：热符号的内存 LRU 缓存
• 关系缓存：预计算的符号关系
• 嵌入缓存：缓存的向量嵌入

基准测试

操作	时间	备注
解析 1K LOC 文件	~50ms	Tree-sitter
提取符号	~10ms	后处理
生成嵌入	~200ms	FastEmbed (CPU)
语义搜索 (10K 符号)	~15ms	余弦相似度
总计（冷）	~275ms	首次查询
总计（热）	~25ms	缓存

使用示例

示例 1：理解代码

用户："解释 AppState::new 如何初始化数据库"

IfAI 检索：

1. AppState 结构体定义
2. new() 方法实现
3. 数据库连接设置代码
4. 相关迁移文件

响应：使用实际代码上下文解释初始化过程。

示例 2：重构

用户："将验证逻辑提取到单独的函数"

IfAI 检索：

1. 带有验证的当前函数
2. 相关的验证函数
3. 正在验证的类型
4. 代码库中的使用模式

响应：使用正确类型生成新函数并更新调用点。

示例 3：调试

用户："为什么 authenticate() 对有效令牌失败？"

IfAI 检索：

1. authenticate() 实现
2. 令牌解析函数
3. 错误定义
4. 测试用例

响应：识别令牌过期检查逻辑中的问题。

限制与未来工作

当前限制

• 动态语言：JavaScript、Python（类型信息有限）
• 生成的代码：宏、元编程
• 跨语言：多语言项目

计划改进

• [ ] 动态语言的类型推断
• [ ] 宏展开支持
• [ ] 跨语言符号桥接
• [ ] 实时索引（文件监视器）

参考

• Tree-sitter 文档
• FastEmbed 仓库
• 原始 RAG 实现

贡献

要改进符号引擎：

1. 添加对新语言的支持
2. 改进符号关系检测
3. 优化嵌入生成
4. 增强语义搜索

详情参见贡献指南。