使用 Vertex AI RAG 引擎将 Vertex AI Search 用作检索后端

本页介绍了 Vertex AI Search 与 Vertex AI RAG Engine 的集成。

Vertex AI Search 提供了一种解决方案，可在 Vertex AI RAG 应用中检索和管理数据。通过使用 Vertex AI Search 作为检索后端，您可以提高性能、可伸缩性和集成便利性。

  • 增强的性能和可伸缩性：Vertex AI Search 旨在以极低的延迟处理大量数据。这意味着，您的 RAG 应用的响应时间更短，性能更出色，尤其是在处理复杂或广泛的知识库时。

  • 简化数据管理：从各种来源（例如网站、BigQuery 数据集和 Cloud Storage 存储桶）导入数据，从而简化数据注入流程。

  • 无缝集成：Vertex AI 提供与 Vertex AI Search 的内置集成，让您可以选择 Vertex AI Search 作为 RAG 应用的语料库后端。这简化了集成过程，并有助于确保组件之间的最佳兼容性。

  • 提高 LLM 输出质量：通过使用 Vertex AI Search 的检索功能，您可以帮助确保 RAG 应用从语料库中检索相关性最高的信息，从而生成更准确、信息更丰富的 LLM 生成输出。

Vertex AI Search 汇集了深度信息检索、自然语言处理和最新的大语言模型 (LLM) 处理功能，有助于理解用户的意图并返回与用户相关度最高的结果。

借助 Vertex AI Search，您可以使用自己控制的数据构建 Google 品质的搜索应用。

如需设置 Vertex AI Search，请执行以下操作：

  1. 创建搜索数据存储区。
  2. 创建搜索应用。

将 Vertex AI Search 用作 Vertex AI RAG Engine 的检索后端

设置 Vertex AI Search 后，请按照以下步骤将其设置为 RAG 应用的检索后端。

将 Vertex AI Search 设置为检索后端以创建 RAG 语料库

以下代码示例展示了如何将 Vertex AI Search 配置为 RAG 语料库的检索后端。

REST

如需使用命令行创建 RAG 语料库，请执行以下操作：

  1. 创建 RAG 语料库

    替换代码示例中使用的以下变量：

    • PROJECT_ID：您的 Google Cloud 项目的 ID。
    • LOCATION：处理请求的区域。
    • DISPLAY_NAME：您要创建的 RAG 语料库的显示名称。
    • ENGINE_NAME：Vertex AI Search 引擎或 Vertex AI Search 数据存储区的完整资源名称。
    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/ragCorpora" \
-d '{
  "display_name" : "DISPLAY_NAME",
  "vertex_ai_search_config" : {
    "serving_config": "ENGINE_NAME/servingConfigs/default_search"
  }
}'

  2. 监控进度

    替换代码示例中使用的以下变量：

    • PROJECT_ID：您的 Google Cloud 项目的 ID。
    • LOCATION：处理请求的区域。
    • OPERATION_ID：RAG 语料库创建操作的 ID。
    curl -X GET \
-H "Authorization: Bearer $(gcloud auth print-access-token)" \
-H "Content-Type: application/json" \
"https://LOCATION-aiplatform.googleapis.com/v1beta1/projects/PROJECT_ID/locations/LOCATION/operations/OPERATION_ID"

Python

在尝试此示例之前，请按照《Vertex AI 快速入门：使用客户端库》中的 Python 设置说明执行操作。 如需了解详情，请参阅 Vertex AI Python API 参考文档。

如需向 Vertex AI 进行身份验证，请设置应用默认凭据。 如需了解详情，请参阅为本地开发环境设置身份验证。 


from vertexai.preview import rag
import vertexai

# TODO(developer): Update and un-comment below lines
# PROJECT_ID = "your-project-id"
# vertex_ai_search_engine_name = "projects/{PROJECT_ID}/locations/{LOCATION}/collections/default_collection/engines/{ENGINE_ID}"
# display_name = "test_corpus"
# description = "Corpus Description"

# Initialize Vertex AI API once per session
vertexai.init(project=PROJECT_ID, location="us-central1")

# Configure Search
vertex_ai_search_config = rag.VertexAiSearchConfig(
    serving_config=f"{vertex_ai_search_engine_name}/servingConfigs/default_search",
)

corpus = rag.create_corpus(
    display_name=display_name,
    description=description,
    vertex_ai_search_config=vertex_ai_search_config,
)
print(corpus)
# Example response:
# RagCorpus(name='projects/1234567890/locations/us-central1/ragCorpora/1234567890',
# display_name='test_corpus', description='Corpus Description'.
# ...

使用 RAG API 检索上下文

创建 RAG 语料库后，您可以通过 RetrieveContexts API 从 Vertex AI Search 检索相关上下文。

REST

此代码示例演示了如何使用 REST 检索上下文。

替换代码示例中使用的以下变量：

  • PROJECT_ID：您的 Google Cloud 项目的 ID。
  • LOCATION：处理请求的区域。
  • RAG_CORPUS_RESOURCE：RAG 语料库资源的名称。

    格式：projects/{project}/locations/{location}/ragCorpora/{rag_corpus}.

  • TEXT：要获取相关上下文的查询文本。
