Quickstart を利用して Snowflake における RAG ベースの LLM アシスタントの構築手順を確認してみる #SnowflakeDB

モダンデータスタック(MDS)

#Snowflake

安原朋紀

2024.06.11

はじめに

2024年5月のリリースで、一部のリージョンではありますが、Snowflake 上でベクトルデータの管理と操作が可能な以下の機能が一般提供になりました。

VECTOR データ型
Snowflake Cortex LLM ベース関数
- EMBED_TEXT_768
Vector similarity functions

この機能により Snowflake 上で RAG（Retrieval-Augmented Generation）アプローチを取る AI アプリケーションの構築を行えるようになります。
また、このためのチュートリアルとして、公式から「Build a Retrieval Augmented Generation (RAG) based LLM assistant using Streamlit and Snowflake Cortex」が提供されています。

本記事ではこのクイックスタートをベースに Snowflake における RAG ベースの AI アシスタントの構築について、基本的な方法や手順についてまとめてみます。
クイックスタートに基づく各要素や概念の基本的な部分を抑えた内容なので、各要素やコマンド等の詳細までは触れていませんのでご注意ください。

RAG の概要

RAG（Retrieval-Augmented Generation）は、大規模言語モデル（LLM）を活用する際、事前に学習したデータに含まれていない情報に基づいて回答を生成できるようにするアプローチの一つです。

例えば、企業が AI アシスタントを利用する場合、その企業独自の規則やポリシーに関する質問に対して、標準的なモデルだけで正確な回答を提供することは困難です。

RAG では、ユーザーからの質問に基づいて関連する文章（ドキュメント）を検索し、それらの文章を LLM の入力（プロンプト）に背景情報として加えます。これにより、モデルが事前に学習していない内容でも、適切な回答を生成することが可能になります。

RAG については以下の記事でも紹介されていますので、あわせてご覧ください。

コサイン類似度でドキュメントの類似性を評価

上述の通り RAG によるアプローチでは、ユーザーからの質問と関連するドキュメントを検索し、プロンプトに付け加えます。つまり、ユーザーからの質問と関連のある（類似度の高い）ドキュメントを何らかの方法で抽出できるような仕組みが必要です。

ここで「ベクトル化」や「コサイン類似度」といったアプローチが使用されます。

ベクトル化とは、自然言語の文章をコンピュータが扱いやすい数値データに変換することを指します。この変換により、文章は数値の配列（ベクトル）として表現され、これを基に文章間の類似度を数値的に計算できるようになります。

RAG ではユーザーからの質問と関連のあるドキュメントそれぞれをベクトル化することで、類似性の評価を試みます。

具体的な類似性の評価方法には、コサイン類似度が使用されます。
ベクトル形式になった文章を利用し、それらがなす角度から類似性を評価します。出力は -1 から 1 の範囲になり 1 に近いほどベクトルが同じ方向、-1 に近い程、逆の方向を向いていることを指します。
コサイン類似度については、以下でも紹介されていますので、あわせてご参照ください。

まとめると、以下のアプローチにより、RAG では事前に学習済みでない内容であっても回答することを可能とします。

社内規則などの関連するドキュメントを事前にベクトル形式で保存
ユーザーからの質問内容をベクトル化
ベクトル化された質問と社内ドキュメントの類似性をコサイン類似度で評価
類似性が高いドキュメントの内容をプロンプトに背景情報として追加し、LLM が適切な回答を生成

クイックスタートからの引用ですが、下図のイメージです。

Snowflake における RAG ベースの AI アシスタントの構築

利用できる機能

上記のアプローチを達成する手段として、Snowflake では現在、以下の機能が使用できます。

文章のベクトル化
- Snowflake Cortex LLM ベース関数
  - EMBED_TEXT_768
  - EMBED_TEXT_1024
ベクトル形式のデータの保存
- VECTOR データ型
コサイン類似度の評価
- Snowflake Cortex LLM ベース関数
  - VECTOR_COSINE_SIMILARITY
LLM による回答の生成（プロンプトを与えると、サポートされている LLM モデルを使用して応答を生成）
- Snowflake Cortex LLM ベース関数
  - COMPLETE

さらに Streamlit in Snowflake を使用することで、Snowflake 内でデータを外部に出すことなく RAG ベースの AI アシスタント（チャットボット）の構築が行えます。

※図はクイックスタートより引用

Snowflake における AI の信頼性と安全性

Snowflake における AI 利用時のポリシーに関係するよくある疑問点が以下にまとめられています。代表的な内容ですと顧客データが LLM のトレーニングデータとして使用されることはありません。

以降でクイックスタートをベースに構築手順を確認してみます。

Snowflake アカウントの作成

