使用 Rust、OpenAI 和 Qdrant 構建 Agentic RAG

在本文中，我們將討論如何使用 Rust 構建代理 RAG 工作流！我們將構建一個代理，它可以獲取 CSV 檔案、對其進行解析並將其嵌入到 Qdrant 中，以及從 Qdrant 中檢索相關嵌入以回答使用者有關 CSV 檔案內容的問題。

有興趣部署還是隻想看看最終的程式碼是什麼樣的？您可以在這裡找到儲存庫。

什麼是 Agentic RAG？
Agentic RAG（即 Agentic Retrieval Augmented Generation）是將 AI 代理與 RAG 相結合的概念，以便能夠生成比代理工作流程更能適應特定用例的工作流程。

本質上，此工作流程與常規代理工作流程之間的區別在於，每個代理都可以單獨訪問來自向量資料庫的嵌入，以便能夠檢索上下文相關的資料 - 從而在 AI 代理工作流程中獲得更準確的答案！

入門
首先，使用cargo shuttle init建立一個新專案。

接下來，我們將使用 shell 程式碼片段新增所需的依賴項：

cargo add anyhow
cargo add async-openai
cargo add qdrant-client
cargo add serde -F derive
cargo add serde-json
cargo add shuttle-qdrant
cargo add uuid -F v4

我們還需要確保擁有 Qdrant URL 和 API 金鑰以及 OpenAI API 金鑰。Shuttle 透過SecretStore主函式中的宏使用環境變數，並且可以儲存在Secrets.toml檔案中：
OPENAI_API_KEY = ""

接下來，我們將更新主函式以包含 Qdrant 宏和 secrets 宏。我們將遍歷每個 secret 並將其設定為環境變數 - 這使我們能夠全域性使用 secret，而無需引用變數SecretStore：

#[shuttle_runtime::main]
async fn main(
    #[shuttle_qdrant::Qdrant] qdrant_client: QdrantClient,
    #[shuttle_runtime::Secrets] secrets: SecretStore,
) -> shuttle_axum::ShuttleAxum {
    secrets.into_iter().for_each(|x| {
        set_var(x.0, x.1);
    });

    let router = Router::new()
        .route(<font>"/", get(hello_world));

    Ok(router.into())
}

構建代理 RAG 工作流程
設定我們的代理
代理本身非常簡單：它擁有一個 OpenAI 客戶端和一個 Qdrant 客戶端，以便能夠搜尋相關的文件嵌入。您還可以在此處新增其他欄位，具體取決於您的代理所需的功能。

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

pub struct MyAgent {
    openai_client: OpenAIClient<OpenAIConfig>,
    qdrant_client: QdrantClient,
}

接下來，我們將建立一個輔助方法來建立代理，以及一個系統訊息，稍後我們會將其輸入到模型提示中。

static SYSTEM_MESSAGE: &str = <font>"
        You are a world-class data analyst, specialising in analysing comma-delimited CSV files.

        Your job is to analyse some CSV snippets and determine what the results are for the question that the user is asking.

        You should aim to be concise. If you don't know something, don't make it up but say 'I don't know.'.
"

impl MyAgent {
    pub fn new(qdrant_client: QdrantClient) -> Self {
        let api_key = std::env::var("OPENAI_API_KEY").unwrap();
        let config = OpenAIConfig::new().with_api_key(api_key);

        let openai_client = OpenAIClient::with_config(config);

        Self {
            openai_client,
            qdrant_client,
        }
    }
}

檔案解析並嵌入到 Qdrant
接下來，我們將實現一個File用於 CSV 檔案解析的結構 - 它應該能夠將檔案路徑、內容以及行儲存為Vec<String>（字串陣列，或更準確地說是字串向量）。我們將行儲存為的原因有幾個Vec<String>：

更小的區塊可以提高檢索的準確性，這是 RAG 面臨的最大挑戰之一。檢索錯誤或不準確的文件會嚴重影響準確性。
提高檢索準確性可增強上下文相關性 - 這對於需要特定問題的複雜查詢非常重要。
處理和索引較小的塊

pub struct File {
    pub path: String,
    pub contents: String,
    pub rows: Vec<String>,
}

