Usar a pesquisa vetorial da Agent Platform com o mecanismo RAG

Nesta página, mostramos como conectar o mecanismo RAG à busca vetorial da Gemini Enterprise Agent Platform.

Você também pode acompanhar usando o notebook Mecanismo de RAG com a busca vetorial da Agent Platform.

Configuração da pesquisa vetorial da Agent Platform

A Pesquisa Vetorial da Agent Platform é baseada na tecnologia de pesquisa vetorial desenvolvida pela Google Research. Com a Pesquisa Vetorial, é possível usar a mesma infraestrutura que fornece uma base para os produtos do Google, como Pesquisa Google, YouTube e Google Play.

Para fazer a integração com o mecanismo RAG, é necessário um índice da Pesquisa vetorial vazio.

Configurar o SDK da Vertex AI

Para configurar o SDK da Vertex AI, consulte Configuração.

Criar um índice da Pesquisa vetorial

Para criar um índice de Pesquisa Vetorial compatível com seu corpus de RAG, o índice precisa atender aos seguintes critérios:

  1. IndexUpdateMethod precisa ser STREAM_UPDATE. Consulte Criar índice de stream.

  2. O tipo de medida de distância precisa ser definido explicitamente como um dos seguintes:

    • DOT_PRODUCT_DISTANCE
    • COSINE_DISTANCE

  3. A dimensão do vetor precisa ser consistente com o modelo de embedding que você planeja usar no corpus de RAG. Outros parâmetros podem ser ajustados com base nas suas escolhas, que determinam se os parâmetros adicionais podem ser ajustados.

Python

Para saber como instalar o SDK da Vertex AI para Python, consulte Instalar o SDK da Vertex AI para Python. Para mais informações, consulte a documentação de referência da API Python. 

def vector_search_create_streaming_index(
    project: str, location: str, display_name: str, gcs_uri: Optional[str] = None
) -> aiplatform.MatchingEngineIndex:
    """Create a vector search index.

    Args:
        project (str): Required. Project ID
        location (str): Required. The region name
        display_name (str): Required. The index display name
        gcs_uri (str): Optional. The Google Cloud Storage uri for index content

    Returns:
        The created MatchingEngineIndex.
    """
    # Initialize the Vertex AI client
    aiplatform.init(project=project, location=location)

    # Create Index
    index = aiplatform.MatchingEngineIndex.create_tree_ah_index(
        display_name=display_name,
        contents_delta_uri=gcs_uri,
        description="Matching Engine Index",
        dimensions=100,
        approximate_neighbors_count=150,
        leaf_node_embedding_count=500,
        leaf_nodes_to_search_percent=7,
        index_update_method="STREAM_UPDATE",  # Options: STREAM_UPDATE, BATCH_UPDATE
        distance_measure_type=aiplatform.matching_engine.matching_engine_index_config.DistanceMeasureType.DOT_PRODUCT_DISTANCE,
    )

    return index

Criar endpoint de índice da Pesquisa Vetorial

Os endpoints públicos são compatíveis com o mecanismo RAG.

Python

Para saber como instalar o SDK da Vertex AI para Python, consulte Instalar o SDK da Vertex AI para Python. Para mais informações, consulte a documentação de referência da API Python. 

def vector_search_create_index_endpoint(
    project: str, location: str, display_name: str
) -> None:
    """Create a vector search index endpoint.

    Args:
        project (str): Required. Project ID
        location (str): Required. The region name
        display_name (str): Required. The index endpoint display name
    """
    # Initialize the Vertex AI client
    aiplatform.init(project=project, location=location)

    # Create Index Endpoint
    index_endpoint = aiplatform.MatchingEngineIndexEndpoint.create(
        display_name=display_name,
        public_endpoint_enabled=True,
        description="Matching Engine Index Endpoint",
    )

    print(index_endpoint.name)

Implantar um índice em um endpoint de índice

Antes de fazer a pesquisa de vizinho mais próximo, o índice precisa ser implantado em um endpoint de índice.

Python

Para saber como instalar o SDK da Vertex AI para Python, consulte Instalar o SDK da Vertex AI para Python. Para mais informações, consulte a documentação de referência da API Python. 

