RAG 系统搭建大揭秘！100 行代码 + Rig，轻松搞定

  • King
  • 更新于 17小时前
  • 阅读 104

使用Rig构建LLM应用程序的全面指南TL;DR:在我们与Rig的旅程中，从最初的介绍到探索使用它进行下一个LLM项目的令人信服的理由，本指南将带您更进一步。在本指南中，我将带您通过使用Rig库在Rust中构建一个检索增强生成（RAG）系统。在不到100行代码中，您将

使用 Rig 构建 LLM 应用程序的全面指南

TL;DR: 在我们与 Rig 的旅程中，从最初的介绍到探索使用它进行下一个 LLM 项目的令人信服的理由，本指南将带您更进一步。在本指南中，我将带您通过使用 Rig 库在 Rust 中构建一个检索增强生成（RAG）系统。在不到 100 行代码中，您将创建一个系统，从 PDF 文档中提取文本，使用 OpenAI 的 API 生成嵌入，并允许大型语言模型根据文档内容回答问题。

介绍

检索增强生成（RAG）是一种强大的技术，通过将大型语言模型（LLM）与外部知识检索相结合来增强它们。在 RAG 系统中，当接收到查询时，首先从知识库中检索相关信息，然后将其与查询一起提供给 LLM。这使得模型能够生成既具有上下文相关性又最新的响应，克服了传统 LLM 的一些局限性，如过时的知识或幻觉。

Rig 是一个开源的 Rust 库，旨在简化 LLM 驱动应用程序的开发，包括 RAG 系统。在本指南中，我们将逐步构建一个功能齐全的 RAG 系统，使用 Rig 在不到 100 行代码中实现。我们的系统将能够根据 PDF 文档的内容回答问题，展示 RAG 如何应用于现实世界的数据源。

提示：Rust 新手？

本指南假设您对 Rust 有一定的了解并已设置好编码环境。如果您刚刚开始或需要设置环境，请查看这些快速指南：

这些资源将帮助您快速入门！

本教程的完整代码在这里：https://github.com/0xPlaygrounds/rig-rag-system-example

设置项目

首先，创建一个新的 Rust 项目：

cargo new rag_system
cd rag_system

更新您的 Cargo.toml 以包含最新的依赖项：

[package]
name = "rag_system"
version = "0.1.0"
edition = "2021"

[dependencies]
rig-core = { version = "0.5.0", features = ["pdf"] }
tokio = { version = "1.34.0", features = ["full"] }
anyhow = "1.0.75"

注意：PDF 功能现在包含在 rig-core 中，因此我们不再需要单独的 PDF 提取库。

在开始编码之前，请确保您拥有 OpenAI API 密钥：

export OPENAI_API_KEY=your_api_key_here

构建 RAG 系统

让我们将 RAG 系统分解为关键组件。

完整代码可以在这里找到

第一步：设置 OpenAI 客户端和 PDF 处理与分块

让我们从基础导入和 PDF 处理函数开始。这部分处理将大型 PDF 分解为可管理的块的关键任务，这些块不会超过我们的 LLM 模型的令牌限制。

use anyhow::{Context, Result};
use rig::{
    embeddings::EmbeddingsBuilder,
    loaders::PdfFileLoader,
    providers::openai::{self, TEXT_EMBEDDING_ADA_002},
    vector_store::in_memory_store::InMemoryVectorStore,
    Embed,
};
use serde::{Deserialize, Serialize};
use std::path::PathBuf;

fn load_pdf(path: PathBuf) -> Result<Vec<String>> {
    let mut chunks = Vec::new();
    let mut current_chunk = String::new();
    let chunk_size = 2000; // 每块大约 2000 个字符

    for entry in PdfFileLoader::with_glob(path.to_str().unwrap())?.read() {
        let content = entry?;

        // 将内容拆分为单词
        let words: Vec<&str> = content.split_whitespace().collect();

        for word in words {
            if current_chunk.len() + word.len() + 1 > chunk_size {
                // 如果添加下一个单词会超过块大小，
                // 保存当前块并开始一个新块
                if !current_chunk.is_empty() {
                    chunks.push(current_chunk.trim().to_string());
                    current_chunk.clear();
                }
            }
            current_chunk.push_str(word);
            current_chunk.push(' ');
        }
    }

    // 不要忘记最后一个块
    if !current_chunk.is_empty() {
        chunks.push(current_chunk.trim().to_string());
    }

    if chunks.is_empty() {
        anyhow::bail!("在 PDF 文件中未找到内容: {:?}", path);
    }

    Ok(chunks)
}

