【Databricks】Databricks Apps上で動的なグラフ出力を実現する
データを紐づけたレポートを作成する際に、ユーザーの入力に合わせて動的にグラフを出力したい場合があると思います。今回は、そういった場面で便利な「vega_lite」×「Streamlit」を紹介します。
はじめに
データを紐づけたレポートを作成する際に、ユーザーの入力に合わせて動的にグラフを出力したい場合があると思います。今回は、そういった場面で便利な「vega_lite」×「Streamlit」を紹介します。
想定シナリオ
事前準備：
今回は他ブログで構築した、Walmartのデータを使用してPDFレポートを作成するアプリを使用します。現在のレポートではLLMと対話してグラフを表示する機能は実装されていませんが、今回は「Vega_Lite」を使用して動的にグラフを作成する機能を追加します。
現状の挙動確認
まずは、現在のアプリの挙動を確認します。アプリを起動すると以下のような画面になります。
ここでグラフを表示ボタンを押すと、
Databricks上のデータを取得し、グラフと合計金額を算出してくれます。さらに、AIが取得したデータに基づいて簡単なインサイトを出力してくれています。
さらにPDFで保存を押すとレポートがPDFとして保存されます。
今回はこのレポートにAIとのチャット機能を作成します。具体的にはこのレポートで表示されていない店舗に関するグラフをユーザーが求めた場合、出力します。
参考：
https://blog.since2020.jp/data_analysis/databricks_streamlit/
https://blog.since2020.jp/data_analysis/databricks-apps_llm/
Vega_liteとは
Vega-Lite は、データの可視化を簡単に行うためのツールです。JSON 形式で可視化の仕様を記述することで、プログラムコードを書かずにグラフやチャートを作成できます。Vega-Lite は、視覚化の対象データとその見せ方を宣言的に記述することを基本にしています。
例えば、以下の内容をJSON形式で記載すると
{
"data": {
"values": [
{"category": "A", "value": 28},
{"category": "B", "value": 55},
{"category": "C", "value": 43},
{"category": "D", "value": 91},
{"category": "E", "value": 81}
]
},
"mark": "bar",
"encoding": {
"x": {"field": "category", "type": "nominal"},
"y": {"field": "value", "type": "quantitative"}
}
}次のようなグラフが記載されます。
上記のJSONをVEGA EDITORに貼り付けてみると実際にグラフが生成されることが分かると思います。
参考：
https://vega.github.io/editor/
今回の想定は、LLMによって動的にJSON形式のグラフ情報を生成してもらい、その内容をStreamlitで可視化します。
実装
本アプリのコードを一部紹介します。
下記はチャット機能を実装したUIの部分です。
handle_chat_message関数でLLMを呼び出しリスポンスを取得します。
handle_chat_message関数は以下です。
この関数は、ユーザーが入力したテキストや対象データなどを受け取り以下を出力します。
{
"mode": "chart" or "text", # 応答の種類
"sql": "SELECT ..." or None, # LLMが生成したSQL
"vega_lite_spec": {...} or None, # Vega-Liteのグラフ定義
"explanation": "説明テキスト", # 日本語の説明
"data": DataFrame or None # SQL実行結果
}詳細を確認していきます。
この関数ではbuild_chat_system_promptでシステムプロンプトを取得し、過去の対話を含めながら、query_llmでLLMにリクエストを投げリスポンスを取得します。
def handle_chat_message(user_message: str, conn, table_name: str, schema_info: str, sample_data: str) -> dict:
"""チャットメッセージを処理してLLM応答を取得・パース・実行"""
if not SERVING_ENDPOINT:
return {
"mode": "text",
"sql": None,
"vega_lite_spec": None,
"explanation": "※ SERVING_ENDPOINTが設定されていないため、応答できません。",
"data": None
}
system_prompt = build_chat_system_prompt(table_name, schema_info, sample_data)
# 会話履歴を構築（直近のやり取りをコンテキストとして含める）
messages = [{"role": "system", "content": system_prompt}]
# 過去の会話履歴を追加（最大10ターン）
chat_history = st.session_state.get("chat_history", [])
for entry in chat_history[-10:]:
messages.append({"role": "user", "content": entry["user"]})
messages.append({"role": "assistant", "content": json.dumps(entry["raw_response"], ensure_ascii=False)})
messages.append({"role": "user", "content": user_message})
try:
raw_response = query_llm(SERVING_ENDPOINT, messages, max_tokens=2000)
except Exception as e:
return {
"mode": "text",
"sql": None,
"vega_lite_spec": None,
"explanation": f"LLM呼び出しエラー: {e}",
"data": None
}
parsed = parse_llm_response(raw_response)
# SQLが含まれている場合は実行
data = None
if parsed.get("sql"):
try:
data = fetch_data_by_query(conn, parsed["sql"])
except Exception as e:
parsed["explanation"] = f"SQLの実行でエラーが発生しました: {e}\n\n元の回答: {parsed.get('explanation', '')}"
parsed["mode"] = "text"
parsed["vega_lite_spec"] = None
parsed["data"] = data
return parsed
システムプロンプトは以下です。利用可能なデータを入力として持たせ、グラフを出力する場合、しない場合に分けて出力形式を指定します。なお、グラフを出力する場合はデータが必要なのでデータ取得用のSQLも同時に生成を依頼します。
続いてLLMを実行します。先ほど作成したプロンプトをLLMに渡してリスポンスを返り値に設定します。
def query_llm(endpoint_name: str, messages: list[dict], max_tokens: int = 1000) -> str:
"""LLMのサービングエンドポイントを呼び出す"""
client = get_deploy_client("databricks")
response = client.predict(
endpoint=endpoint_name,
inputs={
"messages": messages,
"max_tokens": max_tokens
}
)
if "choices" in response:
return response["choices"][0]["message"]["content"]
elif "messages" in response:
return response["messages"][-1]["content"]
else:
raise ValueError(f"Unexpected response format: {response}")その後render_chat_responseで出力された内容を可視化します。Streamlitでは「vega_lite_chart」でVega-Lite形式のグラフを可視化することが出来ます。
これにより、入力からプロンプトの生成、LLMによる回答生成、結果の可視化を実装することが出来ました。実際に作成したアプリを実行時してみます。
デモ
それでは、先ほど作成したアプリをデプロイし、動かしてみます。
まずは、アプリを起動します。
データを取得してグラフを表示を押します。
これまで通り、３店舗の売上情報とデータに基づくインサイトが返ってきました。
それでは以下のようにチャット画面に他店舗を含めた売上の棒グラフを出力するように要求してみます。
このように各店舗ごとの売上と実行に使用したSQLを出力してくれました。
さらにある店舗に注目して適切なグラフを出力できるか確認してみました。
その結果が以下です。
このようにStore6に関する売上の推移を出力してくれました。
終わりに
今回はDatabricks Apps上で動的なグラフ出力を実現するために「vega_lite」×「Streamlit」を使用して実装しました。結果として、ユーザーの要求に対して動的にグラフを作成する機能を実装できました。
今後はより複雑なグラフ生成が可能かどうかを検証していきます。
