Uso de ejemplo del SDK de AI Search Python

Abrir la versión 'notebook' de esta página

En este cuaderno se muestra cómo usar ai Search Python SDK, que proporciona AISearchClient como API principal para trabajar con AI Search.

Como alternativa, puede llamar directamente a la API REST.

Requirements

Este cuaderno de trabajo da por hecho que existe un punto de conexión de Modelo de servicio denominado databricks-gte-large-en. Para crear ese punto de conexión, consulta el cuaderno de trabajo Llamar a un modelo de incrustaciones de GTE mediante Modelo de servicio.

%pip install --upgrade --force-reinstall databricks-ai-search langchain
dbutils.library.restartPython()

from databricks.ai_search.client import AISearchClient
vsc = AISearchClient()

help(AISearchClient)

Carga el conjunto de datos de ejemplo en la tabla Delta de origen

A continuación se crea la tabla Delta de origen.

# Specify the catalog and schema to use. You must have USE_CATALOG privilege on the catalog and USE_SCHEMA and CREATE_TABLE privileges on the schema.
# Change the catalog and schema here if necessary.

catalog_name = "main"
schema_name = "default"

source_table_name = "en_wiki"
source_table_fullname = f"{catalog_name}.{schema_name}.{source_table_name}"

# Uncomment this line to start from scratch.

# spark.sql(f"DROP TABLE {source_table_fullname}")

source_df = spark.read.parquet("/databricks-datasets/wikipedia-datasets/data-001/en_wikipedia/articles-only-parquet").limit(10)
display(source_df)

source_df.write.format("delta").option("delta.enableChangeDataFeed", "true").saveAsTable(source_table_fullname)

display(spark.sql(f"SELECT * FROM {source_table_fullname}"))

Creación de un punto de conexión

ai_search_endpoint_name = "ai-search-demo-endpoint"

vsc.create_endpoint(
    name=ai_search_endpoint_name,
    endpoint_type="STANDARD" # or "STORAGE_OPTIMIZED"
)

endpoint = vsc.get_endpoint(
  name=ai_search_endpoint_name)
endpoint

Crear índice

# AI Search index
vs_index = "en_wiki_index"
vs_index_fullname = f"{catalog_name}.{schema_name}.{vs_index}"

embedding_model_endpoint = "databricks-qwen3-embedding-0-6b"

index = vsc.create_delta_sync_index(
  endpoint_name=ai_search_endpoint_name,
  source_table_name=source_table_fullname,
  index_name=vs_index_fullname,
  pipeline_type='TRIGGERED',
  primary_key="id",
  embedding_source_column="text",
  embedding_model_endpoint_name=embedding_model_endpoint
)
index.describe()

Obtención del índice

Use get_index() para recuperar el objeto de índice de AI Search usando el nombre del índice. También puede usar describe() en el objeto de índice para ver un resumen de la información de configuración del índice.

index = vsc.get_index(endpoint_name=ai_search_endpoint_name, index_name=vs_index_fullname)

index.describe()

# Wait for index to come online. Expect this command to take several minutes.
import time
while not index.describe().get('status').get('detailed_state').startswith('ONLINE'):
  print("Waiting for index to be ONLINE...")
  time.sleep(5)
print("Index is ONLINE")
index.describe()

Búsqueda de similitud

Consulte el índice de búsqueda de IA para buscar documentos similares.

# Returns [col1, col2, ...]
# You can set this to any subset of the columns.
all_columns = spark.table(source_table_fullname).columns

results = index.similarity_search(
  query_text="Greek myths",
  columns=all_columns,
  num_results=2)

results

# Search with a filter. Note that the syntax depends on the endpoint type.

# Standard endpoint syntax
results = index.similarity_search(
  query_text="Greek myths",
  columns=all_columns,
  filters={"id NOT": ("13770", "88231")},
  num_results=2)

# Storage-optimized endpoint syntax
# results = index.similarity_search(
#   query_text="Greek myths",
#   columns=all_columns,
#   filters='id NOT IN ("13770", "88231")',
#   num_results=2)

results

Conversión de resultados en documentos langChain

La primera columna recuperada se carga en page_contenty el resto en metadatos.

from langchain_core.documents import Document
from typing import List

def convert_vector_search_to_documents(results) -> List[Document]:
  column_names = []
  for column in results["manifest"]["columns"]:
      column_names.append(column)

  langchain_docs = []
  for item in results["result"]["data_array"]:
      metadata = {}
      score = item[-1]
      # print(score)
      i = 1
      for field in item[1:-1]:
          # print(field + "--")
          metadata[column_names[i]["name"]] = field
          i = i + 1
      doc = Document(page_content=item[0], metadata=metadata)  # , 9)
      langchain_docs.append(doc)
  return langchain_docs

langchain_docs = convert_vector_search_to_documents(results)

langchain_docs

Elimina el índice.

vsc.delete_index(index_name=vs_index_fullname)

Cuaderno de ejemplo

Uso de ejemplo del SDK de AI Search Python

Obtención del cuaderno