Databricks Apps + Vector Search + AI関数で作る「名言レコメンダー」（第2回：実装編）

はじめに

前回は、名言レコメンダーアプリを構築するためのデータ準備とVector Searchインデックスの作成を行いました。

今回は、いよいよStreamlitアプリの実装に入ります。Vector Search APIの呼び出し方、AI関数（ai_gen）を使った理由生成、セッション管理など、コードの詳細を解説していきます。

全体構成

まず、アプリの全体像を把握しましょう。

app.py              # メインアプリケーション
requirements.txt    # 依存パッケージ
app.yaml           # Databricks Apps用設定（本番）
app-local.yaml     # ローカル開発用設定
databricks.yml     # Asset Bundles設定

今回はapp.pyとrequirements.txtを中心に解説します。

1. 依存パッケージ（requirements.txt）

databricks-sdk==0.76.0
databricks-sql-connector
streamlit
pandas

パッケージ	用途
databricks-sdk	Vector Search API、WorkspaceClient
databricks-sql-connector	AI関数実行用のSQL接続
streamlit	WebアプリUIフレームワーク
pandas	データ整形・表示

2. アプリの基本設定

2.1 インポートと初期設定

import os
import json
from typing import Any, Dict, List, Tuple, Optional

import pandas as pd
import streamlit as st
from databricks.sdk import WorkspaceClient
from databricks.sdk.core import Config
from databricks import sql

Streamlitのページ設定を行います。

st.set_page_config(page_title="Quote Recommender", layout="centered")
st.title("📺 Quote Recommender（Vector Search + おすすめ理由）")

2.2 環境変数の読み込み

アプリで使用する設定値は環境変数から取得します。これにより、ローカル開発と本番環境で設定を切り替えられます。

# Vector Searchインデックス名
INDEX_NAME = os.getenv("VECTOR_SEARCH_INDEX")

# SQL Warehouse ID（AI関数実行用）
WAREHOUSE_ID = os.getenv("WAREHOUSE_ID") or os.getenv("DATABRICKS_WAREHOUSE_ID")
HTTP_PATH = f"/sql/1.0/warehouses/{WAREHOUSE_ID}" if WAREHOUSE_ID else None

# AI関数で参照するソーステーブル
VS_SOURCE_TABLE = os.getenv("VS_SOURCE_TABLE", "media.tv_series.quotes_vs_source")

# 取得するカラム一覧
DEFAULT_COLUMNS = [
    "quote_id",
    "series_name",
    "season",
    "episode",
    "episode_name",
    "quote",
    "helpful",
]

ポイント：

• INDEX_NAMEとWAREHOUSE_IDは必須設定です
• VS_SOURCE_TABLEはデフォルト値を持たせ、変更可能にしています
• 環境変数はapp.yamlまたはapp-local.yamlで設定します（第3回で解説）

3. セッション状態の管理

Streamlitはリクエストごとにスクリプトが再実行されるため、状態を保持するにはst.session_stateを使います。

if "last_query" not in st.session_state:
    st.session_state["last_query"] = ""
if "last_df" not in st.session_state:
    st.session_state["last_df"] = pd.DataFrame()
if "reasons" not in st.session_state:
    # key: (query_text, quote_id) -> reason string
    st.session_state["reasons"] = {}

キー	用途
last_query	最後に検索したクエリ文字列
last_df	検索結果のDataFrame
reasons	生成した理由のキャッシュ（クエリと名言IDのペアをキーに）

reasonsを辞書でキャッシュすることで、同じ名言に対して何度もAI関数を呼び出すことを防いでいます。

4. 認証とクライアント初期化

4.1 WorkspaceClient（Vector Search用）

@st.cache_resource
def get_workspace_client() -> WorkspaceClient:
    return WorkspaceClient()

@st.cache_resourceデコレータにより、クライアントインスタンスがキャッシュされ、毎回の再初期化を防ぎます。

Databricks Apps上での認証：
Databricks Apps内で実行される場合、WorkspaceClient()は自動的に統合認証を使用します。明示的なトークン設定は不要です。

4.2 SQL接続（AI関数用）

AI関数を実行するためのSQL接続は、やや複雑です。Apps環境とローカル環境で認証方法が異なるためです。

