Vertex AI RAG Engine에서 Vertex AI Feature Store 사용

이 페이지에서는 Vertex AI Feature Store를 RAG Engine과 함께 사용할 벡터 데이터베이스로 설정하는 방법을 보여줍니다.

이 노트북 Vertex AI Feature Store를 사용한 RAG Engine을 사용하여 따라할 수도 있습니다.

RAG Engine은 Spanner를 기반으로 하는 기본 제공 벡터 데이터베이스를 사용하여 텍스트 문서의 벡터 표현을 저장하고 관리합니다. 벡터 데이터베이스는 문서가 특정 쿼리와 얼마나 의미상 유사한지를 기반으로 관련 문서를 검색합니다.

Vertex AI Feature Store를 추가 벡터 데이터베이스로 통합하면 RAG Engine에서 Vertex AI Feature Store를 사용하여 짧은 지연 시간으로 대용량 데이터를 처리할 수 있으므로 RAG 애플리케이션의 성능과 확장성을 향상시키는 데 도움이 됩니다.

Vertex AI Feature Store 설정

관리형 클라우드 네이티브 서비스인 Vertex AI Feature Store는 Vertex AI의 필수 구성요소입니다. BigQuery 테이블이나 뷰 내에서 특성 데이터를 관리할 수 있어 머신러닝(ML) 특성 관리 및 온라인 서빙을 간소화합니다. 이를 통해 지연 시간이 짧은 온라인 특성 서빙이 가능해집니다.

최적화된 온라인 서빙으로 만든 FeatureOnlineStore 인스턴스의 경우 벡터 유사성 검색을 활용하여 의미상 유사하거나 관련 있는 항목의 목록(근사 최근접 이웃이라고 함)을 검색할 수 있습니다.

다음 섹션에서는 RAG 애플리케이션용 Vertex AI Feature Store 인스턴스를 설정하는 방법을 보여줍니다.

BigQuery 테이블 스키마 만들기

Google Cloud 콘솔을 사용하여 BigQuery 테이블 스키마를 만듭니다. 데이터 소스로 사용하려면 다음 필드가 포함되어야 합니다.

필드 이름 데이터 유형 상태
corpus_id String 필수
file_id String 필수
chunk_id String 필수
chunk_data_type String Null 허용
chunk_data String Null 허용
file_original_uri String Null 허용
embeddings Float 반복

이 코드 샘플에서는 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
  );

FeatureOnlineStore 인스턴스 프로비저닝

특성 온라인 서빙을 사용 설정하려면 Vertex AI Feature Store CreateFeatureOnlineStore API를 사용하여 FeatureOnlineStore 인스턴스를 설정합니다. FeatureOnlineStore를 처음 프로비저닝하는 경우 작업이 완료되는 데 약 5분이 걸릴 수 있습니다.

REST

온라인 상점 인스턴스를 만들려면 featureOnlineStores.create 메서드를 사용하여 POST 요청을 보냅니다.

요청 데이터를 사용하기 전에 다음을 바꿉니다.

  • LOCATION_ID: FeatureOnlineStore 인스턴스를 만들 리전입니다(예: us-central1).
  • PROJECT_ID: 프로젝트 ID입니다.
  • FEATUREONLINESTORE_NAME: 새 FeatureOnlineStore 인스턴스의 이름입니다.

HTTP 메서드 및 URL:

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

JSON 요청 본문:

{
  "optimized": {}
}

요청을 보내려면 다음 옵션 중 하나를 선택합니다.

curl

요청 본문을 request.json 파일에 저장하고 다음 명령어를 실행합니다.

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

요청 본문을 request.json 파일에 저장하고 다음 명령어를 실행합니다.

$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

다음과 비슷한 JSON 응답이 표시됩니다.