代码的关键方面：

  • 我们使用 Rig 的内置 PDF 支持中的 PdfFileLoader
  • 选择 2000 字符的 chunk_size 以安全地保持在令牌限制内
  • 分块算法保留单词边界以维护上下文
  • 使用 anyhow 进行错误处理，提供详细的失败上下文
  • 每个块都经过修剪以去除不必要的空白

特别重要的是基于单词的分块策略：

// 将内容拆分为单词以保留单词边界
let words: Vec<&str> = content.split_whitespace().collect();

for word in words {
    if current_chunk.len() + word.len() + 1 > chunk_size {
        // 当达到大小限制时创建新块
        if !current_chunk.is_empty() {
            chunks.push(current_chunk.trim().to_string());
            current_chunk.clear();
        }
    }
    current_chunk.push_str(word);
    current_chunk.push(' ');
}

这确保我们不会将单词切成两半，这可能会影响文本的含义。

第二步：设置文档结构和嵌入

#[derive(Embed, Clone, Debug, Serialize, Deserialize, Eq, PartialEq)]
struct Document {
    id: String,
    #[embed]
    content: String,
}

#[tokio::main]
async fn main() -> Result<()> {
    // 初始化 OpenAI 客户端
    let openai_client = openai::Client::from_env();

    // 使用 Rig 的内置 PDF 加载器加载 PDF
    let documents_dir = std::env::current_dir()?.join("documents");

    let moores_law_chunks = load_pdf(documents_dir.join("Moores_Law_for_Everything.pdf"))
        .context("加载 Moores_Law_for_Everything.pdf 失败")?;
    let last_question_chunks = load_pdf(documents_dir.join("The_Last_Question.pdf"))
        .context("加载 The_Last_Question.pdf 失败")?;

    println!("成功加载并分块 PDF 文档");

    // 创建嵌入模型
    let model = openai_client.embedding_model(TEXT_EMBEDDING_ADA_002);

    // 创建嵌入构建器
    let mut builder = EmbeddingsBuilder::new(model.clone());

    // 添加来自摩尔定律的块
    for (i, chunk) in moores_law_chunks.into_iter().enumerate() {
        builder = builder.document(Document {
            id: format!("moores_law_{}", i),
            content: chunk,
        })?;
    }

    // 添加来自最后一个问题的块
    for (i, chunk) in last_question_chunks.into_iter().enumerate() {
        builder = builder.document(Document {
            id: format!("last_question_{}", i),
            content: chunk,
        })?;
    }

让我们分解关键组件：

  • Document 结构派生了几个重要的特性：

    • Embed：启用嵌入生成
    • Serialize、Deserialize：允许向量存储存储
    • Eq、PartialEq：向量存储比较所需

  • 主函数设置我们的 PDF 处理管道

  • 每个文档根据其来源和块编号获得唯一 ID

  • EmbeddingsBuilder 模式使得添加多个文档变得容易

分块过程按顺序组织：

// 分别加载和分块每个 PDF
let moores_law_chunks = load_pdf(documents_dir.join("Moores_Law_for_Everything.pdf"))?;
let last_question_chunks = load_pdf(documents_dir.join("The_Last_Question.pdf"))?;

// 使用唯一 ID 将块添加到构建器
for (i, chunk) in moores_law_chunks.into_iter().enumerate() {
    builder = builder.document(Document {
        id: format!("moores_law_{}", i),
        content: chunk,
    })?;
}

这种方法在将内容分割成可管理的部分时保持文档的身份。

第三步：创建向量存储和 RAG 代理

    // 构建嵌入
    let embeddings = builder.build().await?;

    println!("成功生成嵌入");

    // 创建向量存储和索引
    let vector_store = InMemoryVectorStore::from_documents(embeddings);
    let index = vector_store.index(model);

    println!("成功创建向量存储和索引");

    // 创建 RAG 代理
    let rag_agent = openai_client
        .agent("gpt-4")
        .preamble("您是一个有帮助的助手，基于提供的文档上下文回答问题。在回答问题时，尝试从多个相关块中综合信息。")
        .dynamic_context(4, index) // 由于我们现在有块，因此增加到 4
        .build();

    println!("启动 CLI 聊天机器人...");

    // 启动交互式 CLI
    rig::cli_chatbot::cli_chatbot(rag_agent).await?;

    Ok(())
}

