实战指南：在 Rust 中使用 Qwen 和知识库通过 RAG 构建高效 LLM 应用

一般基于使用大模型与检索增强生成（Retrieval-Augmented Generation, RAG）结合开发应用的调用流程，主要涉及以下几个步骤。这个过程能够确保应用在处理用户查询时，不仅依赖于LLM的知识，还能动态地从外部数据源中检索信息，以提供更加准确和最新的回答。

具体步骤如下：

用户提交查询：用户通过应用界面输入问题或指令。
查询预处理：对用户的查询进行初步处理
检索阶段（RAG部分）：
- 使用嵌入技术（Embedding），将用户查询转化为向量表示。
- 在预先构建好的文档或知识库索引中搜索最相关的条目。这通常涉及到计算查询向量与索引中各文档向量之间的相似度，并选择最匹配的结果。
- 检索到的相关内容被提取出来，准备用于生成回答。
生成阶段（LLM部分）：
- 将检索到的相关内容作为上下文，连同原始用户查询一起传递给LLM。
- LLM基于提供的上下文生成详细且精确的回答。这里，LLM不仅依赖其内部训练的数据集，还利用了实时检索到的最新信息来丰富回答的内容。
答案后处理：
- 对LLM生成的答案进行必要的格式化或简化，使其更易于理解。
- 确保答案符合应用的安全性和合规性要求，例如过滤敏感信息。
返回结果给用户：最终，经过处理的答案通过应用界面展示给用户。

目前这种流程有大量的框架来帮助我们实现简化，那么 Rust 有类似的框架吗？答案是有的。

Swiftide 是一个 Rust 库，用于构建 LLM 应用程序，支持快速数据摄取、转换和索引，以实现有效的查询和提示注入，称为检索增强生成。它为创建各种代理提供了灵活的构建块，允许使用最少的代码从概念快速开发到生产。

不过目前为止，Swiftide 只支持 OpenAI、Groq、Ollama、AWS Bedrock 这四个大模型/平台，Qwen 大模型除非你使用 Ollama 加载本地模型版本。但是本地模型需要GPU 资源，显然使用线上模型更有性价比。幸运的是，Qwen 官方模型也有OpenAI 模式接口调用，因此我们只需要对Swiftide 项目进行一些改动即可兼容。

Qwen 官方文档： https://help.aliyun.com/zh/model-studio/getting-started/models

Swiftide 添加 Qwen 支持

添加功能前我们先看看 Swiftide 有哪些组件：

LLM 大模型：这是最主要的。
LLM Embedding：嵌入模型，因为我们需要把知识库转换为特征向量。
特征向量库：特征向量需要有存储载体。

Qwen 适配 OpenAI 接口

前面说到，Qwen 官方是支持 OpenAI 接口协议的，因此我们只需要做少量的适配。需要实现async_openai::config::Config。

先定义好结构：

const QWEN_API_BASE: &str = "https://dashscope.aliyuncs.com/compatible-mode/v1";

#[derive(Clone, Debug, Deserialize)]
#[serde(default)]
pub struct QwenConfig {
    api_base: String,
    api_key: Secret<String>,
}

impl Default for QwenConfig {
    fn default() -> Self {
        Self {
            api_base: QWEN_API_BASE.to_string(),
            api_key: get_api_key().into(),
        }
    }
}

对于API key,默认从环境变量里获取：

fn get_api_key() -> String {
    std::env::var("QWEN_API_KEY")
        .unwrap_or_else(|_| std::env::var("DASHSCOPE_API_KEY").unwrap_or_default())
}

实现 async_openai::config::Config, 主要参考这个接口的请求头：

curl -X POST https://dashscope.aliyuncs.com/compatible-mode/v1/chat/completions \
-H "Authorization: Bearer $DASHSCOPE_API_KEY" \
-H "Content-Type: application/json" \
-d '{
    "model": "qwen-plus",
    "messages": [
        {
            "role": "system",
            "content": "You are a helpful assistant."
        },
        {
            "role": "user", 
            "content": "你是谁？"
        }
    ]
}'


impl async_openai::config::Config for QwenConfig {
    fn headers(&self) -> HeaderMap {
        let mut headers = HeaderMap::new();

        headers.insert(
            AUTHORIZATION,
            format!("Bearer {}", self.api_key.expose_secret())
                .as_str()
                .parse()
                .unwrap(),
        );

        headers
    }

    fn url(&self, path: &str) -> String {
        format!("{}{}", self.api_base, path)
    }

    fn api_base(&self) -> &str {
        &self.api_base
    }

    fn api_key(&self) -> &Secret<String> {
        &self.api_key
    }

