Google 会使用 AI 技术将内容翻译成您偏好的语言。AI 翻译可能包含错误。

RAG 快速入门

本页面介绍了如何使用 Vertex AI SDK 在 Gemini Enterprise Agent Platform 任务上运行 RAG 引擎。

您也可以使用此笔记本 RAG 引擎简介进行操作。

所需的角色

向您的用户账号授予角色。对以下每个 IAM 角色运行以下命令一次：roles/aiplatform.user

gcloud projects add-iam-policy-binding PROJECT_ID --member="user:USER_IDENTIFIER" --role=ROLE

替换以下内容：

PROJECT_ID：您的项目 ID。
USER_IDENTIFIER：您的用户账号的标识符。账号。如需查看示例，请参阅在 IAM 政策中表示员工池用户。
ROLE：您向用户账号授予的 IAM 角色。

准备您的 Google Cloud 控制台

如需使用 RAG 引擎，请执行以下操作：

安装 Agent Platform SDK for Python。
在 Google Cloud 控制台中运行此命令以设置您的项目。

gcloud config set project {project}
运行此命令以授权您的登录。

gcloud auth application-default login

运行 RAG 引擎

将以下示例代码复制并粘贴到 Google Cloud 控制台中，以运行 RAG 引擎。

Python

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

import agentplatform

from agentplatform import types
from google import genai
from google.genai import types as genai_types


# Create a RAG Corpus, Import Files, and Generate a response

# TODO(developer): Update and un-comment below lines
# PROJECT_ID = "your-project-id"
# MODEL_ID = "gemini-3.5-flash"
# display_name = "test_corpus"
# gcs_path = "gs://my_bucket/my_files_dir/*"
# google_drive_path ="https://drive.google.com/file/d/123"

# Initialize Agent Platform client once per session
client = agentplatform.Client(project=PROJECT_ID, location="us-east4")

# Configure embedding model, for example "text-embedding-005".
embedding_model_config = types.RagEmbeddingModelConfig(
    vertex_prediction_endpoint=types.RagEmbeddingModelConfigVertexPredictionEndpoint(
        endpoint="publishers/google/models/text-embedding-005"
    ),
)

# Create RagCorpus
rag_corpus = client.rag.create_corpus(
    rag_corpus=types.RagCorpus(
        display_name=display_name,
        rag_vector_db_config=types.RagVectorDbConfig(
            rag_embedding_model_config=embedding_model_config
        )
    )
)

# Import Files to the RagCorpus
client.rag.import_files(
    name=rag_corpus.name,
    import_config=types.ImportRagFilesConfig(
        gcs_source=genai_types.GcsSource(uris=[gcs_path]),
        rag_file_transformation_config=types.RagFileTransformationConfig(
            rag_file_chunking_config=types.RagFileChunkingConfig(
                chunk_size=512,
                chunk_overlap=100,
            )
        ), # optional
        max_embedding_requests_per_min=1000, # optional
    )
)

# Direct context retrieval
rag_retrieval_config = genai_types.RagRetrievalConfig(
    top_k=3,  # Optional
    filter=genai_types.RagRetrievalConfigFilter(
        vector_distance_threshold=0.5
    ),  # Optional
)
response = client.rag.retrieve_contexts(
    vertex_rag_store=genai_types.VertexRagStore(
        rag_resources=[
            genai_types.VertexRagStoreRagResource(
                rag_corpus=rag_corpus.name,
            )
        ],
    ),
    query=types.RagQuery(
        text="What is RAG and why it is helpful?",
        rag_retrieval_config=rag_retrieval_config,   
    )
)
print(response)

# Enhance generation
# Create a RAG retrieval tool
rag_retrieval_tool = genai_types.Tool(
    retrieval=genai_types.Retrieval(
        vertex_rag_store=genai_types.VertexRagStore(
            rag_resources=[
                genai_types.VertexRagStoreRagResource(
                    rag_corpus=rag_corpus.name,
                    # Optional: supply IDs from `rag.list_files()`.
                    # rag_file_ids=["rag-file-1", "rag-file-2", ...],
                )
            ],
            rag_retrieval_config=rag_retrieval_config,
        ),
    )
)

# Call generate_content with the tool using the GenAI SDK