執筆時点（2024/6/11）で Snowflake Cortex LLM ベース関数を利用可能なリージョンには制限があります。ここでは AWS US West 2 (Oregon) リージョンにアカウントを作成し、以降の手順を進めました。
Large Language Model (LLM) Functions (Snowflake Cortex) | Snowflake Documentation

ドキュメントの追加

RAG ベースのアプリケーション構築のために、背景情報となるドキュメントに Snowflake アカウントからアクセスできるようにします。
本クイックスタートでは、PDF 形式のドキュメントを使用します。ドキュメントへのアクセス方法はシンプルで、Snowflake ステージに PDF ファイルを追加します。クイックスタートでは以下のコマンドで内部ステージを作成しています。

create or replace stage docs ENCRYPTION = (TYPE = 'SNOWFLAKE_SSE') DIRECTORY = ( ENABLE = true );

内部ステージ作成後は、Snowsight からドキュメントをアップロード可能です。クイックスタート用のサンプルとして、以下のバイクに関するドキュメントが提供されています。

Vector Store の作成

ドキュメントを追加後、ドキュメントの内容をベクトル形式で保存するためのデータベースである Vector Store（ベクトルストア）を用意します。

Snowflake では、ベクトル形式のデータ保存に VECTOR データ型を使用できます。具体的には以下のように、ベクトル形式のデータを格納するカラムを VECTOR データ型として定義します。

ここではCHUNK_VECカラムにベクトル形式のデータが格納されます。

create or replace TABLE DOCS_CHUNKS_TABLE ( 
    RELATIVE_PATH VARCHAR(16777216), -- Relative path to the PDF file
    SIZE NUMBER(38,0), -- Size of the PDF
    FILE_URL VARCHAR(16777216), -- URL for the PDF
    SCOPED_FILE_URL VARCHAR(16777216), -- Scoped url (you can choose which one to keep depending on your use case)
    CHUNK VARCHAR(16777216), -- Piece of text
    CHUNK_VEC VECTOR(FLOAT, 768) );  -- Embedding using the VECTOR data type

テキストのベクトルストアへの格納

次にベクトルストアにドキュメントのデータを格納します。ここでは以下の手順を実施します。

PDF からテキストを抽出
抽出したテキストをチャンク化
各チャンクをベクトル化しテーブルに追加

PDF からテキストを抽出

PDF からテキストを抽出する方法として、クイックスタートではPyPDF2.PdfReaderによるPython UDF を使用しています。

抽出したテキストをチャンク化

テキストデータをテーブルに追加する前に、テキストをチャンクに分割します。
LLM 使用時にはプロンプトのトークン数に制限があるため「チャンク」としてテキストを分割します。ここでのチャンク化にはlangchain.text_splitterのRecursiveCharacterTextSplitterを使用しています。

下図は公式ドキュメントからの引用ですが、COMPLETE 関数であればモデルごとに入力できるトークン数（Context Window）が以下のように異なります。

ここでトークン数が多すぎると、下図のようにエラーとなります。

なお、トークンとは Snowflake Cortex LLM ベース関数使用時に処理されるテキストの最小単位を指します。実際にどの程度のテキストが消費されるかはモデルにより異なる場合がありますが、約 4 文字のテキストに相当するとされています。
Cost Considerations | Snowflake DOCUMENTATION

また、以下はチャンク化のパラメータを指定している箇所の引用ですが、チャンクサイズとあわせて「オーバーラップ」を指定しています。

text_splitter = RecursiveCharacterTextSplitter(
	chunk_size = 4000, #Adjust this as you see fit
	chunk_overlap  = 400, #This let's text have some form of overlap. Useful for keeping chunks contextuallength_function = len
	)

オーバーラップとは、後続のチャンクの先頭に前のチャンクの末尾を追加することを指します。

短いドキュメントであれば問題ないかもしれませんが、単に文字数でテキストをチャンク化すると、チャンク間で文章が途中で切れてしまうことがあります。これにより、本来異なる文脈で意味が解釈されてしまい、検索や応答生成時の性能低下につながる可能性があるため、オーバーラップにより前後の文脈を含むチャンクを作成します。

各チャンクをベクトル化しテーブルに追加

ベクトル化には Snowflake Cortex LLM ベース関数である EMBED_TEXT_768 を使用しています。この関数は以下のように使用する LLM モデルと対象となるテキストを引数に与えることで、768次元のベクトルとして出力します。

SELECT SNOWFLAKE.CORTEX.EMBED_TEXT_768('snowflake-arctic-embed-m', 'hello world');

