Databricks上でLlamaIndexのフルスタックWebアプリケーションを動かしてみる

さらにこちらの続きです。

自分のデータを使ってRAGを構築したら、GUIと組み合わせたくなるのが人情(?)です。

マニュアルにもPutting it All Together(すべてをまとめる)のセクションがあります。

こちらを参考に実装します。

A Guide to Building a Full-Stack Web App with LLamaIndex

これまでのおさらいを兼ねて、コードをウォークスルーします。

ライブラリのインストール

%pip install -U llama-index
%pip install "databricks-sql-connector[sqlalchemy]"
dbutils.library.restartPython()

Unity Catalogとの接続

クラスターのHTTPパス、カタログ名、スキーマ名は適宜変更してください。

from llama_index import download_loader
from llama_index.utilities.sql_wrapper import SQLDatabase

DatabaseReader = download_loader('DatabaseReader')

access_token    = dbutils.notebook.entry_point.getDbutils().notebook().getContext().apiToken().getOrElse(None)
server_hostname = "<Databricksワークスペースのホスト名>"
http_path       = "sql/protocolv1/o/2556758628403379/1118-004519-nu0u899x"
catalog         = "takaakiyayoi_catalog"
schema          = "qiita_2023"

reader = DatabaseReader(
    sql_database = SQLDatabase.from_uri(f"databricks://token:{access_token}@{server_hostname}?" +
        f"http_path={http_path}&catalog={catalog}&schema={schema}")
)

テキストデータのロード

テキストを格納しているテーブルに応じてSQLを変更してください。

query = f"""
SELECT
    body
FROM takaakiyayoi_catalog.qiita_2023.taka_qiita_2023
"""

documents = reader.load_data(query=query)

インデックスの作成

OpenAIのAPIを使っています。Azure OpenAIを使う際にはこちらを参考にしてください。

import os
os.environ["OPENAI_API_KEY"] = dbutils.secrets.get("demo-token-takaaki.yayoi", "openai_api_key")

from llama_index import VectorStoreIndex, SimpleDirectoryReader

index = VectorStoreIndex.from_documents(documents)

RAGへの問い合わせ

動作確認します。ウィジェットを使ってます。

query_engine = index.as_query_engine()
response = query_engine.query(dbutils.widgets.get("prompt"))
print("回答: ", response)

回答: Databricksの利点は、以下のようなものです：

• Apache Sparkを基にしており、大規模データの効率的な処理が可能です。
• クラウドベースのプラットフォームであり、データの処理を容易にします。
• データの共有や協力を容易にするための強力なコラボレーションツールがあります。
• Machine Learningのための強力なツールがあります。
• データの管理や安全性を確保するための強力なデータ管理機能があります。
• 各種データソースとの統合が可能です。

Flask

ここからが新規パートです。事前に以下のポートを決めておきます。

• Flaskのポート: ここでは1501
• streamlitのポート: ここでは1502

動作確認のためFlaskのURLを表示させます。ここではドライバープロキシーの機能を使っています。なお、この時点で表示されているURLにアクセスしても動作しません。

from dbruntime.databricks_repl_context import get_context

def front_url(port):
    """
    フロントエンドを実行するための URL を返す

    Returns
    -------
    proxy_url : str
        フロントエンドのURL
    """
    ctx = get_context()
    proxy_url = f"https://{ctx.browserHostName}/driver-proxy/o/{ctx.workspaceId}/{ctx.clusterId}/{port}/"

    return proxy_url

PORT = 1501

# Driver ProxyのURLを表示
print(front_url(PORT))

# 利便性のためにリンクをHTML表示
displayHTML(f"<a href='{front_url(PORT)}' target='_blank' rel='noopener noreferrer'>別ウインドウで開く</a>")

以下を実行して、Flaskサーバーを起動します。

from flask import Flask, request

app = Flask("llama_index")

@app.route("/", methods=["GET"])
def query_index():
    global index
    query_text = request.args.get("text", None)
    if query_text is None:
        return (
            "テキストが入力されていません。URLに ?text=blah パラメータを含めてください",
            400,
        )
    query_engine = index.as_query_engine()
    response = query_engine.query(query_text)
    return str(response), 200
  

app.run(host="0.0.0.0", port=PORT)

上で表示されたURLにアクセスします。質問が入力されていない旨のメッセージが表示されます。

URLの末尾に以下のようにパラメータを指定します。

.../1501/?text=クラスターのおすすめは

以下のように結果を確認できます。これでサーバーが稼働している状態になりました。

シングルユーザーアクセスモードのクラスターがおすすめです。このクラスターは、使用できる機能に制限がなく、単一のユーザーに割り当てて利用することができます。ただし、ユーザーごとにクラスターが作成されるため、管理の手間が増えることに注意が必要です。

streamlit

あとは画面です。別のファイルとしてpyファイルを作成します。作成 > ファイルを選択して、以下のファイルを作成します。テキストのやり取りなのでシンプルです。

streamlit.py

import streamlit as st 
import numpy as np 
import json
import requests

st.title('Databricks Q&A bot')
#st.header('Databricks Q&A bot')

def generate_answer(question):
  # Driver Proxyと異なるクラスター、ローカルからDriver Proxyにアクセスする際にはパーソナルアクセストークンを設定してください
  token = "" 
  url = f"http://127.0.0.1:1501/?text={question}"

  headers = {
      "Content-Type": "application/json",
      "Authentication": f"Bearer {token}"
  }

  response = requests.get(url, headers=headers)
  if response.status_code != 200:
    raise Exception(
       f"Request failed with status {response.status_code}, {response.text}"
    )
  
  response_text = response.text
  return response_text

question = st.text_input("**質問**")

if question != "":
  response = generate_answer(question)

  st.write(f"**回答:** {response}")

streamlitを稼働させるために、今度は別のノートブックを作成します。

%pip install streamlit
dbutils.library.restartPython()

こちらはstreamlit用のポートを指定します。

from dbruntime.databricks_repl_context import get_context

def front_url(port):
    """
    フロントエンドを実行するための URL を返す

    Returns
    -------
    proxy_url : str
        フロントエンドのURL
    """
    ctx = get_context()
    proxy_url = f"https://{ctx.browserHostName}/driver-proxy/o/{ctx.workspaceId}/{ctx.clusterId}/{port}/"

    return proxy_url

PORT = 1502

# Driver ProxyのURLを表示
print(front_url(PORT))

# 利便性のためにリンクをHTML表示
displayHTML(f"<a href='{front_url(PORT)}' target='_blank' rel='noopener noreferrer'>別ウインドウで開く</a>")

上で作成したstreamlitのファイルを指定して、streamlitを起動します。

streamlit_file = "/Workspace/Users/[e-mail]/20240116_llama_index/streamlit.py"

!streamlit run {streamlit_file} --server.port {PORT}

上で表示されたリンクにアクセスします。

動きました！

当たり前ですが、Databricksの記事に特化しているので、質問にDatabricksを含めなくてもDatabricksに関する答えが返ってきます。

RAGアプリケーションを動かすという目的であれば、LlamaIndexだと手軽に実装できて便利ですね。

はじめてのDatabricks

はじめてのDatabricks

Databricks無料トライアル

Databricks無料トライアル