RAG용 Vertex AI의 LlamaIndex에서 Vertex AI Feature Store 사용

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

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

Vertex AI Feature Store를 추가 벡터 데이터베이스로 통합하면 RAG용 Vertex AI의 LlamaIndex에서 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 테이블 스키마를 정의하는 방법을 보여줍니다.

  # Use this sql query as reference for creating the table
  CREATE TABLE `your-project-id.input_us_central1.rag_source_new` (
    `corpus_id` STRING ,
    `file_id` STRING,
    `chunk_id` STRING,
    `chunk_data_type` STRING,
    `chunk_data` STRING,
    `embeddings` ARRAY<FLOAT64>,
    `file_original_uri` STRING,
  );

FeatureOnlineStore 인스턴스 프로비저닝

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

  # TODO(developer): Update and uncomment the following lines:
  # PROJECT_ID = "your-project-id"
  #
  # Set feature_online_store_id. 
  # Example: "rag_fos_test"
  # FEATURE_ONLINE_STORE_ID="your-feature-online-store-id"

  # Call CreateFeatureOnlineStore to create a FeatureOnlineStore instance
  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=${FEATURE_ONLINE_STORE_ID}   -d '{
      "optimized": {},
  }'

  # TODO(developer): Update and uncomment the following lines:
  # Get operation_id returned in CreateFeatureOnlineStore
  # OPERATION_ID="your-operation-id"

  # Poll Operation status until done = true in the response
  curl -X GET \
  -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/operations/${OPERATION_ID}

FeatureView 리소스 만들기

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

이 코드 샘플에서는 FeatureView 리소스를 만드는 방법을 보여줍니다.

  # 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

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

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

RAG용 Vertex AI의 LlamaIndex에서 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 코퍼스를 만드는 방법을 보여줍니다.

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

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

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

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

  # 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

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

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

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

  # Call Feature Store SyncFeatureView API to run the synchronization process
  curl   "https://us-central1-aiplatform.googleapis.com/v1/${FEATURE_VIEW_RESOURCE_NAME}:sync" \
    -X POST \
    -H "Authorization: Bearer $(gcloud auth print-access-token)" \
    -H "Content-Type: application/json; charset=utf-8"

  # TODO(developer): Update and uncomment the following lines:
  # Call Vertex AI Feature Store GetFeatureViewSync API to check the running synchronization // status
  # FEATURE_VIEW_SYNC_ID = "your-feature-view-sync-id" returned in SyncFeatureView
  curl   "https://us-central1-aiplatform.googleapis.com/v1/${FEATURE_VIEW_RESOURCE_NAME}/featureViewSyncs/${FEATURE_VIEW_SYNC_ID}" \
    -X GET \
    -H "Authorization: Bearer $(gcloud auth print-access-token)" \
    -H "Content-Type: application/json; charset=utf-8"

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

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

  # 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
      }
    }'

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

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

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

다음 단계