出力：どんなインプットも768次元のベクトルに変換される

    [-0.032858,0.036389,0.048158,0.001509,0.032917,-0.001449,0.046142,0.018617,-0.016666,0.013738,-0.050149,-0.007717,-0.011574,-0.021794,0.026656,0.002162,-0.017918,-0.010640,-0.012588,-0.053700,-0.054597,-0.022195,0.009654,-0.013271,-0.021270,0.052335,0.004682,0.020020,-0.072372,-0.080175,-0.000698,0.032825,-0.001189,-0.054464,-0.040156,-0.019773,-0.007758,0.007947,-0.043258,0.024828,-0.029887,-0.028298,0.020229,0.007422,-0.041415,-0.000634,-0.069232,0.041466,-0.038445,-0.035873,0.038067,0.048551,0.056075,0.016130,0.006108,0.011639,0.007035,-0.080729,-0.015572,-0.048236,0.089430,0.017616,0.021103,0.039455,0.008811,-0.057156,-0.012811,-0.049962,-0.016864,-0.018913,-0.032345,-0.002754,0.000006,-0.005927,0.015059,-0.040694,0.033418,0.015427,0.008078,0.027225,-0.022935,-0.014708,0.066962,0.046100,-0.003607,-0.019633,0.017352,0.029912,-0.004652,0.032199,-0.001801,-0.026827,0.056377,0.029610,0.033632,-0.011250,0.045732,0.018888,0.010995,0.012014,-0.072778,-0.028628,-0.000007,-0.025623,-0.075888,-0.025290,0.039146,0.000521,-0.035164,0.070573,-0.037640,-0.026167,-0.008870,-0.035098,-0.020823,0.085145,0.003122,0.023665,0.021641,0.019044,0.029397,0.063050,0.020626,0.059092,0.044762,-0.010749,0.013017,0.012802,-0.042113,-0.052899,0.019535,0.068804,0.042628,-0.040636,0.036740,-0.018702,-0.036101,-0.024203,0.008790,-0.034028,0.036629,0.062777,0.022717,-0.068141,0.042430,0.013881,-0.030856,0.027111,0.025598,0.049491,0.027249,0.039879,-0.048316,-0.050903,-0.001659,0.036691,-0.002667,-0.022170,-0.036229,-0.009680,0.024727,-0.024906,0.014692,0.025015,-0.036671,-0.022536,0.020899,0.027890,-0.007853,-0.011033,-0.047707,-0.069503,0.018603,-0.032965,0.010079,0.003919,0.010399,-0.020688,0.015375,0.011903,-0.009906,0.027769,-0.020321,0.029659,0.009824,-0.006620,0.000038,-0.049348,-0.018708,0.025637,-0.031214,-0.049385,-0.004005,0.004502,0.009184,0.017372,-0.033499,-0.006058,-0.047030,0.061738,0.030114,0.025036,0.019426,-0.046480,-0.034448,0.052054,0.001549,-0.054223,-0.017858,0.098957,-0.026008,0.013030,0.029441,-0.002916,-0.035051,0.001071,0.012341,-0.065667,-0.012416,-0.004425,0.014381,-0.068219,-0.005340,0.015116,-0.012697,0.031564,0.078412,-0.018455,0.020183,0.011334,0.008175,-0.004704,0.043941,-0.042011,-0.012948,0.039800,0.026167,0.022659,0.018743,-0.031375,0.052089,0.107426,-0.026660,-0.001946,-0.056365,0.012070,0.053726,-0.042542,0.006759,-0.006019,-0.010913,-0.075143,-0.026915,-0.006381,0.081104,-0.006681,0.022566,0.045363,0.038600,0.008515,-0.081798,0.063457,-0.022195,-0.047645,-0.064559,-0.042789,-0.047111,-0.067431,0.028864,-0.009412,-0.002030,0.027907,-0.014548,0.025241,-0.011498,-0.064184,0.040182,0.059966,-0.070598,-0.014169,-0.005946,0.065191,-0.010093,-0.016125,0.054851,-0.032352,0.000121,-0.072817,0.020456,-0.021622,0.001159,0.074700,-0.059641,0.000674,0.025977,-0.006790,0.010166,-0.016179,-0.031281,-0.002177,0.029666,0.007822,0.007970,-0.041332,-0.036586,-0.010047,-0.019492,-0.041724,0.013058,0.067281,0.036541,0.000045,-0.000368,0.052646,0.082510,-0.001615,-0.016617,0.008321,-0.061382,0.035367,-0.003814,0.059335,0.019503,0.010717,-0.010472,-0.042571,-0.038445,0.056656,-0.010642,-0.036865,-0.027693,-0.025485,-0.027349,0.017086,0.014756,-0.035073,-0.002554,-0.014352,-0.053727,0.002491,-0.027894,-0.001272,-0.039139,-0.059469,-0.004620,-0.021868,-0.044203,0.016970,-0.060933,0.059201,0.018617,0.017147,-0.032952,-0.043599,-0.019164,-0.022788,0.025518,0.043143,-0.027799,0.076194,0.014546,0.005107,-0.021904,0.019416,0.003609,0.010423,-0.016976,-0.008060,0.039599,-0.049444,-0.019263,-0.038318,0.018919,-0.017908,-0.016810,-0.048235,0.030679,0.011744,-0.027958,-0.006774,0.004862,-0.062531,0.005336,0.055178,0.003131,-0.052110,0.007469,0.022316,-0.078821,0.005919,0.053965,0.057252,-0.051198,-0.014011,0.009770,0.019818,-0.048278,0.035846,0.060350,-0.026849,0.004585,0.021397,-0.006189,0.013810,0.025495,0.048084,0.030157,-0.026026,0.055048,-0.093150,-0.018711,0.041717,0.042309,0.060815,-0.093706,-0.033072,0.045370,0.037914,0.020102,-0.039707,-0.034351,-0.049535,-0.000511,0.007140,-0.016244,-0.024565,0.013175,0.008940,-0.035863,0.019188,-0.045998,-0.021100,-0.016913,0.021939,-0.027042,-0.027040,-0.005715,-0.000597,-0.032854,-0.016977,-0.024760,-0.013980,0.058509,0.006683,0.006590,-0.009875,0.054557,-0.055559,0.018817,-0.008321,0.013835,-0.014729,0.026970,0.020800,-0.047116,0.029564,-0.000768,-0.003098,0.026382,-0.061850,0.038636,0.005210,-0.029467,-0.054451,0.044589,0.027840,-0.009708,0.025598,0.032270,-0.006928,-0.049293,0.044158,0.019610,0.019948,-0.025866,-0.015850,0.034210,0.048011,-0.037846,-0.007523,-0.004416,-0.010112,-0.022988,0.017144,-0.072232,0.029938,0.040650,0.011194,0.033793,-0.037585,-0.035095,0.041628,-0.001569,-0.054865,0.010310,-0.011348,-0.055397,0.012719,0.020510,0.022086,0.013379,0.021168,0.032426,0.015332,0.018244,-0.023512,0.008284,-0.064752,-0.021358,0.003235,-0.103931,0.047833,0.024245,-0.087685,-0.015741,-0.034183,-0.011050,0.017964,0.029842,-0.020753,-0.023489,-0.007431,-0.030349,-0.031201,-0.003824,0.017445,0.008783,0.010470,0.001655,0.009500,-0.044007,0.009033,0.007440,0.038497,-0.065685,0.020682,0.002992,-0.029490,0.024536,0.041165,-0.043204,-0.014845,0.045809,-0.037072,-0.008651,0.032817,0.000657,0.020825,0.018723,0.064041,0.028010,-0.044210,-0.041313,0.023409,-0.021781,-0.023892,0.006784,-0.013820,-0.051877,-0.059424,-0.012378,-0.041178,0.032001,-0.013536,-0.031700,-0.038871,-0.037561,-0.053154,0.035963,0.088215,-0.034972,0.064522,0.039196,0.029848,-0.032507,-0.000271,0.040493,-0.014980,-0.053235,0.016222,-0.004568,0.015411,0.048765,0.020560,-0.028412,0.027806,0.009608,0.017335,0.008623,-0.000666,-0.036620,0.032315,-0.014574,-0.056052,-0.021910,-0.010501,-0.011130,0.056128,-0.053079,0.057192,-0.042287,-0.001254,0.012339,0.022313,0.068580,0.045093,-0.054312,-0.013255,-0.030909,0.056036,0.034661,-0.009781,0.017045,-0.042452,-0.066802,-0.006364,0.016283,0.057360,0.019494,-0.036123,0.014936,0.037369,-0.099000,-0.058392,0.066460,-0.052186,-0.026488,-0.025125,-0.030141,-0.073468,0.069982,0.053295,0.001655,0.002447,-0.019913,0.017231,-0.030827,0.029150,0.041863,0.071274,0.045297,0.025203,-0.037056,0.016575,0.033972,-0.025142,-0.008662,-0.048499,0.070374,0.036593,-0.048089,0.026337,-0.030154,0.031974,-0.022513,-0.048942,-0.059333,0.018051,-0.012454,-0.011369,0.041344,0.007743,0.026571,0.036890,0.023192,-0.023668,0.039622,0.000727,-0.025847,0.017826,0.068080,0.004516,-0.035740,0.029113,0.058188,-0.033431,-0.012088,0.054447,-0.011842,0.061572,0.005824,0.066272,-0.035049,-0.015460,0.019313,-0.042187,0.014442,0.020244,0.029648,-0.065102,-0.002060,-0.005043,0.031330,-0.059832,0.013470,0.044435,-0.050419,-0.023848,0.003812,-0.035683,0.047078,0.005681,-0.029079,-0.034666,-0.032798,0.044730,0.076335,0.003229,0.029985,0.061113,0.001280,-0.013108,0.045857,0.064978,-0.062995,0.014849,0.030959,0.002227,0.052298,0.061020,0.018498,-0.015008,-0.039610,-0.020839,0.045437,-0.012486,0.051835,-0.039826,0.008686,-0.064059,0.042249,0.057247,0.020072,0.030287,-0.020596,0.018325,-0.040275,0.001696,-0.015361,-0.014573,0.051381,0.014755,0.013745,-0.040593,-0.007947,0.039248,0.028329,-0.054030,-0.003811,-0.020322,0.042980,0.017963,-0.020386,-0.046997,0.037754,-0.035768,-0.026911,-0.028785,0.005464,-0.016384,-0.037259,0.029688,0.060227,0.022316,-0.013421,-0.018721,0.020193,0.049298]

