Usar o RagManagedDb com o mecanismo de RAG da Vertex AI

Nesta página, mostramos como o mecanismo de RAG da Vertex AI usa o RagManagedDb, que é um banco de dados de vetores pronto para empresas usado para armazenar e gerenciar representações vetoriais dos seus documentos. O banco de dados de vetores é usado para recuperar documentos relevantes com base na semelhança semântica do documento a uma determinada consulta.

Além disso, esta página mostra como implementar a CMEK.

Gerenciar sua estratégia de recuperação

O RagManagedDb oferece as seguintes estratégias de recuperação para oferecer suporte aos seus casos de uso de RAG:

Estratégia de recuperação Descrição
k-vizinhos mais próximos (KNN) (padrão) Encontra os vizinhos mais próximos exatos comparando todos os pontos de dados no seu corpus RAG. Se você não especificar uma estratégia durante a criação do corpus de RAG, a KNN será a estratégia de recuperação padrão usada.
  • Verifica o recall perfeito (1,0) durante a recuperação.
  • Ótimo para aplicativos sensíveis a recall.
  • Ótimo para corpora RAG de pequeno a médio porte, que armazenam menos de 10.000 arquivos RAG.
  • Exige a pesquisa em todos os pontos de dados. Portanto, a latência aumenta com o número de arquivos RAG no corpus.
Vizinhos mais próximos aproximados (ANN) Usa técnicas de aproximação para encontrar vizinhos semelhantes mais rápido do que a técnica de KNN.
  • Reduz significativamente as latências de consulta em grandes corpora de RAG.
  • O recall diminuiu um pouco devido às técnicas de aproximação usadas.
  • É muito eficaz quando você tem grandes corpora de RAG, que são aproximadamente mais de 10.000 arquivos de RAG.
  • A quantidade de perda de recall aceitável depende do caso de uso, mas, na maioria dos casos em grande escala, perder um pouco de recall em troca de um desempenho de consulta aprimorado é uma troca aceitável.

Criar um corpus RAG com KNN RagManagedDb

Este exemplo de código demonstra como criar um corpus RAG usando KNN RagManagedDb.

Python 

from vertexai import rag
import vertexai

PROJECT_ID = YOUR_PROJECT_ID
LOCATION = YOUR_RAG_ENGINE_LOCATION
DISPLAY_NAME = YOUR_RAG_CORPUS_DISPLAY_NAME

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

vector_db = rag.RagManagedDb(retrieval_strategy=rag.KNN())
rag_corpus = rag.create_corpus(
    display_name=DISPLAY_NAME, backend_config=rag.RagVectorDbConfig(vector_db=vector_db))

REST

Substitua as seguintes variáveis:

  • PROJECT_ID: o ID do projeto.
  • LOCATION: a região para processar a solicitação.
  • CORPUS_DISPLAY_NAME: o nome de exibição do corpus RAG.
PROJECT_ID=PROJECT_ID
LOCATION=LOCATION
CORPUS_DISPLAY_NAME=CORPUS_DISPLAY_NAME

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" : '\""${CORPUS_DISPLAY_NAME}"\"',
      "vector_db_config": {
        "ragManagedDb": {
          "knn": {}
        }
      }
    }'

Criar um corpus RAG com ANN RagManagedDb

Para oferecer o recurso de ANN, o RagManagedDb usa uma estrutura baseada em árvore para particionar dados e facilitar pesquisas mais rápidas. Para ativar o melhor recall e a menor latência, a estrutura dessa árvore precisa ser configurada por testes para se adequar ao tamanho e à distribuição dos seus dados. O RagManagedDb permite configurar o tree_depth e o leaf_count da árvore.

O tree_depth determina o número de camadas ou níveis na árvore. Siga estas diretrizes:

  • Se você tiver aproximadamente 10.000 arquivos RAG no corpus RAG, defina o valor como 2.
  • Se você tiver mais arquivos RAG, defina como 3.
  • Se tree_depth não for especificado, o mecanismo de RAG da Vertex AI vai atribuir um valor padrão de 2 a esse parâmetro.

