Vector Search gebruiken met een OAuth-token

In dit notebook ziet u hoe u een Vector Search-eindpunt aanroept met behulp van de Vector Search SDK of HTTP met een nieuw OAuth-token. In beide gevallen wordt het netwerk-geoptimaliseerde pad gebruikt, zoals wordt aanbevolen voor elke productiebelasting.

De HTTP-aanroepen voor het maken van een token en het aanroepen van het eindpunt kunnen worden geïmplementeerd in een taal van uw keuze. Houd er voor productietoepassingen rekening mee dat het token elke 60 minuten moet worden vernieuwd. Om fouten vanwege een verlopen token te voorkomen, raadt Databricks aan om het te vernieuwen met een interval van minder dan 60 minuten.

Installatie

%pip install databricks-sdk
%pip install databricks-vectorsearch
dbutils.library.restartPython()

import requests
import json
import random
from databricks.sdk import WorkspaceClient
import logging

Configuratieconstanten

Zie de Databricks-documentatie voor meer informatie over het maken van een service-principal.

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

Query Vector Search met behulp van de Python-client

Zie de API-documentatie voor meer informatie.

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)

Vector Search aanroepen met behulp van een HTTP-aanvraag

In deze sectie wordt beschreven hoe u een vectorzoekindex aanroept met behulp van HTTP, die u vervolgens kunt implementeren in een hulpprogramma of taal van uw keuze.

Een OAuth-token genereren met behulp van een API-aanroep

Als u een eindpunt wilt aanroepen met behulp van een geoptimaliseerd netwerkpad, hebt u een OAuth-token nodig. Met de volgende code wordt het token gemaakt.

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

Query Vector Search met behulp van Python-client

Als u JavaScript of een ander hulpprogramma wilt gebruiken om een query uit te voeren op vectorzoekopdrachten, maakt u de query en verzendt u deze naar het eindpunt voor vectorzoekopdrachten.

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))

Voorbeeld van notebook

Vector Search gebruiken met een OAuth-token

Notebook krijgen