ここでは、各チャンクに対してこの関数を適用します。クイックスタートでは、以下の箇所でテキストの抽出からベクトル化までをまとめて行っています。

insert into docs_chunks_table (relative_path, size, file_url,
                            scoped_file_url, chunk, chunk_vec)
    select relative_path, 
            size,
            file_url, 
            build_scoped_file_url(@docs, relative_path) as scoped_file_url, 
            func.chunk as chunk,
            SNOWFLAKE.CORTEX.EMBED_TEXT_768('e5-base-v2',chunk) as chunk_vec   --チャンクのベクトル化
    from 
        directory(@docs),
        TABLE(pdf_text_chunker(build_scoped_file_url(@docs, relative_path))) as func;   --パス内のPDFファイルをチャンク化

この時点でのベクトルストアの中身は下図のようになります。

[CHUNK] カラムがもともと１ファイルだった PDF を指定のサイズとオーバーラップでチャンク化した内容です。[CHUNK_VEC] には、各 [CHUNK] がベクトル化された状態で格納されています。

RAG を使用するシンプルなチャット機能の実装

ここまでの手順で RAG のプロンプトに背景情報として加える際の情報源となるドキュメントの用意が完了しました。

実際に RAG ベースで検索を行う際は、ユーザーからの質問を基にベクトルストアから関連するドキュメントの検索が必要です。そこでドキュメントの検索には、上述のコサイン類似度を使用できます。Snowflake の機能としては Snowflake Cortex LLM ベース関数の VECTOR_COSINE_SIMILARITY を使用できます。