curl -X POST \
  -H "Content-Type: application/json" \
  -H "Authorization: Bearer $(gcloud auth print-access-token)" \
"https://LOCATION-aiplatform.googleapis.com/v1beta1/projects/PROJECT_ID/locations/LOCATION:retrieveContexts" \
  -d '{
    "vertex_rag_store": {
      "rag_resources": {
        "rag_corpus": "RAG_CORPUS_RESOURCE"
      }
    },
    "query": {
      "text": "TEXT"
    }
  }'

Python 版 Vertex AI SDK

如需了解如何安装或更新 Vertex AI SDK for Python，请参阅安装 Vertex AI SDK for Python。 如需了解详情，请参阅 Python 版 Vertex AI SDK API 参考文档。 


from vertexai.preview import rag
import vertexai

# TODO(developer): Update and un-comment below lines
# PROJECT_ID = "your-project-id"
# corpus_name = "projects/[PROJECT_ID]/locations/us-central1/ragCorpora/[rag_corpus_id]"

# Initialize Vertex AI API once per session
vertexai.init(project=PROJECT_ID, location="us-central1")

response = rag.retrieval_query(
    rag_resources=[
        rag.RagResource(
            rag_corpus=corpus_name,
            # Optional: supply IDs from `rag.list_files()`.
            # rag_file_ids=["rag-file-1", "rag-file-2", ...],
        )
    ],
    text="Hello World!",
    similarity_top_k=10,  # Optional
    vector_distance_threshold=0.5,  # Optional
    # vector_search_alpha=0.5, # Optional - Only supported for Weaviate
)
print(response)
# Example response:
# contexts {
#   contexts {
#     source_uri: "gs://your-bucket-name/file.txt"
#     text: "....
#   ....

使用 Vertex AI Gemini API 生成内容

REST

如需使用 Gemini 模型生成内容，请调用 Vertex AI GenerateContent API。通过在请求中指定 RAG_CORPUS_RESOURCE，它会自动从 Vertex AI Search 检索数据。

替换示例代码中使用的以下变量：

  • PROJECT_ID：您的 Google Cloud 项目的 ID。
  • LOCATION：处理请求的区域。
  • MODEL_ID：用于内容生成的 LLM 模型。例如 gemini-1.5-flash-002。
  • GENERATION_METHOD：用于生成内容的 LLM 方法。例如：generateContent、streamGenerateContent。
  • INPUT_PROMPT：发送到 LLM 用于生成内容的文本。尝试使用与 Vertex AI Search 中的文档相关的提示。
  • RAG_CORPUS_RESOURCE：RAG 语料库资源的名称。 格式：projects/{project}/locations/{location}/ragCorpora/{rag_corpus}

  • SIMILARITY_TOP_K（可选）：要检索的热门上下文数量。

    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/publishers/google/models/MODEL_ID:GENERATION_METHOD" \
-d '{
  "contents": {
    "role": "user",
    "parts": {
      "text": "INPUT_PROMPT"
    }
  },
  "tools": {
    "retrieval": {
      "disable_attribution": false,
      "vertex_rag_store": {
        "rag_resources": {
            "rag_corpus": "RAG_CORPUS_RESOURCE"
          },
        "similarity_top_k": SIMILARITY_TOP_K
      }
    }
  }
}'

Python 版 Vertex AI SDK

如需了解如何安装或更新 Vertex AI SDK for Python，请参阅安装 Vertex AI SDK for Python。 如需了解详情，请参阅 Python 版 Vertex AI SDK API 参考文档。 


from vertexai.preview import rag
from vertexai.preview.generative_models import GenerativeModel, Tool
import vertexai

# TODO(developer): Update and un-comment below lines
# PROJECT_ID = "your-project-id"
# corpus_name = "projects/{PROJECT_ID}/locations/us-central1/ragCorpora/{rag_corpus_id}"

# Initialize Vertex AI API once per session
vertexai.init(project=PROJECT_ID, location="us-central1")

rag_retrieval_tool = Tool.from_retrieval(
    retrieval=rag.Retrieval(
        source=rag.VertexRagStore(
            rag_resources=[
                rag.RagResource(
                    rag_corpus=corpus_name,
                    # Optional: supply IDs from `rag.list_files()`.
                    # rag_file_ids=["rag-file-1", "rag-file-2", ...],
                )
            ],
            similarity_top_k=3,  # Optional
            vector_distance_threshold=0.5,  # Optional
        ),
    )
)

rag_model = GenerativeModel(
    model_name="gemini-1.5-flash-001", tools=[rag_retrieval_tool]
)
response = rag_model.generate_content("Why is the sky blue?")
print(response.text)
# Example response:
#   The sky appears blue due to a phenomenon called Rayleigh scattering.
#   Sunlight, which contains all colors of the rainbow, is scattered
#   by the tiny particles in the Earth's atmosphere....
#   ...

后续步骤