{
  "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용 Vertex AI SDK

Python용 Vertex AI SDK를 설치하거나 업데이트하는 방법은 Python용 Vertex AI SDK 설치를 참조하세요. 자세한 내용은 Python용 Vertex AI SDK API 참조 문서를 확인하세요.


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: 프로젝트 ID입니다.
  • location: FeatureOnlineStore 인스턴스를 만들 리전입니다(예: us-central1).
  • feature_online_store_id: 새 FeatureOnlineStore 인스턴스의 이름입니다.

FeatureView 리소스 만들기

특성 데이터 소스를 저장하는 BigQuery 테이블을 FeatureOnlineStore 인스턴스에 연결하려면 CreateFeatureView API를 호출하여 FeatureView 리소스를 만듭니다. FeatureView 리소스를 만들 때는 내적 음수로 정의된 기본 거리 측정항목 DOT_PRODUCT_DISTANCE을 선택합니다(DOT_PRODUCT_DISTANCE가 작을수록 유사성이 높음).

이 코드 샘플에서는 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용 Vertex AI SDK

Python용 Vertex AI SDK를 설치하거나 업데이트하는 방법은 Python용 Vertex AI SDK 설치를 참조하세요. 자세한 내용은 Python용 Vertex AI SDK API 참조 문서를 확인하세요.


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

데이터 업로드 및 온라인 서빙

RAG API는 데이터 업로드와 온라인 서빙을 처리합니다.

RAG Engine에서 Vertex AI Feature Store 사용

Vertex AI Feature Store 인스턴스가 설정된 후 다음 섹션에서는 RAG 애플리케이션과 함께 사용할 벡터 데이터베이스로 설정하는 방법을 보여줍니다.

Vertex AI Feature Store 인스턴스를 벡터 데이터베이스로 사용하여 RAG 코퍼스 만들기

RAG 코퍼스를 만들려면 FEATURE_VIEW_RESOURCE_NAME을 사용해야 합니다. RAG 코퍼스가 생성되고 Vertex AI Feature Store 인스턴스와 자동으로 연결됩니다. RAG API는 생성된 rag_corpus_id를 사용하여 데이터를 Vertex AI Feature Store 인스턴스에 업로드하고 rag_corpus_id에서 관련 컨텍스트를 검색합니다.

이 코드 샘플에서는 Vertex AI Feature Store 인스턴스를 벡터 데이터베이스로 사용하여 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용 Vertex AI SDK

Python용 Vertex AI SDK를 설치하거나 업데이트하는 방법은 Python용 Vertex AI SDK 설치를 참조하세요. 자세한 내용은 Python용 Vertex AI SDK API 참조 문서를 확인하세요.


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

RAG API를 사용하여 BigQuery 테이블에 파일 가져오기

ImportRagFiles API를 사용하여 Google Cloud Storage 또는 Google Drive의 파일을 Vertex AI Feature Store 인스턴스의 BigQuery 테이블로 가져옵니다. 파일은 BigQuery 테이블에 삽입되어 저장됩니다.

이 코드 샘플에서는 RAG API를 사용하여 파일을 BigQuery 테이블로 가져오는 방법을 보여줍니다.

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용 Vertex AI SDK

Python용 Vertex AI SDK를 설치하거나 업데이트하는 방법은 Python용 Vertex AI SDK 설치를 참조하세요. 자세한 내용은 Python용 Vertex AI SDK API 참조 문서를 확인하세요.


from vertexai 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,
    transformation_config=rag.TransformationConfig(
        rag.ChunkingConfig(chunk_size=512, chunk_overlap=100)
    ),
    max_embedding_requests_per_min=900,  # Optional
)
print(f"Imported {response.imported_rag_files_count} files.")
# Example response:
# Imported 2 files.

동기화 프로세스를 실행하여 FeatureOnlineStore 색인 생성

데이터를 BigQuery 테이블에 업로드한 후 동기화 프로세스를 실행하여 데이터를 온라인 서빙에 사용할 수 있게 합니다. FeatureView를 사용하여 FeatureOnlineStore 색인을 생성해야 하며 동기화 프로세스가 완료되는 데 20분이 걸릴 수 있습니다.

이 코드 샘플에서는 동기화 프로세스를 실행하여 FeatureOnlineStore 색인을 생성하는 방법을 보여줍니다.

REST

요청 데이터를 사용하기 전에 다음을 바꿉니다.

  • LOCATION_ID: 온라인 저장소가 있는 리전입니다(예: us-central1).
  • PROJECT_ID: 프로젝트 ID입니다.
  • FEATUREONLINESTORE_NAME: 특성 뷰가 포함된 온라인 스토어의 이름입니다.
  • FEATUREVIEW_NAME: 데이터 동기화를 수동으로 시작하려는 특성 뷰의 이름입니다.

HTTP 메서드 및 URL:

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

요청을 보내려면 다음 옵션 중 하나를 선택합니다.

curl

다음 명령어를 실행합니다.

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

다음 명령어를 실행합니다.

$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

다음과 비슷한 JSON 응답이 표시됩니다.

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

RAG API를 사용하여 관련 컨텍스트 가져오기

동기화 프로세스가 완료되면 RetrieveContexts API를 통해 FeatureOnlineStore 색인에서 관련 컨텍스트를 검색할 수 있습니다.

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용 Vertex AI SDK

Python용 Vertex AI SDK를 설치하거나 업데이트하는 방법은 Python용 Vertex AI SDK 설치를 참조하세요. 자세한 내용은 Python용 Vertex AI SDK API 참조 문서를 확인하세요.


from vertexai 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!",
    rag_retrieval_config=rag.RagRetrievalConfig(
        top_k=10,
        filter=rag.utils.resources.Filter(vector_distance_threshold=0.5),
    ),
)
print(response)
# Example response:
# contexts {
#   contexts {
#     source_uri: "gs://your-bucket-name/file.txt"
#     text: "....
#   ....

Vertex AI Gemini API를 사용하여 콘텐츠 생성

Vertex AI GenerateContent API를 호출하여 Gemini 모델을 사용해 콘텐츠를 생성하고 요청에 RAG_CORPUS_RESOURCE를 지정하여 FeatureOnlineStore 색인에서 데이터를 검색합니다.

REST

# TODO(developer): Update and uncomment the following lines:
# MODEL_ID=gemini-2.0-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_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용 Vertex AI SDK

Python용 Vertex AI SDK를 설치하거나 업데이트하는 방법은 Python용 Vertex AI SDK 설치를 참조하세요. 자세한 내용은 Python용 Vertex AI SDK API 참조 문서를 확인하세요.


from vertexai import rag
from vertexai.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", ...],
                )
            ],
            rag_retrieval_config=rag.RagRetrievalConfig(
                top_k=10,
                filter=rag.utils.resources.Filter(vector_distance_threshold=0.5),
            ),
        ),
    )
)

rag_model = GenerativeModel(
    model_name="gemini-2.0-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....
#   ...

다음 단계