Uso de la búsqueda de vectores con un token de OAuth

En este cuaderno se muestra cómo llamar a un punto de conexión de búsqueda vectorial mediante el SDK de búsqueda vectorial o HTTP con un token de OAuth nuevo. En ambos casos, se usa la ruta de acceso optimizada de la red, como se recomienda para cualquier carga de trabajo de producción.

Las llamadas HTTP para crear un token y llamar al punto de conexión se pueden implementar en un lenguaje de su elección. En el caso de las aplicaciones de producción, tenga en cuenta que el token debe actualizarse cada 60 minutos. Para evitar errores debido a un token obsoleto, Databricks recomienda actualizarlo en un intervalo de menos de 60 minutos.

Configuración

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

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

Constantes de configuración

Para obtener más detalles sobre cómo crear un principal de servicio, consulte la documentación de 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')}"

Búsqueda de vectores de consulta usando el cliente de Python

Para más información, consulte la documentación de api.

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)

Búsqueda de vectores de llamada mediante una solicitud HTTP

En esta sección se muestra cómo llamar a un índice de búsqueda vectorial mediante HTTP, que luego puede implementar en una herramienta o idioma de su elección.

Generación de un token de OAuth mediante una llamada API

Para llamar a un punto de conexión mediante una ruta de acceso optimizada para red, necesita un token de OAuth. El código siguiente crea el 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")

Búsqueda de vectores de consulta mediante cliente de Python

Para usar JavaScript u otra herramienta para consultar la búsqueda de vectores, cree la consulta y envíela al punto de conexión de búsqueda vectorial.

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

Cuaderno de ejemplo

Uso de la búsqueda de vectores con un token de OAuth

Obtener el portátil