RAG 快速入門

本頁面說明如何使用 Vertex AI SDK,在 Gemini Enterprise Agent Platform 執行 RAG Engine 工作。

您也可以參考「RAG Engine 簡介」筆記本。

必要的角色

將角色授予使用者帳戶。針對下列各個 IAM 角色執行下列指令一次:roles/aiplatform.user

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

更改下列內容:

  • PROJECT_ID:專案 ID。
  • USER_IDENTIFIER:使用者帳戶 ID。 例如:myemail@example.com。
  • ROLE:授予使用者帳戶的 IAM 角色。

準備 Google Cloud 控制台

如要使用 RAG Engine,請按照下列步驟操作:

  1. 安裝 Agent Platform SDK for Python。

  2. 在 Google Cloud 控制台執行這項指令,設定專案。

    gcloud config set project {project}

  3. 執行下列指令,授權登入。

    gcloud auth application-default login

執行 RAG Engine

複製這段程式碼範例並貼到 Google Cloud 控制台,即可執行 RAG Engine。

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.8-flash"
# display_name = "test_corpus"
# gcs_path = "gs://my_bucket/my_files_dir/*"

# Initialize Agent Platform client once per session.
# See https://docs.cloud.google.com/gemini-enterprise-agent-platform/build/rag-engine/rag-overview#supported-regions
client = agentplatform.Client(project=PROJECT_ID, location="us-east5")

# 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(vertexai=True, project=PROJECT_ID, location="global")

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

  1. 建立 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 語料庫範例」。

  2. 匯入 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 檔案範例」。

  3. 執行 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」。

  4. 生成內容。

    {
    "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"
    

    詳情請參閱「RAG Engine API」。

後續步驟