# Create a GenAI SDK client
genai_client = genai.Client(enterprise=True, project=PROJECT_ID, location="us-east4")


response = genai_client.models.generate_content(
    model=MODEL_ID,
    contents="What is RAG and why it is helpful?",
    config=genai_types.GenerateContentConfig(
        tools=[rag_retrieval_tool]
    )
)
print(response.text)
# Example response:
#   RAG stands for Retrieval-Augmented Generation.
#   It's a technique used in AI to enhance the quality of responses
# ...

curl

创建 RAG 语料库。

  export LOCATION=LOCATION
  export PROJECT_ID=PROJECT_ID
  export CORPUS_DISPLAY_NAME=CORPUS_DISPLAY_NAME

  // CreateRagCorpus
  // Output: CreateRagCorpusOperationMetadata
  curl -X POST \
  -H "Authorization: Bearer $(gcloud auth print-access-token)" \
  -H "Content-Type: application/json" \
  https://LOCATION-aiplatform.googleapis.com/v1/projects/PROJECT_ID/locations/LOCATION/ragCorpora \
  -d '{
        "display_name" : "'"CORPUS_DISPLAY_NAME"'"
    }'

如需了解详情，请参阅创建 RAG 语料库示例。

导入 RAG 文件。

  // ImportRagFiles
  // Import a single Cloud Storage file or all files in a Cloud Storage bucket.
  // Input: LOCATION, PROJECT_ID, RAG_CORPUS_ID, GCS_URIS
  export RAG_CORPUS_ID=RAG_CORPUS_ID
  export GCS_URIS=GCS_URIS
  export CHUNK_SIZE=CHUNK_SIZE
  export CHUNK_OVERLAP=CHUNK_OVERLAP
  export EMBEDDING_MODEL_QPM_RATE=EMBEDDING_MODEL_QPM_RATE

  // Output: ImportRagFilesOperationMetadataNumber
  // Use ListRagFiles, or import_result_sink to get the correct rag_file_id.
  curl -X POST \
  -H "Authorization: Bearer $(gcloud auth print-access-token)" \
  -H "Content-Type: application/json" \
  https://LOCATION-aiplatform.googleapis.com/v1/projects/PROJECT_ID/locations/LOCATION/ragCorpora/RAG_CORPUS_ID/ragFiles:import \
  -d '{
    "import_rag_files_config": {
      "gcs_source": {
        "uris": "GCS_URIS"
      },
      "rag_file_chunking_config": {
        "chunk_size": CHUNK_SIZE,
        "chunk_overlap": CHUNK_OVERLAP
      },
      "max_embedding_requests_per_min": EMBEDDING_MODEL_QPM_RATE
    }
  }'

如需了解详情，请参阅导入 RAG 文件示例。

运行 RAG 检索查询。

  export RAG_CORPUS_RESOURCE=RAG_CORPUS_RESOURCE
  export VECTOR_DISTANCE_THRESHOLD=VECTOR_DISTANCE_THRESHOLD
  export SIMILARITY_TOP_K=SIMILARITY_TOP_K

  {
  "vertex_rag_store": {
      "rag_resources": {
        "rag_corpus": "RAG_CORPUS_RESOURCE"
      },
      "vector_distance_threshold": VECTOR_DISTANCE_THRESHOLD
    },
    "query": {
    "text": TEXT
    "similarity_top_k": SIMILARITY_TOP_K
    }
  }

  curl -X POST \
      -H "Authorization: Bearer $(gcloud auth print-access-token)" \
      -H "Content-Type: application/json; charset=utf-8" \
      -d @request.json \
      "https://LOCATION-aiplatform.googleapis.com/v1/projects/PROJECT_ID/locations/LOCATION:retrieveContexts"

如需了解详情，请参阅 RAG Engine API。

生成内容。

{
"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",
    "vector_distance_threshold": VECTOR_DISTANCE_THRESHOLD
  }
  }
}
}

curl -X POST \
    -H "Authorization: Bearer $(gcloud auth print-access-token)" \
    -H "Content-Type: application/json; charset=utf-8" \
    -d @request.json \
    "https://LOCATION-aiplatform.googleapis.com/v1/projects/PROJECT_ID/locations/LOCATION/publishers/google/models/MODEL_ID:GENERATION_METHOD"