Rust开发一个最简单的RAG
由于之前本机电脑运行LM studio的效果比Ollama好很多,就来试试使用LM Studio提供的OpenAI兼容API来实现简单Agent功能
现在用的比较多的库是Python的LangChain,但是为了让我学过的rust不会生疏,还是得多用起来
Rust中对AI相关的支持库还是挺多的,比如Rig,今天想从最简单的方式去尝试开发,不用Rig库,这样也知道其中的细节流程
RAG运行步骤
- 参考数据准备,包括数据清洗,分割
- 对分割好的Chuck数据片段向量编码(嵌入)
- 把数据片段和它的向量值存入向量数据库,供以后增强检索
- 用户查询文本向量化后,在向量数据库中检索出k个和这个向量最近邻的相关数据
- 将查询到的相关数据重排后和用户的查询数据一起作为上下文提供给大模型
- 大模型根据额外的上下文知识,进行推理给出最终结果到用户
下面就按上面的基本步骤来实现最简单的RAG
cargo.toml需要添加以下依赖
1 | [dependencies] |
文本分割
src/ingest.rs 中对数据清洗,长文本分割为文本片段,并去调用嵌入模型获取嵌入向量。我这里只是最简单的按长度进行文本分割。
1 | use anyhow::Result; |
数据嵌入向量化
src/embedding.rs 中使用reqwest库直接访问LM Studio提供的API接口,将输入文本通过文本嵌入模型获得对应的嵌入向量的值,这个值就是f32类型的一维数组。
1 | use anyhow::Result; |
量数据库存储和检索
向量数据库有很多,AI推荐的是Qdrant,但是这个需要Docker环境在windows使用有点麻烦,我选择了LanceDB,这是个使用rust实现的开源向量数据库。它支持本地数据文件存储,不需要运行任何服务,和SQLite有点像。虽然这个库是rust实现的内核,但是对rust支持挺一般的。我主要参考了官方的指南的这个代码 https://github.com/lancedb/docs/blob/main/tests/rs/quickstart.rs
src/vectordb.rs 这个是目前整个工程中最长的代码了,虽然也就100多行,主要是我让AI帮我生成代码,始终编译有问题,走了弯路,最后还是参考官方代码正常实现了。
Lancedb需要使用arrow_array的数据结构来往LanceDB中存储数据,因此需要实现records_to_reader()方法来把文本和对应的向量数据转换成arrow_array的RecordBatch。schema是用来告诉数据库这个表的结构是什么样的。具体这个库的使用有很多细节,包括建立索引,查询选择不同的算法,在官方指南有详细介绍算法的实现,这里我只是用了最简单的方法。
1 | use anyhow::{anyhow, Result, Context}; |
实现RAG流程
src/rag.rs 中按RAG的流程逐步调用
1 | use anyhow::Result; |
调用LLM获取返回结果
src/llm.rs负责接收提示词,使用配置的大语言模型进行推理,并获取最终的结果返回。这里主要是调整提示词,用来在不同的使用场景获取更好的效果。
1 | use anyhow::Result; |
Agent应用
RAG只是基于大模型的一种应用,我们可以根据不同的目的开发不同的Agent满足需求。增加了一个Agent层用来管理多个不同的Agent。src/agent.rs目前只有一个rag的功能的agent,它把用户的输入传给rag模块,获取返回的结果。
1 | use anyhow::Result; |
应用程序总入口
src/main.rs 从终端获取用户输入,并将输入给Agent,并将Agent返回结果显示在终端。这里输入了三段背景知识资料。对于复杂系统会把pdf文件转成文本,进行分割,存储到向量数据库中,作为额外的知识库。
1 | mod llm; |
环境配置
项目的根目录下新建.env文件,其中内容为环境变量配置值,用来在程序中获取API和模型配置信息
1 | LLM_API=http://localhost:1234/v1/chat/completions |
另外还要配置LM Studio,在它的开发者界面中打开服务运行,并同时加载千问3.5-9b模型和文本嵌入模型

最终运行效果
因为我运行了多次这个程序,导致背景知识三段话被重复插入到了数据库中,当我询问tell me something about memorywalker时,向量数据库只返回了和memorywalker相关3条记录,rust的记录没有一条返回,的确找到了相关的背景知识。虽然3条记录的内容相同,但是id是不同的,这是因为我重复运行程序,main函数的测试数据被存储了3次。
另外看模型的思考过程中,它发现背景知识中memorywalker is from China and he love studing有语法错误,它做了一些纠结后,最后还是以一个专业助手的角度把语法错误改正,并给出了英文结论Based on the provided context, memorywalker is from China and he loves studying.。
当我再次问模型Is he a good guy?,模型改为了用中文思考,并用中文给出了回答。
1 | Agent start! |
第二次换了一个问题,不知道为什么这个模型开始用中文回复了
1 | > Is he a good guy? |