Skip to content

AI Query

AI Query is a transformation component that uses the Databricks ai_query() function to obtain an answer to a natural-language question. This function uses a Databricks chat model serving endpoint made available by Databricks Foundation Model APIs. The component takes one or more input columns from your source table, combines the inputs with a user prompt, and sends this data to the Databricks chat model for processing.

The output is a string containing the chat model's response to the question.


Make sure you have read and understand the Requirements set out by Databricks before using this component.


  • 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.


Name = string

A human-readable name for the component.

Model = drop-down

Select the Databricks model serving endpoint that will be used to answer the query. The following models are currently supported:

  • DBRX Instruct
  • Meta-Llama-3-70B-Instruct
  • Meta-Llama-2-70B-Chat
  • Mixtral-8x7B Instruct

User Prompt = text editor

Use the text editor to write a question for the chat model to respond to.

Columns = column editor

Select the source columns to feed as input to the chat model.

  • Column Name: A column from the input table.
  • Descriptive Name: An alternate descriptive name to better contextualize the column. Recommended if your column names are low-context.

Include Input Columns = boolean

  • Yes: Outputs both your source input columns and the query response. This will also include those input columns not selected in Columns.
  • No: Only includes the query response.

Snowflake Databricks Amazon Redshift