AI Search Python SDK-Beispielverwendung

Notizbuchversion dieser Seite öffnen

Dieses Notebook zeigt, wie Sie das AI Search Python SDK verwenden, das AISearchClient als primäre API für die Arbeit mit AI Search bereitstellt.

Alternativ können Sie die REST-API direkt aufrufen.

Requirements

In diesem Notizbuch wird davon ausgegangen, dass ein benannter databricks-gte-large-en Model Serving-Endpunkt vorhanden ist. Informationen zum Erstellen dieses Endpunkts finden Sie im Notizbuch "Aufrufen eines GTE-Einbettungsmodells mithilfe von Model Serving".

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

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

help(AISearchClient)

Laden des Toy-Datensatzes in die Delta-Quelltabelle

Im Folgenden wird die Delta-Quelltabelle erstellt.

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

Endpunkt erstellen

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

Erstellen eines Index

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

Abrufen des Indexes

Verwenden Sie get_index(), um das AI Search-Indexobjekt unter Verwendung des Indexnamens abzurufen. Sie können auch für das Indexobjekt verwenden describe() , um eine Zusammenfassung der Konfigurationsinformationen des Indexes anzuzeigen.

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

Ähnlichkeitssuche

Fragen Sie den AI-Suchindex ab, um ähnliche Dokumente zu finden.

# 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

Konvertieren von Ergebnissen in LangChain-Dokumente

Die erste abgerufene Spalte wird in page_content geladen, und der Rest wird in Metadaten geladen.

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

Löschen des Indexes

vsc.delete_index(index_name=vs_index_fullname)

Beispiel-Notebook

AI Search Python SDK-Beispielverwendung

Notebook abrufen