Share via

Query a served model with ai_query()

  • Article

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

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