AI Query
Public preview
Editions
Production use of this feature is available for specific editions only. Contact our sales team for more information.
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.
Note
Make sure you have read and understand the Requirements set out by Databricks before using this component.
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.
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.
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 |
|---|---|---|
| ❌ | ✅ | ❌ |