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

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

Query a served model with ai_query()

What is ai_query()?