impl File {
    pub fn new(path: PathBuf) -> Result<Self> {
        let contents = std::fs::read_to_string(&path)?;

        let path_as_str = format!(<font>"{}", path.display());

        let rows = contents
            .lines()
            .map(|x| x.to_owned())
            .collect::<Vec<String>>();

        Ok(Self {
            path: path_as_str,
            contents,
            rows
        })
    }
}

雖然上述解析方法很有用（將所有行收集到中Vec<String>），但請注意，這是一種簡單的實現。根據您的 CSV 檔案的分隔方式和/或是否有需要清理的髒資料，您可能需要準備資料以使其已經準備就緒，或者包括某種形式的資料清理或驗證。這方面的一些示例可能是：

unicode-segmentation-用於拆分句子的庫箱
csv_log_cleaner-用於清理 CSV 的二進位制包
validator-用於驗證結構/列舉欄位的庫包

接下來，我們將回到我們的代理並實現一種將文件嵌入到 Qdrant 中的方法，該方法將採用File我們定義的結構。

為此，我們需要執行以下操作：

獲取我們之前建立的行並將其新增為嵌入請求的輸入。
建立嵌入（使用 openAI）並建立有效載荷，用於與 Qdrant 中的嵌入一起儲存。

請注意，雖然我們使用了uuid::Uuid唯一儲存，但您也可以透過在結構中新增數字計數器並在插入嵌入後將其加 1 來輕鬆使用數字。
假設沒有錯誤，返回Ok(())

use async_openai::types::{ CreateEmbeddingRequest, EmbeddingInput };
use async_openai::Embeddings;
use qdrant_client::prelude::{Payload, PointStruct};

static COLLECTION: &str = <font>"my-collection";

// text-embedding-ada-002 is the model name from OpenAI that deals with embeddings<i>
static EMBED_MODEL: &str = "text-embedding-ada-002";

impl MyAgent {
pub async fn embed_document(&self, file: File) -> Result<()> {
        if file.rows.is_empty() {
            return Err(anyhow::anyhow!("There's no rows to embed!"));
        }

        let request = CreateEmbeddingRequest {
            model: EMBED_MODEL.to_string(),
            input: EmbeddingInput::StringArray(file.rows.clone()),
            user: None,
            dimensions: Some(1536),
            ..Default::default()
        };

        let embeddings_result = Embeddings::new(&self.openai_client).create(request).await?;

        for embedding in embeddings_result.data {
            let payload: Payload = serde_json::json!({
                "id": file.path.clone(),
                "content": file.contents,
                "rows": file.rows
            })
            .try_into()
            .unwrap();

            println!("Embedded: {}", file.path);

            let vec = embedding.embedding;

            let points = vec![PointStruct::new(
                uuid::Uuid::new_v4().to_string(),
                vec,
                payload,
            )];
            self.qdrant_client
                .upsert_points(COLLECTION, None, points, None)
                .await?;
        }
        Ok(())
    }
}

文件搜尋
現在我們已經嵌入了文件，我們需要一種方法來檢查我們的嵌入是否與使用者給出的提示具有上下文相關性。為此，我們將建立一個search_document執行以下操作的函式：

使用嵌入提示CreateEmbeddingRequest並從結果中獲取嵌入。我們將在文件搜尋中使用此嵌入。因為我們只在這裡新增了一個句子來嵌入（提示），所以它只會返回一個句子 - 因此我們可以從向量中建立一個迭代器並嘗試找到第一個結果。
透過結構體為我們的文件搜尋建立一個引數列表SearchPoints（見下文）。在這裡，我們需要設定集合名稱、我們要搜尋的向量（即輸入）、如果有匹配項，我們希望返回多少個結果，以及有效載荷選擇器。
在資料庫中搜尋結果 - 如果沒有結果，則返回錯誤；如果有結果，則返回結果。

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

