Usa la ricerca basata sull'IA con un token OAuth

Questo notebook illustra come chiamare un endpoint di ricerca di intelligenza artificiale usando AI Search SDK o HTTP con un nuovo token OAuth. In entrambi i casi viene usato il percorso ottimizzato per la rete, come consigliato per qualsiasi carico di lavoro di produzione.

Le chiamate HTTP per la creazione di un token e la chiamata dell'endpoint possono essere implementate in un linguaggio di propria scelta. Per le applicazioni di produzione, tenere presente che il token deve essere aggiornato ogni 60 minuti. Per evitare errori dovuti a un token non aggiornato, Databricks consiglia di aggiornarlo a intervalli inferiori a 60 minuti.

Configurazione

%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

Costanti di configurazione

Per informazioni dettagliate su come creare un Principal del Servizio, consultare la documentazione di 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')}"

Eseguire query sulla ricerca di intelligenza artificiale usando il client Python

Per informazioni dettagliate, vedere la documentazione dell'API.

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)

Chiamare ricerca di intelligenza artificiale usando una richiesta HTTP

Questa sezione illustra come chiamare un indice di ricerca di intelligenza artificiale usando HTTP, che è quindi possibile implementare in uno strumento o in un linguaggio di propria scelta.

Generare un token OAuth usando una chiamata API

Per chiamare un endpoint usando un percorso ottimizzato per la rete, è necessario un token OAuth. Il codice seguente crea il 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")

Eseguire query sulla ricerca di intelligenza artificiale usando Python client

Per usare JavaScript o un altro strumento per eseguire una query in AI Search, creare la query e inviarla all'endpoint di AI Search.

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

Notebook di esempio

Usa la ricerca basata sull'IA con un token OAuth

Ottieni il notebook