无服务器语义搜索

Andre Bogus

2023 年 7 月 12 日

你想在你的网站或在线应用中插入语义搜索功能吗？现在你可以做到，而且无需花费任何费用！在这个示例中，你将学习如何为你的非商业目的创建一个免费的原型搜索引擎。

所需材料

一套 Rust 工具链
cargo lambda （通过包管理器安装，下载二进制文件或运行 cargo install cargo-lambda）
AWS CLI
Qdrant 实例（免费层级可用）
选择一个你喜欢的 Embedding 提供商服务（请参阅我们的Embedding 文档。你可能会从 AI Grant 获得积分，Cohere 也提供一个速率限制的非商业免费层级）
AWS Lambda 账户（提供 12 个月的免费层级）

你将要构建什么

你将把 Embedding 提供商和 Qdrant 实例结合起来，构建一个精巧的语义搜索功能，通过一个小的 Lambda 函数调用这两个服务。

lambda integration diagram

现在让我们看看在连接它们之前如何使用每种材料。

Rust 和 cargo-lambda

你希望你的函数快速、精简且安全，因此使用 Rust 是显而易见的最佳选择。为了编译 Rust 代码以在 Lambda 函数中使用，构建了 cargo-lambda 子命令。cargo-lambda 可以将你的 Rust 代码打包到一个 zip 文件中，然后 AWS Lambda 可以在精简的 provided.al2 运行时上部署它。

要与 AWS Lambda 交互，你的 Rust 项目需要在 Cargo.toml 中包含以下依赖项

[dependencies]
tokio = { version = "1", features = ["macros"] }
lambda_http = { version = "0.8", default-features = false, features = ["apigw_http"] }
lambda_runtime = "0.8"

这提供了一个接口，包含启动 Lambda 运行时的入口点以及注册 HTTP 调用处理程序的方法。将以下代码片段放入 src/helloworld.rs

use lambda_http::{run, service_fn, Body, Error, Request, RequestExt, Response};

/// This is your callback function for responding to requests at your URL
async fn function_handler(_req: Request) -> Result<Response<Body>, Error> {
    Response::from_text("Hello, Lambda!")
}

#[tokio::main]
async fn main() {
    run(service_fn(function_handler)).await
}

你还可以使用闭包将其他参数绑定到你的函数处理程序（service_fn 调用将变为 service_fn(|req| function_handler(req, ...))）。此外，如果你想从请求中提取参数，可以使用 Request 方法（例如 query_string_parameters 或 query_string_parameters_ref）来实现。

将以下内容添加到你的 Cargo.toml 中以定义二进制文件

[[bin]]
name = "helloworld"
path = "src/helloworld.rs"

在 AWS 方面，你需要设置一个 Lambda 和 IAM 角色以与你的函数一起使用。

create lambda web page

选择你的函数名称，选择“在 Amazon Linux 2 上提供你自己的引导程序”。架构选择 arm64。你还需要激活一个函数 URL。这里你可以选择通过 IAM 保护它，或者将其保持开放，但请注意，开放的端点可以被任何人访问，如果流量过大，可能会产生费用。

默认情况下，这还会创建一个基本角色。要查找该角色，你可以进入函数概览

function overview

点击靠近“▸ Function overview”标题的“信息”链接，然后在左侧选择“权限”选项卡。

你会在执行角色下方直接找到“角色名称”。记下它以便稍后使用。

function overview

要测试你的“Hello, Lambda”服务是否工作，你可以编译并上传函数

$ export LAMBDA_FUNCTION_NAME=hello
$ export LAMBDA_ROLE=<role name from lambda web ui>
$ export LAMBDA_REGION=us-east-1
$ cargo lambda build --release --arm --bin helloworld --output-format zip
  Downloaded libc v0.2.137
# [..] output omitted for brevity
    Finished release [optimized] target(s) in 1m 27s
$ # Delete the old empty definition
$ aws lambda delete-function-url-config --region $LAMBDA_REGION --function-name $LAMBDA_FUNCTION_NAME
$ aws lambda delete-function --region $LAMBDA_REGION --function-name $LAMBDA_FUNCTION_NAME
$ # Upload the function
$ aws lambda create-function --function-name $LAMBDA_FUNCTION_NAME \
    --handler bootstrap \
    --architectures arm64 \
    --zip-file fileb://./target/lambda/helloworld/bootstrap.zip \
    --runtime provided.al2 \
    --region $LAMBDA_REGION \
    --role $LAMBDA_ROLE \
    --tracing-config Mode=Active
