Skip to main content
Last updated on

Use AI Search with an OAuth token

This notebook shows how to call an AI Search endpoint using the AI Search SDK or HTTP with a fresh OAuth token. In both cases, the network optimized path is used, as is recommended for any production workload.

The HTTP calls for creating a token and calling the endpoint can be implemented in a language of your choice. For production applications, keep in mind that the token must be refreshed every 60 minutes. To prevent errors due to a stale token, Databricks recommends refreshing it at an interval of less than 60 minutes.

Setup​

Python
%pip install databricks-sdk
%pip install databricks-vectorsearch
dbutils.library.restartPython()
Python
import requests
import json
import random
from databricks.sdk import WorkspaceClient
import logging

Configuration Constants​

For details on how to create a service principal, see the Databricks documentation.

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

Query AI Search using the Python client​

For details, see the API documentation.

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

Call AI Search using an HTTP request​

This section shows how to call an AI Search index using HTTP, which you can then implement in a tool or language of your choice.

Generate an OAuth token using an API call​

To call an endpoint using a network optimized path, you need an OAuth token. The following code creates the token.

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

Query AI Search using Python client​

To use JavaScript or some other tool to query AI Search, create the query and send it to the AI Search endpoint.

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

Example notebook​

Use AI Search with an OAuth token

Open notebook in new tab
On this page