impl MyAgent {
    async fn search_document(&self, prompt: String) -> Result<String> {
        let request = CreateEmbeddingRequest {
            model: EMBED_MODEL.to_string(),
            input: EmbeddingInput::String(prompt),
            user: None,
            dimensions: Some(1536),
            ..Default::default()
        };

        let embeddings_result = Embeddings::new(&self.openai_client).create(request).await?;

        let embedding = &embeddings_result.data.first().unwrap().embedding;

        let payload_selector = WithPayloadSelector {
            selector_options: Some(SelectorOptions::Enable(true)),
        };

        <font>// set parameters for search<i>
        let search_points = SearchPoints {
            collection_name: COLLECTION.to_string(),
            vector: embedding.to_owned(),
            limit: 1,
            with_payload: Some(payload_selector),
            ..Default::default()
        };

        // if the search is successful<i>
        // attempt to iterate through the results vector and find a result<i>
        let search_result = self.qdrant_client.search_points(&search_points).await?;
        let result = search_result.result.into_iter().next();

        match result {
            Some(res) => Ok(res.payload.get("contents").unwrap().to_string()),
            None => Err(anyhow::anyhow!("There were no results that matched :(")),
        }
    }
}

現在我們已經設定好了有效使用代理所需的一切，我們可以設定一個提示功能！

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

static PROMPT_MODEL: &str = <font>"gpt-4o";

impl MyAgent {
    pub async fn prompt(&self, prompt: &str) -> anyhow::Result<String> {
        let context = self.search_document(prompt.to_owned()).await?;
        let input = format!(
            "{prompt}

            Provided context:
            {}
            ",
            context // this is the payload from Qdrant<i>
        );

        let res = self
            .openai_client
            .chat()
            .create(
                CreateChatCompletionRequestArgs::default()
                    .model(PROMPT_MODEL)
                    .messages(vec![
                        //First we add the system message to define what the Agent does<i>
                        ChatCompletionRequestMessage::System(
                            ChatCompletionRequestSystemMessageArgs::default()
                                .content(SYSTEM_MESSAGE)
                                .build()?,
                        ),
                        //Then we add our prompt<i>
                        ChatCompletionRequestMessage::User(
                            ChatCompletionRequestUserMessageArgs::default()
                                .content(input)
                                .build()?,
                        ),
                    ])
                    .build()?,
            )
            .await
            .map(|res| {
                //We extract the first one<i>
                res.choices[0].message.content.clone().unwrap()
            })?;

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

        Ok(res)
    }
}

將代理連線到我們的 Web 服務
因為我們將代理邏輯與 Web 服務邏輯分開了，所以我們只需要將各個部分連線在一起就可以完成了！

首先，我們將建立幾個結構 -Prompt採用 JSON 提示的結構和AppState充當 Axum Web 伺服器中的共享應用程式狀態的函式。

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

#[derive(Clone)]
pub struct AppState {
    agent: MyAgent,
}

我們還將在這裡介紹我們的提示處理程式端點：

async fn prompt(
    State(state): State<AppState>,
    Json(json): Json<Prompt>,
) -> Result<impl IntoResponse> {
    let prompt_response = state.agent.prompt(&json.prompt).await?;

    Ok((StatusCode::OK, prompt_response))
}

然後我們需要在主函式中解析我們的 CSV 檔案，建立AppState並嵌入 CSV，以及設定我們的路由器：

#[shuttle_runtime::main]
async fn main(
    #[shuttle_qdrant::Qdrant] qdrant_client: QdrantClient,
    #[shuttle_runtime::Secrets] secrets: SecretStore,
) -> shuttle_axum::ShuttleAxum {
    secrets.into_iter().for_each(|x| {
        set_var(x.0, x.1);
    });

    <font>// note that this already assumes you have a file called "test.csv"<i>
    // in your project root<i>
    let file = File::new("test.csv".into())?;

    let state = AppState {
        agent: MyAgent::new(qdrant_client),
    };

    state.agent.embed_document(file).await?;

    let router = Router::new()
        .route("/", get(hello_world))
        .route("/prompt", post(prompt))
        .with_state(state);

    Ok(router.into())
}

部署
要部署，您需要做的就是使用cargo shuttle deploy（--ad如果在具有未提交更改的 Git 分支上則使用標誌），然後坐下來觀看奇蹟發生！

總結
透過結合 AI 代理和 RAG 的力量，我們可以建立強大的工作流程，以滿足許多不同的用例。藉助 Rust，我們可以利用效能優勢，以較低的記憶體佔用安全地執行我們的工作流程。

使用 Rust、OpenAI 和 Qdrant 構建 Agentic RAG

相關文章