BedrockとAurora Serverless PostgreSQL pgvectorでRAGする - Amazon Titan Embeddingsでベクトル入門③

Aurora Serverless PostgreSQLのpgvectorをVector DBとして使用したRAGを構成します。

# 概念図
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3573242/0225f260-2917-f334-ed02-3633d98e4a03.png)

# 構成図
VPCで閉じたかったのでCloud9を使っています。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3573242/63dde7f3-6d2f-d233-be11-5faf7ae587cd.png)

# Aurora Serverless PostgreSQLの構成
エンジンタイプ Aurora PostgreSQL Compatible
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3573242/623fa115-1a7b-684b-ba4b-a411f9286c9c.png)

バージョン 15.4 開発/テスト クラスター識別子`vector-test`
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3573242/4028b627-bc89-2d68-2db6-ec7757158bf4.png)

パスワードは適当に（覚えておいてください）
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3573242/91678b69-9a8a-8cf2-39b1-261035e049bb.png)

Serverless v2
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3573242/f747ee14-50f3-7227-f27d-4793dccc7f05.png)

「接続」はとりあえずデフォルトで（defaultのVPCとセキュリティグループ）
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3573242/7e088095-340f-ad4d-dd7d-ce38baa2fe14.png)

データベースの作成
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3573242/16977207-386c-660d-067b-e7ee2e6412bd.png)

作成が終わったらライターのエンドポイントをメモしておきます
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3573242/76c714ef-e47a-471d-4866-e474d765c81d.png)

## 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の有効化は完了です。
```:psqlターミナル
postgres=> CREATE EXTENSION vector;
CREATE EXTENSION
```
テーブル等はLangChainが作成してくれるのでここで作成する必要はありません。
後は確認して抜けます。
```:psqlターミナル
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に格納します。

https://docs.aws.amazon.com/ja_jp/bedrock/latest/userguide/bedrock-ug.pdf

まず上記ファイルをCloud9にUploadします。
今回の例では`/home/ec2-user/environment/bedrock-ug.pdf`に配置しています。

以下のプログラムを作成・実行します。
`CONNECTION_STRING`の`{password}`はpostgreSQLのマスターパスワードを設定してください。`vector-test.cluster-xxxxxxxxxxxx.us-east-1.rds.amazonaws.com`はライターエンドポイントを設定してください。

```Python:loader.py
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は以下です。

https://api.python.langchain.com/en/latest/vectorstores/langchain.vectorstores.pgvector.PGVector.html

# 格納内容の確認
psqlでリーダーエンドポイントに接続します。
```:ターミナル
psql -h vector-test.cluster-xx-xxxxxxxxxxxx.us-east-1.rds.amazonaws.com -U postgres -p 5432
```
## LangChainにより作成されたテーブルの確認

```:psqlターミナル
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`を確認します。
```:psqlターミナル
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`を確認します。

```:psqlターミナル
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=> 
```
件数を確認します。

```:psqlターミナル
postgres=> select count(*) from langchain_pg_embedding;
 count 
-------
   156
(1 row)

postgres=> 
```
PDFファイルのページ数相当のレコードが作成されている事が分かります。
少し中を見てみます。embeddingはベクトル値が格納されている巨大なカラムなのでまずは外して検索します。
```:psqlターミナル
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に格納されている内容を確認します。
```:psqlターミナル
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`はリーダーエンドポイントを設定してください。

```python:pgv_search.py
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`からプレビュー画面を表示します。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3573242/0e5706f5-cbba-a050-4d84-0aa1bfd66f85.png)
素のClaude2が知らないBedrockの情報を聞いてみます。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3573242/d72ab4f0-1e64-65e2-41ed-78b7d94ad403.png)
素のClaude2が知らない情報について答えてくれたので、PostgreSQLに格納した情報を元に回答してくれている事が分かります。また150ページ中の15ページを取得して回答を組み立てているので、質問に近い文章が取得されてそうです。

以上のように、Aurora Serverless PostgreSQLのpgvectorを有効化すれば、Vector DB（Vector Store）として使用してRAGを実装出来る事が分かりました。

今回はLangChainがいろいろ良い感じに処理してくれましたが、Agents for Amazon Bedrockが始まれば同様に良い感じにやってくれる感じでしょうか。

https://aws.amazon.com/jp/blogs/news/preview-enable-foundation-models-to-complete-tasks-with-agents-for-amazon-bedrock/

https://qiita.com/cyberBOSE/items/fe66c2049aaa04568d61