手把手构建TinyCodeRAG：轻量级代码知识库解决方案

在上一篇文章中，我们拆解了RAG系统的核心组件。今天，我们来点更酷的——亲自构建一个专为代码优化的TinyCodeRAG！

💡 快速科普：RAG（Retrieval-Augmented Generation，检索增强生成）技术通过结合外部知识库和AI生成能力，有效缓解大模型的"幻觉"问题。

你可能会好奇："已有TinyRAG珠玉在前，为何再造轮子？" 原因有二：

首先，造轮子是最扎实的学习路径。

其次，在造轮子的过程中，可以想办法把它造得更好看一点，这是一个创造的过程，也令人心旷神怡。

于是，TinyCodeRAG诞生了！它带来四大核心升级：
✅ 代码智能分块：专门解析代码数据结构，构建精准向量集
✅ 开箱即用：提供测试API key（用完我会定期续费）
✅ 模块化测试：每个组件都有独立测试用例
✅ 对话体验优化：完整支持多轮上下文对话

话不多说，让我们一起动起来！

0. 项目文件速览 🗂

先快速过一遍项目结构（完整代码已开源）：
(https://github.com/codemilestones/TinyCodeRAG)

TinyCodeRAG
├── RAG
│   ├── embeddings.py       # 向量化功能封装
│   ├── chunker_text.py     # 通用文本分割器
│   ├── chunker_code.py     # 专用代码分割器 👈 关键创新点！
│   ├── vector_base.py      # 轻量向量数据库
│   ├── llm.py              # LLM接口封装
│   ├── test_*.py           # 各模块测试脚本（共4个）
│   ├── tiny_code_rag.py    # RAG系统整合入口

下面分别按照，向量化、文本和代码的拆分、向量库的实现、LLM的封装、TinyCodeRAG的封装应用进行介绍。每个模块均有对应的测试文件，方便大家进行理解。

1.向量化引擎 🔢

我们首先来实现RAG系统的核心基础：向量化处理。在 embeddings.py 文件中，我们定义了一个专门的类，主要负责两个关键功能：

  1. get_embeddings：将文本（或代码片段）转化为对应的向量表示。
  2. cosine_similarity：计算两个向量之间的余弦相似度得分。

向量生成引擎：OpenAI 我们选择了OpenAI的 text-embedding-3-small 模型来驱动向量生成。这个模型非常灵活，能够将任意长度的输入文本（哪怕超级长），都高效地转换成一个固定长度的1536维向量。这为后续的相似度计算奠定了坚实基础。

验证：不仅仅是转换，更要看效果 为了确保这套向量化逻辑真正奏效，我们设计了测试用例，重点验证两个能力：

    # 对照文本
    test_text_1 = "Hello, world! This is a test."

    # 与test_text_1 高度相关
    test_text_2 = "Hello, world! This is a embedding test."

    # 与test_text_1 主题不同
    test_text_3 = "I want to study how to use the embedding model."     

    # 超长文本测试
    test_text_long = "... repeat long text ..." * 100

测试结果符合预期：

这验证了我们的向量化和相似度计算在识别语义关联性方面是可靠和有效的。

2. 文本分割模块 ✂️

在构建RAG（检索增强生成）系统时，分块（Chunking）是确保系统性能的关键环节，其核心在于将文本分割为长度适中且语义完整的片段。

为满足多格式文本处理需求，本项目实现了基于文件的文本分块模块（chunker_text.py）。该模块继承自TinyRAG项目，支持解析并分块PDF、Markdown及TXT三种常见格式的文档内容。调用示例如下：

# 读取指定路径下的文档，分块最大长度600，相邻块重叠150个字符
docs = ReadFiles('~/workspace/tiny-universe').get_content(max_token_len=600, cover_content=150)

进一步地，为适配代码库的特殊结构，我们新增了专用代码分割器（chunker_code.py）。与通用文本分块不同，该模块专门优化了代码处理逻辑：

  1. 函数完整性保留：优先将同一函数代码保持在同一块中
  2. 智能文件过滤：根据后缀名自动识别并处理源代码文件
  3. 目录级处理：支持直接输入文件夹路径进行批处理
# 拆分指定目录下的源代码，保留150字符的块间重叠
code_docs = split_to_segment("~/workspace/tiny-universe", cover_content=150)

参数说明：

通过test_chunker.py对tiny-universe项目进行分块测试。可以查看分块后的文本内容。

3. 向量数据库模块 🗄️

对于向量数据库，现在有非常多的选择，比如milvus、Pinecone、Weaviate等。甚至ES也支持向量检索。

虽然当前向量数据库选择丰富，但完整部署方案存在较高运维成本，且不利于理解RAG系统的核心运作逻辑。为此我们基于TinyRAG项目实现了轻量级向量存储模块。


class VectorStore:
    # 初始化文本存储容器
    def __init__(self, document: List[str] = ['']) -> None:  

    # 调用嵌入模型将文本转化为向量（注意消耗token）
    def get_vector(self, EmbeddingModel: BaseEmbeddings) -> List[List[float]]:  

    # 向量数据持久化存储
    def persist(self, path: str = 'storage'):  

    # 从存储路径加载预处理向量
    def load_vector(self, path: str = 'storage') -> bool:  

    # 计算两个向量的相似度
    def get_similarity(self, vector1: List[float], vector2: List[float]) -> float:  

    # 语义检索：输入查询文本，返回最相关的k个结果
    def query(self, query: str, EmbeddingModel: BaseEmbeddings, k: int = 1) -> List[str]:

通过load_vector()重用预存向量数据避免重复计算，显著降低token消耗。query()函数则封装了从文本编码到相似度匹配的全流程。

如test_vector_base.py的示例所示，只需三步即可完成业务整合：初始化文档容器→加载/生成向量→执行语义查询：

    vector_store = VectorStore(document=doc_contents)

    # 优先加载已有向量数据，不存在时实时生成
    if not vector_store.load_vector():
        vector_store.get_vector(OpenAIEmbedding())
        vector_store.persist()

    # 执行语义检索并打印结果
    for doc in vector_store.query("RAG 的组成部分是那些?", OpenAIEmbedding(), 3):
        print(doc)
        print("-"*100)

4. 大模型的封装 🧠

我们基于 OpenAI API 封装了一套简洁的 LLM 调用模块（位于 llm.py 中）。为了让大家轻松上手测试，项目里默认集成了一个免费的测试密钥。

考虑到成本限制（token消耗），目前限定使用 Doubao-1.5-lite-32k 模型。同时想聊聊 RAG 系统的一个关键点：模型的规模不是越大越好。

实际上，模型参数的选择更应该取决于你的具体需求复杂度。如果你的逻辑设计本身并不复杂，使用大模型反而可能引入更多不确定性——小模型有时反而更稳。

5. TinyCodeRAG系统整合 🤖

在tiny_code_rag.py文件中，我们封装了TinyCodeRAG的调用逻辑，直接实现了RAG系统的核心功能。来看一下具体怎么玩转它吧：


if __name__ == "__main__":
    ## 1. 准备代码语料和文本语料
    code_docs = split_to_segmenmt("~/workspace/tiny-universe", cover_content=50)
    text_docs = ReadFiles('~/workspace/tiny-universe').get_content(max_token_len=600, cover_content=150)

    # 将代码语料和文本语料合并
    doc_contents = [doc.content for doc in code_docs] + text_docs

    vector_store = VectorStore(document=doc_contents)

    # 向量化并做持久化
    if not vector_store.load_vector():
        vector_store.get_vector(OpenAIEmbedding())
        vector_store.persist()

    # 2. 获取大模型
    model = DoubaoLiteModel()

    # 3. 写入用户输入，待大模型输出后，将其加入历史数据，并不断循环。
    history = []
    while True:
        user_input = input("请输入问题: ")
        contents = vector_store.query(user_input, OpenAIEmbedding(), 3)
        response = model.chat(user_input, history, "\n".join(contents))
        print("\n")
        history.append({'role': 'user', 'content': user_input})
        history.append({'role': 'assistant', 'content': response})

当你准备好RAG的语料数据库后，该文件可以直接运行，并进行多轮对话。每轮对话都会包含历史数据。

6. 总结

本项目主要实现了RAG系统的核心逻辑，包括向量化、文本和代码的拆分、向量库的实现、LLM的封装、TinyCodeRAG的封装应用。

通过本项目，我们完整实现了：
🔧 代码敏感型知识库构建
🔁 检索-生成闭环系统
💬 可扩展的多轮对话框架

立即体验：👉

TinyCodeRAG

✨ 欢迎Star/Fork/Issue三连！你的反馈是我持续优化的动力~