语义缓存实战：Qdrant + Rust 优化 LLM 成本

你好！今天我们将学习如何使用 Qdrant 在 Rust 中实现语义缓存。本教程完成后，你将拥有一个可以执行以下操作的 Rust 应用程序：

• 摄取 CSV 文件，在 LLM 的帮助下将其转换为 embedding，并将其插入 Qdrant

• 在 Qdrant 中创建两个集合 - 一个用于常规用途，一个用于缓存

• 利用语义缓存实现更快速的访问

有兴趣部署或者迷路了想找到代码存储库？你可以在这里找到。

什么是语义缓存，为什么要使用它？

在常规数据缓存中，我们存储信息以便稍后更快地检索。例如，你可能有一个在 Nginx 后面提供的网络服务。我们可以让 Nginx 缓存所有响应，或者只缓存最常访问的端点。这改进了性能并降低了网络服务器的负载。

从这个角度来说，语义缓存相当相似。使用向量数据库，我们可以创建存储查询本身的数据库集合。例如，这两个问题在语义上表达相同的意思：

• Rust 编程语言的最佳实践是什么？

• Rustlang 的最佳实践是什么？

我们可以在缓存集合中存储查询的副本，并将答案作为 JSON 负载存储。如果用户随后提出类似的问题，我们可以检索 embedding 并从负载中获取答案。这避免了我们必须使用 LLM 来获取答案。

语义缓存有几个好处：

• 需要长响应的提示词可以实现显著的成本节省。

• 它相当容易实现且成本相对便宜 - 唯一的成本是存储和使用 embedding 模型

• 你可以使用比常规 embedding 更便宜的模型

语义缓存通常与 RAG（检索增强生成）一起使用。RAG 是一个允许从预先嵌入的材料中进行上下文检索的框架。例如，CSV 文件或文档可以使用模型转换为 embedding 并存储在数据库中。每当用户想要查找与给定提示词相似的文档时，他们会嵌入该提示词并在给定数据库中进行搜索。

当然，也有充分的理由不使用语义缓存。需要不同、多样答案的提示词不会从语义缓存中获益。这在生成式 AI 使用中特别相关。获取存储的查询将降低响应的创意度。无论如何，如果你的管道的某个部分能够利用语义缓存，那么这样做是个好主意。

首先，别忘了使用 shuttle init 和 Axum 框架。我们将使用下面的 shell 代码片段安装依赖项：

cargo add qdrant-client@1.7.0 anyhow async-openai serde serde-json \
shuttle-qdrant uuid -F uuid/v4,serde/derive

你可以在这里找到我们的快速开始文档。

为了设置我们的密钥，我们将使用位于项目根目录的 Secrets.toml 文件（你需要手动创建此文件）。然后你可以使用下面的格式添加你需要的任何密钥：

OPENAI_API_KEY = ""
QDRANT_URL = ""
QDRANT_API_KEY = ""

创建集合

既然我们可以开始了，我们将添加一些更通用的方法来创建常规集合以及缓存集合，以模拟真实场景（以及一个 new() 函数来创建 RAGSystem 结构体本身）。我们首先创建结构体：注意这里虽然我们使用 1536 维的向量，但你需要的维度数可能取决于你使用的模型。

use qdrant_client::prelude::QdrantClient;
use async_openai::{config::OpenAIConfig, Client};

struct RagSystem {
    qdrant_client: QdrantClient,
    openai_client: Client<OpenAIConfig>
}

static REGULAR_COLLECTION_NAME: &str = "my-collection";
static CACHE_COLLECTION_NAME: &str = "my-collection-cached";

impl RAGSystem {
    fn new(qdrant_client: QdrantClient) -> Self {
        let openai_api_key = env::var("OPENAI_API_KEY").unwrap();

let openai_config = OpenAIConfig::new()
            .with_api_key(openai_api_key)
            .with_org_id("qdrant-shuttle-semantic-cache");

let openai_client = Client::with_config(openai_config);

Self {
            openai_client,
            qdrant_client,
        }
    }
}

