データ抽出から可視化まで担う AI 分析アシスタントの試作
概要
冬休みの自由研究で AI 分析アシスタントを試作したので、日記の代わりに記事を書きます。
こんな感じで、自然言語で分析内容を指定すると、データの抽出と可視化をやってくれます。
ちなみに利用しているデータは theLook eCommerce という BigQuery の public dataset です。
コードは GitHub に置いてあります。
構成
- 画面は streamlit 製
- データ抽出は Cube の GraphQL API を叩く
- 可視化方法は Veta-Lite で指定する
- AI の呼び出しは Pydantic AI を介して行う
以下、それぞれについてもう少し詳しい構成や選定理由を書きます。
Streamlit
データに関するアプリケーションを試作するならこれでしょ、という安定感がありますよね。
久しぶりに公式ドキュメントを見たら LLM を利用したチャットアプリのチュートリアルが追加されていたので、基本的にはそれを真似て実装しています。
Cube
Semantic Layer を実装するためのツールです。
少し前から AI x Semantic Layer の実験をしてみたかったので採用しました。
AI に SQL を生成させてデータを取得するよりも、Semantic Layer 経由でデータを取得させる方が簡単だと期待しています。
また Cube の REST API ではなく GraphQL API を採用したのは、dimension や measure の名前も AI に正確に理解させたいからです。以下のような GraphQL の schema だけ与えれば、AI は dimension や measure 名を正確に理解して API を叩けるはずです。
# 抜粋
type UsersMembers {
count: Float
count_male: Float
count_female: Float
ratio_male: Float
ratio_female: Float
email: String
gender: String
first_name: String
last_name: String
full_name: String
count_order: Float
}
type Result {
users(where: UsersWhereInput, orderBy: UsersOrderByInput): UsersMembers!
}
type Query {
cube(where: RootWhereInput, limit: Int, offset: Int, timezone: String, renewQuery: Boolean, ungrouped: Boolean, orderBy: RootOrderByInput): [Result!]!
}
一方で、REST API からデータ抽出する場合は openspec.yml を参考に API を叩くことになりますが、これだけでは dimension や measure の名前が分かりません。meta データ用の API を叩いたりして追加の情報と組み合わせる必要があり、難易度が一段上がってしまう印象です。
Vega-Lite
Python の有名な可視化ツールの一つに Altair があります。Vega-Lite はそのベースにもなっている可視化の仕様です。JSON 形式で指定できるので AI で生成しやすいと考えました。
JSON Schema が公開されているので、そこから Pydantic のコードを自動生成できると期待したのですがうまくいかず[1]、結局以下のように手作業で実装しています。手作業で実装したものは Vega-Lite の仕様のごく一部ですが、これでもシンプルなグラフなら記述できます。
ちなみに streamlit で Vega-Lite に基づくグラフを表示するするのは st.vega_lite_chart を実行するだけなのでとても簡単です。
Pydantic AI
streamlit を使うと決めたので、Python の AI フレームワークを探していたら最初に見つりました。validation の分野で信頼されている Pydantic チーム製というのも心強く、即採用です。
ちなみに今回は Agent を三つに分けて実装しています。
- ルーティング Agent
- ユーザーの要望を理解し、他の Agent に指示を出します
- ユーザーの要望が曖昧な場合は、詳細を特定するために質問を返します
- そもそも分析が不可能な場合は拒否します（Cube の Model 定義を instruction として与えることで、そこに記述がないものは拒否させます）
- ユーザーとの会話履歴の管理も担います
- GraphQL Query 生成 Agent
- ルーティング Agent の指示に基づいて、Cube の GraphQL API を叩くための GraphQL Query を生成します
- instructions として GraphQL の Schema を与えることで、誤った Query を返さないように validation もできるようになっています
- Vega-Lite の JSON 生成 Agent
- ルーティング Agent の指示に基づいて Vega-Lite の JSON を生成します
- instruction として、実際に Cube から取得したデータを数行与えることで、精度向上を期待しています
所感
Good
- AI に生成させる部分を GraphQL Query と Vega-Lite の JSON に限定したことで、最低限の指示でも AI が働いてくれている気がする
Bad
- データ抽出処理と可視化処理で Agent を分けたことにより、そこで文脈が途切れている気がする（例えばデータ抽出時は細かく group by されたデータを持ってくるが、結局は可視化処理でまとまるから無駄、みたいなケースがちらほら）
もうちょい工夫すれば Bad であげたポイントも多少改善できそうな気がしますが、今回は飽きたのでここまでです。
Discussion