def normalize_server_hostname(host: str) -> str:
    """ホスト名からプロトコルとスラッシュを除去"""
    return host.replace("https://", "").replace("http://", "").rstrip("/")


def connect_sql(http_path: str):
    """
    Apps: 統合認証（Config + credentials_provider）
    Local: DATABRICKS_HOST + DATABRICKS_TOKEN
    """
    try:
        # Apps環境: Configから認証情報を取得
        cfg = Config()
        return sql.connect(
            server_hostname=normalize_server_hostname(cfg.host),
            http_path=http_path,
            credentials_provider=lambda: cfg.authenticate,
        )
    except Exception:
        # ローカル環境: 環境変数から取得
        host = os.getenv("DATABRICKS_HOST", "")
        token = os.getenv("DATABRICKS_TOKEN", "")
        if not host or not token:
            raise RuntimeError(
                "Databricks認証情報が不足しています。"
                "ローカル実行なら DATABRICKS_HOST / DATABRICKS_TOKEN が必要です。"
            )
        return sql.connect(
            server_hostname=normalize_server_hostname(host),
            http_path=http_path,
            access_token=token,
        )

認証フローの解説：

1. まずConfig()で統合認証を試みる（Apps環境で有効）
2. 失敗した場合は環境変数DATABRICKS_HOSTとDATABRICKS_TOKENを使用
3. どちらもない場合はエラー

この二段構えにより、同じコードがApps環境でもローカル開発でも動作します。

5. Vector Search クエリ

5.1 クエリ関数

def query_vector_search(
    w: WorkspaceClient,
    index_name: str,
    query_text: str,
    num_results: int,
    columns: List[str],
    query_type: str,
    filters_json: Optional[str] = None,
):
    kwargs: Dict[str, Any] = {
        "index_name": index_name,
        "columns": columns,
        "num_results": int(num_results),
        "query_text": query_text,
        "query_type": query_type,
    }
    if filters_json:
        kwargs["filters_json"] = filters_json
    return w.vector_search_indexes.query_index(**kwargs)

パラメータ解説：

パラメータ	説明
index_name	検索対象のインデックス名
query_text	ユーザーの検索クエリ
num_results	取得件数
columns	取得するカラム一覧
query_type	検索タイプ（ann/hybrid/FULL_TEXT）
filters_json	オプションのフィルタ条件（JSON文字列）

5.2 検索タイプの違い

タイプ	説明	適したケース
ann	ベクトル類似度のみで検索	意味的な類似性を重視
hybrid	ベクトル + キーワード検索の組み合わせ	バランスの取れた検索
FULL_TEXT	キーワードベースの全文検索	特定の単語を含む結果が欲しい

5.3 レスポンスのDataFrame変換

Vector Search APIのレスポンスをpandas DataFrameに変換します。

def _safe_get(obj: Any, path: List[str], default=None):
    """ネストしたオブジェクトから安全に値を取得"""
    cur = obj
    for p in path:
        if cur is None:
            return default
        if isinstance(cur, dict):
            cur = cur.get(p)
        else:
            cur = getattr(cur, p, None)
    return cur if cur is not None else default


def response_to_df(resp: Any) -> pd.DataFrame:
    # カラム情報の取得
    col_infos = _safe_get(resp, ["manifest", "columns"], default=[]) or []
    col_names: List[str] = []
    for c in col_infos:
        name = c.get("name") if isinstance(c, dict) else getattr(c, "name", None)
        if name:
            col_names.append(name)

    # データ配列の取得
    data = _safe_get(resp, ["result", "data_array"], default=[]) or []
    if not data:
        return pd.DataFrame()

    # scoreカラムの補完（環境差に対応）
    if col_names and len(data[0]) == len(col_names) + 1 and "score" not in col_names:
        col_names = col_names + ["score"]

    df = pd.DataFrame(data, columns=col_names if col_names else None)

    # 数値カラムの型変換
    for c in ["season", "episode", "helpful", "score"]:
        if c in df.columns:
            df[c] = pd.to_numeric(df[c], errors="coerce")

    return df

なぜこの処理が必要か：

