RAG 引擎跨语料库检索

重要提示：创建 RAG 语料库时，请设置语料库 description。description 创建后无法修改。高质量的说明对于有效的跨语料库检索至关重要，因为系统会根据这些说明为您的查询选择合适的语料库。

本页介绍了两个 API：AsyncRetrieveContexts 和 AskContexts，它们支持 RAG 跨语料库检索，由后端的 Agentic Retrieval 提供支持。

AsyncRetrieveContexts API：这是一个异步 API，可让用户从多个 RAG 管理的语料库中检索相关上下文。这是一个长时间运行的 API。用户可以使用操作 ID 获取查询的状态和结果。

AskContexts API：这是一个同步 API，可通过搜索多个 RAG 管理的语料库直接生成查询的答案。

重要提示：如需使用此功能，您必须向 RAG Engine 服务账号 service-[Your project's automatically generated project number]@gcp-sa-vertex-rag.iam.gserviceaccount.com 授予项目中的 Vertex AI User 角色。

重要提示：此功能仅支持 us-central1。

系统架构和语料库选择

本部分介绍了 RAG 跨语料库检索系统的架构，以及该系统如何选择合适的语料库进行检索。

系统架构

跨语料库检索系统采用代理方法来协调跨多个语料库的检索。该架构的概要概述包括以下组件：

  • 编排器/路由器：接收用户查询并协调检索流程。
  • 规划代理：分析查询，并根据语料库的 descriptions 确定哪些语料库最相关。
  • 检索引擎：对所选 RAG 语料库执行语义搜索。
  • 推理代理：评估检索到的上下文是否足以回答查询。如果不足，则生成反馈，以根据“足够的上下文感知”(SCA) 触发另一个检索循环。如需了解详情，请参阅 SCA 博文和研究论文。
  • LLM 生成器（适用于 AskContexts）：综合检索到的上下文以生成最终答案。
跨语料库检索架构

语料库的选择方式

当查询在多个 RAG 管理的语料库中执行时，系统必须确定哪个语料库包含回答查询所需的信息。这是通过智能体检索实现的：

  1. 语料库映射：系统会维护可用 RAG 语料库的映射以及它们的技术说明（在 RAG 语料库的 description 字段中指定）。
  2. 语义匹配：战略规划智能体（由 Gemini 模型提供支持）会根据可用语料库的 descriptions 评估用户查询。
  3. 有针对性的路由：规划器会生成一个检索计划，将查询的特定部分映射到最相关的语料库，确保检索有重点且高效，而不是盲目地搜索所有语料库。

代码示例

以下代码示例演示了如何使用 RAG 跨语料库检索 API。

AsyncRetrieveContexts API

此代码示例演示了如何从多个 RAG 语料库异步检索上下文。

Python

from vertexai.preview import rag
import vertexai

PROJECT_ID = "PROJECT_ID"
LOCATION = "LOCATION"
RAG_CORPUS_1_ID = "RAG_CORPUS_1_ID"
RAG_CORPUS_2_ID = "RAG_CORPUS_2_ID"

# Initialize API once per session
vertexai.init(project=PROJECT_ID, location=LOCATION)

# Async retrieve contexts from multiple corpora
# This can be run in an environment like a Colab notebook.
response = await rag.async_retrieve_contexts(
    text="Why is the sky blue?",
    rag_resources=[
        rag.RagResource(
            rag_corpus=f"projects/{PROJECT_ID}/locations/{LOCATION}/ragCorpora/{RAG_CORPUS_1_ID}",
        ),
        rag.RagResource(
            rag_corpus=f"projects/{PROJECT_ID}/locations/{LOCATION}/ragCorpora/{RAG_CORPUS_2_ID}",
        ),
    ],
)
print(response)

REST

PROJECT_ID="PROJECT_ID"
LOCATION="LOCATION"
RAG_CORPUS_1_ID="RAG_CORPUS_1_ID"
RAG_CORPUS_2_ID="RAG_CORPUS_2_ID"
QUERY="Why is the sky blue?"

curl -X POST \
  -H "Authorization: Bearer $(gcloud auth print-access-token)" \
  -H "Content-Type: application/json" \
  "https://${LOCATION}-aiplatform.googleapis.com/v1beta1/projects/${PROJECT_ID}/locations/${LOCATION}:asyncRetrieveContexts" \
  -d '{
    "query": {
      "text": "'"$QUERY"'"
    },
    "tools": {
      "retrieval": {
        "disable_attribution": false,
        "vertex_rag_store": {
          "rag_resources": {
            "rag_corpus": "projects/'"${PROJECT_ID}"'/locations/'"${LOCATION}"'/ragCorpora/'"${RAG_CORPUS_1_ID}"'"
          },
          "rag_resources": {
            "rag_corpus": "projects/'"${PROJECT_ID}"'/locations/'"${LOCATION}"'/ragCorpora/'"${RAG_CORPUS_2_ID}"'"
          }
        }
      }
    }
  }'

