ChatGPT × MySQL Shell：SQLをAIに考えてもらうプラグイン

4.ChatGPT × MySQL Shell

SQLを自動生成するプラグインには，以下の要素が必要です。

• ユーザからプロンプトを受け取る
• スキーマ，テーブル，カラムの情報を接続中のMySQLから取得する
• 上記の情報をあわせてプロンプトを構築する
• OpenAIのAPIを呼び出し，結果を表示する

プロンプトは引数で受け取れば良いでしょう。カラムの情報は，INFORMATION_SCHEMAをクエリするのが良さそうです。プロンプトの構築は，C言語ならともかくPythonなら朝飯前です。OpenAIのAPIは，SDKから一発で呼び出すことが可能です。びっくりするほど簡単に作れそうです。今回はMySQL 8.0.32を使って実装してみます。

プラグイン自体は，オフィシャルのドキュメントで紹介されている例をベースに簡単に作ることができます。まずは，プラグインに引数を渡しましょう。add_extension_object_memberのparametersで，プラグインに必要な引数を指定します。今回は，クエリを生成したいMySQLにつながっているsessionと，プロンプトのstringを受け取ることにします。これで，「プラグイン名.generate(session, “プロンプト”)」という呼び出しができるようになります。

def generate(session, query):
    pass
# 略
shell.add_extension_object_member(\
    plugin_obj,\
    "generate",\
    generate,\
    {"brief": "Generate SQL",
     "parameters": [
        {"name": "session", "type": "object", "class": "Session", "brief": "session"},
        {"name": "query", "type": "string", "brief": "query string"}
    ]}\
)

カラムの情報を取るために，INFORMATION_SCHEMA.COLUMNSをSELECTします。sessionからのメソッドチェーンで簡単にクエリを実行できます。

i_s = session.get_schema("information_schema")
columns = i_s.get_table("COLUMNS")\
    .select("table_schema", "table_name", "column_name", "column_type")\
    .where(f"table_schema NOT IN ({','.join(SYS_SCHEMAS)})")\
    .execute()\
    .fetch_all()

MySQL ShellからOpenAIのSDKを利用するために，pipでインストールします。これは通常の手順と少々異なりますが，コマンド数回で簡単にインストールが可能です。

wget https://bootstrap.pypa.io/get-pip.py
mysqlsh --py -f get-pip.py
mysqlsh --pym pip install --user openai

OpenAIのChat completion APIの呼び出しは，もはやこの記事での解説は不要でしょう。プロンプトは，図3のものをそのまま使います。これらをまとめると，最終的に以下のようなコードになります。

from mysqlsh import mysqlx
import openai


SYS_SCHEMAS = list(map(lambda x: f"'{x}'", ["sys", "mysql", "information_schema", "performance_schema"]))

PROMPT_SYSTEM = "あなたはMySQLに精通したDBAであり，問い合わせに対して最適なSQLを提示することができます。"
PROMPT_TABLE = "次のようなテーブルがあるMySQLを仮定します。\n"
PROMPT_MAIN = "問い合わせは「{q}」です。問い合わせへの回答となるSQLのみを出力してください。説明や装飾は不要です。コードブロックで囲まないでください。SQLを出力できない場合は，何も出力しないでください。"

MODEL = "gpt-3.5-turbo"
openai.api_key = "XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX"

def generate(session, query):
    # INFORMATION_SCHEMA.COLUMNSからカラムの情報を取得する
    i_s = session.get_schema("information_schema")
    columns = i_s.get_table("COLUMNS")\
        .select("table_schema", "table_name", "column_name", "column_type")\
        .where(f"table_schema NOT IN ({','.join(SYS_SCHEMAS)})")\
        .execute()\
        .fetch_all()

    # カラムの情報をCSV形式にしてプロンプトに含める
    csv = "table_schema,table_name,column_name,column_type\n" + "\n".join(list(map(lambda x: ",".join(x), columns)))
    messages = [
        {"role": "system", "content": PROMPT_SYSTEM},
        {"role": "user", "content": PROMPT_TABLE + csv},
        {"role": "user", "content": PROMPT_MAIN.format(q=query)}
    ]

    # APIを呼び出して結果が正常に取得できていたら出力する
    completion = openai.ChatCompletion.create(
        model=MODEL,
        messages=messages
    )
    if ("choices" in completion and len(completion["choices"]) > 0 and
        "message" in completion["choices"][0] and "content" in completion["choices"][0]["message"]):
        print(completion["choices"][0]["message"]["content"])


plugin_obj = shell.create_extension_object()

shell.add_extension_object_member(\
    plugin_obj,\
    "generate",\
    generate,\
    {"brief": "Generate SQL",
     "parameters": [
        {"name": "session", "type": "object", "class": "Session", "brief": "session"},
        {"name": "query", "type": "string", "brief": "query string"}
    ]}\
)

shell.register_global(\
    "mc2q",\
    plugin_obj,\
    {"brief": "my_Chat2Query"}\
)

今回作成したプラグインは，Chat2Queryに敬意を込めて，「my_Chat2Query」と命名しました。このコードを「~/.mysqlsh/plugins/my_chat2query/init.py」として配置すれば，準備完了です（API_KEYの設定をお忘れなく）。

早速使ってみましょう。オフィシャルのExample Databaseのworldsをインポートして，このデータベースに対するクエリを生成してもらいます。worldsデータベースは，国，都市，人口，言語などの情報が格納されています。図5にmy_Chat2Queryで生成したSQLを示します。

シンプルなクエリも，2つのJOINが必要なちょっと複雑なクエリも，完璧にこなしてくれました。注目ポイントは，レコードの情報がないのに，それっぽいWHERE句を付けてくれている点です。Language=’English’だけでなく，enum(‘T’, ‘F’)もきちんと理解していて偉いですね。どちらのSQLも，引数で入力したプロンプトよりも圧倒的に長いため，人間の作業量を減らすことができた，と言えるでしょう。今回のデータベースの規模は3テーブル，合計24カラムで，total_tokensは500弱でした。

生成されたSQLを実際に実行すると，図6のような結果が得られます。

良さそうです。ツッコミどころは，私がMySQL Shellを使わずにクエリを実行していることくらいしかありません。Chat2Queryほどの性能はないかもしれませんが，なかなか使えそうなものができたのではないでしょうか？

5.まとめ

今回はOpenAI Chat completions APIを用いてSQLを自動生成するMySQL Shellプラグインを実装しました。自分で記述したコード量は非常に少ないですが，生成されるSQLは十分使えそうなレベルに達しています。人間でも簡単にかける程度の複雑さのものを，人間よりも高速に生成してくれる，という点に価値があるのではないでしょうか。

一方，今回のようなシンプルな実装の場合プロンプトが長くなるため，大規模なデータベースではプロンプトの長さの制限にひっかかったり，課金額が大きくなったりする，という問題点もあります。これらの問題をクリアすれば，頼れるアシスタントプラグインになりそうです。

次世代システム研究室では，グループ全体のインテグレーションを支援してくれるアーキテクトを募集しています。アプリケーション開発者の方，次世代システム研究室にご興味を持って頂ける方がいらっしゃいましたら，ぜひ募集職種一覧からご応募をお願いします。

参考

PingCAP – ChatGPT搭載の革新的なSQLジェネレータ「Chat2Query」とは
lefred’s blog – MySQL Shell and extra Python modules
MySQL Shell API 8.0.32
MySQL – 7.2.4 MySQL Shell 拡張オブジェクトの例
Other MySQL Documentation (Example Databases)
OpenAI API Documentation