• Vector Search APIのレスポンス構造は、SDKバージョンや環境によって微妙に異なることがある
• scoreカラムがメタデータに含まれないケースがある
• _safe_getでNoneチェックを一元化し、堅牢なパースを実現

6. AI関数による理由生成

アプリの目玉機能である「おすすめ理由の生成」を実装します。

def generate_reason_with_ai_gen(user_input: str, quote_id: str) -> str:
    if not HTTP_PATH:
        raise RuntimeError("WAREHOUSE_IDが未設定です（理由生成にはSQL Warehouseが必要）。")

    sql_text = f"""
    SELECT
      ai_gen(
        concat(
          'あなたはTVドラマの名言レコメンダです。ユーザー入力に対して、なぜこの名言が合うかを日本語で2文で説明してください。',
          '名言全文の繰り返しは避け、要点だけ述べてください。\\n',
          'ユーザー入力: ', ?, '\\n',
          'シリーズ: ', series_name, '\\n',
          'エピソード: ', coalesce(episode_name, ''), '\\n',
          '名言: ', quote
        )
      ) AS reason
    FROM {VS_SOURCE_TABLE}
    WHERE quote_id = ?
    """

    with connect_sql(HTTP_PATH) as conn:
        with conn.cursor() as cur:
            cur.execute(sql_text, (user_input, quote_id))
            row = cur.fetchone()

    return (row[0] if row and row[0] else "").strip()

処理の流れ：

1. ソーステーブルからquote_idに該当する名言データを取得
2. ユーザー入力と名言データを組み合わせてプロンプトを構築
3. ai_gen関数でLLMに理由を生成させる
4. 結果を返却

プロンプト設計のポイント：

• 「2文で」と指定し、簡潔な出力を促す
• 「名言全文の繰り返しは避け」で冗長な応答を防止
• coalesce(episode_name, '')でNULL対策

パラメータバインディング：
SQLインジェクション対策として、?プレースホルダとパラメータタプルを使用しています。

7. サイドバー（検索設定UI）

ユーザーが検索条件を調整できるサイドバーを実装します。

with st.sidebar:
    st.subheader("設定")

    # 必須設定のバリデーション
    if not INDEX_NAME:
        st.error("VECTOR_SEARCH_INDEXが未設定です。")
    if not WAREHOUSE_ID:
        st.warning("WAREHOUSE_IDが未設定です（理由生成に必要）。")

    # 検索タイプ
    query_type = st.selectbox(
        "検索タイプ",
        ["hybrid", "ann", "FULL_TEXT"],
        index=1
    )

    # 表示件数
    num_results = st.slider("表示件数", 1, 50, 1)

    # 取得カラム（カスタマイズ可能）
    columns_csv = st.text_input(
        "取得列（カンマ区切り）",
        value=",".join(DEFAULT_COLUMNS)
    )
    columns = [c.strip() for c in columns_csv.split(",") if c.strip()]

    # フィルタ条件
    min_helpful = st.number_input(
        "helpful下限（この数以上）",
        min_value=0,
        value=0,
        step=1
    )

    # 内部取得件数（類似度上位が薄い時の保険）
    fetch_k = st.slider("内部取得件数（保険）", 1, 200, 50)

    # 表示オプション
    show_table = st.toggle("結果を表で表示", value=False)
    show_raw = st.toggle("生レスポンス表示（デバッグ）", value=False)

    st.divider()
    st.caption(f"Index: {INDEX_NAME or '(not set)'}")
    st.caption(f"Warehouse: {WAREHOUSE_ID or '(not set)'}")

UIの工夫：

• 必須設定が欠けている場合は警告を表示
• fetch_kで内部的に多めに取得し、フィルタ後に件数が足りなくなることを防止
• デバッグ用の生レスポンス表示オプション

8. メイン検索フォーム

with st.form("search_form"):
    st.text_input(
        "入力文字列（例: 気分/状況/キーワード）",
        key="q_input",
        value=st.session_state.get("last_query", ""),
    )
    submitted = st.form_submit_button(
        "検索する",
        type="primary",
        disabled=(not INDEX_NAME)
    )