この関数では、以下のように比較したいベクトルデータ型のデータを与えます。

SELECT
    VECTOR_COSINE_SIMILARITY(
        SNOWFLAKE.CORTEX.EMBED_TEXT_768('snowflake-arctic-embed-m', 'hello world'),
        SNOWFLAKE.CORTEX.EMBED_TEXT_768('snowflake-arctic-embed-m', 'hello world')
     ) as similarity;

出力：上の例では同じインプットなので類似度は１

SIMILARITY
1

RAG ベースの AI アシスタントとして使用する際は、ユーザーからの質問を基にドキュメントを検索することになるので、以下のようにユーザーからの質問とベクトル化した各チャンクを比較し、類似度が高い（出力が１に近い）チャンクを関連性が高いドキュメントとして抽出することになります。

SELECT
    VECTOR_COSINE_SIMILARITY(
        SNOWFLAKE.CORTEX.EMBED_TEXT_768('snowflake-arctic-embed-m', '<ユーザーからの質問>'),
        SNOWFLAKE.CORTEX.EMBED_TEXT_768('snowflake-arctic-embed-m', '<ベクトル化したチャンク>')
     ) as similarity;

クイックスタートでは Streamlit in Snowflake を使用してはじめに以下のコードで RAG ベースの回答を生成するアプリケーションを作成しています。

クイックスタートのコード

import streamlit as st # Import python packages
from snowflake.snowpark.context import get_active_session
session = get_active_session() # Get the current credentials

import pandas as pd

pd.set_option("max_colwidth",None)
num_chunks = 3 # Num-chunks provided as context. Play with this to check how it affects your accuracy

def create_prompt (myquestion, rag):

    if rag == 1:    

        cmd = """
            with results as
            (SELECT RELATIVE_PATH,
            VECTOR_COSINE_SIMILARITY(docs_chunks_table.chunk_vec,
                    SNOWFLAKE.CORTEX.EMBED_TEXT_768('e5-base-v2', ?)) as similarity,
            chunk
            from docs_chunks_table
            order by similarity desc
            limit ?)
            select chunk, relative_path from results 
            """
    
        df_context = session.sql(cmd, params=[myquestion, num_chunks]).to_pandas()      
        
        context_lenght = len(df_context) -1

        prompt_context = ""
        for i in range (0, context_lenght):
            prompt_context += df_context._get_value(i, 'CHUNK')

        prompt_context = prompt_context.replace("'", "")
        relative_path =  df_context._get_value(0,'RELATIVE_PATH')
    
        prompt = f"""
            'You are an expert assistance extracting information from context provided. 
            Answer the question based on the context. Be concise and do not hallucinate. 
            If you don´t have the information just say so.
            Context: {prompt_context}
            Question:  
            {myquestion} 
            Answer: '
            """
        cmd2 = f"select GET_PRESIGNED_URL(@docs, '{relative_path}', 360) as URL_LINK from directory(@docs)"
        df_url_link = session.sql(cmd2).to_pandas()
        url_link = df_url_link._get_value(0,'URL_LINK')

    else:
        prompt = f"""
            'Question:  
            {myquestion} 
            Answer: '
            """
        url_link = "None"
        relative_path = "None"
        
    return prompt, url_link, relative_path

