掘金 人工智能 10小时前
手把手构建TinyCodeRAG
index_new5.html
../../../zaker_core/zaker_tpl_static/wap/tpl_guoji1.html

 

本文详细介绍了如何构建一个专为代码优化的轻量级RAG(检索增强生成)系统——TinyCodeRAG。通过代码智能分块、模块化测试、多轮对话优化等核心升级,实现了代码敏感型知识库的构建。文章涵盖了向量化、文本/代码分割、向量数据库、大模型封装以及TinyCodeRAG系统整合等关键环节,并提供了开源代码和测试API Key,方便读者实践和理解RAG系统的工作原理。

🔢 向量化引擎:TinyCodeRAG使用OpenAI的text-embedding-3-small模型进行向量化,将文本或代码片段转化为1536维向量,并通过测试验证了其文本长度适应性和语义相似度判别的有效性。

✂️ 文本分割模块:项目包含基于文件的文本分块模块(chunker_text.py),支持解析PDF、Markdown和TXT等多种格式。同时,为适配代码库,新增了专用代码分割器(chunker_code.py),该分割器优化了代码处理逻辑,保留函数完整性、智能过滤文件并支持目录级处理。

🗄️ 向量数据库模块:TinyCodeRAG基于TinyRAG项目实现轻量级向量存储模块,通过load_vector()重用预存向量数据,降低token消耗。query()函数封装了从文本编码到相似度匹配的全流程,方便业务整合。

🧠 大模型封装:项目基于OpenAI API 封装了一套简洁的LLM 调用模块,并默认集成测试密钥。考虑到成本,目前限定使用 Doubao-1.5-lite-32k 模型。文章强调了模型规模的选择应根据需求复杂度进行,小模型在某些情况下可能更稳定。

🤖 TinyCodeRAG系统整合:在tiny_code_rag.py文件中,封装了TinyCodeRAG的调用逻辑,实现了RAG系统的核心功能,支持多轮对话,并提供了代码语料和文本语料的准备、向量化、大模型调用等步骤。

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

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

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

你可能会好奇:"已有TinyRAG珠玉在前,为何再造轮子?" 原因有二:

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

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

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

话不多说,让我们一起动起来!

0. 项目文件速览 🗂

先快速过一遍项目结构(完整代码已开源):
(github.com/codemilesto…)

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 文件中,我们定义了一个专门的类,主要负责两个关键功能:

    get_embeddings:将文本(或代码片段)转化为对应的向量表示。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)。与通用文本分块不同,该模块专门优化了代码处理逻辑:

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

参数说明

通过test_chunker.pytiny-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三连!你的反馈是我持续优化的动力~

Fish AI Reader

Fish AI Reader

AI辅助创作,多种专业模板,深度分析,高质量内容生成。从观点提取到深度思考,FishAI为您提供全方位的创作支持。新版本引入自定义参数,让您的创作更加个性化和精准。

FishAI

FishAI

鱼阅,AI 时代的下一个智能信息助手,助你摆脱信息焦虑

联系邮箱 441953276@qq.com

相关标签

RAG 代码知识库 TinyCodeRAG 向量化 LLM
相关文章