现在我们将创建用于初始化常规集合的方法。注意我们只需要使用这些方法一次。在集合已经创建之后，如果我们尝试再次初始化它们，我们会得到一个错误。

use qdrant_client::prelude::CreateCollection;
use qdrant_client::qdrant::{
    vectors_config::Config,VectorParams,
    VectorsConfig, WithPayloadSelector,
};

impl RagSystem {
    async fn create_regular_collection(&self) -> Result<()> {
        self.qdrant_client
            .create_collection(&CreateCollection {
                collection_name: REGULAR_COLLECTION_NAME.to_string(),
                vectors_config: Some(VectorsConfig {
                    config: Some(Config::Params(VectorParams {
                        size: 1536,
                        distance: Distance::Cosine.into(),
                        ..Default::default()
                    })),
                }),
                ..Default::default()
            })
            .await?;

Ok(())
    }
}

接下来，我们将创建缓存集合。在创建此集合时，注意我们使用 Distance::Euclid 而不是 Distance::Cosine。这两者可以定义如下：

Distance::Cosine（或"余弦相似性"）测量两个向量指向相同方向的紧密程度。例如，如果我们在图表上绘制两个向量，位于 [2,1] 的向量会比 [-1, -2] 更接近 [1,1]。余弦相似性在文本分析中测量文档相似性时被广泛使用。

Distance::Euclid（或"欧氏距离"）测量两个向量相互之间的接近程度 - 即从 A 到 B 的距离，其中 A 和 B 是图表上的两个点。这里我们不是试图确定相似性，而是想确定某个东西在大多数情况下或完全相同。

impl RagSystem {
    async fn create_cache_collection(&self) -> Result<()> {
        self.qdrant_client
            .create_collection(&CreateCollection {
                collection_name: CACHE_COLLECTION_NAME.to_string(),
                vectors_config: Some(VectorsConfig {
                    config: Some(Config::Params(VectorParams {
                        size: 1536,
                        distance: Distance::Euclid.into(),
                        hnsw_config: None,
                        quantization_config: None,
                        on_disk: None,
                        ..Default::default()
                    })),
                }),
                ..Default::default()
            })
            .await?;

Ok(())
    }
}

接下来，我们需要从文件输入创建 embedding - 以 CSV 文件为例。为此，我们需要执行以下操作：

• 读取文件输入并将其解析为字符串（std::fs::read_to_string() 自动解析为 String 类型）

• 将文件内容分块为适当的数量（这里为了说明，我们简单地按行处理）

• 批量嵌入 embedding 并将其添加到 Qdrant

这里我们使用 async-openai 库来创建 embedding - 但如果你不想使用 OpenAI，你总是可以使用 fastembed-rs 作为替代方案，或者选择任何其他允许创建 embedding 的 crate。

use std::path::PathBuf;
use async_openai::types::{CreateEmbeddingRequest, EmbeddingInput};
use anyhow::Result;

impl RagSystem {
    async fn embed_and_upsert_csv_file(&self, file_path: PathBuf) -> Result<()> {
        let file_contents = std::fs::read_to_string(&file_path)?;

// 注意这里我们跳过 1 因为 CSV 文件通常有标题
        // 如果你没有任何标题，你可以删除它
        let chunked_file_contents: Vec<String> =
            file_contents.lines().skip(1).map(|x| x.to_owned()).collect();

let embedding_request = CreateEmbeddingRequest {
            model: "text-embedding-ada-002".to_string(),
            input: EmbeddingInput::StringArray(chunked_file_contents.to_owned()),
            encoding_format: None, // 默认为 f32
            user: None,
            dimensions: Some(1536),
        };

let embeddings = Embeddings::new(&self.openai_client)
            .create(embedding_request)
            .await?;

if embeddings.data.is_empty() {
            return Err(anyhow::anyhow!(
                "There were no embeddings returned by OpenAI!"
            ));
        }

let embeddings_vec: Vec<Vec<f32>> =
            embeddings.data.into_iter().map(|x| x.embedding).collect();

// 注意我们稍后创建 upsert_embedding 函数
        for embedding in embeddings_vec {
            self.upsert_embedding(embedding, file_contents.clone())
                .await?;
        }

Ok(())
    }
}