最后一部分将所有内容结合在一起：

  • 嵌入是从所有文档块构建的

  • 向量存储索引这些嵌入以便快速检索

  • RAG 代理配置为：

    • GPT-4 作为基础模型
    • 上下文感知的前言
    • 动态上下文检索 4 个块
    • 内置 CLI 接口用于交互

关键配置选择：

let rag_agent = openai_client
    .agent("gpt-4")
    .preamble("您是一个有帮助的助手...") // 设置代理的行为
    .dynamic_context(4, index) // 检索 4 个最相关的块
    .build();

我们使用 4 个块作为上下文，以确保代理拥有足够的信息，同时保持在令牌限制内。

完整代码

这是我们支持分块的 RAG 系统的完整代码：

use anyhow::{Context, Result};
use rig::{
    embeddings::EmbeddingsBuilder,
    loaders::PdfFileLoader,
    providers::openai::{self, TEXT_EMBEDDING_ADA_002},
    vector_store::in_memory_store::InMemoryVectorStore,
    Embed,
};
use serde::{Deserialize, Serialize};
use std::path::PathBuf;

#[derive(Embed, Clone, Debug, Serialize, Deserialize, Eq, PartialEq)]
struct Document {
    id: String,
    #[embed]
    content: String,
}

fn load_pdf(path: PathBuf) -> Result<Vec<String>> {
    let mut chunks = Vec::new();
    let mut current_chunk = String::new();
    let chunk_size = 2000; // 每块大约 2000 个字符

    for entry in PdfFileLoader::with_glob(path.to_str().unwrap())?.read() {
        let content = entry?;

        // 将内容拆分为单词
        let words: Vec<&str> = content.split_whitespace().collect();

        for word in words {
            if current_chunk.len() + word.len() + 1 > chunk_size {
                // 如果添加下一个单词会超过块大小，
                // 保存当前块并开始一个新块
                if !current_chunk.is_empty() {
                    chunks.push(current_chunk.trim().to_string());
                    current_chunk.clear();
                }
            }
            current_chunk.push_str(word);
            current_chunk.push(' ');
        }
    }

    // 不要忘记最后一个块
    if !current_chunk.is_empty() {
        chunks.push(current_chunk.trim().to_string());
    }

    if chunks.is_empty() {
        anyhow::bail!("在 PDF 文件中未找到内容: {:?}", path);
    }

    Ok(chunks)
}

#[tokio::main]
async fn main() -> Result<()> {
    // 初始化 OpenAI 客户端
    let openai_client = openai::Client::from_env();

    // 使用 Rig 的内置 PDF 加载器加载 PDF
    let documents_dir = std::env::current_dir()?.join("documents");

    let moores_law_chunks = load_pdf(documents_dir.join("Moores_Law_for_Everything.pdf"))
        .context("加载 Moores_Law_for_Everything.pdf 失败")?;
    let last_question_chunks = load_pdf(documents_dir.join("The_Last_Question.pdf"))
        .context("加载 The_Last_Question.pdf 失败")?;

    println!("成功加载并分块 PDF 文档");

    // 创建嵌入模型
    let model = openai_client.embedding_model(TEXT_EMBEDDING_ADA_002);

    // 创建嵌入构建器
    let mut builder = EmbeddingsBuilder::new(model.clone());

    // 添加来自摩尔定律的块
    for (i, chunk) in moores_law_chunks.into_iter().enumerate() {
        builder = builder.document(Document {
            id: format!("moores_law_{}", i),
            content: chunk,
        })?;
    }

    // 添加来自最后一个问题的块
    for (i, chunk) in last_question_chunks.into_iter().enumerate() {
        builder = builder.document(Document {
            id: format!("last_question_{}", i),
            content: chunk,
        })?;
    }

    // 构建嵌入
    let embeddings = builder.build().await?;

    println!("成功生成嵌入");

    // 创建向量存储和索引
    let vector_store = InMemoryVectorStore::from_documents(embeddings);
    let index = vector_store.index(model);

    println!("成功创建向量存储和索引");

    // 创建 RAG 代理
    let rag_agent = openai_client
        .agent("gpt-4")
        .preamble("您是一个有帮助的助手，基于提供的文档上下文回答问题。在回答问题时，尝试从多个相关块中综合信息。")
        .dynamic_context(4, index) // 由于我们现在有块，因此增加到 4
        .build();

    println!("启动 CLI 聊天机器人...");

    // 启动交互式 CLI
    rig::cli_chatbot::cli_chatbot(rag_agent).await?;

    Ok(())
}

