SQL Database Agent example#
This example demonstrates the use of the SQL Database Agent for answering questions over a Databricks database.
from langchain.agents import create_sql_agent
from langchain.agents.agent_toolkits import SQLDatabaseToolkit
toolkit = SQLDatabaseToolkit(db=db, llm=llm)
agent = create_sql_agent(
    llm=llm,
    toolkit=toolkit,
    verbose=True
)
agent.run("What is the longest trip distance and how long did it take?")
> Entering new AgentExecutor chain...
Action: list_tables_sql_db
Action Input: 
Observation: trips
Thought:I should check the schema of the trips table to see if it has the necessary columns for trip distance and duration.
Action: schema_sql_db
Action Input: trips
Observation: 
CREATE TABLE trips (
	tpep_pickup_datetime TIMESTAMP, 
	tpep_dropoff_datetime TIMESTAMP, 
	trip_distance FLOAT, 
	fare_amount FLOAT, 
	pickup_zip INT, 
	dropoff_zip INT
) USING DELTA
/*
3 rows from trips table:
tpep_pickup_datetime	tpep_dropoff_datetime	trip_distance	fare_amount	pickup_zip	dropoff_zip
2016-02-14 16:52:13+00:00	2016-02-14 17:16:04+00:00	4.94	19.0	10282	10171
2016-02-04 18:44:19+00:00	2016-02-04 18:46:00+00:00	0.28	3.5	10110	10110
rtdocs_stable/api.python.langchain.com/en/stable/integrations/databricks/databricks.html