嵌入提示

接下来，我们需要嵌入任何进一步的输入，以搜索任何匹配的嵌入。embed_prompt 函数看起来与我们 embed_and_upsert_csv_file 函数的嵌入部分非常相似。但是，它会返回一个 Vec<f32>，因为我们之后想用它来搜索我们的集合。

impl RagSystem {
    pub async fn embed_prompt(&self, prompt: &str) -> Result<Vec<f32>> {
        let embedding_request = CreateEmbeddingRequest {
            model: "text-embedding-ada-002".to_string(),
            input: EmbeddingInput::String(prompt.to_owned()),
            encoding_format: None, // defaults to f32
            user: None,
            dimensions: Some(1536),
        };

let embeddings = Embeddings::new(&self.openai_client)
            .create(embedding_request)
            .await?;

if embeddings.data.is_empty() {
            return Err(anyhow::anyhow!(
                "There were no embeddings returned by OpenAI!"
            ));
        }

Ok(embeddings.data.into_iter().next().unwrap().embedding)
    }
}

更新（插入）嵌入

一旦我们创建了嵌入，我们将创建一个方法来向 Qdrant 中添加嵌入，称为 upsert_embedding。这将处理为我们的嵌入创建有效负载并将其插入数据库。一旦添加到集合中，我们稍后可以搜索我们的集合，并获得与嵌入相关的 JSON 有效负载！

该函数如下所示：

use qdrant_client::prelude::PointStruct;

impl RAGSystem {
    async fn upsert_embedding(&self, embedding: Vec<f32>, file_contents: String) -> Result<()> {
        let payload = serde_json::json!({
           "document": file_contents
        })
        .try_into()
        .map_err(|x| anyhow::anyhow!("Ran into an error when converting the payload: {x}"))?;

let points = vec![PointStruct::new(
            uuid::Uuid::new_v4().to_string(),
            embedding,
            payload,
        )];

self.qdrant_client
            .upsert_points(REGULAR_COLLECTION_NAME.to_owned(), None, points, None)
            .await?;

Ok(())
    }
}

这里，我们使用 uuid::Uuid 作为嵌入的唯一标识符。你也可以通过拥有一个随着每个嵌入而增加的 u64 计数器来实现相同的功能。但是，你需要确保不要意外覆盖你自己的嵌入！使用与集合中现有嵌入相同的 ID 插入新嵌入将覆盖该嵌入。

当然，我们还需要创建一个方法来向缓存添加内容。请注意，这里的有效负载是不同的。我们使用 answer 而不是 document 有效负载字段，因为有效负载将包含对问题的预生成答案。

impl RagSystem {
    pub async fn add_to_cache(&self, embedding: Vec<f32>, answer: String) -> Result<()> {
        let payload = serde_json::json!({
           "answer": answer
        })
        .try_into()
        .map_err(|x| anyhow::anyhow!("Ran into an error when converting the payload: {x}"))?;

let points = vec![PointStruct::new(
            uuid::Uuid::new_v4().to_string(),
            embedding,
            payload,
        )];

self.qdrant_client
            .upsert_points(CACHE_COLLECTION_NAME.to_owned(), None, points, None)
            .await?;

Ok(())
    }
}

搜索 Qdrant 集合

既然我们在 Qdrant 中创建了可以搜索的内容，我们需要为 VectorDB 实现一些方法。我们将其分为两个方法：

• search_regular_collection
• search_cache_collection

在搜索嵌入时，我们应该尝试使用 search_cache_collection 搜索我们的语义缓存 - 如果找不到任何内容，我们应该使用常规的 search_regular_collection 方法来获取文档，用它提示 OpenAI，然后返回结果。