$ # Add the function URL
$ aws lambda add-permission \
    --function-name $LAMBDA_FUNCTION_NAME \
    --action lambda:InvokeFunctionUrl \
    --principal "*" \
    --function-url-auth-type "NONE" \
    --region $LAMBDA_REGION \
    --statement-id url
$ # Here for simplicity unauthenticated URL access. Beware!
$ aws lambda create-function-url-config \
    --function-name $LAMBDA_FUNCTION_NAME \
    --region $LAMBDA_REGION \
    --cors "AllowOrigins=*,AllowMethods=*,AllowHeaders=*" \
    --auth-type NONE

现在你可以前往你的函数概览并点击函数 URL。你应该会看到类似这样的内容

Hello, Lambda!

Bearer ！你已经在 Rust 中设置了一个 Lambda 函数。接着看下一个材料

Embedding

大多数提供商提供简单的 HTTPS GET 或 POST 接口，你可以使用 API 密钥进行调用，你需要将密钥放在认证头部中。如果你将其用于非商业目的，只需几次点击即可获得 Cohere 提供的速率限制试用密钥。前往他们的欢迎页面，注册后即可进入控制面板，其中有一个“API keys”菜单项，将带你到以下页面：Cohere 控制面板

在那里，你可以点击 API 密钥旁边的 ⎘ 符号将其复制到剪贴板。不要将你的 API 密钥放在代码中！而是从 Lambda 环境中设置的环境变量中读取它。这可以避免意外将你的密钥放入公共仓库。现在，获取 Embedding 所需的只是少量代码。首先，你需要通过 reqwest 扩展你的依赖项，并添加 anyhow 以便更轻松地进行错误处理

anyhow = "1.0"
reqwest =  { version = "0.11.18", default-features = false, features = ["json", "rustls-tls"] }
serde = "1.0"

现在，有了上面的 API 密钥，你就可以调用来获取 Embedding 向量了

use anyhow::Result;
use serde::Deserialize;
use reqwest::Client;

#[derive(Deserialize)]
struct CohereResponse { outputs: Vec<Vec<f32>> }

pub async fn embed(client: &Client, text: &str, api_key: &str) -> Result<Vec<Vec<f32>>> {
    let CohereResponse { outputs } = client
        .post("https://api.cohere.ai/embed")
        .header("Authorization", &format!("Bearer {api_key}"))
        .header("Content-Type", "application/json")
        .header("Cohere-Version", "2021-11-08")
        .body(format!("{{\"text\":[\"{text}\"],\"model\":\"small\"}}"))
        .send()
        .await?
        .json()
        .await?;
    Ok(outputs)
}

请注意，如果文本超出输入维度，这可能会返回多个向量。Cohere 的 small 模型具有 1024 个输出维度。

其他提供商也有类似的接口。有关更多信息，请查阅我们的Embedding 文档。看看获取 Embedding 只需这么少的代码？

顺便一提，写一个小测试来检查 Embedding 是否工作以及向量是否是预期的大小是个好主意

#[tokio::test]
async fn check_embedding() {
    // ignore this test if API_KEY isn't set
    let Ok(api_key) = &std::env::var("API_KEY") else { return; }
    let embedding = crate::embed("What is semantic search?", api_key).unwrap()[0];
    // Cohere's `small` model has 1024 output dimensions.
    assert_eq!(1024, embedding.len());
}

在运行此代码时，设置 API_KEY 环境变量，以检查 Embedding 是否工作。

Qdrant 搜索

现在你已经有了 Embedding，是时候将它们放入 Qdrant 了。你当然可以使用 curl 或 python 来设置你的集合并上传点，但既然你已经有了 Rust，并且包含一些获取 Embedding 的代码，你可以继续使用 Rust，并添加 qdrant-client。

use anyhow::Result;
use qdrant_client::prelude::*;
use qdrant_client::qdrant::{VectorsConfig, VectorParams};
use qdrant_client::qdrant::vectors_config::Config;
use std::collections::HashMap;

