こちらで触れたSQLDatabaseチェーン向けDatabricksコネクターを試してみます。
詳細説明ページへのリンクが切れているようですが、ソースコードを見ながら動かしてみました。
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ノートブックからアクセスします。
%pip install langchain==0.0.205
%pip install databricks-sql-connector
dbutils.library.restartPython()
LLMはOpenAI APIを使うので、APIキーを設定します。
import os
# OpenAI APIキーの設定
os.environ["OPENAI_API_KEY"] = dbutils.secrets.get("demo-token-takaaki.yayoi", "openai_api_key")
カタログとスキーマを指定してSQLDatabaseインスタンスを作成します。
重要！
OpenAI APIがアクセスしても問題のない、そして、誤操作しても問題がないスキーマにアクセスして試すようにしてください。
from langchain import OpenAI, SQLDatabase, SQLDatabaseChain
# カタログ名とスキーマ(データベース)名を指定して接続
db = SQLDatabase.from_databricks(catalog="takaakiyayoi_catalog", schema="covid")
# LLM
llm = OpenAI(temperature=0, verbose=True)
ちなみにこのようなテーブルが格納されています。日本のCOVID-19感染者数のデータです。
テーブルの中身はこんな感じ。
都道府県の数を聞いてみます。
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.'
なんと、ちゃんと答えが返ってくる。
ならば、日本語かつ少し複雑な問い合わせを。
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]: '最も感染者数の合計の多い都道府県は東京都です。'
恐ろしい...そして、すごい。
更新処理をやってみる
参照系だけではなく、更新系もやってみます。ドキドキ。
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テーブルとしてコピーされました。'
ちゃんとコピーされる。すごい。
もう少しトリッキーなものを。
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（地域）です。'
カラムのコメントを更新します。デフォルトだと構文エラーが出るので、構文の例を含めます。
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 '都道府県名'."
更新されました！複数カラムのコメント更新もやりたいところです。
