LlamaIndexとOpenAI APIで、テーブル定義から分析に使いたいテーブルの検索をできるか流れを試してみた

機械学習チームの鈴木です。

LLMを使って自然言語でデータベースに対して検索をすることで、データ分析の効率化を模索する取り組みが行われています。

私もやってみたいなと思い、分析対象の候補となるテーブルの選定がOpenAI APIでできそうかやってみました。

とても簡単な例ですが、実際やってみると実現方法が頭の中で整理されてよかったので、試したことをご紹介します。

本記事の内容

なにか分析したい目的があった場合に、データカタログなどからテーブルに関する情報を取得し、どのテーブルを見に行くか効率よく知りたいです。

このような目的のシステムによって嬉しいユースケースは以下の2つがあると考えていて、今回は1点目を試してみました。

1. データカタログから候補となるテーブルを探してくれたり、分析例を提案してくれる。
2. 自然言語から直接SQL文に変換・実行して結果例を返してくれる。

1点目の例として、以下のブログを参考にしました。Glue Data Catalogなどに応用することを見据えつつ、今回は流れをまず掴みたかったので、自分が利用しやすかったGoogle Colabにて試してみました。

2点目の例は、上記も記載がありますが、LangChainおよびLlamaIndexに実装面での案内がありました。

今回はLlamaIndexのガイドを参考にしつつ、Google Colab上で作ったDuckDBのテーブルの定義を踏まえて、分析したい内容をLLMに問い合わせし、期待した案内を生成してくれるか確認してみました。

クライアント実行環境

実行環境はGoogle Colaboratoryを使いました。ハードウェアアクセラレータ無し、ランタイム仕様は標準としました。

Pythonのバージョンは以下でした。

!python --version
# Python 3.10.12

また、ライブラリは以下のようにインストールしました。

# インデックス作成で使用
!pip install llama-index
!pip install python-dotenv
!pip install duckdb duckdb-engine

インストールされたライブラリのバージョンは以下でした。

!pip freeze | grep -e "openai" -e "llama-index" -e "langchain" -e "SQLAlchemy" -e "duckdb-engine"
# duckdb-engine==0.9.0
# langchain==0.0.225
# langchainplus-sdk==0.0.20
# llama-index==0.7.1
# openai==0.27.8
# SQLAlchemy==2.0.16

環境変数の設定

まず、必要な値を.envファイルに書き込みます。

!echo 'OPENAI_API_KEY="<トークン>"' >.env

load_dotenv()で環境変数を読み込みました。

from dotenv import load_dotenv
load_dotenv()

やってみる

1. ライブラリをインポートする

以下のようにライブラリをインポートしました。

import itertools
import logging
import os
import sys

logging.basicConfig(stream=sys.stdout, level=logging.INFO)
logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))

from llama_index import SimpleDirectoryReader, Document, GPTListIndex
from sqlalchemy import create_engine, MetaData, Table, Column, String, Integer, Float, Date, column

2. テーブルを作成する

DuckDBにテーブルを作成しました。今回は以下の3テーブルを作成しました。

• city_stats： 都市の人口を管理するテーブル
• store_information： 店舗の売上高を管理するテーブル
• weather_information: 天候の情報を管理するテーブル

テーブルは定義だけでデータは入れません。また、あくまで実際にテーブルSQLAlchemyでテーブル定義を取得するところも含めてやってみたいというだけなので、検証の中心となるのは次のテーブルの情報を持った文書の作成からになります。

engine = create_engine("duckdb:///:memory:")
metadata_obj = MetaData()

table_name = "city_stats"
city_stats_table = Table(
    table_name,
    metadata_obj,
    Column("city_name", String(16), primary_key=True),
    Column("population", Integer),
    Column("country", String(16), nullable=False),
)

table_name2 = "store_information"
store_information_table = Table(
    table_name2,
    metadata_obj,
    Column("store_name", String(16), primary_key=True),
    Column("amount_of_sales", Integer),
    Column("city", String(16), nullable=False),
)

table_name3 = "weather_information"
weather_information_table = Table(
    table_name3,
    metadata_obj,
    Column("date", Date),
    Column("weather", String(8)),
    Column("tempelature", Float),  
    Column("city", String(16), nullable=False),
)

metadata_obj.create_all(engine)

3. テーブルの情報を持った文書を作成する

ここからが本題で、LLMへの問い合わせをする際にプロンプトに含めるテーブルの情報を作成しました。

先ほど作成したテーブル定義から、テーブル名およびカラム名を取得して、ローカルにファイル出力しました。

data.txtへは、テーブル名|カラム名のようにパイプ繋ぎにして、カラムごとに書き出す形としました。

# テーブル名とカラム名の取得
table_names = [t.name for t in [city_stats_table, store_information_table, weather_information_table]]
table_cols = [t.c.keys() for t in [city_stats_table, store_information_table, weather_information_table]]

# テーブル名とカラム名を出力するための整理
table_names_names_tmp = [[n, k] for n, k in zip(table_names, table_cols)]
table_names_names = []
for lst in table_names_names_tmp:
  new_lst = [f"{lst[0]}|{i}" for i in lst[1]]
  table_names_names.append(new_lst)

