Utilize a Pesquisa Vetorial com um token OAuth

Abra a versão de caderno desta página

Este notebook mostra como acessar um endpoint de Pesquisa Vetorial usando o SDK de Pesquisa Vetorial ou HTTP com um novo token OAuth. Em ambos os casos, usa-se o caminho otimizado pela rede, conforme recomendado para qualquer carga de trabalho de produção.

As chamadas HTTP para criar um token e chamar o endpoint podem ser implementadas em uma linguagem de sua escolha. Para aplicativos de produção, tenha em mente que o token deve ser atualizado a cada 60 minutos. Para evitar erros devido a um token obsoleto, o Databricks recomenda atualizá-lo em um intervalo de menos de 60 minutos.

Configuração

%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 configuração

Para obter detalhes sobre como criar uma entidade de serviço, consulte a documentação do 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')}"

Consultar a Pesquisa de Vetores usando o cliente Python

Para obter detalhes, consulte a documentação da 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)

Chamar a Pesquisa de Vetor usando uma solicitação HTTP

Esta seção mostra como chamar um índice de pesquisa de vetor usando HTTP, que você pode implementar em uma ferramenta ou linguagem de sua escolha.

Gerar um token OAuth usando uma chamada à API

Para chamar um endpoint usando um caminho otimizado para a rede, é necessário um token OAuth. O código a seguir cria o 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")

Pesquisa de Vetor de Consulta usando cliente Python

Para usar JavaScript ou alguma outra ferramenta para consultar o vetor de pesquisa, crie a consulta e envie-a para o endpoint vetor de pesquisa.

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

Caderno de exemplo

Utilize a Pesquisa Vetorial com um token OAuth

Obter laptop

  • Last updated on