Usar o Vertex AI Feature Store no RAG Engine

Esta página mostra como configurar o Feature Store da Vertex AI como o banco de dados de vetores a ser usado com o RAG Engine.

Você também pode seguir este notebook Mecanismo de RAG com o Feature Store da Vertex AI.

O RAG Engine usa um banco de dados vetorial integrado com suporte do Spanner para armazenar e gerenciar representações vetoriais de documentos de texto. O banco de dados de vetores recupera documentos relevantes com base na semelhança semântica deles com uma determinada consulta.

Ao integrar o Feature Store da Vertex AI como um banco de dados de vetores adicional, o RAG Engine pode usar o Feature Store da Vertex AI para processar grandes volumes de dados com baixa latência, o que ajuda a melhorar o desempenho e a escalabilidade dos seus aplicativos RAG.

Configurar um Vertex AI Feature Store

O Feature Store da Vertex AI, um serviço gerenciado nativo da nuvem, é um componente essencial da Vertex AI. Ele simplifica o gerenciamento de recursos de machine learning (ML) e a exibição on-line, permitindo que você gerencie dados de recursos em uma tabela ou visualização do BigQuery. Isso permite a veiculação de recursos on-line de baixa latência.

Para instâncias de FeatureOnlineStore criadas com a exibição on-line otimizada, é possível aproveitar uma pesquisa de similaridade vetorial para recuperar uma lista de entidades semanticamente parecidas ou relacionadas, conhecidas como vizinhos mais próximos aproximados.

As seções a seguir mostram como configurar uma instância do Feature Store da Vertex AI para seu aplicativo RAG.

Criar um esquema de tabela do BigQuery

Use o console do Google Cloud para criar um esquema de tabela do BigQuery. Ela precisa conter os seguintes campos para servir como a fonte de dados.

Nome do campo Tipo de dado Status
corpus_id String Obrigatório
file_id String Obrigatório
chunk_id String Obrigatório
chunk_data_type String Anulável
chunk_data String Anulável
file_original_uri String Anulável
embeddings Float Repetida

Este exemplo de código demonstra como definir o esquema da tabela do BigQuery.

SQL

  CREATE TABLE `PROJECT_ID.input_us_central1.rag_source_new` (
    `corpus_id` STRING NOT NULL,
    `file_id` STRING NOT NULL,
    `chunk_id` STRING NOT NULL,
    `chunk_data_type` STRING,
    `chunk_data` STRING,
    `embeddings` ARRAY<FLOAT64>,
    `file_original_uri` STRING
  );

Provisionar uma instância FeatureOnlineStore

Para ativar a veiculação on-line de recursos, use a API CreateFeatureOnlineStore do Feature Store da Vertex AI para configurar uma instância FeatureOnlineStore. Se você estiver provisionando um FeatureOnlineStore pela primeira vez, a operação poderá levar aproximadamente cinco minutos para ser concluída.

REST

Para criar uma instância de loja on-line, envie uma solicitação POST usando o método featureOnlineStores.create.

Antes de usar os dados da solicitação abaixo, faça as substituições a seguir:

  • LOCATION_ID: região onde você quer criar a instância de FeatureOnlineStore, como us-central1.
  • PROJECT_ID: o ID do projeto.
  • FEATUREONLINESTORE_NAME: o nome da nova instância FeatureOnlineStore.

Método HTTP e URL: 

POST https://LOCATION_ID-aiplatform.googleapis.com/v1/projects/PROJECT_ID/locations/LOCATION_ID/featureOnlineStores?feature_online_store_id=FEATUREONLINESTORE_NAME

Corpo JSON da solicitação:

{
  "optimized": {}
}

Para enviar a solicitação, escolha uma destas opções:

curl

Salve o corpo da solicitação em um arquivo com o nome request.json e execute o comando a seguir: 

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

PowerShell

Salve o corpo da solicitação em um arquivo com o nome request.json e execute o comando a seguir: 

$cred = gcloud auth print-access-token
$headers = @{ "Authorization" = "Bearer $cred" }

