Query a served model with ai_query
Important
This feature is in Public Preview.
This article describes how to query a model serving endpoint from SQL with ai_query()
.
What is ai_query()
?
The ai_query()
function is a built-in Azure Databricks SQL function, part of AI functions. It allows these types of models to be accessible from SQL queries:
- Custom models hosted by a model serving endpoint.
- Models hosted by Databricks Foundation Model APIs.
- External models (third-party models hosted outside of Databricks).
For syntax and design patterns, see ai_query function.
When this function is used to query a model serving endpoint, it is only available in workspaces and regions where Model Serving is available and enabled.
Requirements
- See Requirements.
Query the endpoint with ai_query()
You can query the model behind the endpoint using ai_query()
on serverless or pro SQL warehouses. For scoring request and response formats see Query generative AI models.
Note
- For Databricks Runtime 14.2 and above, this function is supported in notebook environments including Databricks notebooks and jobs.
- For Databricks Runtime 14.1 and below, this function is not supported in notebook environments, including Databricks notebooks.
Example: Query a large language model
The following example queries the model behind the sentiment-analysis
endpoint with the text
dataset and specifies the return type of the request.
SELECT text, ai_query(
"sentiment-analysis",
text,
returnType => "STRUCT<label:STRING, score:DOUBLE>"
) AS predict
FROM
catalog.schema.customer_reviews
Example: Query a predictive model
The following example queries a classification model behind the spam-classification
endpoint to batch predict whether the text
is spam in inbox_messages
table. The model takes 3 input features: timestamp, sender, text. The model returns a boolean array.
SELECT text, ai_query(
endpoint => "spam-classification",
request => named_struct(
"timestamp", timestamp,
"sender", from_number,
"text", text),
returnType => "BOOLEAN") AS is_spam
FROM catalog.schema.inbox_messages