    fn query(&self) -> Vec<(&str, &str)> {
        vec![]
    }
}

我们需要限制 Qwen 模型名称和嵌入模型名称，因此需要定义2个枚举：

Qwen 模型版本：

#[derive(Debug, Default, Clone, PartialEq)]
pub enum QwenModel {
    #[default]
    Max,
    Plus,
    Turbo,
    Long,
}
impl Display for QwenModel {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            QwenModel::Max => write!(f, "qwen-max"),
            QwenModel::Plus => write!(f, "qwen-plus"),
            QwenModel::Turbo => write!(f, "qwen-turbo"),
            QwenModel::Long => write!(f, "qwen-long"),
        }
    }
}

Qwen Embedding 版本：

#[derive(Debug, Default, Clone, PartialEq)]
pub enum QwenEmbedding {
    #[default]
    TextEmbeddingV1,
    TextEmbeddingV2,
    TextEmbeddingV3,
    TextEmbeddingAsyncV1,
    TextEmbeddingAsyncV2,
}


impl Display for QwenEmbedding {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            QwenEmbedding::TextEmbeddingV1 => write!(f, "text-embedding-v1"),
            QwenEmbedding::TextEmbeddingV2 => write!(f, "text-embedding-v2"),
            QwenEmbedding::TextEmbeddingV3 => write!(f, "text-embedding-v3"),
            QwenEmbedding::TextEmbeddingAsyncV1 => write!(
                f,
                "text-embedding-async-v1

"
            ),
            QwenEmbedding::TextEmbeddingAsyncV2 => write!(
                f,
                "text-embedding-async-v2

"
            ),
        }
    }
}

impl From<&String> for QwenEmbedding {
    fn from(value: &String) -> Self {
        match value.as_str() {
            "text-embedding-v1" => QwenEmbedding::TextEmbeddingV1,
            "text-embedding-v2" => QwenEmbedding::TextEmbeddingV2,
            "text-embedding-v3" => QwenEmbedding::TextEmbeddingV3,
            "text-embedding-async-v1" => QwenEmbedding::TextEmbeddingAsyncV1,
            "text-embedding-async-v2" => QwenEmbedding::TextEmbeddingAsyncV2,
            _ => panic!("Invalid embedding model"),
        }
    }
}

定义 Qwen 结构体

因为我们需要的是 OpenAI 兼容接口，所以 client 要用的OpenAI 库里的 Client。

#[derive(Debug, Builder, Clone)]
#[builder(setter(into, strip_option))]
pub struct Qwen {
    #[builder(default = "default_client()", setter(custom))]
    client: Arc<async_openai::Client<QwenConfig>>,
    /// Default options for prompt models.
    #[builder(default)]
    default_options: Options,
}

Options 可以参考 OpenAI 的实现或者 Groq 的实现：

#[derive(Debug, Default, Clone, Builder)]
#[builder(setter(into, strip_option))]
pub struct Options {
    /// The default prompt model to use, if specified.
    #[builder(default)]
    pub prompt_model: Option<String>,
    #[builder(default)]
    pub embed_model: Option<String>,
    #[builder(default)]
    pub dimensions: u16,
}

这里需要添加dimensions 字段来保存嵌入模型的维度，因为Qwen 的Embedding 模型的维度不太统一：


模型名称	向量维度
text-embedding-v3	1,024（默认）、768或512
text-embedding-v2	1,536
text-embedding-v1	1,536
text-embedding-async-v2	1,536
text-embedding-async-v1	1,536
因此我们在对外提供传参的维度参数时需要做相应的验证:

/// 设置或验证默认维度
/// 
/// 此方法主要用于设置或验证与嵌入模型相关的维度参数。如果提供了嵌入模型，则根据模型类型验证传入的维度是否符合预期。
/// 如果维度不符合预期，程序将断言失败并抛出错误信息。如果default_options尚未初始化，或者没有提供嵌入模型，则会初始化或更新default_options中的dimensions字段。
/// 
/// # 参数
/// - `dimensions`: u16类型，表示维度值。
/// 
/// # 返回
/// 返回`&mut Self`，允许进行方法链式调用。
pub fn default_dimensions(&mut self, dimensions: u16) -> &mut Self {
    // 尝试获取可变引用到default_options
    if let Some(options) = self.default_options.as_mut() {
        // 如果提供了嵌入模型，根据模型类型验证维度
        if let Some(model) = &options.embed_model {
            let embed_model: QwenEmbedding = model.into();
            match embed_model {
                QwenEmbedding::TextEmbeddingV1 => assert_eq!(
                    dimensions, 1536,
                    "Dimensions must be 1536 for this embedding model"
                ),
                QwenEmbedding::TextEmbeddingV2 => assert_eq!(
                    dimensions, 1536,
                    "Dimensions must be 1536 for this embedding model"
                ),
                QwenEmbedding::TextEmbeddingV3 => assert!(
                    matches!(dimensions, 1024 | 768 | 512),
                    "Dimensions must be one of [1024, 768, 512] for TextEmbeddingV3"
                ),
                QwenEmbedding::TextEmbeddingAsyncV1 => assert_eq!(
                    dimensions, 1536,
                    "Dimensions must be 1536 for this embedding model"
                ),
                QwenEmbedding::TextEmbeddingAsyncV2 => assert_eq!(
                    dimensions, 1536,
                    "Dimensions must be 1536 for this embedding model"
                ),
            }
        }

        // 更新options中的dimensions字段
        options.dimensions = dimensions;
    } else {
        // 如果default_options尚未初始化，进行初始化
        self.default_options = Some(Options {
            dimensions,
            ..Default::default()
        });
    }
    self
}