O leaf_count determina o número de nós folha na estrutura baseada em árvore. Cada nó folha contém grupos de vetores estreitamente relacionados e o centroide correspondente. Siga estas diretrizes:

  • O valor recomendado é 10 * sqrt(num of RAG files in your RAG corpus).

  • Se não for especificado, o mecanismo de RAG da Vertex AI vai atribuir um valor padrão de 500 a esse parâmetro.

Python 

from vertexai import rag
import vertexai

PROJECT_ID = YOUR_PROJECT_ID
LOCATION = YOUR_RAG_ENGINE_LOCATION
DISPLAY_NAME = YOUR_RAG_CORPUS_DISPLAY_NAME
TREE_DEPTH = YOUR_TREE_DEPTH # Optional: Acceptable values are 2 or 3. Default is 2.
LEAF_COUNT = YOUR_LEAF_COUNT # Optional: Default is 500.

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

ann_config = rag.ANN(tree_depth=TREE_DEPTH, leaf_count=LEAF_COUNT)
vector_db = rag.RagManagedDb(retrieval_strategy=ann_config)
rag_corpus = rag.create_corpus(
    display_name=DISPLAY_NAME, backend_config=rag.RagVectorDbConfig(vector_db=vector_db))

REST

Substitua as seguintes variáveis:

  • PROJECT_ID: o ID do projeto.
  • LOCATION: a região para processar a solicitação.
  • CORPUS_DISPLAY_NAME: o nome de exibição do corpus RAG.
  • TREE_DEPTH: a profundidade da árvore.
  • LEAF_COUNT: sua contagem de folhas.
PROJECT_ID=PROJECT_ID
LOCATION=LOCATION
CORPUS_DISPLAY_NAME=CORPUS_DISPLAY_NAME
TREE_DEPTH=TREE_DEPTH
LEAF_COUNT=LEAF_COUNT

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" : '\""${CORPUS_DISPLAY_NAME}"\"',
      "vector_db_config": {
        "ragManagedDb": {
          "ann": {
            "tree_depth": '"${TREE_DEPTH}"',
            "leaf_count": '"${LEAF_COUNT}"'
          }
        }
      }
    }'

Importar seus dados para a ANN RagManagedDb

É possível usar a API ImportRagFiles ou UploadRagFile para importar seus dados para a ANN RagManagedDb. No entanto, ao contrário da estratégia de recuperação do KNN, a abordagem da ANN exige que o índice baseado em árvore seja reconstruído pelo menos uma vez e, opcionalmente, após a importação de quantidades significativas de dados para um recall ideal. Para que o mecanismo de RAG da Vertex AI recrie seu índice ANN, defina rebuild_ann_index como "true" na solicitação da API ImportRagFiles.

O seguinte é importante:

  1. Antes de consultar o corpus de RAG, é preciso recriar o índice ANN pelo menos uma vez.
  2. Apenas uma recriação de índice simultânea é aceita em um projeto em cada local.

Para fazer upload do seu arquivo local no corpus de RAG, consulte Fazer upload de um arquivo RAG. Para importar dados para seu corpus de RAG e acionar uma recriação do índice ANN, consulte o exemplo de código a seguir, que demonstra como importar do Cloud Storage. Para saber mais sobre as fontes de dados compatíveis, consulte Fontes de dados compatíveis com RAG.

Python 

from vertexai import rag
import vertexai

PROJECT_ID = YOUR_PROJECT_ID
LOCATION = YOUR_RAG_ENGINE_LOCATION
CORPUS_ID = YOUR_CORPUS_ID
PATHS = ["gs://my_bucket/my_files_dir"]
REBUILD_ANN_INDEX = REBUILD_ANN_INDEX # Choose true or false.

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

