Verwenden der KI-Suche mit einem OAuth-Token

Dieses Notizbuch zeigt, wie Sie einen KI-Suchendpunkt mithilfe des AI Search 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-ai-search
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')}"

Abfragen der KI-Suche mithilfe des Python-Clients

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

from databricks.ai_search.client import AISearchClient

vsc_dp = AISearchClient(
  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 KI-Suche mithilfe einer HTTP-Anforderung

In diesem Abschnitt wird gezeigt, wie Sie einen KI-Suchindex 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")

Abfragen der KI-Suche mit Python-Client

Um JavaScript oder ein anderes Tool zum Abfragen der KI-Suche zu verwenden, erstellen Sie die Abfrage, und senden Sie sie an den AI Search-Endpunkt.

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 KI-Suche mit einem OAuth-Token

Notebook abrufen