为了使我们的方法更具容错能力，我们在结果上使用了 .into_iter().next()。这试图通过仅遍历向量中的第一个项来查找向量中的第一个项。这有效是因为我们只在寻找一个单一的嵌入，但你可以根据需要增加或减少限制。

找到匹配后，我们需要从与嵌入匹配相关的 JSON 有效负载中获取文档密钥并返回它。我们稍后会在我们的 RAG 提示中将其用作上下文！

use qdrant_client::qdrant::{
    with_payload_selector::SelectorOptions, SearchPoints, WithPayloadSelector
};

impl RagSystem {
        pub async fn search(&self, embedding: Vec<f32>) -> Result<String> {
        let payload_selector = WithPayloadSelector {
            selector_options: Some(SelectorOptions::Enable(true)),
        };

let search_points = SearchPoints {
            collection_name: REGULAR_COLLECTION_NAME.to_owned(),
            vector: embedding,
            limit: 1,
            with_payload: Some(payload_selector),
            score_threshold: Some(0.35f32),
            ..Default::default()
        };

let search_result = self
            .qdrant_client
            .search_points(&search_points)
            .await
            .inspect_err(|x| println!("An error occurred while searching for points: {x}"))
            .unwrap();

let result = search_result.result.into_iter().next();

let Some(result) = result else {
            return Err(anyhow::anyhow!("There's nothing matching."));
        };

Ok(result.payload.get("document").unwrap().to_string())
    }
}

当然，你也可能想实现一个函数来搜索你的缓存集合。请注意，虽然函数大多数相同，但我们从有效负载中获取 answer 字段而不是 document 以获得语义。

impl RagSystem {
    pub async fn search_cache(&self, embedding: Vec<f32>) -> Result<String> {
        let payload_selector = WithPayloadSelector {
            selector_options: Some(SelectorOptions::Enable(true)),
        };

let search_points = SearchPoints {
            collection_name: CACHE_COLLECTION_NAME.to_owned(),
            vector: embedding,
            limit: 1,
            with_payload: Some(payload_selector),
            ..Default::default()
        };

let search_result = self
            .qdrant_client
            .search_points(&search_points)
            .await
            .inspect_err(|x| println!("An error occurred while searching for points: {x}"))?;

let result = search_result.result.into_iter().next();

let Some(result) = result else {
            return Err(anyhow::anyhow!("There's nothing matching."));
        };

Ok(result.payload.get("answer").unwrap().to_string())
    }
}

当然，既然一切都完成了，最后要做的是提示！在这里，你可以看到下面我们生成一个提示，基本上由我们想要的提示以及提供的上下文组成。然后我们从 OpenAI 获取第一个结果并返回消息内容。

use async_openai::types::{
    ChatCompletionRequestMessage, ChatCompletionRequestSystemMessageArgs,
    ChatCompletionRequestUserMessageArgs, CreateChatCompletionRequestArgs
};

impl RagSystem {
    pub async fn prompt(&self, prompt: &str, context: &str) -> Result<String> {
        let input = format!(
            "{prompt}

Provided context:
            {context}
            "
        );

let res = self
            .openai_client
            .chat()
            .create(
                CreateChatCompletionRequestArgs::default()
                    .model("gpt-4o")
                    .messages(vec![
                        ChatCompletionRequestMessage::User(
                            ChatCompletionRequestUserMessageArgs::default()
                                .content(input)
                                .build()?,
                        ),
                    ])
                    .build()?,
            )
            .await
            .map(|res| {
                // We extract the first result
                match res.choices[0].message.content.clone() {
                    Some(res) => Ok(res),
                    None => Err(anyhow::anyhow!("There was no result from OpenAI")),
                }
            })??;

println!("Retrieved result from prompt: {res}");

Ok(res)
    }
}

在 Rust Web 服务中使用 Qdrant

让我们快速查看一个真实世界的例子。下面是 Axum 框架的一个 HTTP 端点，它将我们的 RAGSystem 作为应用程序状态。它将嵌入提示并尝试搜索缓存。如果没有结果，它会在常规集合中搜索匹配。结果文档有效负载被添加到增强提示中，问题和答案被添加到缓存中。最后，一个响应从端点返回。

