LangchainのSQLDatabaseチェーンを使ってDatabricksのデータベースを操作する

こちらで触れた**SQLDatabaseチェーン向けDatabricksコネクター**を試してみます。

https://qiita.com/taka_yayoi/items/073057f869a0c5cecd78

詳細説明ページへのリンクが切れているようですが、ソースコードを見ながら動かしてみました。

https://github.com/hwchase17/langchain/blob/master/langchain/sql_database.py#L127

# `SQLDatabase.from_databricks`メソッド

`SQLDatabase.from_databricks`でDatabricksのカタログ、スキーマ、テーブルにアクセスできるようになります。

メソッドの説明は以下の通り。

> Databricks接続からSQLDatabaseインスタンスを作成するクラスメソッド。このメソッドには`databricks-sql-connector`パッケージが必要です。インストールされていない場合、`pip install databricks-sql-connector`で追加できます。


引数は以下の通りです。

- `catalog`: Databricksのカタログ名。
- `schema`: カタログ内にあるスキーマ(データベース)名。
- `host`(オプション): Databricksワークスペースのホスト名から`https://`を除いたもの。指定されない場合、環境変数`DATABRICKS_HOST`を参照しようとします。それも指定されておらず、Databricksノートブックで実行される際、現在のワークスペースのホスト名がデフォルト値となります。デフォルト値はNoneです。
- `api_token`(オプション): DatabricksのSQLウェアハウスやクラスターにアクセスするためのDatabricksパーソナルアクセストークン。指定されない場合、`DATABRICKS_TOKEN`を取得しようとします。それも指定されておらず、Databricksノートブックで実行される際、現在のユーザー向けの一時トークンが生成されます。デフォルト値はNoneです。
- `warehouse_id`(オプション): Databricks SQLのウェアハウスID。指定された場合、このメソッドはこのウェアハウスを使用するように接続を設定します。`cluster_id`と一緒には使用できません。デフォルト値はNoneです。
- `cluster_id`(オプション): DatabricksランタイムにおけるクラスターID。指定されると、このメソッドはこのクラスターを使用するように接続を設定します。`warehouse_id`と一緒には使用できません。Databricksで実行しており、`warehouse_id`と`cluster_id`がNoneの場合、ノートブックがアタッチされているクラスターIDを使用します。デフォルト値はNoneです。
- `engine_args`(オプション): Databricksに接続する際に使用される引数です。デフォルト値はNoneです。
- `**kwargs`(任意): `from_uri`メソッドにおける追加のキーワード引数です。

# Databricksにアクセスする

Databricksノートブックからアクセスします。

```py
%pip install langchain==0.0.205
%pip install databricks-sql-connector
```
```py
dbutils.library.restartPython()
```

LLMはOpenAI APIを使うので、APIキーを設定します。

```py
import os

# OpenAI APIキーの設定
os.environ["OPENAI_API_KEY"] = dbutils.secrets.get("demo-token-takaaki.yayoi", "openai_api_key")
```

カタログとスキーマを指定してSQLDatabaseインスタンスを作成します。

:::note warn
**重要！**
OpenAI APIがアクセスしても問題のない、そして、誤操作しても問題がないスキーマにアクセスして試すようにしてください。
:::

```py
from langchain import OpenAI, SQLDatabase, SQLDatabaseChain

# カタログ名とスキーマ(データベース)名を指定して接続
db = SQLDatabase.from_databricks(catalog="takaakiyayoi_catalog", schema="covid")
# LLM
llm = OpenAI(temperature=0, verbose=True)
```

ちなみにこのようなテーブルが格納されています。日本のCOVID-19感染者数のデータです。
![Screenshot 2023-06-19 at 21.04.25.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/a932be5d-19f5-1eed-0c38-ef65487bebf4.png)
テーブルの中身はこんな感じ。
![Screenshot 2023-06-19 at 21.10.16.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/81577346-2db0-d6c9-75ba-bd05b6cd8d66.png)

都道府県の数を聞いてみます。

```py
db_chain = SQLDatabaseChain.from_llm(llm, db, verbose=True)
db_chain.run("How many Prefecture are there?")
```
```
> Entering new  chain...
How many Prefecture are there?
SQLQuery:SELECT COUNT(DISTINCT Prefecture) FROM covid_cases
SQLResult: [(47,)]
Answer:There are 47 Prefectures.
> Finished chain.
Out[3]: 'There are 47 Prefectures.'
```