def vector_search_deploy_index(
    project: str,
    location: str,
    index_name: str,
    index_endpoint_name: str,
    deployed_index_id: str,
) -> None:
    """Deploy a vector search index to a vector search index endpoint.

    Args:
        project (str): Required. Project ID
        location (str): Required. The region name
        index_name (str): Required. The index to update. A fully-qualified index
          resource name or a index ID.  Example:
          "projects/123/locations/us-central1/indexes/my_index_id" or
          "my_index_id".
        index_endpoint_name (str): Required. Index endpoint to deploy the index
          to.
        deployed_index_id (str): Required. The user specified ID of the
          DeployedIndex.
    """
    # Initialize the Vertex AI client
    aiplatform.init(project=project, location=location)

    # Create the index instance from an existing index
    index = aiplatform.MatchingEngineIndex(index_name=index_name)

    # Create the index endpoint instance from an existing endpoint.
    index_endpoint = aiplatform.MatchingEngineIndexEndpoint(
        index_endpoint_name=index_endpoint_name
    )

    # Deploy Index to Endpoint
    index_endpoint = index_endpoint.deploy_index(
        index=index, deployed_index_id=deployed_index_id
    )

    print(index_endpoint.deployed_indexes)

Se for a primeira vez que você implanta um índice em um endpoint de índice, leva aproximadamente 30 minutos para criar e iniciar automaticamente o back-end antes que o índice possa ser armazenado. Após a primeira implantação, o índice fica pronto em segundos. Para ver o status da implantação do índice, abra o Console da Pesquisa Vetorial, selecione a guia Endpoints de índice e escolha seu endpoint de índice.

Identifique o nome do recurso do índice e do endpoint de índice, que têm os seguintes formatos:

  • projects/${PROJECT_NUMBER}/locations/${LOCATION_ID}/indexes/${INDEX_ID}
  • projects/${PROJECT_NUMBER}/locations/${LOCATION_ID}/indexEndpoints/${INDEX_ENDPOINT_ID}.

Usar a Pesquisa Vetorial da Agent Platform no mecanismo de RAG

Depois que a instância da Pesquisa Vetorial estiver configurada, siga as etapas nesta seção para defini-la como o banco de dados vetorial do aplicativo RAG.

Definir o banco de dados de vetores para criar um corpus RAG

Ao criar o corpus de RAG, especifique apenas os valores completos de INDEX_ENDPOINT_NAME e INDEX_NAME. Use o ID numérico para os nomes de recursos do índice e do endpoint de índice. O corpus de RAG é criado e associado automaticamente ao índice da Pesquisa Vetorial. As validações são realizadas nos critérios. Se algum dos requisitos não for atendido, a solicitação será rejeitada.

Python

Antes de testar este exemplo, siga as instruções de configuração do Python no Guia de início rápido da plataforma de agentes: como usar bibliotecas de cliente.

Para autenticar na Plataforma de Agentes, configure o Application Default Credentials. Para mais informações, consulte Configurar a autenticação para um ambiente de desenvolvimento local. 


import agentplatform
from agentplatform import types

# TODO(developer): Update and un-comment below lines
# PROJECT_ID = "your-project-id"
# vector_search_index_name = "projects/{PROJECT_ID}/locations/{LOCATION}/indexes/{INDEX_ID}"
# vector_search_index_endpoint_name = "projects/{PROJECT_ID}/locations/{LOCATION}/indexEndpoints/{INDEX_ENDPOINT_ID}"
# display_name = "test_corpus"
# description = "Corpus Description"

# 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 (Optional)
embedding_model_config = types.RagEmbeddingModelConfig(
    vertex_prediction_endpoint=types.RagEmbeddingModelConfigVertexPredictionEndpoint(
        endpoint="publishers/google/models/text-embedding-005"
    )
)

# Configure Vector DB
vector_db = types.RagVectorDbConfigVertexVectorSearch(
    index=vector_search_index_name,
    index_endpoint=vector_search_index_endpoint_name
)

corpus = client.rag.create_corpus(
    rag_corpus=types.RagCorpus(
        display_name=display_name,
        description=description,
        rag_vector_db_config=types.RagVectorDbConfig(
            rag_embedding_model_config=embedding_model_config,
            vertex_vector_search=vector_db,
        ),
    )
)
print(corpus)
# Example response:
# RagCorpus(name='projects/1234567890/locations/us-central1/ragCorpora/1234567890',
# display_name='test_corpus', description='Corpus Description', embedding_model_config=...
# ...