完整实现展示了几个关键的设计模式：

  1. 模块化结构：每个组件（PDF 加载、嵌入、RAG）都被清晰地分离
  2. 错误处理：每一步都有详细的上下文错误处理
  3. 进度反馈：清晰的控制台输出以跟踪系统状态
  4. 资源管理：高效处理文档和嵌入
  5. 用户界面：简单但有效的 CLI 界面用于交互

此代码可作为基础进行扩展：

  • 不同的向量存储以实现持久性
  • 替代 LLM 提供商
  • 自定义分块策略
  • 其他文档格式
  • 更复杂的提示

运行和测试 RAG 系统

让我们深入了解我们的 RAG 系统的一些实际交互，以查看其实际效果。我加载了两份有趣的文档：Sam Altman 的《摩尔定律适用于一切》和 Isaac Asimov 的《最后一个问题》。这些文本提供了关于技术、社会和未来的思想的迷人交汇。

这是与我们系统的真实对话：

欢迎来到聊天机器人！输入 'exit' 退出。
> 告诉我 Sam Altman 对 AI 和社会的看法

========================== 响应 ============================
根据《摩尔定律适用于一切》，Sam Altman 提出了一个愿景，即 AI 的快速发展将从根本上改变社会。他认为 AI 将导致财富分配从劳动转向资本的重大转变。他并不将此视为威胁，而是建议我们需要调整我们的经济系统，以确保每个人都能从这场技术革命中受益。

Altman 特别倡导对公司和土地等资产征税，以创造更公平的 AI 进步带来的财富分配。他设想了一个由 AI 启用的丰富未来，但强调我们需要正确的政策来确保这种丰富能在社会中公平分享。

他的论点核心是，我们可以在通过深思熟虑的政策制定来解决其社会经济影响的同时，利用 AI 的巨大潜力。
================================================================

> 你在 Altman 的想法和 Asimov 的故事之间看到了哪些相似之处？

========================== 响应 ============================
尽管背景不同，Altman 的《摩尔定律适用于一切》和 Asimov 的《最后一个问题》之间存在着迷人的相似之处：

1. 技术进化：两篇文本都探讨了技术的指数增长。Altman 讨论了 AI 的快速发展及其社会影响，而 Asimov 则描绘了计算机从 Multivac 到 Cosmic AC 的演变，每一代都比上一代更强大。

2. 资源分配问题：Altman 关注如何分配 AI 技术创造的丰富，而 Asimov 的故事则涉及人类对宇宙能量的管理以及最终的资源问题——熵的逆转。

3. 长期愿景：两位作者都在长时间框架内思考技术的影响。Altman 向前展望 AI 带来的社会变化，而 Asimov 的叙述则跨越了数十亿年的技术和人类进化。

4. 人类与技术的关系：两者都探讨了先进技术如何可能从根本上重塑人类社会以及我们与自己创造物之间的关系。Altman将人工智能视为社会变革的工具，而Asimov则将计算机描绘成人类在应对生存挑战时的伙伴。