なんと、ちゃんと答えが返ってくる。

ならば、日本語かつ少し複雑な問い合わせを。

```py
db_chain.run("最も感染者数の合計の多い都道府県は")
```
```
> Entering new  chain...
最も感染者数の合計の多い都道府県は
SQLQuery:SELECT Prefecture, SUM(Cases) AS Total_Cases FROM covid_cases GROUP BY Prefecture ORDER BY Total_Cases DESC LIMIT 5;
SQLResult: [('Tokyo', 3984044), ('Osaka', 2604677), ('Kanagawa', 2028176), ('Aichi', 2019228), ('Saitama', 1705754)]
Answer:最も感染者数の合計の多い都道府県は東京都です。
> Finished chain.
Out[6]: '最も感染者数の合計の多い都道府県は東京都です。'
```

恐ろしい...そして、すごい。

# 更新処理をやってみる

参照系だけではなく、更新系もやってみます。ドキドキ。

```py
db_chain.run("covid_casesをcovid_cases_copyとしてコピーしてください")
```
```
> Entering new  chain...
covid_casesをcovid_cases_copyとしてコピーしてください
SQLQuery:CREATE TABLE covid_cases_copy AS SELECT * FROM covid_cases
SQLResult: []
Answer:covid_casesテーブルがcovid_cases_copyテーブルとしてコピーされました。
> Finished chain.
Out[6]: 'covid_casesテーブルがcovid_cases_copyテーブルとしてコピーされました。'
```

ちゃんとコピーされる。すごい。

もう少しトリッキーなものを。

```py
db_chain.run("covid_cases_copyの各カラムの説明文をください")
```
```
> Entering new  chain...
covid_cases_copyの各カラムの説明文をください
SQLQuery:SELECT `Prefecture`, `Cases`, date_timestamp, pref_no, `Area` FROM covid_cases_for_forecast LIMIT 5
SQLResult: [('Tokyo', 86, datetime.datetime(2022, 1, 1, 0, 0), 13, '関東地方'), ('Tokyo', 65, datetime.datetime(2022, 1, 2, 0, 0), 13, '関東地方'), ('Tokyo', 130, datetime.datetime(2022, 1, 3, 0, 0), 13, '関東地方'), ('Tokyo', 151, datetime.datetime(2022, 1, 4, 0, 0), 13, '関東地方'), ('Tokyo', 430, datetime.datetime(2022, 1, 5, 0, 0), 13, '関東地方')]
Answer:covid_cases_for_forecastテーブルのカラムは、Prefecture（都道府県）、Cases（感染者数）、date_timestamp（日付）、pref_no（都道府県番号）、Area（地域）です。
> Finished chain.
Out[7]: 'covid_cases_for_forecastテーブルのカラムは、Prefecture（都道府県）、Cases（感染者数）、date_timestamp（日付）、pref_no（都道府県番号）、Area（地域）です。'
```

カラムのコメントを更新します。デフォルトだと構文エラーが出るので、構文の例を含めます。

```py
db_chain.run("covid_cases_copyのカラムの日本語の説明文でcovid_cases_copyのカラムのコメントを更新してください。その際には構文「ALTER TABLE StudentInfo ALTER COLUMN name COMMENT 'new comment'」を使ってください。")
```
```
> Entering new  chain...
covid_cases_copyのカラムPrefectureの日本語の説明文でcovid_cases_copyのカラムのコメントを更新してください。その際には構文「ALTER TABLE StudentInfo ALTER COLUMN name COMMENT 'new comment'」を使ってください。
SQLQuery:ALTER TABLE covid_cases_copy ALTER COLUMN Prefecture COMMENT '都道府県名'
SQLResult: []
Answer:The comment for the column Prefecture in the table covid_cases_copy has been updated to '都道府県名'.
> Finished chain.
Out[22]: "The comment for the column Prefecture in the table covid_cases_copy has been updated to '都道府県名'."
```

更新されました！複数カラムのコメント更新もやりたいところです。
![Screenshot 2023-06-20 at 9.31.13.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/b73846f9-bb4b-75fc-ecfb-f04972eb5393.png)

### Databricksクイックスタートガイド

[Databricksクイックスタートガイド](https://www.amazon.co.jp/dp/B09V1YXFVQ/)


### Databricks無料トライアル

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