Använda vektorsökning med en OAuth-token

Den här notebook-filen visar hur du anropar en slutpunkt för vektorsökning med SDK för vektorsökning eller HTTP med en ny OAuth-token. I båda fallen används den optimerade nätverksvägen, vilket rekommenderas för varje produktionsarbetsbelastning.

HTTP-anropen för att skapa en token och anropa slutpunkten kan implementeras på ett valfritt språk. Tänk på att token måste uppdateras var 60:e minut för produktionsprogram. För att förhindra fel på grund av en inaktuell token rekommenderar Databricks att du uppdaterar den med ett intervall på mindre än 60 minuter.

Inställningar

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

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

Konfigurationskonstanter

Mer information om hur du skapar ett huvudnamn för tjänsten finns i Databricks-dokumentationen.

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

Frågevektorsökning med hjälp av Python-klienten

Mer information finns i API-dokumentationen.

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)

Anropa vektorsökning med hjälp av en HTTP-begäran

Det här avsnittet visar hur du anropar ett vektorsökningsindex med HJÄLP av HTTP, som du sedan kan implementera i ett verktyg eller valfritt språk.

Generera en OAuth-token med ett API-anrop

Om du vill anropa en slutpunkt med en nätverksoptimerad sökväg behöver du en OAuth-token. Följande kod skapar 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")

Frågevektorsökning med hjälp av Python klient

Om du vill använda JavaScript eller något annat verktyg för att fråga vektorsökning skapar du frågan och skickar den till slutpunkten för vektorsökning.

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

Exempelanteckningsbok

Använda vektorsökning med en OAuth-token

Hämta anteckningsbok