OpenAI APIとLangChainで自然言語を元にSQL文を生成＆実行する

環境

Google Colaboratory

参考にした記事

以下の記事のコードをベースに、一部改変を加えています。 zenn.dev

Code(テーブルの作成まで)

!pip install langchain==0.0.145
!pip install openai==0.27.4

import os

#TODO: APIキーの登録が必要
os.environ["OPENAI_API_KEY"] = "..."

from sqlalchemy import text, create_engine, MetaData, Table
import pandas as pd

# サンプルデータの定義
df = pd.DataFrame({
    "name": ["佐藤", "鈴木", "吉田","松崎"],
    "age": ["20", "30", "40","35"],
    "sex": ["男", "女", "男","男"],
})

# 
sql_uri = "sqlite:///sample.db"
engine = create_engine(sql_uri, echo=False)

# メタデータのインスタンス作成
metadata = MetaData()

# テーブルが存在する場合、それをドロップする
with engine.connect() as conn:
    users = Table('users', metadata, autoload_with=engine)
    if users.exists(conn):
        users.drop(conn)

# Pandas DFをテーブル化
df.to_sql("users", con=engine)

# 作成されたテーブルの確認
with engine.connect() as conn:
    print(conn.execute(text("SELECT * FROM users")).fetchall())

Code(プロンプトと実行結果)

プロンプト

db_chain.run("sexが男のユーザーの平均年齢は？")

実行結果

> Entering new SQLDatabaseChain chain...
sexが男のユーザーの平均年齢は？
SQLQuery: SELECT AVG("age") FROM users WHERE sex = '男';
SQLResult: [(31.666666666666668,)]
Answer: 男のユーザーの平均年齢は31.67歳です。
> Finished chain.
 男のユーザーの平均年齢は31.67歳です。

プロンプト

db_chain.run("sexが男かつ、ageが25歳以上のユーザーの一覧を出力して")

実行結果

> Entering new SQLDatabaseChain chain...
sexが男かつ、ageが25歳以上のユーザーの一覧を出力して
SQLQuery: SELECT "name" FROM users WHERE "sex" = '男' AND "age" >= 25 LIMIT 5;
SQLResult: [('吉田',), ('松崎',)]
Answer: 吉田さんと松崎さんがsexが男かつ、ageが25歳以上のユーザーです。
> Finished chain.
 吉田さんと松崎さんがsexが男かつ、ageが25歳以上のユーザーです。