corpus_name = f"projects/{PROJECT_ID}/locations/{LOCATION}/ragCorpora/{CORPUS_ID}"
# This is a non blocking call.
response = await rag.import_files_async(
    corpus_name=corpus_name,
    paths=PATHS,
    rebuild_ann_index=REBUILD_ANN_INDEX
)

# Wait for the import to complete.
await response.result()

REST 

GCS_URI=GCS_URI
REBUILD_ANN_INDEX=<true/false>

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/${CORPUS_ID}/ragFiles:import \
-d '{
  "import_rag_files_config": {
    "gcs_source": {
      "uris": '\""${GCS_URI}"\"',
      },
    "rebuild_ann_index": '${REBUILD_ANN_INDEX}'
  }
}'

Gerenciar sua criptografia

O mecanismo de RAG da Vertex AI oferece opções robustas para gerenciar a criptografia dos seus dados em repouso. Por padrão, todos os dados do usuário no RagManagedDb são criptografados usando um Google-owned and Google-managed encryption key, que é a configuração padrão. Essa configuração padrão ajuda a verificar se os dados estão seguros sem exigir nenhuma configuração específica.

Se você precisar de mais controle sobre as chaves usadas para criptografia, o mecanismo RAG da Vertex AI oferece suporte à chave de criptografia gerenciada pelo cliente (CMEK). Com a CMEK, você pode usar suas chaves criptográficas, gerenciadas no Cloud Key Management Service (KMS), para proteger os dados do seu corpus de RAG.

Para informações sobre limitações da CMEK para corpora de RAG, consulte Limitações da CMEK para o mecanismo de RAG da Vertex AI.

Configurar a chave do KMS e conceder permissões

Antes de criar um corpus de RAG criptografado com CMEK, é necessário configurar uma chave criptográfica no Google Cloud KMS e conceder à conta de serviço do mecanismo de RAG da Vertex AI as permissões necessárias para usar essa chave.

Pré-requisitos

Para realizar as etapas de configuração a seguir, verifique se sua conta de usuário tem as permissões adequadas do Identity and Access Management (IAM) no projeto Google Cloud em que você pretende criar a chave do KMS e o corpus de RAG. Normalmente, é necessário um papel como o de administrador do Cloud KMS (roles/cloudkms.admin).

Ativar a API

Para ativar a API Cloud KMS, faça o seguinte:

  1. Acesse o console do Google Cloud .
  2. Selecione o projeto em que você quer gerenciar suas chaves e crie o corpus de RAG.
  3. Na barra de pesquisa, digite "Gerenciamento de chaves" e selecione o serviço "Gerenciamento de chaves".
  4. Se a API não estiver ativada, clique em Ativar. Talvez seja necessário aguardar alguns minutos para que a API seja totalmente provisionada.

Criar um keyring e uma chave do KMS

Para criar um keyring, faça o seguinte:

  1. Na seção Gerenciamento de chaves, clique em Criar keyring.

    Digite o seguinte:

    • Nome do keyring: insira um nome exclusivo para o keyring, como rag-engine-cmek-keys.
    • Tipo de local: selecione "Região".
    • Local: escolha a região selecionada, como us-central1. Essa região precisa corresponder à região em que os recursos do RAG Engine vão residir.

  2. Clique em Criar.

Para criar uma chave no keyring, faça o seguinte:

  1. Depois que o keyring for criado, você vai receber uma solicitação ou poderá navegar até Criar chave.

    Digite o seguinte:

    • Nome da chave: insira um nome exclusivo para a chave, como my-rag-corpus-key.
    • Nível de proteção: escolha um nível de proteção (Software ou HSM). Se você precisar de chaves protegidas por hardware, selecione HSM.
    • Finalidade: selecione Criptografia/descriptografia simétrica. Isso é necessário para a CMEK.
    • Origem do material da chave: selecione Chave gerada.
    • Período de rotação: opcional. Recomendado. Configure uma programação de rotação de chaves de acordo com as políticas de segurança da sua organização, como a cada 90 dias.

  2. Clique em Criar.