def complete(myquestion, model_name, rag = 1):

    prompt, url_link, relative_path =create_prompt (myquestion, rag)
    cmd = f"""
                select SNOWFLAKE.CORTEX.COMPLETE(?,?) as response
            """
    
    df_response = session.sql(cmd, params=[model_name, prompt]).collect()
    return df_response, url_link, relative_path

def display_response (question, model, rag=0):
    response, url_link, relative_path = complete(question, model, rag)
    res_text = response[0].RESPONSE
    st.markdown(res_text)
    if rag == 1:
        display_url = f"Link to [{relative_path}]({url_link}) that may be useful"
        st.markdown(display_url)

#Main code

st.title("Asking Questions to Your Own Documents with Snowflake Cortex:")
st.write("""You can ask questions and decide if you want to use your documents for context or allow the model to create their own response.""")
st.write("This is the list of documents you already have:")
docs_available = session.sql("ls @docs").collect()
list_docs = []
for doc in docs_available:
    list_docs.append(doc["name"])
st.dataframe(list_docs)

#Here you can choose what LLM to use. Please note that they will have different cost & performance
model = st.sidebar.selectbox('Select your model:',(
                                    'mixtral-8x7b',
                                    'snowflake-arctic',
                                    'mistral-large',
                                    'llama3-8b',
                                    'llama3-70b',
                                    'reka-flash',
                                        'mistral-7b',
                                        'llama2-70b-chat',
                                        'gemma-7b'))

question = st.text_input("Enter question", placeholder="Is there any special lubricant to be used with the premium bike?", label_visibility="collapsed")

rag = st.sidebar.checkbox('Use your own documents as context?')

print (rag)

if rag:
    use_rag = 1
else:
    use_rag = 0

if question:
    display_response (question, model, use_rag)

以降で、このコードのポイントとなる部分を見ていきます。

はじめの関数create_promptでは以下の操作を行っています。

ユーザーからの質問を受け取り
ユーザーからの質問をベクトル化し、ベクトル化済みのチャンクとの類似度を評価
類似度の高いチャンクを抽出
抽出したチャンクのテキストを結合しプロンプトに追加するための背景情報を作成

以下の部分で、ユーザーから受け取った質問をベクトル化しつつ、ベクトルストアのカラムとの類似度を評価しています。計算された類似度を基準に降順に並び替えることで、類似度の高いドキュメントを抽出するという仕組みです。

cmd = """
    with results as
    (SELECT RELATIVE_PATH,
    VECTOR_COSINE_SIMILARITY(docs_chunks_table.chunk_vec,
            SNOWFLAKE.CORTEX.EMBED_TEXT_768('e5-base-v2', ?)) as similarity,
    chunk
    from docs_chunks_table
    order by similarity desc
    limit ?)
    select chunk, relative_path from results 
    """

df_context = session.sql(cmd, params=[myquestion, num_chunks]).to_pandas()

また、limit句ではnum_chunks変数で上位いくつまでのチャンクをコンテキストに追加するかをコントロールしています。
※SQL に含まれる?は、session.sqlメソッドのparams引数で指定されたmyquestionとnum_chunksによって置き換えられます。

実際にプロンプトに指定する文章は、以下の箇所で作成しています。

prompt_context = ""
    for i in range (0, context_lenght):
        prompt_context += df_context._get_value(i, 'CHUNK')

    prompt_context = prompt_context.replace("'", "")
    relative_path =  df_context._get_value(0,'RELATIVE_PATH')

    prompt = f"""
    'You are an expert assistance extracting information from context provided. 
    Answer the question based on the context. Be concise and do not hallucinate. 
    If you don´t have the information just say so.
    Context: {prompt_context}
    Question:  
    {myquestion} 
    Answer: '
    """

prompt_contextとしてfor 文で指定のチャンク数文取得した背景情報のテキストを結合し、prompt 変数に挿入しています。これにより、背景情報を提供しつつ、ユーザーからの質問に回答するようにしています。

このコードでは RAG ベースでない質問も可能なように条件分岐されており、その場合は以下のように、ユーザーからの質問のみからなるシンプルなプロンプトが生成されます。

prompt = f"""
         'Question:  
           {myquestion} 
           Answer: '
           """

