Aurora Serverless PostgreSQLのpgvectorをVector DBとして使用したRAGを構成します。
概念図
構成図
VPCで閉じたかったのでCloud9を使っています。
Aurora Serverless PostgreSQLの構成
エンジンタイプ Aurora PostgreSQL Compatible
バージョン 15.4 開発/テスト クラスター識別子vector-test
パスワードは適当に（覚えておいてください）
Serverless v2
「接続」はとりあえずデフォルトで（defaultのVPCとセキュリティグループ）
データベースの作成
作成が終わったらライターのエンドポイントをメモしておきます
Cloud 9の構成
Cloud9を作成します。デフォルトでOKです。
PostgreSQLのセキュリティグループ（default）のインバウンドルールにCloud9のセキュリティグループを追加します。とりあえずならすべてのトラフィック
PostgreSQLクライアントのインストール
sudo amazon-linux-extras install postgresql14
pgvectorの有効化
psqlでライターエンドポイントに接続します。
psql -h vector-test.cluster-xxxxxxxxxxxx.us-east-1.rds.amazonaws.com -U postgres -p 5432
パスワードの入力を求められるので入力してEnter
Password for user postgres:
psql (14.8, server 15.4)
WARNING: psql major version 14, server major version 15.
Some psql features might not work.
SSL connection (protocol: TLSv1.2, cipher: AES128-SHA256, bits: 128, compression: off)
Type "help" for help.
postgres=>
今回の用途であればワーニングは気にしなくて大丈夫です。
後は以下のように入力してpgvectorの有効化は完了です。
postgres=> CREATE EXTENSION vector;
CREATE EXTENSION
テーブル等はLangChainが作成してくれるのでここで作成する必要はありません。
後は確認して抜けます。
postgres=> \dx
List of installed extensions
Name | Version | Schema | Description
---------+---------+------------+--------------------------------------------
plpgsql | 1.0 | pg_catalog | PL/pgSQL procedural language
vector | 0.5.0 | public | vector data type and ivfflat access method
(2 rows)
postgres=> exit
Python環境の構成
必要ライブラリのインストール
pip install boto3
pip install langchain
pip install pgvector
pip install psycopg2-binary
pip install pypdf
pip install streamlit
Vector DBの作成
以下のPDFをPostgreSQLに格納します。
まず上記ファイルをCloud9にUploadします。
今回の例では/home/ec2-user/environment/bedrock-ug.pdfに配置しています。
以下のプログラムを作成・実行します。
CONNECTION_STRINGの{password}はpostgreSQLのマスターパスワードを設定してください。vector-test.cluster-xxxxxxxxxxxx.us-east-1.rds.amazonaws.comはライターエンドポイントを設定してください。
from langchain.embeddings import BedrockEmbeddings
from langchain.vectorstores.pgvector import PGVector
from langchain.document_loaders import PyPDFLoader
# 接続文字列の定義
CONNECTION_STRING = "postgresql+psycopg2://postgres:{password}@vector-test.cluster-xxxxxxxxxxxx.us-east-1.rds.amazonaws.com:5432"
# Embeddingsの定義
embeddings = BedrockEmbeddings(model_id="amazon.titan-embed-text-v1")
# PostgreSQLの定義
db = PGVector(
connection_string=CONNECTION_STRING,
embedding_function=embeddings
)
# PDFファイルをDBに追加
loader = PyPDFLoader(r"/home/ec2-user/environment/bedrock-ug.pdf")
pages = loader.load_and_split()
db.add_documents(documents=pages)
これで上記のPDFファイルが1ページ単位に分割、Titan EmbeddingsによってEmbeddingされて、PostgreSQLに格納されているはずです。
ちなみにPGVectorのI/Fは以下です。
格納内容の確認
psqlでリーダーエンドポイントに接続します。
psql -h vector-test.cluster-xx-xxxxxxxxxxxx.us-east-1.rds.amazonaws.com -U postgres -p 5432
LangChainにより作成されたテーブルの確認
postgres=> \d
List of relations
Schema | Name | Type | Owner
--------+-------------------------+-------+----------
public | langchain_pg_collection | table | postgres
public | langchain_pg_embedding | table | postgres
(2 rows)
postgres=>
LangChainによってlangchain_pg_collectionとlangchain_pg_embeddingの2テーブルが作成されています。
コレクションの確認
langchain_pg_collectionを確認します。
postgres=> \d langchain_pg_collection;
Table "public.langchain_pg_collection"
Column | Type | Collation | Nullable | Default
-----------+-------------------+-----------+----------+---------
name | character varying | | |
cmetadata | json | | |
uuid | uuid | | not null |
Indexes:
"langchain_pg_collection_pkey" PRIMARY KEY, btree (uuid)
Referenced by:
TABLE "langchain_pg_embedding" CONSTRAINT "langchain_pg_embedding_collection_id_fkey" FOREIGN KEY (collection_id) REFERENCES langchain_pg_collection(uuid) ON DELETE CASCADE
postgres=> select * from langchain_pg_collection;
name | cmetadata | uuid
-----------+-----------+--------------------------------------
langchain | null | 79292fae-56a5-4934-ad94-6e1f2b7a9185
(1 row)
1レコード作成されています。
nameのlangchainはLangChainのデフォルト値です。今回はプログラムからの作成時に指定していないのでデフォルト値になりましたが、引数で指定して作成する事も出来ます。またuuidとしてコレクションとやらにIDが振られています。
Embeddingテーブルの確認
langchain_pg_embeddingを確認します。
postgres=> \d langchain_pg_embedding;
Table "public.langchain_pg_embedding"
Column | Type | Collation | Nullable | Default
---------------+-------------------+-----------+----------+---------
collection_id | uuid | | |
embedding | vector | | |
document | character varying | | |
cmetadata | json | | |
custom_id | character varying | | |
uuid | uuid | | not null |
Indexes:
"langchain_pg_embedding_pkey" PRIMARY KEY, btree (uuid)
Foreign-key constraints:
"langchain_pg_embedding_collection_id_fkey" FOREIGN KEY (collection_id) REFERENCES langchain_pg_collection(uuid) ON DELETE CASCADE
postgres=>
件数を確認します。
postgres=> select count(*) from langchain_pg_embedding;
count
-------
156
(1 row)
postgres=>
PDFファイルのページ数相当のレコードが作成されている事が分かります。
少し中を見てみます。embeddingはベクトル値が格納されている巨大なカラムなのでまずは外して検索します。
postgres=> select collection_id,document,custom_id,uuid from langchain_pg_embedding;
collection_id | document
| custom_id | uuid
--------------------------------------+--------------------------------------------------------------------------------------------------------------------------------
---------------------------+--------------------------------------+--------------------------------------
79292fae-56a5-4934-ad94-6e1f2b7a9185 | Amazon Bedrock
+| d5f0a7e4-7bc8-11ee-8460-06a674c2adb7 | 28d92844-45e8-4880-a6f8-79379a981e62
| ユーザーガイド
| |
79292fae-56a5-4934-ad94-6e1f2b7a9185 | Amazon Bedrock ユーザーガイド
+| d5f0a992-7bc8-11ee-8460-06a674c2adb7 | 90199c78-cd20-4143-adc7-6d11e474a972
| Amazon Bedrock: ユーザーガイド
+| |
| Copyright © 2023 Amazon Web Services, Inc. and/or its affiliates. All rights reserved.
+| |
| Amazon の商標およびトレードドレスは、顧客に混乱を招く可能性がある態様、または Amazon の信用を傷つけ
+| |
| たり、失わせたりする態様において、Amazon のものではない製品またはサービスに関連して使用してはなりませ
+| |
| ん。Amazon が所有しない商標はすべてそれぞれの所有者に所属します。所有者は必ずしも Amazon との提携や関連
+| |
| があるわけではありません。また、Amazon の支援を受けているとは限りません。
| |
79292fae-56a5-4934-ad94-6e1f2b7a9185 | Amazon Bedrock ユーザーガイド
+| d5f0aa50-7bc8-11ee-8460-06a674c2adb7 | ee5cde36-2e92-4b2a-9f7d-415032315a6b
| Table of Contents
+| |
| アマゾン・ベッドロックとは？ ..................................................................................................
...........1 +| |
| Bedrock モデルにアクセスしてください。 .....................................................................................1
+| |
文字数が多くてちょっとわかりにくいですが、collection_idにはコレクションのuuidが格納されています。
documentには各ページのテキストが格納されています。
uuidにはレコード毎に異なるIDが振られています。
embeddingに格納されている内容を確認します。
postgres=> select embedding from langchain_pg_embedding;
[0.18457031,0.38671875,0.484375,-0.022338867,0.7578125,0.63671875,-0.18652344,-0.00037384033,0.012207031,-0.3515625,-0.095703125,0.69140625,0.4765625,0.140625,-0.3710
9375,-0.083984375,-0.47460938,0.29492188,-0.30078125,-0.029541016,-0.76171875,0.98828125,0.053466797,0.9453125,-0.42773438,0.625,0.08691406,-0.37890625,-0.609375,-0.31
054688,-0.75,0.65234375,-0.2578125,-0.56640625,-0.20019531,-0.12890625,0.19726562,-1.1015625,-0.65234375,0.69921875,-0.0859375,-0.5546875,-0.123046875,0.21386719,0.279
29688,-0.36132812,0.07128906,0.29101562,1.171875,0.67578125,-0.03149414,0.0026855469,0.21484375,0.30273438,0.16796875,0.26367188,0.50390625,-0.47265625,0.65625,0.26757
812,-0.024169922,0.35546875,-0.13085938,0.25390625,0.4296875,0.0859375,0.80078125,-0.30859375,-0.52734375,0.92578125,-0.671875,0.94140625,0.27148438,-0.2265625,-0.0986
3281,0.48632812,-0.41015625,0.47070312,0.828125,0.056152344,-0.21972656,0.5703125,0.20214844,0.29101562,-0.39257812,-0.51953125,-0.45898438,0.29296875,0.00013828278,
意味が無いので途中で切りますが、embeddingにはベクトル値が格納されている事が分かります。
つまりベクトル検索を行うと、embeddingが近いレコードから取得され、実際の文章としてはdocumentが取得できそうだなという事が分かりました。
RAGの実装
以下のプログラムを作成・実行します。
前回同様に検索内容のEmbeddingとしてはVectorDBと同じTitan Embeddingsを使用し、最後の文章整理についてはClaude2を使用します。
実行しやすいようにStreamlitでガワを被せています。
プログラム上部の定義部分以外は前回のChromaDBのプログラムと変わりありません。
CONNECTION_STRINGの{password}はpostgreSQLのマスターパスワードを設定してください。vector-test.cluster-xx-xxxxxxxxxxxx.us-east-1.rds.amazonaws.comはリーダーエンドポイントを設定してください。
from langchain.embeddings import BedrockEmbeddings
from langchain.vectorstores.pgvector import PGVector
from langchain.llms import Bedrock
from langchain.chains import RetrievalQA
from langchain.prompts import PromptTemplate
import streamlit as st
# 接続文字列の定義
CONNECTION_STRING = "postgresql+psycopg2://postgres:{password}@vector-test.cluster-xx-xxxxxxxxxxxx.us-east-1.rds.amazonaws.com:5432"
# Embeddingsの定義
embeddings = BedrockEmbeddings(model_id="amazon.titan-embed-text-v1")
# PostgreSQLの定義
db = PGVector(
connection_string=CONNECTION_STRING,
embedding_function=embeddings
)
# LLMの定義
llm = Bedrock(
model_id="anthropic.claude-v2",
model_kwargs={"max_tokens_to_sample": 1000},
)
# promptの定義
prompt_template = """
<documents>{context}</documents>
\n\nHuman: 上記の内容を参考文書として、質問の内容に対して詳しく説明してください。言語の指定が無い場合は日本語で答えてください。
もし質問の内容が参考文書に無かった場合は「文書にありません」と答えてください。回答内容には質問自体やタグは含めないでください。
Take a deep breath and work on this problem step-by-step.
<question>{question}</question>
\n\nAssistant:"""
PROMPT = PromptTemplate(template=prompt_template, input_variables=["context", "question"])
chain_type_kwargs = {"prompt": PROMPT}
# Chainの定義。検索結果の上位15件を使用
qa = RetrievalQA.from_chain_type(llm,retriever=db.as_retriever(search_kwargs={"k": 15}),chain_type_kwargs=chain_type_kwargs)
# Streamlit
st.title("PostgreSQL PGVectoRAG")
input_text = st.text_input("入力された文字列でVectorDBを検索し回答します")
send_button = st.button("送信")
if send_button:
# 実行
st.write(qa.run(input_text))
以下のように実行します。
python -m streamlit run pgv_search.py --server.port 8080
Preview→Preview Running Applicationからプレビュー画面を表示します。
素のClaude2が知らないBedrockの情報を聞いてみます。
素のClaude2が知らない情報について答えてくれたので、PostgreSQLに格納した情報を元に回答してくれている事が分かります。また150ページ中の15ページを取得して回答を組み立てているので、質問に近い文章が取得されてそうです。
以上のように、Aurora Serverless PostgreSQLのpgvectorを有効化すれば、Vector DB（Vector Store）として使用してRAGを実装出来る事が分かりました。
今回はLangChainがいろいろ良い感じに処理してくれましたが、Agents for Amazon Bedrockが始まれば同様に良い感じにやってくれる感じでしょうか。