REST

  # TODO(developer): Update and un-comment the following lines:
  # CORPUS_DISPLAY_NAME = "YOUR_CORPUS_DISPLAY_NAME"
  # Full index/indexEndpoint resource name
  # Index: projects/${PROJECT_NUMBER}/locations/${LOCATION_ID}/indexes/${INDEX_ID}
  # IndexEndpoint: projects/${PROJECT_NUMBER}/locations/${LOCATION_ID}/indexEndpoints/${INDEX_ENDPOINT_ID}
  # INDEX_RESOURCE_NAME = "YOUR_INDEX_ENDPOINT_RESOURCE_NAME"
  # INDEX_NAME = "YOUR_INDEX_RESOURCE_NAME"
  # Call CreateRagCorpus API to create a new RagCorpus
  curl -X POST -H "Authorization: Bearer $(gcloud auth print-access-token)" -H "Content-Type: application/json" https://${LOCATION_ID}-aiplatform.googleapis.com/v1beta1/projects/${PROJECT_NUMBER}/locations/${LOCATION_ID}/ragCorpora -d '{
        "display_name" : '\""${CORPUS_DISPLAY_NAME}"\"',
        "rag_vector_db_config" : {
                "vertex_vector_search": {
                  "index":'\""${INDEX_NAME}"\"'
              "index_endpoint":'\""${INDEX_ENDPOINT_NAME}"\"'
                }
          }
    }'

  # Call ListRagCorpora API to verify the RagCorpus is created successfully
  curl -sS -X GET \
  -H "Content-Type: application/json" \
  -H "Authorization: Bearer $(gcloud auth print-access-token)" \
  "https://${LOCATION_ID}-aiplatform.googleapis.com/v1beta1/projects/${PROJECT_NUMBER}/locations/${LOCATION_ID}/ragCorpora"

Importar arquivos usando a API RAG

Use o método da API ragFiles.import para importar arquivos do Cloud Storage ou do Google Drive para o índice da Pesquisa Vetorial. Os arquivos são incorporados e armazenados no índice da Pesquisa vetorial.

REST

# TODO(developer): Update and uncomment the following lines:
# RAG_CORPUS_ID = "your-rag-corpus-id"
#
# Google Cloud Storage bucket/file location.
# For example, "gs://rag-fos-test/"
# GCS_URIS= "your-gcs-uris"

# Call ImportRagFiles API to embed files and store in the BigQuery table
curl -X POST \
-H "Authorization: Bearer $(gcloud auth print-access-token)" \
-H "Content-Type: application/json" \
https://us-central1-aiplatform.googleapis.com/v1beta1/projects/${PROJECT_NUMBER}/locations/us-central1/ragCorpora/${RAG_CORPUS_ID}/ragFiles:import \
-d '{
  "import_rag_files_config": {
    "gcs_source": {
      "uris": '\""${GCS_URIS}"\"'
    },
    "rag_file_chunking_config": {
      "chunk_size": 512
    }
  }
}'

# Call ListRagFiles API to verify the files are imported successfully
curl -X GET \
-H "Content-Type: application/json" \
-H "Authorization: Bearer $(gcloud auth print-access-token)" \
https://us-central1-aiplatform.googleapis.com/v1beta1/projects/${PROJECT_NUMBER}/locations/us-central1/ragCorpora/${RAG_CORPUS_ID}/ragFiles

Python

Para saber como instalar o SDK da Vertex AI para Python, consulte Instalar o SDK da Vertex AI para Python. Para mais informações, consulte a documentação de referência da API Python. 


import agentplatform
from agentplatform import types

from google.genai import types as genai_types

# TODO(developer): Update and un-comment below lines
# PROJECT_ID = "your-project-id"
# corpus_name = "projects/{PROJECT_ID}/locations/us-east5/ragCorpora/{rag_corpus_id}"
# paths = ["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")

response = client.rag.import_files(
    name=corpus_name,
    import_config=types.ImportRagFilesConfig(
        gcs_source=genai_types.GcsSource(uris=paths),
        rag_file_transformation_config=types.RagFileTransformationConfig(
            rag_file_chunking_config=types.RagFileChunkingConfig(
                chunk_size=512,
                chunk_overlap=100,
            )
        ),  # optional
        max_embedding_requests_per_min=900,  # optional
    ),
)

print(f"Imported {response.imported_rag_files_count} files.")
# Example response:
# Imported 2 files.

Recuperar contextos relevantes usando a API RAG

Após a conclusão das importações de arquivos, o contexto relevante pode ser recuperado do índice da Pesquisa Vetorial usando a API RetrieveContexts.

REST

