Utiliser la recherche AI avec un jeton OAuth

Ouvrir dans Databricks

Ce notebook montre comment appeler un Endpoint de recherche IA en utilisant le SDK de recherche IA ou HTTP avec un jeton OAuth frais. Dans les deux cas, le chemin optimisé par le réseau est utilisé, comme cela est recommandé pour toute workload de production.

Les appels HTTP pour créer un jeton et appeler l'Endpoint peuvent être implémentés dans la langue de votre choix. Pour les applications de production, gardez à l'esprit que le jeton doit être refreshé toutes les 60 minutes. Pour éviter les erreurs dues à un jeton obsolète, Databricks recommande de le refresher à un intervalle de moins de 60 minutes.

Configuration

Python

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

Python

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

Constantes de configuration

Pour plus de détails sur la création d'un Service Principal, consultez la documentation Databricks.

Python

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

Interroger AI Search à l'aide du client Python

Pour plus de détails, consultez la documentation de l'API.

Python

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)

Appeler la Recherche IA à l'aide d'une requête HTTP

Cette section montre comment appeler un index de recherche IA à l’aide de HTTP, que vous pouvez ensuite implémenter dans un outil ou un langage de votre choix.

Générer un jeton OAuth à l’aide d’un appel d’API

Pour appeler un Endpoint à l'aide d'un chemin optimisé pour le réseau, vous avez besoin d'un jeton OAuth. Le code suivant crée le jeton.

Python

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

Interroger la recherche IA à l'aide du client Python

Pour utiliser JavaScript ou un autre outil pour query AI Search, créez la query et envoyez-la à l'endpoint AI Search.

Python

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

Exemple de Notebook

Utilisez AI Search avec un jeton OAuth

Ouvrir le notebook dans un nouvel ongletOuvrir dans Databricks