📖

【LangChain】データベースを扱う①　LLMを用いたクエリ生成

2024/04/07に公開

LangChain

tech

会話アプリや社内データ検索アプリを作成するにおいて、データベースを参照する機会が必ず出てくると思います。

この章と次の章では、LangChainでデータベースを扱う方法を学びます。

まずこの章では簡単な問いかけを想定して、langchainがクエリを作成できることを学びます。
次の章では、SQLエージェントを使ってより複雑なデータ抽出、集計ができるようにします。

実践

Google colabを実行環境にしています

SQLite3とサンプルデータのインストール

サンプルデータとして、公式チュートリアルでも用いられているchinook.dbを使います。
chinookの説明はこちら

音楽テーマが中心で、主に音楽トラックの販売に関連する情報を含んでいます

!apt-get install sqlite3

!wget https://raw.githubusercontent.com/lerocha/chinook-database/master/ChinookDatabase/DataSources/Chinook_Sqlite.sql

!sqlite3 Chinook.db < Chinook_Sqlite.sql

これでChinook.dbが作成されましたので、中身を確認しましょう

import sqlite3
conn = sqlite3.connect('Chinook.db')
cursor = conn.cursor()

# 例としてアーティストテーブルからデータを取得
cursor.execute("SELECT * FROM Artist LIMIT 5;")
rows = cursor.fetchall()

for row in rows:
    print(row)

# 接続を閉じる
cursor.close()
conn.close()

結果

(1, 'AC/DC')
(2, 'Accept')
(3, 'Aerosmith')
(4, 'Alanis Morissette')
(5, 'Alice In Chains')

LangChainを使ってクエリを作成する

データベースをlangchainで使えるようにする

先ほどのステップでcolabにChinook.dbを作成しましたので、これをlangchainで扱えるようにします。

from langchain_community.utilities import SQLDatabase

db = SQLDatabase.from_uri("sqlite:///Chinook.db")
print(db.dialect)
print('------------')
print(db.get_usable_table_names())
print('------------')
db.run("SELECT * FROM Artist LIMIT 10;")

以下のようにクエリ実行結果が得られます

[(1, 'AC/DC'), (2, 'Accept'), (3, 'Aerosmith'), (4, 'Alanis Morissette'), (5, 'Alice In Chains'), (6, 'Antônio Carlos Jobim'), (7, 'Apocalyptica'), (8, 'Audioslave'), (9, 'BackBeat'), (10, 'Billy Cobham')]

LLMを使ってクエリを作成する

create_sql_query_chainを使うと、人間の言葉からクエリを作成することができます

from langchain.chains import create_sql_query_chain
from langchain_openai import ChatOpenAI

llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0)
chain = create_sql_query_chain(llm, db)
response = chain.invoke({"question": "Artistテーブルに含まれるデータの数はいくつですか"})
response

SELECT COUNT("ArtistId") AS TotalArtists FROM "Artist"
LIMIT 1;

私はあまりSQLに詳しくありませんので、人間の言葉をクエリに翻訳してくれるだけでもすごいと思ってしまいました。

作成したクエリ（ここではresponse）をdb.run()にぶち込めば、先ほどと同様にクエリ実行結果を得られます。
※まだ回答を生成していませんので、実行結果だけの素っ気ないレスポンスです

db.run(response)

#-> [(275,)]

chainの中で何が行われているかは、以下のコードで確認できます

chain.get_prompts()[0].pretty_print()

中身は英語なので、chatGPTさんに日本語に翻訳してもらいました。
けっこうシンプルなプロンプトテンプレートです

あなたはSQLiteのエキスパートです。入力された質問に対して、まず文法的に正しいSQLiteクエリを作成し、そのクエリの結果を見てから入力質問の回答を返してください。ユーザーが質問で特定の数を指定していない限り、SQLiteのLIMIT句を用いて最大5つの結果を取得してください。データベース内の情報が最も豊富なデータを返すように結果を順序付けしてください。
テーブルからすべての列を問い合わせることはしないでください。質問に答えるのに必要な列のみを問い合わせてください。区切り識別子として列名を二重引用符(")で囲んでください。
以下のテーブルで見られる列名のみを使用し、存在しない列を問い合わせないように注意してください。また、どの列がどのテーブルにあるかも注意してください。
質問に「今日」という言葉が含まれている場合は、現在の日付を取得するためにdate('now')関数を使用してください。

データベース接続権限を可能な限り最小限に抑えます

上記の方法で、人間の言葉をクエリに翻訳してデータベースから検索できました。
ところで、実際はデータベースには誰でも彼でもアクセスしていいわけではなく、アクセスして良いテーブルを制限する必要があります。

やることとしては、QuerySQLDataBaseToolを加えるだけです。
（実際はもっと書く必要があるのでしょうが、クイックスタートでは詳細は書かれていませんでした。必要に応じて調べると良いでしょう）

from langchain_community.tools.sql_database.tool import QuerySQLDataBaseTool

execute_query = QuerySQLDataBaseTool(db=db) # 追加
write_query = create_sql_query_chain(llm, db)
chain = write_query | execute_query # 追加
chain.invoke({"question": "employeesの数はいくつですか"})

回答を生成させる

先ほどまではクエリを作成して実行する部分だけでしたので、ここから会話として成り立たせるようにします。

といっても、これまで学んできたprompt、modelを使うだけですので、以下コードを見れば何を渡して何をしているのかわかるはずです。

from operator import itemgetter

from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import PromptTemplate
from langchain_core.runnables import RunnablePassthrough

answer_prompt = PromptTemplate.from_template(
    """Given the following user question, corresponding SQL query, and SQL result, answer the user question.

Question: {question}
SQL Query: {query}
SQL Result: {result}
Answer: """
)

answer = answer_prompt | llm | StrOutputParser()
chain = (
    RunnablePassthrough.assign(query=write_query).assign(
        result=itemgetter("query") | execute_query
    )
    | answer
)

chain.invoke({"question": "How many employees are there"})

結果

There are a total of 8 employees.

この章で、LLMでクエリを作成し、データベースからデータを引っ張ってくる方法がひと通りわかったと思います。
次の章ではSQLエージェントを使って、検索してみましょう。