Vertex AI RAG 引擎

从托管的 Vertex AI RAG 引擎语料库中检索有依据的上下文 — 无需 自行托管向量存储、嵌入提供商或摄取管道。


功能介绍

Vertex AI RAG Engine 是 Google Cloud 托管的 RAG 后端:您可以将 文档导入 RAG 语料库,平台则负责分块、嵌入和向量搜索。adk-ragvertex-rag 功能提供:

  • VertexRagEngineClient — 经 ADC 身份验证、只读的数据面客户端:获取/列出语料库、列出已导入的文件,以及检索上下文
  • VertexAiRagRetrievalTool — 将检索作为 adk_core::Tool,即 adk-python 的 VertexAiRagRetrieval 的 Rust 对应实现

**范围:**仅限检索。语料库创建和文件导入属于配置工作 — 请使用 Vertex AI 控制台RagCorpora/RagFiles 管理 APIs。


安装

[dependencies]
adk-rag = { version = "2.1.0", features = ["vertex-rag"] }

身份验证使用应用程序默认凭据:

gcloud auth application-default login

客户端

use adk_rag::vertex_rag::{RetrieveContextsRequest, VertexRagConfig, VertexRagEngineClient};

#[tokio::main]
async fn main() -> adk_core::Result<()> {
    // Or VertexRagConfig::from_env() reading GOOGLE_CLOUD_PROJECT / GOOGLE_CLOUD_LOCATION
    let config = VertexRagConfig::new("my-project", "us-central1");
    let client = VertexRagEngineClient::new_with_adc(config)?;

    // Verify the corpus exists and has imported files; fails with
    // actionable guidance when it is missing, empty, or in ERROR state.
    let corpus = client.ensure_corpus_ready("1234567890").await?;
    println!("corpus: {:?} ({:?} files)", corpus.display_name, corpus.rag_files_count);

    // Enumerate what's in the project and the corpus.
    let corpora = client.list_corpora().await?;
    let files = client.list_rag_files("1234567890").await?;
    println!("{} corpora, {} files", corpora.len(), files.len());

    // Retrieve the most relevant passages for a query.
    let request = RetrieveContextsRequest::new("what is the refund policy?", ["1234567890"])
        .similarity_top_k(5)
        .vector_distance_threshold(0.7);
    for context in client.retrieve_contexts(&request).await? {
        println!(
            "[{:.3}] {} — {}",
            context.score.unwrap_or_default(),
            context.source_display_name.as_deref().unwrap_or("<unknown>"),
            context.text.as_deref().unwrap_or(""),
        );
    }
    Ok(())
}

语料库可以以裸 ID 传递(根据客户端的项目和位置解析),也可以传递完整的 projects/*/locations/*/ragCorpora/* 资源名称。

注意:similarity_top_kvector_distance_threshold 保留 adk-python 的名称,但会通过当前线路路径发送 — query.ragRetrievalConfig.topKquery.ragRetrievalConfig.filter.vectorDistanceThreshold。已弃用的 v1beta1 拼写形式(query.similarityTopKvertexRagStore.vectorDistanceThreshold)已从 v1 中移除,且永远不会 生成。vector_similarity_threshold 是过滤器的另一个互斥分支。


检索工具

VertexAiRagRetrievalTool 接受一个必需的 query 字符串,并返回 一个包含 {text, sourceUri, sourceDisplayName, score} 对象的 JSON 数组。它 声明自身为只读且线程安全,因此 ToolExecutionStrategy::Auto 可以将其与其他读取操作并行调度。

use std::sync::Arc;
use adk_agent::LlmAgentBuilder;
use adk_model::GeminiModel;
use adk_rag::vertex_rag::{VertexAiRagRetrievalTool, VertexRagConfig, VertexRagEngineClient};

fn main() -> anyhow::Result<()> {
    let config = VertexRagConfig::new("my-project", "us-central1");
    let client = Arc::new(VertexRagEngineClient::new_with_adc(config)?);

    let retrieval = VertexAiRagRetrievalTool::new(client, vec!["1234567890".into()])
        .similarity_top_k(5)
        .vector_distance_threshold(0.7);

    let api_key = std::env::var("GOOGLE_API_KEY")?;
    let agent = LlmAgentBuilder::new("rag-assistant")
        .description("Answers questions grounded in a Vertex AI RAG Engine corpus")
        .model(Arc::new(GeminiModel::new(&api_key, "gemini-3.7-flash")?))
        .tool(Arc::new(retrieval))
        .instruction(
            "Answer using the vertex_rag_retrieval tool. Retrieve first, then \
             answer strictly from the retrieved passages, citing sourceDisplayName.",
        )
        .build()?;
    let _ = agent;
    Ok(())
}

请参阅 examples/vertex_rag 获取完整的可运行代理:

cargo run --manifest-path examples/vertex_rag/Cargo.toml

API 参考

操作端点返回值
get_corpus(corpus)GET v1beta1/{corpus}RagCorpus(404 会转换为可操作的未找到错误)
ensure_corpus_ready(corpus)GET v1beta1/{corpus}RagCorpus;缺失、为空或处于 ERROR 状态时返回错误
list_corpora()GET v1beta1/{parent}/ragCorporaVec<RagCorpus>,后续分页
list_rag_files(corpus)GET v1beta1/{corpus}/ragFilesVec<RagFile>,后续分页
retrieve_contexts(&request)POST v1beta1/{parent}:retrieveContextsVec<RagContext>

响应类型采用宽松反序列化——每个字段都是可选的,未知字段会被忽略,因此新增的服务器字段不会导致解析失败。错误包含组件 Memory(检索属于记忆域;AdkError 没有专用的 RAG 组件)、提供商 vertex_ai 以及机器可读的 rag.vertex.* 代码。


环境变量

变量使用者描述
GOOGLE_CLOUD_PROJECTVertexRagConfig::from_env拥有语料库的项目
GOOGLE_CLOUD_LOCATIONVertexRagConfig::from_env区域,例如 us-central1
VERTEX_RAG_CORPUS示例 / 实时测试语料库 ID 或完整资源名称

改用自托管的 RAG?

对于自行运行的管道——支持可插拔的分块器、嵌入提供商和向量存储(Qdrant、LanceDB、pgvector、SurrealDB)——请参阅 RAG