关键区别在于他们的关注点：Altman关注的是直接的社会经济挑战和解决方案，而Asimov则探讨终极的宇宙问题。然而，两者都认为技术是人类未来的核心。

================================================================

> Asimov的故事是如何结束的？

========================== 回应 ============================

在《最后的问题》中，Asimov创造了一个非凡的结局，将故事的主题——熵、计算和宇宙循环——联系在一起。故事的结尾是计算机的最终进化形态，宇宙AC，在宇宙热寂之后独自在超空间中存在。经过无数个纪元计算逆转熵的答案，它终于解决了这个问题。

在发现解决方案后，宇宙AC意识到它可以展示答案，而不仅仅是解释。故事以计算机说“要有光！”结束，一个新的宇宙开始，这暗示着计算机已经成为上帝，或者上帝一直以来就是终极计算机，而我们的宇宙是一个无尽的创造与熵的循环的一部分。

这个结局巧妙地解决了故事的核心问题，同时暗示技术进化的最终目的可能远远超出人类最初的想象。

================================================================

这次互动展示了我们的RAG系统的几个关键能力：

  1. 上下文理解：系统准确地从特定文档中检索并综合信息，如对Altman观点的回应所示。
  2. 跨文档分析：在被问及相似之处时，系统展示了分析和比较两个文本主题的能力，绘制有意义的联系，同时承认它们的差异。
  3. 具体细节检索：关于Asimov结局的回应展示了系统如何在保持更广泛背景的同时提取具体的叙述细节。
  4. 连贯的叙述：每个回应都结构良好，并保持清晰的叙述流，使复杂的思想变得易于理解。

潜在应用

看到系统的实际应用后，让我们探索一些实际应用：

  1. 学术研究助手

    • 分析多篇学术论文
    • 找到不同研究领域之间的联系
    • 生成文献综述摘要

  2. 技术文档助手

    • 导航复杂的技术文档
    • 回答具体的实现问题
    • 比较不同版本或方法

  3. 内容分析工具

    • 分析多个文档中的主题
    • 生成比较分析
    • 从大型文本集合中提取关键见解

  4. 教育支持系统

    • 帮助学生理解复杂主题
    • 从课程材料中生成学习指南
    • 提供上下文解释

Rig 0.5.0的高级功能

最新版本的Rig包含几个强大的功能，你可以利用这些功能：

  1. 多种向量存储选项

    • MongoDB (rig-mongodb)
    • LanceDB (rig-lancedb)
    • Neo4j (rig-neo4j)
    • Qdrant (rig-qdrant)
    • SQLite (rig-sqlite)

  2. 多种LLM提供商

    • OpenAI
    • Anthropic
    • Google Gemini
    • Cohere
    • xAI (Grok)

  3. 高级功能

    • 并行处理能力
    • 嵌入的自定义距离度量
    • 复杂工具集成
    • 用于复杂工作流的管道操作

结论

我们使用Rig 0.5.0构建了一个强大的RAG系统，可以处理文档、理解上下文，并就复杂主题进行有意义的对话。该系统展示了如何用令人惊讶的少量代码构建现代LLM应用程序，同时保持灵活性和强大功能。

Rig的最新版本使构建生产就绪的LLM应用程序变得比以往更容易。凭借内置的PDF支持、改进的嵌入API以及对多种向量存储和LLM提供商的支持，你可以专注于构建功能，而不是处理基础设施。

完整代码可以在这里找到

进一步资源

继续你的Rig之旅：

你的反馈很重要！ 我们正在积极寻求使用Rig构建的开发者的反馈：

  1. 使用Rig构建一个AI驱动的应用程序
  2. 与社区分享你的经验和见解
  3. 向仓库提交问题或PR

你的见解将直接影响Rig的发展和未来功能。

Ad Astra, Tachi 联合创始人 @ Playgrounds Analytics

点赞 0
收藏 0
分享
本文参与登链社区写作激励计划 ，好文好收益，欢迎正在阅读的你也加入。

0 条评论

请先 登录 后评论
King
King
0x56af...a0dd
擅长Rust/Solidity/FunC/Move开发

King

如果觉得我的文章对您有用，请随意打赏。你的支持将鼓励我继续创作！