Skip to main content

Setting up Agents to query your Databricks Database

This guide will walk you through configuring Agents to query your Databricks Database. Before following this guide, ensure that you have set up your Databricks tool already by following the instructions here .

To configure this agent, you will require the following components added to your project

  1. Databricks SQL Query Tool Action
  2. LLM model

SQL Query Agent

The SQL Query Agent can be used to writd SQL Queries against a Databricks databricks.

Agent Setups

SQL Query Agent

The Agent setup should have a LLM and SQL Query tool. Databricks SQL Query agent

  1. Add the Databricks SQL Query tool to the Agent
  2. Add the LLM to the Agent

The following connections are needed:

  1. Input line to SQL Query Tool
  2. Tool line to LLM
  3. LLM line to Output

Update the SQL Query Tool by clicking on it, and setting it to receive a value from Input Step Result, as shown below. Databricks SQL Query Action

With this Agent setup, you can pass in an SQL Query as an input, and it's results will be processed by the LLM to show any natural language response.

Databricks query results

The Action can also be updated to accept a parameterized query, to allow the input to set only certain parts of the SQL Statement, such as the table name.

In the following example, we will select the SQL Action again, and type in the exact SQL Query we want the action to execute. However, we will use Agent Variables to set a conditional value dynamically.

Databricks SQL Query Action

Once you save the parameter, it should show up highlighed as a tag in blue: Databricks SQL Query Action

Now, you can run this agent Databricks query results

You can also build more complex agents with this tool as an Action. It is recommended that you provide a prompt to your Model with formatting and summarization instructions that would help it better provide meaningful insights on your Databricks queries.