LangChainを用いたDatabricks SQLエージェントを動かしてみる

以前こちらの記事を書きました。

https://qiita.com/taka_yayoi/items/daab11a3926e16802aa3

それから、Databricks上のマニュアルも充実してきています。

https://docs.databricks.com/ja/large-language-models/langchain.html

その中の[Databricks SQLエージェント](https://docs.databricks.com/ja/large-language-models/langchain.html#databricks-sql-agent)を動かしてみます。

# Langchainを使用してSQLデータベースと対話する

以下のコードは、Databricks SQLエージェントの例を示しています。Databricks SQLエージェントを使用すると、任意のDatabricksユーザーがDatabrick Unity Catalog内の指定されたスキーマと対話し、データに関する洞察を生成できます。

## 要件

- このノートブックを使用するには、OpenAI APIトークンを提供してください。
- Databricks Runtime 13.3 ML以上

## インポート

Databricksは最新バージョンの`langchain`と`databricks-sql-connector`を推奨しています。

```py
%pip install --upgrade langchain databricks-sql-connector langchain-community databricks-sqlalchemy openai langchain-openai
dbutils.library.restartPython()
```

```py
import os
os.environ["OPENAI_API_KEY"] = dbutils.secrets.get(scope="demo-token-takaaki.yayoi", key="openai_api_key")
```

## SQLデータベースエージェント

これはUnity Catalog内の特定のスキーマと対話する方法の例です。エージェントは新しいテーブルを作成したり、テーブルを削除したりすることはできません。テーブルをクエリすることしかできません。

データベースインスタンスは次のように作成されます:

```py
db = SQLDatabase.from_databricks(catalog="...", schema="...")
```

そして、エージェント(と必要なツール)は以下によって作成されます:

```py
toolkit = SQLDatabaseToolkit(db=db, llm=llm)
agent = create_sql_agent(llm=llm, toolkit=toolkit, **kwargs)
```

実際のコードは以下のようになります。

```py
from langchain.agents import create_sql_agent
from langchain.agents.agent_toolkits import SQLDatabaseToolkit
from langchain.sql_database import SQLDatabase
from langchain import OpenAI

db = SQLDatabase.from_databricks(catalog="samples", schema="nyctaxi")
llm = OpenAI(temperature=.7)
toolkit = SQLDatabaseToolkit(db=db, llm=llm)
agent = create_sql_agent(llm=llm, toolkit=toolkit, verbose=True)
```

```py
agent.run("What is the longest trip distance and how long did it take?")
```
```


> Entering new SQL Agent Executor chain...
Action: sql_db_list_tables
Action Input: trips I think the trips table looks like the most relevant one.
Action: sql_db_schema
Action Input: trips
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
TBLPROPERTIES('delta.feature.allowColumnDefaults' = 'enabled')

/*
3 rows from trips table:
tpep_pickup_datetime	tpep_dropoff_datetime	trip_distance	fare_amount	pickup_zip	dropoff_zip
2016-02-13 21:47:53+00:00	2016-02-13 21:57:15+00:00	1.4	8.0	10103	10110
2016-02-13 18:29:09+00:00	2016-02-13 18:37:23+00:00	1.31	7.5	10023	10023
2016-02-06 19:40:58+00:00	2016-02-06 19:52:32+00:00	1.8	9.5	10001	10018
*/ I need to query for the longest trip distance and the time it took using the trip_distance and tpep_dropoff_datetime columns.
Action: sql_db_query
Action Input: SELECT trip_distance, tpep_dropoff_datetime FROM trips ORDER BY trip_distance DESC LIMIT 1[(30.6, datetime.datetime(2016, 2, 22, 22, 0, 58))] I now know the final answer.
Final Answer: The longest trip distance was 30.6 miles and it took 10 minutes.

> Finished chain.
```
```
'The longest trip distance was 30.6 miles and it took 10 minutes.'
```

MLflow Tracingも動きます。

![Screenshot 2025-02-17 at 9.24.22.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/2cb6d07a-a594-44c0-bb38-12bd11a32421.png)



### はじめてのDatabricks

[はじめてのDatabricks](https://qiita.com/taka_yayoi/items/8dc72d083edb879a5e5d)

### Databricks無料トライアル

[Databricks無料トライアル](https://databricks.com/jp/try-databricks)