获取操作结果

由于 AsyncRetrieveContexts 是长时间运行的操作，因此您可以使用 GetOperations API 通过操作 ID 轮询结果。

Python

from google.cloud import aiplatform_v1beta1

PROJECT_ID = "PROJECT_ID"
LOCATION = "LOCATION"
OPERATION_ID = "OPERATION_ID"

client = aiplatform_v1beta1.VertexRagServiceClient(
    client_options={"api_endpoint": f"{LOCATION}-aiplatform.googleapis.com"}
)

operation = client.get_operation(
    request={"name": f"projects/{PROJECT_ID}/locations/{LOCATION}/operations/{OPERATION_ID}"}
)

if operation.done:
    print(operation.response)
else:
    print("Operation is still running")

REST

PROJECT_ID="PROJECT_ID"
LOCATION="LOCATION"
OPERATION_ID="OPERATION_ID"

curl -X GET \
  -H "Authorization: Bearer $(gcloud auth print-access-token)" \
  "https://${LOCATION}-aiplatform.googleapis.com/v1beta1/projects/${PROJECT_ID}/locations/${LOCATION}/operations/${OPERATION_ID}"

AskContexts API

此代码示例演示了如何通过同步搜索多个 RAG 语料库来直接生成答案。

Python

from vertexai.preview import rag
import vertexai

PROJECT_ID = "PROJECT_ID"
LOCATION = "LOCATION"
RAG_CORPUS_1_ID = "RAG_CORPUS_1_ID"
RAG_CORPUS_2_ID = "RAG_CORPUS_2_ID"

# Initialize API once per session
vertexai.init(project=PROJECT_ID, location=LOCATION)

# Ask contexts from multiple corpora
response = rag.ask_contexts(
    text="Why is the sky blue?",
    rag_resources=[
        rag.RagResource(
            rag_corpus=f"projects/{PROJECT_ID}/locations/{LOCATION}/ragCorpora/{RAG_CORPUS_1_ID}",
        ),
        rag.RagResource(
            rag_corpus=f"projects/{PROJECT_ID}/locations/{LOCATION}/ragCorpora/{RAG_CORPUS_2_ID}",
        ),
    ],
)
print(response)

REST

PROJECT_ID="PROJECT_ID"
LOCATION="LOCATION"
RAG_CORPUS_1_ID="RAG_CORPUS_1_ID"
RAG_CORPUS_2_ID="RAG_CORPUS_2_ID"
QUERY="Why is the sky blue?"

curl -X POST \
  -H "Authorization: Bearer $(gcloud auth print-access-token)" \
  -H "Content-Type: application/json" \
  "https://${LOCATION}-aiplatform.googleapis.com/v1beta1/projects/${PROJECT_ID}/locations/${LOCATION}:askContexts" \
  -d '{
    "query": {
      "text": "'"$QUERY"'"
    },
    "tools": {
      "retrieval": {
        "disable_attribution": false,
        "vertex_rag_store": {
          "rag_resources": {
            "rag_corpus": "projects/'"${PROJECT_ID}"'/locations/'"${LOCATION}"'/ragCorpora/'"${RAG_CORPUS_1_ID}"'"
          },
          "rag_resources": {
            "rag_corpus": "projects/'"${PROJECT_ID}"'/locations/'"${LOCATION}"'/ragCorpora/'"${RAG_CORPUS_2_ID}"'"
          }
        }
      }
    }
  }'

后续步骤