AI Query
Public preview
Editions
Production use of this feature is available for specific editions only. Contact our sales team for more information.
The AI Query transformation component 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.
Note
- 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.
Use case
You can use the AI Query component to ask questions of a source text in plain English. A typical use case might be analyzing a table of sales data with questions such as:
- What were the total sales last quarter?
- Show me the top 5 products by revenue in 2024.
- How many users signed up each month in 2023?
- Show orders from California over $1000 in the last 30 days.
Properties
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.
To use variables in this field, type the name of the variable prefixed by the dollar symbol and surrounded by { } brackets, as follows: ${variable}. Once you type ${, a drop-down list of autocompleted suggested variables will appear. This list updates as you type; for example, if you type ${date, functions and variables containing date will be listed.
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.
Got feedback or spotted something we can improve?
We'd love to hear from you. Join the conversation in the Documentation forum!