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

4.ChatGPT × MySQL Shell

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

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"}
    ]}\
)

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()

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

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"}\
)

シンプルなクエリも，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