# TODO(developer): Update and uncomment the following lines:
# RETRIEVAL_QUERY="your-retrieval-query"
#
# Full RAG corpus resource name
# Format:
# "projects/${PROJECT_NUMBER}/locations/us-central1/ragCorpora/${RAG_CORPUS_ID}"
# RAG_CORPUS_RESOURCE="your-rag-corpus-resource"

# Call RetrieveContexts API to retrieve relevant contexts
curl -X POST \
  -H "Content-Type: application/json" \
  -H "Authorization: Bearer $(gcloud auth print-access-token)" \
https://us-central1-aiplatform.googleapis.com/v1beta1/projects/${PROJECT_NUMBER}/locations/us-central1:retrieveContexts \
  -d '{
    "vertex_rag_store": {
      "rag_resources": {
          "rag_corpus": '\""${RAG_CORPUS_RESOURCE}"\"',
        },
    },
    "query": {
      "text": '\""${RETRIEVAL_QUERY}"\"',
      "similarity_top_k": 10
    }
  }'

Python

Para saber como instalar o SDK da Vertex AI para Python, consulte Instalar o SDK da Vertex AI para Python. Para mais informações, consulte a documentação de referência da API Python. 


import agentplatform

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

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

# 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")

response = client.rag.retrieve_contexts(
    vertex_rag_store=genai_types.VertexRagStore(
        rag_resources=[
            genai_types.VertexRagStoreRagResource(
                rag_corpus=corpus_name,
                # Optional: supply IDs from `rag.list_files()`.
                # rag_file_ids=["rag-file-1", "rag-file-2", ...],
            )
        ],
    ),
    query=types.RagQuery(
        text="Hello World!",
        rag_retrieval_config=genai_types.RagRetrievalConfig(
            top_k=10,
            filter=genai_types.RagRetrievalConfigFilter(
                vector_distance_threshold=0.5
            ),
        ),
    )
)
print(response)
# Example response:
# contexts {
#   contexts {
#     source_uri: "gs://your-bucket-name/file.txt"
#     text: "....
#   ....

Gerar conteúdo usando a API Gemini da Agent Platform

Para gerar conteúdo usando modelos do Gemini, faça uma chamada para a API GenerateContent da Agent Platform. Ao especificar o RAG_CORPUS_RESOURCE na solicitação, a API recupera automaticamente os dados do índice da Pesquisa Vetorial.

REST

# TODO(developer): Update and uncomment the following lines:
# MODEL_ID=gemini-2.5-flash
# GENERATE_CONTENT_PROMPT="your-generate-content-prompt"

# GenerateContent with contexts retrieved from the FeatureStoreOnline index
curl -X POST \
-H "Authorization: Bearer $(gcloud auth print-access-token)" \
-H "Content-Type: application/json"  https://us-central1-aiplatform.googleapis.com/v1beta1/projects/${PROJECT_NUMBER}/locations/us-central1/publishers/google/models/${MODEL_ID}:generateContent \
-d '{
  "contents": {
    "role": "user",
    "parts": {
      "text": '\""${GENERATE_CONTENT_PROMPT}"\"'
    }
  },
  "tools": {
    "retrieval": {
      "vertex_rag_store": {
        "rag_resources": {
            "rag_corpus": '\""${RAG_CORPUS_RESOURCE}"\"',
          },
        "similarity_top_k": 8,
        "vector_distance_threshold": 0.32
      }
    }
  }
}'

Python

Para saber como instalar o SDK da Vertex AI para Python, consulte Instalar o SDK da Vertex AI para Python. Para mais informações, consulte a documentação de referência da API Python. 


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

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

rag_retrieval_tool = genai_types.Tool(
    retrieval=genai_types.Retrieval(
        vertex_rag_store=genai_types.VertexRagStore(
            rag_resources=[
                genai_types.VertexRagStoreRagResource(
                    rag_corpus=corpus_name
                )
            ],
            rag_retrieval_config=genai_types.RagRetrievalConfig(
                top_k=10,
                filter=genai_types.RagRetrievalConfigFilter(
                    vector_distance_threshold=0.5
                ),
            ),
        ),
    )
)

# Create a GenAI SDK client to make a generate_content request
genai_client = genai.Client(vertexai=True, project=PROJECT_ID, location="global")

response = genai_client.models.generate_content(
    model="gemini-3.8-flash",
    contents="Why is the sky blue?",
    config=genai_types.GenerateContentConfig(
        tools=[rag_retrieval_tool]
    ),
)
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....
#   ...

A seguir