将矢量搜索与 OAuth 令牌配合使用

打开此页面的笔记本版本

此笔记本演示如何使用矢量搜索 SDK 或 HTTP 和新的 OAuth 令牌调用矢量搜索终结点。 在这两种情况下，都使用网络优化路径，这也是任何生产负载的推荐做法。

可以采用所选语言实现用于创建令牌和调用终结点的 HTTP 调用。 对于生产应用程序，请记住，令牌必须每隔 60 分钟刷新一次。 为了防止由于令牌失效而导致的错误，Databricks 建议在小于 60 分钟的间隔内刷新令牌。

Setup

%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 令牌配合使用

获取笔记本

  • Last updated on