为了方便调用，QwenBuilder 再加上下面的方法：

pub fn default_prompt_model(&mut self, model: &QwenModel) -> &mut Self {
        if let Some(options) = self.default_options.as_mut() {
            options.prompt_model = Some(model.to_string());
        } else {
            self.default_options = Some(Options {
                prompt_model: Some(model.to_string()),
                ..Default::default()
            });
        }
        self
    }

    pub fn default_embed_model(&mut self, model: &QwenEmbedding) -> &mut Self {
        if let Some(options) = self.default_options.as_mut() {
            options.embed_model = Some(model.to_string());
        } else {
            self.default_options = Some(Options {
                embed_model: Some(model.to_string()),
                ..Default::default()
            });
        }
        self
    }

Qwen LLM 适配 Swiftide

Swiftide 里只需要实现 SimplePrompt trait 即可。

#[async_trait]
impl SimplePrompt for Qwen {
	// 使用给定的提示与语言模型进行交互的异步函数
	async fn prompt(&self, prompt: Prompt) -> Result<String> {
	    // 获取用于处理提示的默认模型，确保已设置
	    let model = self
	        .default_options
	        .prompt_model
	        .as_ref()
	        .context("模型未设置")?
	        .to_string();
	
	    // 构造创建聊天补全的请求，包括模型和提示消息
	    let request = CreateChatCompletionRequestArgs::default()
	        .model(model)
	        .messages(vec![ChatCompletionRequestUserMessageArgs::default()
	            .content(prompt.render().await?) // 渲染并添加提示内容
	            .build()?
	            .into()])
	        .build()?;
	
	    // 为了调试目的记录构造的请求
	    tracing::debug!(
	        messages = serde_json::to_string_pretty(&request)?,
	        "[SimplePrompt] 请求发送到qwen"
	    );
	
	    // 发送请求并等待响应
	    let mut response = self.client.chat().create(request).await?;
	
	    // 为了调试目的记录接收到的响应
	    tracing::debug!(
	        response = serde_json::to_string_pretty(&response)?,
	        "[SimplePrompt] 从qwen接收到的响应"
	    );
	
	    // 提取并返回响应中第一个选择的内容，确保其存在
	    response
	        .choices
	        .remove(0)
	        .message
	        .content
	        .take()
	        .context("响应内容错误")
	}

}

Qwen Embedding 适配

Swiftide 适配需要实现 swiftide_core::EmbeddingModel

原始接口如下：

curl --location 'https://dashscope.aliyuncs.com/compatible-mode/v1/embeddings' \
--header "Authorization: Bearer $DASHSCOPE_API_KEY" \
--header 'Content-Type: application/json' \
--data '{
    "model": "text-embedding-v3",
    "input": "衣服的质量杠杠的，很漂亮，不枉我等了这么久啊，喜欢，以后还来这里买",  
    "dimension": "1024",  
    "encoding_format": "float"
}'