Invoke-WebRequest `
    -Method POST `
    -Headers $headers `
    -ContentType: "application/json; charset=utf-8" `
    -InFile request.json `
    -Uri "https://LOCATION_ID-aiplatform.googleapis.com/v1/projects/PROJECT_ID/locations/LOCATION_ID/featureOnlineStores?feature_online_store_id=FEATUREONLINESTORE_NAME" | Select-Object -Expand Content

Você receberá uma resposta JSON semelhante a esta:

{
  "name": "projects/PROJECT_NUMBER/locations/LOCATION_ID/featureOnlineStores/FEATUREONLINESTORE_NAME/operations/OPERATION_ID",
  "metadata": {
    "@type": "type.googleapis.com/google.cloud.aiplatform.v1.CreateFeatureOnlineStoreOperationMetadata",
    "genericMetadata": {
      "createTime": "2023-09-18T17:49:23.847496Z",
      "updateTime": "2023-09-18T17:49:23.847496Z"
    }
  }
}

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.cloud import aiplatform
from vertexai.resources.preview import feature_store


def create_optimized_public_feature_online_store_sample(
    project: str,
    location: str,
    feature_online_store_id: str,
):
    aiplatform.init(project=project, location=location)
    fos = feature_store.FeatureOnlineStore.create_optimized_store(
        feature_online_store_id
    )
    return fos
  • project: o ID do projeto.
  • location: região onde você quer criar a instância de FeatureOnlineStore, como us-central1.
  • feature_online_store_id: o nome da nova instância FeatureOnlineStore.

Criar um recurso FeatureView

Para conectar a tabela do BigQuery, que armazena a origem de dados do recurso, à instância FeatureOnlineStore, chame a API CreateFeatureView para criar um recurso FeatureView. Ao criar um recurso FeatureView, escolha a métrica de distância padrão DOT_PRODUCT_DISTANCE, que é definida como o negativo do produto escalar. Um DOT_PRODUCT_DISTANCE menor indica maior semelhança.

Este exemplo de código demonstra como criar um recurso FeatureView.

REST

  # TODO(developer): Update and uncomment the following lines:
  # Set feature_view_id
  # Example: "feature_view_test"
  # FEATURE_VIEW_ID = "your-feature-view-id"
  #
  # The big_query_uri generated in the above BigQuery table schema creation step
  # The format should be "bq://" + BigQuery table ID
  # Example: "bq://tester.ragtest1.rag_testdata"
  # BIG_QUERY_URI=YOUR_BIG_QUERY_URI

  # Call CreateFeatureView API to create a FeatureView
  curl -X POST -H "Authorization: Bearer $(gcloud auth print-access-token)" -H "Content-Type: application/json" \
  https://us-central1-aiplatform.googleapis.com/v1/projects/${PROJECT_ID}/locations/us-central1/featureOnlineStores/${FEATURE_ONLINE_STORE_ID}/featureViews?feature_view_id=${FEATURE_VIEW_ID} \
    -d '{
          "vertex_rag_source": {
            "uri": '\""${BIG_QUERY_URI}"\"'
          }
      }'

  # Call ListFeatureViews API to verify the FeatureView is created successfully
  curl -X GET -H "Authorization: Bearer $(gcloud auth application-default print-access-token)" -H "Content-Type: application/json" https://us-central1-aiplatform.googleapis.com/v1/projects/${PROJECT_ID}/locations/us-central1/featureOnlineStores/${FEATURE_ONLINE_STORE_ID}/featureViews

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.cloud import aiplatform
from vertexai.resources.preview import feature_store


def create_feature_view_from_rag_source(
    project: str,
    location: str,
    existing_feature_online_store_id: str,
    feature_view_id: str,
    bq_table_uri: str,
):
    aiplatform.init(project=project, location=location)
    fos = feature_store.FeatureOnlineStore(existing_feature_online_store_id)
    fv = fos.create_feature_view(
        name=feature_view_id,
        source=feature_store.utils.FeatureViewVertexRagSource(uri=bq_table_uri),
    )
    return fv

Fazer upload de dados e veiculação on-line

A API RAG processa o upload de dados e a exibição on-line.

Usar o Vertex AI Feature Store no RAG Engine

Depois que a instância do Feature Store da Vertex AI for configurada, as seções a seguir vão mostrar como configurá-la como o banco de dados vetorial a ser usado com o aplicativo RAG.

Use a instância do Feature Store da Vertex AI como o banco de dados de vetores para criar um corpus de RAG

Para criar o corpus RAG, use FEATURE_VIEW_RESOURCE_NAME. O corpus de RAG é criado e associado automaticamente à instância do Feature Store da Vertex AI. As APIs RAG usam o rag_corpus_id gerado para processar o upload de dados para a instância do Vertex AI Feature Store e recuperar contextos relevantes do rag_corpus_id.

Este exemplo de código demonstra como usar a instância do Feature Store da Vertex AI como o banco de dados vetorial para criar um corpus RAG.

REST

# TODO(developer): Update and uncomment the following lines:
# CORPUS_DISPLAY_NAME = "your-corpus-display-name"
#
# Full feature view resource name
# Format: projects/${PROJECT_ID}/locations/us-central1/featureOnlineStores/${FEATURE_ONLINE_STORE_ID}/featureViews/${FEATURE_VIEW_ID}
# FEATURE_VIEW_RESOURCE_NAME = "your-feature-view-resource-name"

# Call CreateRagCorpus API to create a new RAG corpus
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_ID}/locations/us-central1/ragCorpora -d '{
    "display_name" : '\""${CORPUS_DISPLAY_NAME}"\"',
    "rag_vector_db_config" : {
      "vertex_feature_store": {
        "feature_view_resource_name":'\""${FEATURE_VIEW_RESOURCE_NAME}"\"'
      }
    }
  }'

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

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 vertexai.preview import rag
import vertexai

# TODO(developer): Update and un-comment below lines
# PROJECT_ID = "your-project-id"
# feature_view_name = "projects/{PROJECT_ID}/locations/{LOCATION}/featureOnlineStores/{FEATURE_ONLINE_STORE_ID}/featureViews/{FEATURE_VIEW_ID}"
# display_name = "test_corpus"
# description = "Corpus Description"

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

# Configure embedding model (Optional)
embedding_model_config = rag.EmbeddingModelConfig(
    publisher_model="publishers/google/models/text-embedding-004"
)

# Configure Vector DB
vector_db = rag.VertexFeatureStore(resource_name=feature_view_name)

corpus = rag.create_corpus(
    display_name=display_name,
    description=description,
    embedding_model_config=embedding_model_config,
    vector_db=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=...
# ...

Importar arquivos para a tabela do BigQuery usando a API RAG

Use a API ImportRagFiles para importar arquivos do Google Cloud Storage ou do Google Drive para a tabela do BigQuery da instância do Vertex AI Feature Store. Os arquivos são incorporados e armazenados na tabela do BigQuery.

Este exemplo de código demonstra como importar arquivos para a tabela do BigQuery usando a API RAG.

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_ID}/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_ID}/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. 


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}"
# paths = ["https://drive.google.com/file/123", "gs://my_bucket/my_files_dir"]  # Supports Google Cloud Storage and Google Drive Links

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

response = rag.import_files(
    corpus_name=corpus_name,
    paths=paths,
    chunk_size=512,  # Optional
    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.

Executar um processo de sincronização para criar um índice FeatureOnlineStore

Depois de fazer upload dos dados na tabela do BigQuery, execute um processo de sincronização para disponibilizar os dados para veiculação on-line. É necessário gerar um índice FeatureOnlineStore usando o FeatureView. O processo de sincronização pode levar 20 minutos para ser concluído.

Este exemplo de código demonstra como executar um processo de sincronização para criar um índice FeatureOnlineStore.

REST

Antes de usar os dados da solicitação abaixo, faça as substituições a seguir:

  • LOCATION_ID: região em que a loja on-line está localizada, como us-central1.
  • PROJECT_ID: o ID do projeto.
  • FEATUREONLINESTORE_NAME: o nome da loja on-line que contém a visualização do recurso.
  • FEATUREVIEW_NAME: o nome da vista do elemento em que você quer iniciar manualmente a sincronização de dados.

Método HTTP e URL: 

POST https://LOCATION_ID-aiplatform.googleapis.com/v1/projects/PROJECT_ID/locations/LOCATION_ID/featureOnlineStores/FEATUREONLINESTORE_NAME/featureViews/FEATUREVIEW_NAME:sync

Para enviar a solicitação, escolha uma destas opções:

curl

execute o seguinte comando: 

curl -X POST \
     -H "Authorization: Bearer $(gcloud auth print-access-token)" \
     -H "Content-Type: application/json; charset=utf-8" \
     -d "" \
     "https://LOCATION_ID-aiplatform.googleapis.com/v1/projects/PROJECT_ID/locations/LOCATION_ID/featureOnlineStores/FEATUREONLINESTORE_NAME/featureViews/FEATUREVIEW_NAME:sync"

PowerShell

execute o seguinte comando: 

$cred = gcloud auth print-access-token
$headers = @{ "Authorization" = "Bearer $cred" }

Invoke-WebRequest `
    -Method POST `
    -Headers $headers `
    -Uri "https://LOCATION_ID-aiplatform.googleapis.com/v1/projects/PROJECT_ID/locations/LOCATION_ID/featureOnlineStores/FEATUREONLINESTORE_NAME/featureViews/FEATUREVIEW_NAME:sync" | Select-Object -Expand Content

