OAuthトークンでAI Searchを使う
このノートブックでは、新しいOAuthトークンを使用して、AI Search SDKまたはHTTPを使用してAI Searchエンドポイントを呼び出す方法を示します。いずれの場合も、あらゆる本番運用ワークロードに推奨されているため、ネットワーク最適化されたパスが使用されます。
トークンの作成とエンドポイントの呼び出しを行うHTTP呼び出しは、お好みの言語で実装できます。本番環境で使用する場合は、トークンを60分ごとに更新する必要があることに注意してください。古いトークンによるエラーを防ぐため、Databricksでは60分未満の間隔で更新することをお勧めします。
セットアップ
Python
%pip install databricks-sdk
%pip install databricks-vectorsearch
dbutils.library.restartPython()
Python
import requests
import json
import random
from databricks.sdk import WorkspaceClient
import logging
構成定数
サービスプリンシパルを作成する方法の詳細については、「Databricksドキュメント」を参照してください。
Python
# 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 クライアントを使用して AI 検索をクエリします
詳細については、API ドキュメントを参照してください。
Python
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 リクエストを使用して AI 検索を呼び出す
このセクションでは、HTTPを使用してAI検索インデックスを呼び出す方法を説明します。この方法は、選択したツールまたは言語で実装できます。
APIコールを使用してOAuthトークンを生成します
ネットワーク最適化されたパスを使用してエンドポイントを呼び出すには、OAuthトークンが必要です。次のコードでトークンを作成します。
Python
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クライアントを使用してAI検索をクエリします
JavaScript やその他のツールを使用して AI Search をクエリする場合は、クエリを作成し、AI Search エンドポイントに送信します。
Python
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))