Query a served model with ai_query
Preview
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 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