Você receberá uma resposta JSON semelhante a esta:

{
  "featureViewSync": "projects/PROJECT_ID/locations/LOCATION_ID/featureOnlineStores/FEATUREONLINESTORE_NAME/featureViews/FEATUREVIEW_NAME/featureViewSyncs/OPERATION_ID"
}

Extrair contextos relevantes usando a API RAG

Depois que o processo de sincronização for concluído, você poderá extrair contextos relevantes do índice FeatureOnlineStore pela API RetrieveContexts.

REST

# TODO(developer): Update and uncomment the following lines:
# RETRIEVAL_QUERY="your-retrieval-query"
#
# Full RAG corpus resource name
# Format:
# "projects/${PROJECT_ID}/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_ID}/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. 


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
)
print(response)
# Example response:
# contexts {
#   contexts {
#     source_uri: "gs://your-bucket-name/file.txt"
#     text: "....
#   ....

Gerar conteúdo usando a API Gemini da Vertex AI

Chame a API GenerateContent da Vertex AI para usar modelos do Gemini para gerar conteúdo e especifique RAG_CORPUS_RESOURCE na solicitação para extrair dados do índice FeatureOnlineStore.

REST

# TODO(developer): Update and uncomment the following lines:
# MODEL_ID=gemini-1.5-flash-001
# 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_ID}/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,
      }
    }
  }
}'

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 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....
#   ...

A seguir