fn setup<'i>(
    embed_client: &reqwest::Client,
    embed_api_key: &str,
    qdrant_url: &str,
    api_key: Option<&str>,
    collection_name: &str,
    data: impl Iterator<Item = (&'i str, HashMap<String, Value>)>,
) -> Result<()> {
    let mut config = QdrantClientConfig::from_url(qdrant_url);
    config.api_key = api_key;
    let client = QdrantClient::new(Some(config))?;

    // create the collections
    if !client.has_collection(collection_name).await? {
        client
            .create_collection(&CreateCollection {
                collection_name: collection_name.into(),
                vectors_config: Some(VectorsConfig {
                    config: Some(Config::Params(VectorParams {
                        size: 1024, // output dimensions from above
                        distance: Distance::Cosine as i32,
                        ..Default::default()
                    })),
                }),
                ..Default::default()
            })
            .await?;
    }
    let mut id_counter = 0_u64;
    let points = data.map(|(text, payload)| {
        let id = std::mem::replace(&mut id_counter, *id_counter + 1);
        let vectors = Some(embed(embed_client, text, embed_api_key).unwrap());
        PointStruct { id, vectors, payload }
    }).collect();
    client.upsert_points(collection_name, points, None).await?;
    Ok(())
}

根据你是否想有效地过滤数据，你还可以添加一些索引。为了简洁起见，我在此省略了这部分。此外，这也没有实现分块（将数据分割成多个请求进行 upsert，以避免超时错误）。

添加一个合适的 main 方法，你就可以运行此代码插入点（或直接使用示例中的二进制文件）。请确保在 qdrant_url 中包含端口。

现在点已插入，你可以通过 Embedding 进行搜索

use anyhow::Result;
use qdrant_client::prelude::*;
pub async fn search(
    text: &str,
    collection_name: String,
    client: &Client,
    api_key: &str,
    qdrant: &QdrantClient,
) -> Result<Vec<ScoredPoint>> {
    Ok(qdrant.search_points(&SearchPoints {
        collection_name,
        limit: 5, // use what fits your use case here
        with_payload: Some(true.into()),
        vector: embed(client, text, api_key)?,
        ..Default::default()
    }).await?.result)
}

你还可以通过在 SearchPoints 中添加一个 filter: ... 字段进行过滤，并且你很可能想要进一步处理结果，但示例代码已经实现了这一点，所以如果你需要此功能，可以从那里开始。

将所有部分整合起来

现在你已经拥有了所有部分，是时候将它们连接起来了。复制并连接以上代码片段作为留给读者的练习。

你可能需要稍微扩展 main 方法，在启动时连接一次客户端，同时从环境中获取 API 密钥，这样就不需要将它们编译到代码中。为此，你可以在 Rust 代码中使用 std::env::var(_) 获取它们，并在 AWS 控制台中设置环境。

$ export QDRANT_URI=<qour Qdrant instance URI including port>
$ export QDRANT_API_KEY=<your Qdrant API key>
$ export COHERE_API_KEY=<your Cohere API key>
$ export COLLECTION_NAME=site-cohere
$ aws lambda update-function-configuration \
    --function-name $LAMBDA_FUNCTION_NAME \
    --environment "Variables={QDRANT_URI=$QDRANT_URI,\
        QDRANT_API_KEY=$QDRANT_API_KEY,COHERE_API_KEY=${COHERE_API_KEY},\
        COLLECTION_NAME=${COLLECTION_NAME}"`

总之，你最终会得到一个用于插入数据的命令行程序和一个 Lambda 函数。前者只需运行 cargo run 来设置集合。对于后者，你可以再次调用 cargo lambda 和 AWS 控制台。

$ export LAMBDA_FUNCTION_NAME=search
$ export LAMBDA_REGION=us-east-1
$ cargo lambda build --release --arm --output-format zip
  Downloaded libc v0.2.137
# [..] output omitted for brevity
    Finished release [optimized] target(s) in 1m 27s
$ # Update the function
$ aws lambda update-function-code --function-name $LAMBDA_FUNCTION_NAME \
     --zip-file fileb://./target/lambda/page-search/bootstrap.zip \
     --region $LAMBDA_REGION

讨论

Lambda 的工作原理是在 URL 被调用时启动你的函数，因此除非实际使用，否则不需要保留计算资源。这意味着第一次调用会有大约 1-2 秒的函数加载延迟，后续调用会更快。当然，调用 Embedding 提供商和 Qdrant 也会有延迟。另一方面，免费层级无需支付任何费用，所以你得到的肯定物超所值。对于许多用例来说，在一两秒内获得结果是可以接受的。

Rust 最大限度地减少了函数的开销，无论是在文件大小还是运行时方面。使用 Embedding 服务意味着你无需关心细节。知道 URL、API 密钥和 Embedding 大小就足够了。最后，有了 Lambda 和 Qdrant 的免费层级以及 Embedding 提供商的免费积分，唯一的成本就是你设置所有这些所需的时间。免费的东西谁能拒绝呢？

学习