Databricks SQLでのRAGの構築

SQLでRAGが組めるようになったとは。しみじみ。

本日パブリックプレビューになったこちらの`vector_search`関数。

https://qiita.com/taka_yayoi/items/8f11fd1fdebca515942f

これによって、RAGの構成要素であるベクトルストアにSQLからアクセスできるようになりました。

もともと、SQLからLLMを呼び出す機能は`ai_query`関数として提供されていました。

https://docs.databricks.com/ja/sql/language-manual/functions/ai_query.html

これらを組み合わせれば、RAGが作れるわけです。試してみます。

# 準備

## ベクトルストア

[こちら](https://docs.databricks.com/ja/generative-ai/create-query-vector-search.html)の手順に従って、ベクトル検索インデックスを作成しておきます。

![Screenshot 2024-10-02 at 9.21.11.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/8b8c3fea-2da9-df11-fbd9-76d5d7bd4517.png)

## 外部モデルのサービングエンドポイント

`ai_query`関数で呼び出すLLMのサービングエンドポイントを、[こちら](https://docs.databricks.com/ja/generative-ai/external-models/index.html)の手順に従って作成します。gpt-4o-miniを使います。

![Screenshot 2024-10-02 at 9.22.57.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/8d312206-acc4-9644-54f1-bc14fc9812fb.png)

# SQLでのRAGの構築

## 文脈の取得

まずは、`VECTOR_SEARCH`でベクトル検索インデックスを検索して、文脈を構成する文字列を構築します。

```sql
SELECT
    "Databricksとは何か" || "という質問に対して" || (SELECT
      body_sub
    FROM
      VECTOR_SEARCH(
        index => "takaakiyayoi_catalog.vector_seach.taka_qiita_vs",
        query => "Databricksとは何か",
        num_results => 1
      )
    ) || "という文脈が得られています。簡潔かつ適切な回答を生成してください" AS context
```

動いてますね。

![Screenshot 2024-10-02 at 9.25.30.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/a91daec8-1d92-b7a4-a622-de4f5d09a8dc.png)

## 回答の生成

得られた文脈と質問に基づいて回答を生成してみます。

```sql
SELECT
  ai_query(
    'taka-openai-o4-mini',
    "Databricksとは何か" || "という質問に対して" || (
      SELECT
        body_sub
      FROM
        VECTOR_SEARCH(
          index => "takaakiyayoi_catalog.vector_seach.taka_qiita_vs",
          query => "Databricksとは何か",
          num_results => 1
        )
    ) || "という文脈が得られています。質問と得られた文脈に基づいて簡潔かつ適切な回答を生成してください"
  ) AS response
```

いい感じですね。
![Screenshot 2024-10-02 at 9.26.29.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/3aae8e4b-f2b3-79ac-5379-4e0094756129.png)

# まとめ上げる

これだけでは勿体無いので、ダッシュボードに組み込みます。これで簡易のRAGチャットbotが作れるはずです。

プロンプトを洗練して、パラメーター`:query`を埋め込みます。

```sql
SELECT
  ai_query(
    'taka-openai-o4-mini',
    "あなたはDatabricksの専門家です。" || :query || "という質問に対して" || (
      SELECT
        body_sub
      FROM
        VECTOR_SEARCH(
          index => "takaakiyayoi_catalog.vector_seach.taka_qiita_vs",
          query => :query,
          num_results => 1
        )
    ) || "という文脈が得られています。質問と得られた文脈に基づいて、簡潔かつ適切な回答を生成してください。文脈が得られない場合には「わからない」と回答してください。"
  ) AS response
```

こんな感じでダッシュボードを作成します。
![Screenshot 2024-10-02 at 9.28.39.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/7eb92b1d-da2d-ab6a-de86-585db3a312af.png)
![Screenshot 2024-10-02 at 9.28.29.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/e512bbac-6430-1d8c-0817-3d68a38ea8ad.png)

完成です！
![Screenshot 2024-10-02 at 9.29.19.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/1f832c1b-2965-d22f-47ac-f5b53ad25bc8.png)

関係のない質問には「わからない」と返すようになっています。
![Screenshot 2024-10-02 at 9.30.35.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/af6a2b61-e832-9b44-13c3-7ca0add569ae.png)

プロトタイピングでは有用なアプローチになるのではと思いました。AWS東京リージョンでも動きますので、是非試してみてください！

### はじめてのDatabricks

[はじめてのDatabricks](https://qiita.com/taka_yayoi/items/8dc72d083edb879a5e5d)

### Databricks無料トライアル

[Databricks無料トライアル](https://databricks.com/jp/try-databricks)