それぞれ RAG を使用するかどうかで、具体的には以下のようなプロンプトが生成されます。

以降のcomplete 関数では生成したプロンプトを基に Snowflake Cortex の COMPLETE 関数を実行します。

def complete(myquestion, model_name, rag = 1):

    prompt, url_link, relative_path =create_prompt (myquestion, rag)
    cmd = f"""
             select SNOWFLAKE.CORTEX.COMPLETE(?,?) as response
           """
    
    df_response = session.sql(cmd, params=[model_name, prompt]).collect()
    return df_response, url_link, relative_path

例えば、上記の関数ロジックのみを実行すると下図のように、COMPLETE 関数の実行結果を得られます。（その他にも、ここでは類似度の高いドキュメントのリンクも表示するアプリとして作成しているため、その URL も出力されています。）

上図では RAG = 1 としているので、内部的にはユーザーからの質問をベクトル化し、ベクトルストアで類似度の高いドキュメントを抽出し、プロンプトに付け加えるといった処理が行われています。
その他の箇所では、Streamlit による UI の構築などを行っています。

RAG を使用するチャットボット機能の実装

1つ前の手順で、シンプルな RAG ベースの回答を生成することができましたが、前提として LLM はステートレスです。つまり、LLM を複数回呼び出しても、前回の呼び出しの内容は記憶されません。Snowflake Cortex も同様で各 COMPLETE 関数の呼び出しは独立しています。

例えば以下のように続けて関数を実行しても前回の内容に基づく回答は得られません。

会話の流れや文脈を維持するためには、会話の履歴を保持し、LLM への各呼び出しでそれを提供する必要があります。

会話履歴の保持

このクイックスタートでは、チャット履歴の保持のために、Streamlit ライブラリのst.seession_stateを使用しています。ユーザーがチャットに入力を行ったり LLM が回答を生成すると、その内容がst.session_state.messagesに追加されます。これにより、質問や応答がセッション間で保持される仕組みです。
Session State | Streamlit Documentation

st.session_state.messagesリストに追加される各メッセージは、それぞれroleとcontentというキーを持っています。roleキーで、以下の通りそのメッセージがユーザーからのものなのか、LLMからのものなのかを識別できます。

"role": "user"
- そのメッセージがユーザーによる入力（質問など）であることを示す
"role": "assistant"
- そのメッセージが LLM による応答であることを示す

メッセージが追加された時のオブジェクト内のイメージは以下の通りです。

履歴の追加はmain関数の以下の部分で制御しています。

クイックスタートのコード

# Accept user input
if question := st.chat_input("What do you want to know about your products?"):
    # Add user message to chat history
    st.session_state.messages.append({"role": "user", "content": question})
    # Display user message in chat message container
    with st.chat_message("user"):
        st.markdown(question)
    # Display assistant response in chat message container
    with st.chat_message("assistant"):
        message_placeholder = st.empty()

        question = question.replace("'","")

        with st.spinner(f"{st.session_state.model_name} thinking..."):
            response = complete(question)
            res_text = response[0].RESPONSE     
        
            res_text = res_text.replace("'", "")
            message_placeholder.markdown(res_text)
    
    st.session_state.messages.append({"role": "assistant", "content": res_text})

ベクトルストア検索の性能向上

また、ここでは会話履歴を使用しユーザーからの質問に基づく、関連ドキュメント検索時の精度を上げるような仕組みが追加されています。
以下の箇所でget_chat_historyにより会話履歴を取得し、すぐ後に定義されているsummarize_question_with_history関数のインプットとして使用することで、会話履歴とユーザーからの質問を組み合わせ、COMPLETE 関数によって質問を拡張する操作を行っています。

クイックスタートのコード

def get_chat_history():
#Get the history from the st.session_stage.messages according to the slide window parameter
    
    chat_history = []
    
    start_index = max(0, len(st.session_state.messages) - slide_window)
    for i in range (start_index , len(st.session_state.messages) -1):
         chat_history.append(st.session_state.messages[i])

    return chat_history

    
def summarize_question_with_history(chat_history, question):
# To get the right context, use the LLM to first summarize the previous conversation
# This will be used to get embeddings and find similar chunks in the docs for context

    prompt = f"""
        Based on the chat history below and the question, generate a query that extend the question
        with the chat history provided. The query should be in natual language. 
        Answer with only the query. Do not add any explanation.
        
        <chat_history>
        {chat_history}
        </chat_history>
        <question>
        {question}
        </question>
        """
    
    cmd = """
            select snowflake.cortex.complete(?, ?) as response
          """
    df_response = session.sql(cmd, params=[st.session_state.model_name, prompt]).collect()
    sumary = df_response[0].RESPONSE     

    if st.session_state.debug:
        st.sidebar.text("Summary to be used to find similar chunks in the docs:")
        st.sidebar.caption(sumary)

    sumary = sumary.replace("'", "")

    return sumary