use axum::{Json, extract::State, response::IntoResponse, http::StatusCode};
use serde::Deserialize;

#[derive(Deserialize)]
struct Prompt {
    prompt: String,
}

async fn prompt(
    State(state): State<RAGSystem>,
    Json(prompt): Json<Prompt>,
) -> Result<impl IntoResponse, impl IntoResponse> {
    let embedding = match state.embed_prompt(&prompt.prompt).await {
        Ok(embedding) => embedding,
        Err(e) => {
            return Err((
                StatusCode::INTERNAL_SERVER_ERROR,
                format!("An error occurred while embedding the prompt: {e}"),
            ))
        }
    };

if let Ok(answer) = state.search_cache(embedding.clone()).await {
        return Ok(answer);
    }

let search_result = match state.search(embedding.clone()).await {
        Ok(res) => res,
        Err(e) => {
            return Err((
                StatusCode::INTERNAL_SERVER_ERROR,
                format!("An error occurred while prompting: {e}"),
            ))
        }
    };

let llm_response = match state.prompt(&prompt.prompt, &search_result).await {
        Ok(prompt_result) => prompt_result,
        Err(e) => {
            return Err((
                StatusCode::INTERNAL_SERVER_ERROR,
                format!("Something went wrong while prompting: {e}"),
            ))
        }
    };

if let Err(e) = state.add_to_cache(embedding, &llm_response).await {
        return Err((
            StatusCode::INTERNAL_SERVER_ERROR,
            format!("Something went wrong while adding item to the cache: {e}"),
        ));
    };

Ok(llm_response)
}

最后一步是设置主函数。注意我们为主函数添加了 shuttle_qdrant::Qdrant 注解，这允许在本地运行时通过 Docker 自动配置一个本地 Qdrant 实例。但在生产环境中，我们需要填写 cloud_url 和 api_key 密钥。

#[shuttle_runtime::main]
async fn main(
    #[shuttle_qdrant::Qdrant(
        cloud_url = "{secrets.QDRANT_URL}",
        api_key = "{secrets.QDRANT_API_KEY}"
    )]
    qdrant: QdrantClient,
    #[shuttle_runtime::Secrets] secrets: SecretStore,
) -> shuttle_axum::ShuttleAxum {
    secrets.into_iter().for_each(|x| env::set_var(x.0, x.1));

let rag = RAGSystem::new(qdrant);

let setup_required = true;

if setup_required {
        rag.create_regular_collection().await?;
        rag.create_cache_collection().await?;

rag.embed_csv_file("test.csv".into()).await?;
    }

let rtr = Router::new().route("/prompt", post(prompt)).with_state(rag);

Ok(rtr.into())
}

部署时，你只需使用 shuttle deploy（如果在有未提交更改的 Git 分支上，则加上 --ad 标志），然后等待部署完成。部署完成后，任何后续部署只需重新编译应用程序（以及任何新增的依赖项），速度会快得多。

扩展这个示例

想扩展这个示例吗？下面是几种方法。

为语义缓存使用更便宜的模型

虽然使用高性能模型很好，但我们特别想节省的是成本。我们可以使用更便宜的模型，并让它判断一个问题在语义上是否与另一个问题相同。以下是你可以使用的 prompt：

Are these two questions semantically the same? Answer either 'Yes' or 'No'. Do not answer with anything else. If you don't know the answer, say 'I don't know'.

Question 1: <question 1 goes here>
Question 2: <question 2 goes here>

更小的有效载荷索引

需要注意的是，虽然我们的示例有效，但你可能需要考虑有效载荷索引或与特定嵌入相关的关联数据。如果你为大文件中的每个嵌入都插入整个文件内容作为有效载荷，资源使用会迅速增加。你可以通过只为每个嵌入插入文件的相关部分来缓解这个问题（例如本例中，可能是相应的行）。

感谢阅读！通过使用语义缓存，我们可以创建一个性能更好的 RAG 系统，既节省时间也节省成本。