Para copiar o nome do recurso da chave, faça o seguinte:

  1. Depois que a chave for criada, acesse a página de detalhes dela.

  2. Localize o nome do recurso. O formato é projects/YOUR_PROJECT_ID/locations/YOUR_REGION/keyRings/YOUR_KEY_RING_NAME/cryptoKeys/YOUR_KEY_NAME/cryptoKeyVersions/1.

  3. Copie o nome do recurso e remova a parte /cryptoKeyVersions/VERSION_NUMBER. O nome do recurso formatado corretamente é projects/YOUR_PROJECT_ID/locations/YOUR_REGION/keyRings/YOUR_KEY_RING_NAME/cryptoKeys/YOUR_KEY_NAME.

Conceder permissões ao agente de serviço do mecanismo de RAG da Vertex AI

Para que o mecanismo RAG da Vertex AI criptografe e descriptografe dados usando sua chave do KMS, o agente de serviço precisa ter as permissões adequadas nessa chave específica.

Para identificar seu agente de serviço do mecanismo de RAG da Vertex AI, faça o seguinte:

  1. Acesse a página IAM e administrador > IAM no Google Cloud console do seu projeto.

  2. Na página "Identity and Access Management", marque a caixa de seleção Incluir concessões de papel fornecidas pelo Google.

  3. Na barra de filtro ou pesquisa da lista de principais, pesquise o agente de serviço do mecanismo RAG da Vertex AI. Ele segue o padrão service-YOUR_PROJECT_NUMBER@gcp-sa-vertex-rag.iam.gserviceaccount.com.

Para conceder permissões na chave do KMS, faça o seguinte:

  1. Volte ao serviço de gerenciamento de chaves no console Google Cloud .

  2. Selecione o keyring que contém a chave criada.

  3. Selecione a chave específica que você criou.

  4. Na página de detalhes da chave, acesse a guia Permissões.

  5. Clique em Adicionar principal.

  6. No campo Novos principais, digite o endereço de e-mail do agente de serviço do mecanismo RAG da Vertex AI.

  7. No menu suspenso Selecionar um papel, escolha o papel de criptografador/descriptografador de CryptoKey do Cloud KMS (roles/cloudkms.cryptoKeyEncrypterDecrypter). Esse papel concede ao agente de serviço as permissões necessárias para usar a chave em operações de criptografia e descriptografia.

  8. Clique em Salvar.

Criar um corpus RAG com criptografia gerenciada pelo cliente

Este exemplo de código demonstra como criar um corpus RAG criptografado com uma chave de criptografia gerenciada pelo cliente (CMEK).

Substitua as variáveis nos exemplos de código a seguir:

Python 

import vertexai
from google.cloud import aiplatform
from vertexai import rag
from google.cloud.aiplatform_v1.types.encryption_spec import EncryptionSpec

PROJECT_ID = YOUR_PROJECT_ID
LOCATION = YOUR_RAG_ENGINE_LOCATION
DISPLAY_NAME = YOUR_RAG_CORPUS_DISPLAY_NAME
KMS_KEY_NAME = YOUR_KMS_KEY_NAME

vertexai.init(project=PROJECT_ID)

rag_corpus = rag.create_corpus(display_name=DISPLAY_NAME, encryption_spec=EncryptionSpec(kms_key_name=KMS_KEY_NAME))

REST 

PROJECT_ID = YOUR_PROJECT_ID
LOCATION = YOUR_RAG_ENGINE_LOCATION
DISPLAY_NAME = YOUR_RAG_CORPUS_DISPLAY_NAME
KMS_KEY_NAME = YOUR_KMS_KEY_NAME

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}"\"',
      "encryption_spec" : {
        "kms_key_name" : '\""${KMS_KEY_NAME}"\"'
      }
    }'

A seguir