Verwenden der Vektorsuche mit einem OAuth-Token

Notizbuchversion dieser Seite öffnen

Dieses Notizbuch zeigt, wie Sie einen Vektorsuchendpunkt mithilfe des Vektorsuch-SDK oder HTTP mit einem neuen OAuth-Token aufrufen. In beiden Fällen wird der netzwerkoptimierte Pfad verwendet, wie für jede Produktionsauslastung empfohlen.

Die HTTP-Aufrufe zum Erstellen eines Tokens und das Aufrufen des Endpunkts können in einer Sprache Ihrer Wahl implementiert werden. Beachten Sie bei Produktionsanwendungen, dass das Token alle 60 Minuten aktualisiert werden muss. Um Fehler aufgrund eines veralteten Tokens zu vermeiden, empfiehlt Databricks, sie in einem Intervall von weniger als 60 Minuten zu aktualisieren.

Konfiguration

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

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

Konfigurationskonstanten

Ausführliche Informationen zum Erstellen eines Dienstprinzipals finden Sie in der Databricks-Dokumentation.

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

Abfragevektorsuche mit dem Python-Client

Ausführliche Informationen finden Sie in der API-Dokumentation.

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)

Aufrufen der Vektorsuche mithilfe einer HTTP-Anforderung

In diesem Abschnitt wird gezeigt, wie Sie einen Vektorsuchindex mit HTTP aufrufen, den Sie dann in einem Tool oder einer Sprache Ihrer Wahl implementieren können.

Generieren eines OAuth-Tokens mithilfe eines API-Aufrufs

Um einen Endpunkt mithilfe eines netzwerkoptimierten Pfads aufzurufen, benötigen Sie ein OAuth-Token. Der folgende Code erstellt das Token.

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

Abfragevektorsuche mit Python-Client

Um JavaScript oder ein anderes Tool zum Abfragen der Vektorsuche zu verwenden, erstellen Sie die Abfrage, und senden Sie sie an den Vektorsuchendpunkt.

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

Beispiel-Notebook

Verwenden der Vektorsuche mit einem OAuth-Token

Notebook abrufen