Query a served model with `ai_query()`

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

This function is not available on Databricks SQL Classic.
Querying Foundation Model APIs is enabled by default. To query endpoints that serve custom models or external models, you must enroll in the public preview. Please populate and submit the AI Functions Public Preview enrollment form.
An existing model serving endpoint. See Create custom model serving endpoints.
You must enable AWS PrivateLink to use this feature on pro SQL warehouses.

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 foundation models.

Note

For Databricks Runtime 14.2 and above, this function is supported in notebook environments including Databricks notebooks and workflows.
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

Query a served model with ai_query()

What is ai_query()?