🌟

【LangChain】データベースを扱う③ postgreSQLを使う

2025/01/05に公開

LangChain

tech

LangChainでのSQLに関して、2024年は2本記事を書いていました。

今回は新たにpostgreSQLを使う方法です
（LangChain v0.3 公式ドキュメント参照）

※ なお、書いているデータは全て架空のものです。

事前準備

postgreSQLのインストール

私はmac m2なので、brewを使ってインストールしました

データの準備

自慢ではありませんが私はSQLをほとんど書けませんので、基本的にGUIで編集しています。

GUIアプリについては以下記事に色々紹介されています。

スプレッドシートやAppSheetをよく使っていますので、DBeaverを採用しました。

以下のような形式で「ダミー_平均気温」というデータを作成しました。

postgreSQLとDBeaverの簡単な使い方

「スキーマ」がスプレッドシートのファイルに該当（以下の例では"public20250105"という名前のスキーマを作成しています）
「テーブル」がスプレッドシートのシートに該当
テーブルを右クリックでメニューを表示し、「新しく作る 表」で新規作成できます
テーブルを選択し、データタブを選ぶと表で表示される

テーブルを選択し、データのインポートができる

エクスポートする時、列名に日本語を設定している場合は「BOM」にチェックを入れること

コマンドでデータ確認

// postgresの立ち上げ
% psql -h localhost -p 5432 -U <username> -d postgres

// スキーマの確認
postgres=# \dn
      List of schemas
      Name      |  Owner   
----------------+----------
 public20250105 | <username>
(1 row)

// スキーマ下のテーブル確認
postgres=# \dt public20250105.*
                  List of relations
     Schema     |      Name       | Type  |  Owner   
----------------+-----------------+-------+----------
 public20250105 | ダミー_平均気温 | table | <username>
 public20250105 | テスト作成      | table | <username>
(2 rows)

// テーブルのデータを確認
postgres=# SELECT * FROM public20250105."テスト作成";
 名前  |   誕生日   |  id   
-------+------------+-------
 Aimer | 1994-12-21 | aimer
 milet | 1998-05-08 | milet
 ikura | 2002-03-04 | ikura
(3 rows)

// 終了
postgres=# \q

LangChainでの実装

postgreSQLからデータ抽出を試す

SQLDatabaseを使うと、langchainでクエリを実行できます

from langchain_community.utilities import SQLDatabase

# SQLDatabaseを使って、PostgreSQLに接続
db = SQLDatabase.from_uri("postgresql://username:@localhost:5432/postgres", schema="public20250105")
print(db.dialect)
print('------------')
print(db.get_usable_table_names())
print('------------')
print(db.get_table_info())
print('------------')
result = db.run('SELECT * FROM "public20250105"."ダミー_平均気温" LIMIT 10;')
print(result)

ChatOpenAIのチェックと、query-system-promptの内容確認

query-system-promptは、ユーザーが投げた言葉からいい感じにクエリを作成してくれるツール

import getpass
import os

from langchain import hub
from langchain_openai import ChatOpenAI

if not os.environ.get("OPENAI_API_KEY"):
  os.environ["OPENAI_API_KEY"] = getpass.getpass("Enter API key for OpenAI: ")

llm = ChatOpenAI(model="gpt-4o-mini")

query_prompt_template = hub.pull("langchain-ai/sql-query-system-prompt")

assert len(query_prompt_template.messages) == 1
query_prompt_template.messages[0].pretty_print()

実行結果を日本語に訳すと以下の通り

================================ System Message ================================

入力された質問に対して、答えを見つけるために実行する正しい構文の {dialect} クエリを作成してください。ユーザーが特定の数の例を取得するように質問で指定しない限り、常にクエリの結果は最大で {top_k} 件に制限してください。関連するカラムで結果を並べ替えて、データベース内で最も興味深い例を返すことができます。

特定のテーブルからすべてのカラムをクエリしないでください。質問に関連するいくつかのカラムだけを取得してください。

スキーマの説明で確認できるカラム名のみを使用するように注意してください。存在しないカラムをクエリしないように注意してください。また、どのカラムがどのテーブルにあるのかにも注意してください。

以下のテーブルのみを使用してください：
{table_info}

質問: {input}

要は、dialect、top_k、table_info、そしてユーザーの要求であるinputの4つのパラメーターを入れると、いい感じにクエリを作成してくれるってことです。

型の設定とクエリビルダーの作成

LangGraphで使うことを見越しています。

Stateにはユーザーの質問やLLMが作成したクエリ、クエリ実行結果などを設定しています。（本来StateにはFieldで詳細を設定したほうが良い）
QueryOutputでは、この後作成するクエリビルダー（write_query）が出力する型を設定します
write_queryはLangGraphのノードで使うための関数です。
内容としては、先ほどインスタンス化したquery_prompt_templateを使って、プロンプトを作成しています。

※ 今回はLangGraphの作成はやりません

from typing_extensions import Annotated, TypedDict

class State(TypedDict):
    question: str
    query: str
    result: str
    answer: str

state = State(question="", query="", result="", answer="", top_k=10)

class QueryOutput(TypedDict):
    """Generated SQL query."""
    query: Annotated[str, ..., "Syntactically valid SQL query."]

def write_query(state: State):
    """Generate SQL query to fetch information."""
    prompt = query_prompt_template.invoke(
        {
            "dialect": db.dialect,
            "top_k": state["top_k"],
            "table_info": db.get_table_info(),
            "input": state["question"],
        }
    )
    structured_llm = llm.with_structured_output(QueryOutput)
    result = structured_llm.invoke(prompt)
    return {"query": result["query"]}

実行

write_queryに要望を出してみましょう。良い感じにクエリを作成してくれます。

q = write_query({
    "question": "2020年以降のデータを抽出してください？",
    "query": "",
    "result": "",
    "answer": "",
    "top_k": 10})

print(q)

このクエリを使って、dbからデータの抽出を実行してみましょう。

res = db.run(q["query"])

print(res)