# テーブル情報の書き出し
os.mkdir("./data") 
with open("./data/data.txt", "w") as f:
  f.write(",\n".join(list(itertools.chain.from_iterable(table_names_names))))

このように出力できました。

./data/data.txt

city_stats|city_name,
city_stats|population,
city_stats|country,
store_information|store_name,
store_information|amount_of_sales,
store_information|city,
weather_information|date,
weather_information|weather,
weather_information|tempelature,
weather_information|city

4. 質問をしてみる

作成したテーブル情報のファイルを読み込んで、GPTListIndexを作成しました。

documents = SimpleDirectoryReader(input_dir="./data").load_data()
list_index = GPTListIndex.from_documents(documents)

query_engine = list_index.as_query_engine()

まず以下の内容を聞いてみます。これはテーブル名から明らかなのでできて欲しいところです。

テーブル物理名|カラム名のフォーマットの情報を渡します。この中で天気に関する情報が分かるテーブルがどれか知りたいです。テーブル物理名を教えてください。

response = query_engine.query("テーブル物理名|カラム名のフォーマットの情報を渡します。この中で天気に関する情報が分かるテーブルがどれか知りたいです。テーブル物理名を教えてください。")

for i in response.response.split("。"):
    print(i + "。")

以下のように回答されました。

weather_information。

次に、分析対象の組み合わせについても聞いてみました。

テーブル物理名|カラム名の情報を渡しました。この中で天気と店舗の売上に関する情報を分析したいです。どのテーブルを組み合わせるといいか日本語で教えてください。テーブル物理名はそのままにしてください。

response = query_engine.query("テーブル物理名|カラム名の情報を渡しました。この中で天気と店舗の売上に関する情報を分析したいです。どのテーブルを組み合わせるといいか日本語で教えてください。テーブル物理名はそのままにしてください。")

for i in response.response.split("。"):
    print(i + "。")

store_informationテーブルとweather_informationテーブルを組み合わせるといいです。

ということでこの質問についても期待通りの回答を得ることができました。

なお、今回の例では、ノードは一つに収まっているので、最初のプロンプトで全てのテーブル情報を踏まえて回答を生成してくれています。

for doc_id, node in list_index.storage_context.docstore.docs.items():
    get_node_info = node.get_node_info()
    print(f'{doc_id=}, start={get_node_info["start"]}, end={get_node_info["end"]}')

# doc_id='0eed347f-5388-429e-bdb6-ee5c76deabfb', start=0, end=266

考察

結構期待した通りの返事をしてくれたのでとても良かったのですが、とはいえ今回は小さいテーブルが3つかない場合なので、実際のシステム化には考えないといけないことがいくつかあるなと考えています。

インデックスの作り方

今回はテーブル情報が小さいテキストファイルに収まったのでGPTListIndexで回答の合成に全てのテキストを使うことができました。テーブルが非常に多い場合は、これだと時間がかかってしまうため、例えばGPTVectorStoreIndexなどで、クエリと類似度の高いノードだけに絞ることになりますが、テーブル定義とクエリの埋め込みベクトルを使って適切に検索対象のテーブルの類似性を検索できるのかは検証するべきだなと思いました。

テーブル情報の取得方法

テーブル定義に変更や追加がある場合は、テーブル情報はスケジュールしたバッチ処理で定期的に更新する方法がありそうです。テーブル数が多い場合は検索パフォーマンスのためにもデータベースに登録すると思いますが、どのタイミングで入れ替えをするかはテーブル定義の変更頻度によります。データカタログへ変更があったタイミングで入れ替えをするような仕組みを入れるとよいのかもしれません。

プロンプトの改良

プロンプトの内容からお気づきな方もいらっしゃると思いますが、このバージョンでLlamaIndexから問い合わせに使う標準のプロンプトが英語ベースであることもあり、回答が英文になる場合もありました。また、city_stats|city_nameをテーブル名として回答してしまうこともあったので、プロンプトとテーブル情報のファイルの作り方は改善の余地がありました。

上記のような課題は、テーブル情報が個々のシステムの固有の情報であるため、LLMが固有の情報を踏まえて回答を生成する必要があるために発生します。

この課題へのアプローチの一つとして、今回のようにテーブル情報を含めたプロンプトを作成するRetrieval Augmented Generation（RAG）という手法を使って解決することができます。

この記事で紹介した簡単な例でも、システム側の定義情報の更新をインデックスやLLM用のデータベースにどのように反映させるかといったシステム的な設計をよく考えないといけないことが分かります。

なお、RAGついては以下のブログを参考にしてください。

最後に

Google Colaboratory上でデータベースにあるテーブルの定義を取得し、その情報を使ってデータ分析を想定した質問に対して分析の案内ができそうか試してみました。

テーブル数やデータベース数が増えた際に上手く機能しそうかについては引き続き検証してみたいと思います。

参考になりましたら幸いです。