OAuth 토큰과 함께 벡터 검색 사용

이 Notebook은 새 OAuth 토큰이 있는 벡터 검색 SDK 또는 HTTP를 사용하여 벡터 검색 엔드포인트를 호출하는 방법을 보여줍니다. 두 경우 모두 프로덕션 워크로드에 권장되는 것처럼 네트워크 최적화 경로가 사용됩니다.

토큰을 만들고 엔드포인트를 호출하는 HTTP 호출은 원하는 언어로 구현할 수 있습니다. 프로덕션 애플리케이션의 경우 60분마다 토큰을 새로 고쳐야 합니다. 부실 토큰으로 인한 오류를 방지하기 위해 Databricks는 60분 미만의 간격으로 새로 고치는 것이 좋습니다.

설치

%pip install databricks-sdk
%pip install databricks-vectorsearch
dbutils.library.restartPython()
import requests
import json
import random
from databricks.sdk import WorkspaceClient
import logging

구성 상수

서비스 주체를 만드는 방법에 대한 자세한 내용은 Databricks 설명서를 참조하세요.

# Define the secret ID and secret for the service principal
CLIENT_ID = dbutils.secrets.get(scope="scope", key="service_principal_client_id")
CLIENT_SECRET = dbutils.secrets.get(scope="scope", key="service_principal_client_secret")

# You can get it by clicking the copy button next to the index name
INDEX_NAME="UC_CATALOG_INDEX_NAME"
ENDPOINT_NAME="ENDPOINT_NAME"

workspace_url = f"https://{spark.conf.get('spark.databricks.workspaceUrl')}"

Python 클라이언트를 사용하여 벡터 검색 쿼리

자세한 내용은 API 설명서를 참조하세요.

from databricks.vector_search.client import VectorSearchClient

vsc_dp = VectorSearchClient(
  service_principal_client_id=CLIENT_ID,
  service_principal_client_secret=CLIENT_SECRET,
  workspace_url=workspace_url)

index = vsc_dp.get_index(endpoint_name=ENDPOINT_NAME, index_name=INDEX_NAME)
index.similarity_search(["text"], query_vector=[0]*2560, num_results=5, debug_level=0)

HTTP 요청을 사용하여 벡터 검색 호출

이 섹션에서는 HTTP를 사용하여 벡터 검색 인덱스 호출 방법을 보여 줍니다. 그러면 선택한 도구 또는 언어로 구현할 수 있습니다.

API 호출을 사용하여 OAuth 토큰 생성

네트워크 최적화 경로를 사용하여 엔드포인트를 호출하려면 OAuth 토큰이 필요합니다. 다음 코드는 토큰을 만듭니다.

url = f"{workspace_url}/oidc/v1/token"
deets = json.dumps([
                {
                    "type": "unity_catalog_permission",
                    "securable_type": "table",
                    "securable_object_name": INDEX_NAME,
                    "operation": ("ReadVectorIndex"),
                },
            ])
payload = { 'grant_type': 'client_credentials', 'scope': 'all-apis', 'authorization_details': deets}

response = requests.post(
    url=url,
    auth=(CLIENT_ID, CLIENT_SECRET),
    headers={"Content-Type": "application/x-www-form-urlencoded"},
    data=payload,
)
if response.status_code != 200:
    logging.error(f"OAuth token request failed: {response.status_code} - {response.text}")
    response.raise_for_status()

token_data = response.json()
access_token = token_data.get("access_token")
if not access_token:
    raise ValueError("Failed to get access token")

Python 클라이언트를 사용하여 벡터 검색 쿼리

JavaScript 또는 다른 도구를 사용하여 벡터 검색을 쿼리하려면 쿼리를 만들고 벡터 검색 엔드포인트로 보냅니다.

index = vsc_dp.get_index(endpoint_name=ENDPOINT_NAME, index_name=INDEX_NAME)
index_url = index.index_url
print(index_url)

headers = {"Authorization": f"Bearer {access_token}"}
payload = {
    "query_vector": [0]*2560,
    "num_results": 5,
    "columns": ["text"]}

response = requests.post(index_url+"/query", headers=headers, data=json.dumps(payload))
print(response.status_code)
print(json.dumps(response.json(), indent=2))

예제 노트

OAuth 토큰과 함께 벡터 검색 사용

노트북 받기