// 异步处理输入文本以生成相应的嵌入向量。
//
// 参数：
// - `input`: 包含要嵌入的文本的字符串向量。
//
// 返回值：
// - `Result<Embeddings>`: 如果成功，则返回包含生成的嵌入向量的结果类型。
//
// 该函数从默认选项中获取嵌入模型和维度，构造嵌入请求，并将其发送到 Qwen API。然后处理响应以提取嵌入向量并返回结果。
#[async_trait]
impl EmbeddingModel for Qwen {
    async fn embed(&self, input: Vec<String>) -> Result<Embeddings> {
        // 获取嵌入模型，确保其已设置。
        let model = self
            .default_options
            .embed_model
            .as_ref()
            .context("模型未设置")?;
        
        // 从默认选项中获取嵌入维度。
        let dimensions = self.default_options.dimensions;
        
        // 构造嵌入请求参数。
        let request = CreateEmbeddingRequestArgs::default()
            .model(model)
            .dimensions(dimensions)
            .input(&input)
            .build()?;
        
        // 记录嵌入请求的详细信息。
        tracing::debug!(
            num_chunks = input.len(),
            model = &model,
            "[Embed] 请求发送至 qwen"
        );
        
        // 发送嵌入请求并等待响应。
        let response = self.client.embeddings().create(request).await?;
        
        // 记录接收到的嵌入向量数量。
        let num_embeddings = response.data.len();
        tracing::debug!(num_embeddings = num_embeddings, "[Embed] 接收到的响应");

        // 处理响应以提取嵌入向量，假设顺序保持不变。
        // 警告：此假设可能并不总是成立。
        Ok(response.data.into_iter().map(|d| d.embedding).collect())
    }
}

实现文档查询 LLM 应用

为了演示上面成果，我这里以这个页面内容（tauri 框架文档）为例子：

https://tauri.app/zh-cn/develop/resources/

初始化 Qwen 客户端

let client = QwenBuilder::default()
        .default_embed_model(&swiftide::integrations::qwen::QwenEmbedding::TextEmbeddingV2)
        .default_prompt_model(&swiftide::integrations::qwen::QwenModel::Long)
        .default_dimensions(1536)
        .build()?;

初始化向量库

我们使用LanceDB 作为向量库，LanceDB 也是 Rust 实现的，当然你也可以用 Qdrant，Qdrant 也是Rust 实现的，但是 Qdrant 是一个应用，使用比较麻烦。 LanceDB 初始化：

let tempdir = TempDir::new().unwrap();
let lancedb = LanceDB::builder()
        .uri(tempdir.child("lancedb").to_str().unwrap())
        .vector_size(1536)
        .with_vector(EmbeddedField::Combined)
        .with_metadata(metadata_qa_text::NAME)
        .table_name("swiftide_test")
        .build()
        .unwrap();

这里的 vector_size 就是向量维度，需要和你调用的 Embedding 模型维度一致。

构建向量索引管道

这个操作调用嵌入模型生成知识库内容向量，再把向量存入向量库。

// 构建索引管道，从文件加载器开始，经过分块、元数据处理、嵌入，最后存储到LanceDB中
    indexing::Pipeline::from_loader(FileLoader::new(".").with_extensions(&["md"]))
        .with_default_llm_client(client.clone())
        .then_chunk(ChunkMarkdown::from_chunk_range(10..2048))
        .then(MetadataQAText::new(client.clone()))
        .then_in_batch(Embed::new(client.clone()).with_batch_size(10))
        .then_store_with(lancedb.clone())
        .run()
        .await?;

构建查询管道


let pipeline = query::Pipeline::default()
        .then_transform_query(query_transformers::GenerateSubquestions::from_client(
            client.clone(),
        ))
        .then_transform_query(query_transformers::Embed::from_client(client.clone()))
        .then_retrieve(lancedb.clone())
        .then_transform_response(response_transformers::Summary::from_client(client.clone()))
        .then_answer(answers::Simple::from_client(client.clone()));

查询

我们查询一个问题看看

let result = pipeline
        .query("tauri 如何访问文件")
        .await?;
    
    println!("====");
    println!("{:?}", result.answer());

输出如下：

在 Tauri 应用程序中访问文件可以通过 Rust 和 JavaScript 两种方式进行。以下是具体的步骤和示例代码：\n\n### 在 Rust 中访问文件\n\n1. **配置 `resources`**：首先，确保你已经在 `tauri.conf.json` 文件中的 `bundle` 对象中添加了 `resources` 属性，以包含你需要的文件。\n\n    ```json title=tauri.conf.json\n    {\n      \"bundle\": {\n        \"resources\": [\n          \"lang/*\"\n        ]\n      }\n    }\n    ```\n\n2. **使用 `PathResolver` 访问文件**：在 Rust 代码中，你可以使用 `PathResolver` 实例来访问这些文件。\n\n ...................

最后

事实上现在很多大型文档大概就是这么实现AI 查询的，例如 Milvus 向量库的文档，右正角就有一个Ask AI: https://milvus.io/docs 。 lancedb 反而没有自己提供类似的功能。

Milvus 也是一个向量库，理论上也可以在 Swiftide 实现 Milvus 的 embedding 存储。

不使用框架实现也是可以的，但是Swiftide 还提供一些 AI Agents 的一些功能，支持一些Tools 的输入输出，这极大的省了很多开发工作。