Vertex AI Search Retriever

Vertex AI Search 데이터스토어에서 데이터를 가져올 수 있는 Retriever를 설정하는 노트북이에요. 구조화·비구조화·웹사이트 데이터스토어를 각각 테스트하는 방법을 보여드릴게요.

출처: 문서

본문

이 노트북은 Vertex AI Search 데이터스토어에서 가져올 수 있는 Retriever를 설정하는 과정을 안내합니다.

사전 요구 사항 (Pre-requirements)

Google Cloud 프로젝트 설정
Vertex AI Search 데이터스토어 설정
Vertex AI API 활성화

라이브러리 설치

%pip install llama-index-retrievers-vertexai-search

현재 런타임 재시작

새로 설치한 패키지를 이 Jupyter 런타임에서 사용하려면 런타임을 재시작해야 합니다. 아래 셀을 실행하면 현재 커널이 재시작됩니다.

# Colab only
# Automatically restart kernel after installs so that your environment can access the new packages
import IPython


app = IPython.Application.instance()
app.kernel.do_shutdown(True)

노트북 환경 인증 (Colab 전용)

이 노트북을 Google Colab에서 실행한다면 환경을 인증해야 합니다. 이를 위해 아래 새 셀을 실행하세요. Vertex AI Workbench를 사용한다면 이 단계는 필요하지 않습니다.

# Colab only
import sys


if "google.colab" in sys.modules:
    from google.colab import auth


    auth.authenticate_user()

# If you're using JupyterLab instance, uncomment and run the below code.
#!gcloud auth login

from llama_index.retrievers.vertexai_search import VertexAISearchRetriever


# Please note it's underscore '_' in vertexai_search

Google Cloud 프로젝트 정보 설정 및 Vertex AI SDK 초기화

Vertex AI를 시작하려면 기존 Google Cloud 프로젝트가 있어야 하고 Vertex AI API를 활성화해야 합니다.

프로젝트 및 개발 환경 설정에 대해 더 자세히 알아보세요.

PROJECT_ID = "{your project id}"  # @param {type:"string"}
LOCATION = "us-central1"  # @param {type:"string"}
import vertexai


vertexai.init(project=PROJECT_ID, location=LOCATION)

구조화(Structured) 데이터스토어 테스트

DATA_STORE_ID = "{your id}"  # @param {type:"string"}
LOCATION_ID = "global"

struct_retriever = VertexAISearchRetriever(
    project_id=PROJECT_ID,
    data_store_id=DATA_STORE_ID,
    location_id=LOCATION_ID,
    engine_data_type=1,
)

query = "harry potter"
retrieved_results = struct_retriever.retrieve(query)

print(retrieved_results[0])

구조화 데이터스토어는 engine_data_type=1로 지정합니다. 이는 검색 가능한 구조화 데이터(예: 지식 그래프)를 대상으로 합니다.

비구조화(Unstructured) 데이터스토어 테스트

DATA_STORE_ID = "{your id}"
LOCATION_ID = "global"

unstruct_retriever = VertexAISearchRetriever(
    project_id=PROJECT_ID,
    data_store_id=DATA_STORE_ID,
    location_id=LOCATION_ID,
    engine_data_type=0,
)

query = "alphabet 2018 earning"
retrieved_results2 = unstruct_retriever.retrieve(query)

print(retrieved_results2[0])

비구조화 데이터스토어는 engine_data_type=0으로 지정합니다. 이는 문서/PDF 같은 비구조화 데이터를 대상으로 합니다.

웹사이트(Website) 데이터스토어 테스트

DATA_STORE_ID = "{your id}"
LOCATION_ID = "global"
website_retriever = VertexAISearchRetriever(
    project_id=PROJECT_ID,
    data_store_id=DATA_STORE_ID,
    location_id=LOCATION_ID,
    engine_data_type=2,
)

query = "what's diamaxol"
retrieved_results3 = website_retriever.retrieve(query)

print(retrieved_results3[0])

웹사이트 데이터스토어는 engine_data_type=2로 지정합니다. 이는 웹 콘텐츠를 대상으로 합니다.

Query Engine에서 사용하기

# import modules needed
from llama_index.core import Settings
from llama_index.llms.vertex import Vertex
from llama_index.embeddings.vertex import VertexTextEmbedding

vertex_gemini = Vertex(
    model="gemini-1.5-pro",
    temperature=0,
    context_window=100000,
    additional_kwargs={},
)
# setup the index/query llm
Settings.llm = vertex_gemini

from llama_index.core.query_engine import RetrieverQueryEngine


query_engine = RetrieverQueryEngine.from_args(struct_retriever)

response = query_engine.query("Tell me about harry potter")
print(str(response))

Vertex AI Search Retriever