Exemple d’utilisation du SDK Python pour AI Search
Ce Notebook montre comment utiliser le SDK Python AI Search, qui fournit AISearchClient comme API principale pour travailler avec AI Search.
Alternativement, vous pouvez appeler l'API REST directement.
Exigences
Ce notebook suppose qu'un endpoint de Model Serving nommé databricks-gte-large-en existe. Pour créer cet endpoint, consultez le notebook Appeler un modèle d'embeddings GTE à l'aide de Model Serving.
Python
%pip install --upgrade --force-reinstall databricks-ai-search langchain
dbutils.library.restartPython()
Python
from databricks.ai_search.client import AISearchClient
vsc = AISearchClient()
Python
help(AISearchClient)
Chargez un dataset de démonstration dans la table Delta source.
Ce qui suit crée la table Delta source.
Python
# 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"
Python
source_table_name = "en_wiki"
source_table_fullname = f"{catalog_name}.{schema_name}.{source_table_name}"
Python
# Uncomment this line to start from scratch.
# spark.sql(f"DROP TABLE {source_table_fullname}")
Python
source_df = spark.read.parquet("/databricks-datasets/wikipedia-datasets/data-001/en_wikipedia/articles-only-parquet").limit(10)
display(source_df)
Python
source_df.write.format("delta").option("delta.enableChangeDataFeed", "true").saveAsTable(source_table_fullname)
Python
display(spark.sql(f"SELECT * FROM {source_table_fullname}"))
Créer un Endpoint
Python
ai_search_endpoint_name = "ai-search-demo-endpoint"
Python
vsc.create_endpoint(
name=ai_search_endpoint_name,
endpoint_type="STANDARD" # or "STORAGE_OPTIMIZED"
)
Python
endpoint = vsc.get_endpoint(
name=ai_search_endpoint_name)
endpoint
Créer un index
Python
# 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"
Python
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()
Obtenir l’indice
Utilisez get_index() pour récupérer l'objet d'index de recherche IA en utilisant le nom de l'index. Vous pouvez également utiliser describe() sur l'objet d'index pour voir un résumé des informations de configuration de l'index.
Python
index = vsc.get_index(endpoint_name=ai_search_endpoint_name, index_name=vs_index_fullname)
index.describe()
Python
# 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()
Recherche de similarité
Interrogez l'index de recherche IA pour trouver des documents similaires.
Python
# 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
Python
# 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
Convertir les résultats en documents LangChain
La première colonne récupérée est chargée dans page_content, et le reste dans les métadonnées.
Python
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
Supprimer l'index
Python
vsc.delete_index(index_name=vs_index_fullname)