もう少し具体的には、以下のようなことを行っています。

例として、LLM との次のような会話履歴があるとします。

ここで、続けて「AIが健康管理にどのように貢献できるか、もう少し詳しく教えてください。」と質問するとします。summarize_question_with_history関数では、指定のスライドウィンドウを基に、はじめに以下のようなプロンプトを生成します。

Based on the chat history below and the question, generate a query that extend the question
with the chat history provided. The query should be in natual language. 
Answer with only the query. Do not add any explanation.

<chat_history>
[{'role': 'user', 'content': 'こんにちは'}, {'role': 'assistant', 'content': 'こんにちは！お手伝いできることがあれば何でも聞いてくださいね。'}, {'role': 'user', 'content': '最近のニュースについて教えてください。'}, {'role': 'assistant', 'content': 'もちろんです。最近では、AI技術の進展に関するニュースが多く報じられています。特に自動運転車とAIによる健康管理が注目されています。'}, {'role': 'user', 'content': 'AIによる健康管理とは具体的にどういうことですか？'}]
</chat_history>
<question>
AIが健康管理にどのように貢献できるか、もう少し詳しく教えてください。
</question>

続けて、このプロンプトを使用し COMPLETE 関数を実行し、以下のような回答を得ます。

出力は、もともと質問しようとしていた内容である「AIが健康管理にどのように貢献できるか、もう少し詳しく教えてください。」を会話履歴から「最近のAI技術の進歩により、自動車の運転や健康管理などにもAIが活用されています。AIが健康管理にどのように貢献できるのか、もう少し詳しく教えていただけますか？」としてブラッシュアップした内容となっています。この操作を追加することで、LLM の応答をより良い結果とすることが期待できます。

ここで得られたプロンプトは、以下の部分でベクトルストアに対する検索用のクエリとして使用されています。

if chat_history != []: #There is chat_history, so not first question
    question_summary = summarize_question_with_history(chat_history, myquestion)
    prompt_context =  get_similar_chunks(question_summary)

なお、スライドウィンドウについては、以下が参考になります。これにより、長い会話であっても、トークン制限に達することを防ぎます。

最終的なプロンプトの作成

さいごにcreate_prompt 関数でユーザーからの質問をインプットに、類似度の高いドキュメントの抽出による背景情報や会話履歴を付け加えた、プロンプトを作成します。以下のベースが動的に更新されます。

prompt = f"""
        You are an expert chat assistance that extracs information from the CONTEXT provided
        between <context> and </context> tags.
        You offer a chat experience considering the information included in the CHAT HISTORY
        provided between <chat_history> and </chat_history> tags..
        When ansering the question contained between <question> and </question> tags
        be concise and do not hallucinate. 
        If you don´t have the information just say so.
        
        Do not mention the CONTEXT used in your answer.
        Do not mention the CHAT HISTORY used in your asnwer.
        
        <chat_history>
        {chat_history}
        </chat_history>
        <context>          
        {prompt_context}
        </context>
        <question>  
        {myquestion}
        </question>
        Answer: 
        """

ここで得られたプロンプトを COMPLETE 関数に渡すことで得られた回答がユーザーに表示されます。ユーザー自身の質問と、ここでの応答がst.session_state.messagesに随時履歴として加えられます。

UI の構成

UI の主な構成要素である、チャットボットとの会話部分は、以下の部分で定義されています。

# Display chat messages from history on app rerun
for message in st.session_state.messages:
	with st.chat_message(message["role"]):
		st.markdown(message["content"])

st.chat_message ではチャットメッセージコンテナを挿入することができます。role 名がuserやassistantであれば、それぞれあらかじめ用意されている下図のようなアイコンが使用されます。

st.chat_message | Streamlit Documentation

その他の要素では、使用する LLM モデルの選択や RAG を使用するかなどを UI で制御できるように構成しています。

さいごに

公式のクイックスタートを使用し、Snowflake における RAG ベースの AI アシスタントの構築手順を追ってみました。
Snowflake 内の機能で基本的な RAG ベースの AI アシスタントを構築できることが具体的な実装手順とともに理解できたと思います。

背景情報となるドキュメントの追加については、すでに Snowflake 内に情報があればそのまま利用できますし、PDF であれば Python 等の UDF や今後は Document AI を利用したテキストのを抽出も可能と思います。
外部のアプリケーション内のデータであっても外部ネットワークアクセス経由で取得も可能です。

今回であればテキスト抽出後の、チャンクサイズやオーバーラップ、コンテキストとして使用するチャンク数、スライドウィンドウの長さなど、各データに基づいて適切な回答が得られるかの調整が必要です。しかし、各機能が Snowflake 内で利用できるため、迅速に検証を進めることができそうな印象を受けました。

こちらの内容が何かの参考になれば幸いです。