はじめに
今回は、Streamlitを使って、Vertex AIを利用したチャットアプリを作成する方法を紹介します。このアプリでは、自然言語での質問に対してGeminiがSQLクエリを生成し、BigQueryに対してデータベース操作を行ってくれます。また、取得した結果に対してGeminiがデータの傾向などの考察も行ってくれるため、ユーザーはSQLクエリを記述しなくても、自然言語でチャット画面からデータ分析をインタラクティブに進めることができるようになります。
本記事では、今回使用するStreamlitやVertex AIなどの技術に触れつつ、生成AIによってBigQueryのデータに対して分析を行えることの価値を説明して、具体的なアプリケーションのコードとその実装方法を詳しく解説します。
Streamlitとは？
Streamlitは、データサイエンティストやエンジニアがわずかなコードで迅速にインタラクティブなデータアプリケーションを構築できるPythonフレームワークです。
Pythonで簡潔に記述できるため、複雑なフロントエンドの知識がなくても、洗練されたウェブアプリケーションを作成できます。ユーザーはHTMLやJavaScriptを意識せずに、直感的にインタラクティブなダッシュボードを作成できます。
コードが即座に反映されるリアルタイムのプレビュー機能も特徴の一つで、データサイエンティストやエンジニアがプロトタイピングなどの開発を迅速に行える環境を提供しています。開発元はStreamlit社で、現在はSnowflakeに買収されて、Snowflakeから提供されています。
Vertex AIとは？
Vertex AIは、Google Cloudが提供する統合型のAIプラットフォームで、機械学習モデルの構築、トレーニング、デプロイなどを簡単に行えます。Vertex AIでは、事前にトレーニングされた大規模な自然言語処理モデルや画像認識モデルなども利用可能で、複雑なAIシステムを短期間で構築できます。
Vertex AIではモデルガーデンなどとも呼ばれる様々なベンダーの提供する生成AIモデルを選択して利用する機能があり、今回はGoogleの提供するGeminiを利用します。
生成AIによってBigQueryのデータに対して分析を行えることの価値
今回紹介するコードは、StreamlitとVertex AIを使って、Google Cloudが提供するデータエージェントであるBigQuery データキャンバスと同様のアイデアを実現しています。生成AIによって自然言語入力をSQLクエリに変換し、BigQuery上のデータにアクセスするというコンセプトが同等です。
先般開催されたGoogle Cloud Next Tokyo '24の基調講演で説明やデモがありましたが、マネージドサービスであるBigQuery データキャンバスを利用することで、データアナリストなどのユーザーが自然言語を使って、BigQuery上のデータの探索、変換、クエリ、可視化を行うことができるサービスとなっています。
また、Google Cloud Next Tokyo '24の基調講演では、日本テレビの事例が取り上げられていましたが、こちらの事例でも同様に、マネージドサービスであるVertex AI Agent BuilderやLangchain on Vertex AIを用いて自然言語でBigQueryの番組情報などを入手するチャットボットを作成していました。
Google Cloud自身がこうしたマネージドサービスを提供していることからも分かる通り、構造化データとして企業が従来蓄積してきた企業データに対して、生成AIを通してユーザーが分析を直接行えることには、大きなビジネス上の価値があります。
現在生成AIで主流のRAGなどのグラウンディング技術では、PDFやHTMLなどの企業ドキュメントの非構造化された外部ファイルをモデルに対するプロンプトに注入することが良く説明されますが、注入するデータをマニュアルなどから更に拡大して、BigQueryなどに格納された実際の業務データを対象にする点が今後更に重要度を増してくるでしょう。
今回のチャットアプリは、そうしたGoogle Cloudが提供するマネージドサービスではなく、まずは既存の技術で構築可能なコードを開発することで、基本的な設計を理解できるようにします。
コード全体の構成
ここからは、具体的なアプリのコードの解説を行います。今回のコードは、以下のような機能を実現しています：
- StreamlitでのUI設定：アプリケーションの見た目をカスタマイズ。
- Vertex AIとの連携：自然言語の質問に応答する生成AI機能を提供。
- PromptTemplateの活用：AI応答の生成をテンプレート化して柔軟な応答を提供。
- BigQueryとSQLAlchemyによるデータベース操作：SQLクエリによるデータ検索を実現。
- クエリの実行と生成AIによる応答の生成
- ユーザーのチャット履歴管理：ユーザーが行った質問とAIの応答を保持し表示。
1. Streamlitの設定
まずは、Streamlitを使ってウェブアプリの基本的なレイアウトを設定します。このアプリでは、ページ全体を広く使うためにlayout="wide"を指定しています。また、アプリのヘッダーとサイドバーの設定も行い、企業で利用するためのブランド感も演出しています。
import streamlit as st
st.set_page_config(
page_title="Vertex AI Chat for BigQuery",
layout="wide"
)
st.header("Vertex AI Chat for BigQuery", divider="red")
st.set_page_configでページのタイトルやレイアウトを設定し、st.headerでページのヘッダーを作成します。ここでは赤い線を使ってヘッダーの見た目を整え、視覚的に印象的なデザインを作成しています。
カスタムCSSの適用
次に、st.markdownを使って、背景画像やカラー設定をカスタマイズするCSSを適用しています。これにより、ブランドやテーマに合わせたアプリの外観を簡単にカスタマイズできます。
st.markdown(
"""
<style>
.stApp {
background-image: url("app/static/red.png");
background-size: 30% 60%;
background-repeat: no-repeat;
background-position: right;
background-color:rgba(255,255,255,0.8);
background-blend-mode:lighten;
}
</style>
""",
unsafe_allow_html=True
)
サイドバーの設定
アプリのサイドバーにヘッダーを追加し、見た目を整えています。
header_container = st.sidebar
header_container.markdown(
"<h1 style='color: #FFFFFF;'>Vertex AI Chat for BigQuery</h1>",
unsafe_allow_html=True
)
2. Google Vertex AIとの統合
次に、Vertex AIを使用して、自然言語の入力に対して生成AIが応答する機能を実装します。オープンソースの生成AIフレームワークであるLangChainのlangchain_google_vertexaiライブラリを使い、Vertex AIのGeminiモデルを利用します。
from langchain_google_vertexai import ChatVertexAI
from google.cloud import bigquery
from sqlalchemy import create_engine
PROJECT_ID = os.environ.get("GCP_PROJECT") # Your Google Cloud Project ID
LOCATION = os.environ.get("GCP_REGION") # Your Google Cloud Project Region
DATASET_ID = os.environ.get("GCP_DATASETID") # Your Google Cloud Dataset ID
vertexai.init(project=PROJECT_ID, location=LOCATION)
# Vertex AIのモデルをインスタンス化
LLM = VertexAI(
# model_name="text-bison@001",
# model_name="gemini-1.5-pro-001",
model_name=st.session_state.selected_model, # Use selected model from session state
max_output_tokens=1024,
temperature=0.1,
top_p=0.8,
top_k=10,
verbose=True,
)
この部分で、Google CloudのプロジェクトIDを指定してVertex AIを初期化して、Vertex AIの事前学習済みモデルをインスタンス化し、自然言語入力に対する応答を生成する準備を行います。
3. PromptTemplateの活用
PromptTemplateは、Vertex AIが生成する応答に対してテンプレートを使用するための機能です。これにより、ユーザーの質問に応じた構造化された応答を生成できます。
from langchain import PromptTemplate
TEMPLATE = \
"""
Given an input question, first create a syntactically correct, ***BigQuery-compatible*** {dialect} query to run, then look at the results of the query and return the answer based on the query result.
Please note that when you create sql code, make it one line and ***NEVER*** use backticks (``). **DO NOT use backticks (``) in the query ALWAYS.**"
Also when you create sql code ALWAYS add LIMIT 100; to the query if not already specified.
There are two tables Train and Test , use Train table.
Train table has these columns:
"int64_field_0" INT64,
"trans_date_trans_time" TIMESTAMP,
"cc_num" INT64,
"merchant" STRING,
"category" STRING,
Use the following format and strictly follow the order:
Question: "Question here"
SQLQuery: "SQL Query to run"
SQLResult: "Result of the SQLQuery"
Answer: "Final answer here"
You must answer in Japanese.
Only use the following tables: {table_info}
Question: {input}
"""
# SQLDatabaseChainに用いるプロンプトの定義
PROMPT = PromptTemplate(
input_variables = ["input", "table_info","dialect"], # プロンプトのテンプレートで指定した入力変数
template = TEMPLATE, # プロンプトのテンプレート
)
-
TEMPLATE: テンプレートの定義です。
{input}の部分に、ユーザーが実際に入力した質問が差し込まれます。 - PromptTemplate: このクラスを使って、テンプレートに変数を埋め込み、フォーマットに従った応答を生成します。
このテンプレートを用いることで、AIの応答を統一したフォーマットに整え、適切な回答を生成するように指示します。
4. BigQueryとSQLAlchemyを使ったデータベース操作
次に、BigQueryとSQLAlchemyを使用して、SQLクエリを実行し、データベースから結果を取得する仕組みを実装します。SQLAlchemyを使ってBigQueryに接続し、create_engineを使ってクエリを実行します。
from google.cloud import bigquery
from sqlalchemy import create_engine
database_uri = f"bigquery://{PROJECT_ID}/{DATASET_ID}"
# SQLAlchemy エンジンを初期化
tables = ['Train'] # 今回参照するテーブル名
# SQLDatabaseオブジェクト
SQLDatabase = SQLDatabase.from_uri(
database_uri, # 読み込むデータベース
)
- SQLDatabase: このオブジェクトは、データベースに対してSQLクエリを実行したり、データを操作したりするために利用されます。SQLAlchemyを介してBigQueryと連携する際に便利なインターフェースを提供します。
- BigQueryとSQLAlchemyの準備: 最初にBigQueryとSQLAlchemyをインポートし、BigQueryに接続するためのURIを設定します。
- データベースの接続設定: database_uriを使って、BigQueryのデータセットに接続し、参照するテーブル（Train）を指定しています。
- SQLDatabaseオブジェクトの生成: SQLDatabase.from_uriを使って、BigQueryのデータベースに接続し、操作するためのオブジェクトを生成しています。このオブジェクトを使って、SQLクエリを実行したり、データを取得したりすることが可能になります。
5. クエリの実行と生成AIによる応答の生成
# Memoryオブジェクト（会話履歴をdb_chainに持たせるため）
memory = ConversationBufferMemory(output_key="result")
-
ConversationBufferMemory: これは、会話の履歴を保持するためのメモリオブジェクトです。自然言語による会話を記憶し、LLMが前後の文脈を理解できるようにするために使われます。ここでは、出力結果のキーを"result"に設定しています。
db_chain = SQLDatabaseChain(llm = LLM, # LLM
database = SQLDatabase, # SQLDatabaseオブジェクト
prompt = PROMPT, # プロンプトテンプレート
verbose = True, # プロンプトの動的表示有無
return_intermediate_steps = True, # 出力結果に中間処理も含めるか否か
return_direct = True, # Geminiだとなぜか結果でQuestionを繰り返すため、return_directをTrueにする
memory = memory, # Memoryオブジェクト
)
-
SQLDatabaseChain: これは、SQLデータベースとLLMを統合するためのオブジェクトです。これにより、自然言語で与えられた質問をSQLクエリに変換し、そのクエリを実行して得られた結果を再び自然言語でユーザーに返すことができます。-
llm: ここでは、LLM（大規模言語モデル）を指定しています。LLMは、自然言語の処理や生成に使用されます。 -
database:SQLDatabaseオブジェクトは、SQLAlchemy経由でデータベースに接続するためのオブジェクトです。これにより、SQLクエリがデータベース上で実行されます。 -
prompt: プロンプトテンプレートです。これにより、LLMがどのように質問を解釈し、どの形式で結果を返すかを指定できます。 -
verbose:Trueに設定すると、実行されるプロンプトや中間の処理が詳細に表示されます。デバッグやプロンプトの確認に便利です。 -
return_intermediate_steps:Trueに設定すると、SQLクエリの実行プロセスやその結果の中間ステップも返されます。結果がどのように生成されたかを確認するために有効です。 -
return_direct:Trueに設定すると、GeminiというAIモデルの特性に合わせて、LLMが不要な質問を繰り返さないように処理します。 -
memory: 先ほど作成したConversationBufferMemoryオブジェクトを使用し、会話のコンテキストを保持します。
-
データベースのクエリ結果を取得するための関数
def search_db(text):
# db_chain実行
result = db_chain(text)
print(result)
-
search_db関数: この関数は、与えられたテキスト（自然言語の質問）をdb_chainに渡して実行し、結果を取得します。-
db_chain(text): これにより、SQLDatabaseChainオブジェクトが実行され、textに含まれる自然言語の質問がSQLクエリに変換されてデータベースで実行され、結果が返されます。 -
print(result): クエリ結果をコンソールに出力しています。ここで返されるresultには、SQLクエリの実行結果やLLMが生成した応答が含まれます。
-
SQLクエリ結果の加工とLLMを使用した最終応答の生成
# SQL結果の取り出し
sql_result = result # db_chainの実行結果をそのまま使う
# sql_result = result["result"] # resultだけを抜き取ると、LLMが良い回答をしてくれない
-
sql_result:db_chainの実行結果全体を変数sql_resultに代入しています。resultには、SQLクエリの結果やその生成プロセスが含まれています。
# LLMを使って最終的な答えを生成
answer = LLM.predict(f"The SQL query returned the following results: {sql_result}. Structure the result in a best format and provide an answer in *{selected_language}* retaining all the information in original result and add insights. Always explain the SQL query used.")
-
LLM.predict(): ここでは、LLMを使って最終的な回答を生成しています。SQLクエリの結果を自然言語で説明し、回答を生成するために使用されます。selected_languageという変数は、応答がどの言語で生成されるかを指定しています。- インサイトの付加: SQLのクエリ結果を元に、LLMが結果を再構成し、情報を保持しつつ、追加の洞察を提供するように指示されています。
- SQLクエリの説明: また、LLMは使用したSQLクエリについても説明を行います。
6. ユーザーのチャット履歴管理
Streamlitのst.session_stateを使い、ユーザーが行ったチャット履歴を保存し、過去のメッセージを表示できるようにしています。これにより、アプリケーションがインタラクティブかつユーザーにとって使いやすくなります。
def init_messages():
if "messages" not in st.session_state:
st.session_state.messages = []
if __name__ == '__main__':
init_messages()
# 過去のメッセージを表示
for message in st.session_state.messages:
with st.chat_message(message["role"]):
st.markdown(message["content"])
# ユーザーの入力を受け付ける
if user_input := st.chat_input("質問をどうぞ"):
st.session_state.messages.append({"role": "user", "content": user_input})
# Vertex AIまたはSQLエージェントで応答を生成
if user_input.startswith("SQL:"):
query = user_input[4:]
response = search_db(query)
else:
response = LLM.predict(user_input)
# 応答を保存し、表示
st.session_state.messages.append({"role": "assistant", "content": response})
チャット履歴の処理
-
init_messagesでセッションの初期化を行い、st.session_stateに過去のメッセージを保存します。 -
user_inputにユーザーの質問を入力すると、自然言語の質問かSQLクエリかを判別し、それぞれに応じた処理を行います。 - 最終的に、ユーザーの入力とAIの応答をセッションに保存し、過去のメッセージとともに画面に表示します。
まとめ
今回のアプリでは、Streamlitを使って手軽にWebインターフェースを構築し、Vertex AIを活用して自然言語の質問に応答する機能を実装しました。また、BigQueryとSQLAlchemyを用いてデータベースクエリの処理も実現しています。このアプリケーションを通じて、ユーザーが簡単に生成AIと対話し、データベース操作もできる柔軟なツールを提供できます。
今後は、冒頭で説明したデータエージェントであるBigQuery データキャンバスやマネージドサービスであるVertex AI Agent BuilderとLangchain on Vertex AIを用いて、このアプリをリファクタリングできればと考えています。