if submitted:
    q = st.session_state["q_input"].strip()
    if not q:
        st.warning("入力してください")
        st.stop()

    try:
        w = get_workspace_client()

        # helpfulフィルタの構築
        filters_json = None
        if int(min_helpful) > 0:
            filters_json = json.dumps({"helpful >=": int(min_helpful)})

        # 内部取得件数は表示件数とfetch_kの大きい方
        k = max(int(num_results), int(fetch_k))

        resp = query_vector_search(
            w=w,
            index_name=INDEX_NAME,
            query_text=q,
            num_results=k,
            columns=columns,
            query_type=query_type,
            filters_json=filters_json,
        )

        if show_raw:
            st.write(resp.as_dict() if hasattr(resp, "as_dict") else resp)

        df = response_to_df(resp)

        # スコアとhelpfulでソート
        sort_cols = []
        if "score" in df.columns:
            sort_cols.append("score")
        if "helpful" in df.columns:
            sort_cols.append("helpful")
        if sort_cols:
            df = df.sort_values(by=sort_cols, ascending=[False] * len(sort_cols))

        # セッションに保存
        st.session_state["last_query"] = q
        st.session_state["last_df"] = df

    except Exception as e:
        st.error(f"検索に失敗しました: {e}")
        st.stop()

フィルタ機能の活用：
Vector Searchのfilters_jsonを使うと、ベクトル検索にメタデータ条件を追加できます。

# 例: helpful >= 50 の結果のみ取得
filters_json = json.dumps({"helpful >=": 50})

9. 検索結果の表示

df = st.session_state.get("last_df", pd.DataFrame())
query_text = st.session_state.get("last_query", "").strip()

if df is None or df.empty:
    st.info("検索してください（結果がここに表示されます）")
    st.stop()

# 表示件数に丸める
df = df.head(int(num_results))

# テーブル表示（オプション）
if show_table:
    st.dataframe(df, use_container_width=True)

# カード形式で各名言を表示
for _, r in df.iterrows():
    quote_id = str(r.get("quote_id", ""))
    quote = str(r.get("quote", ""))
    series = str(r.get("series_name", ""))
    season = r.get("season", "")
    episode = r.get("episode", "")
    ep_name = str(r.get("episode_name", ""))

    st.markdown(f"### 💬 {quote}")
    st.caption(f"📺 {series} - S{season}E{episode} / {ep_name}")

    if "score" in df.columns:
        st.caption(f"score: {r.get('score', '')} / helpful: {r.get('helpful', '')}")

    # --- 理由生成UI ---
    reason_key: Tuple[str, str] = (query_text, quote_id)

    cols = st.columns([1, 3])
    with cols[0]:
        disabled = (not WAREHOUSE_ID) or (not quote_id) or (not query_text)
        if st.button("理由を生成", key=f"gen_{query_text}_{quote_id}", disabled=disabled):
            try:
                with st.spinner("おすすめ理由を生成中..."):
                    reason = generate_reason_with_ai_gen(query_text, quote_id)
                st.session_state["reasons"][reason_key] = reason if reason else "(生成結果が空でした)"
            except Exception as e:
                st.session_state["reasons"][reason_key] = f"(生成に失敗: {e})"

    with cols[1]:
        if reason_key in st.session_state["reasons"]:
            st.info(st.session_state["reasons"][reason_key])
        else:
            st.caption("（理由は未生成）")

    st.divider()

UI設計のポイント：

• 各名言をカード形式で表示し、視認性を向上
• 「理由を生成」ボタンは名言ごとに独立
• 生成中はst.spinnerでローディング表示
• 生成済みの理由はst.session_state["reasons"]にキャッシュ

まとめ

今回は、名言レコメンダーアプリの実装を詳しく解説しました。

主なポイント：

• WorkspaceClientでVector Search APIを呼び出し
• Apps環境とローカル環境の両方に対応した認証処理
• ai_gen関数でLLMによる動的な理由生成
• st.session_stateによる状態管理とキャッシュ

次回は、このアプリをDatabricks Appsにデプロイする方法を解説します。app.yamlの設定やAsset Bundlesでのリソース管理について詳しく見ていきます。

コード全文

本記事で解説したコードの全文は、以下のリポジトリで公開しています。
（リポジトリURL）

---

前回：第1回 準備編
次回：第3回 デプロイ編