{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Action": "bedrock:*",
            "Resource": "*"
        }
    ]
}

コードを動かす環境の EC2 に GitHub のコードを展開

コードは Python 3.10 で動作します。

ローカルリポジトリにcloneしたら、

pip install -r requirements.txt

でパッケージをインストールします。

venvを使って環境を切り替えられるようにしておくと良いでしょう。

PolarDB for PostgreSQL の準備

クラスタを作成します。

事前に VPC・vSwitch を作成しておきます（途中で新規作成することも可能）。

プロダクト環境で使うときは TDE を有効化します。

クラスタ作成後、プライマリエンドポイント（またはクラスターのエンドポイント）のパブリックネットワークアドレスを申請します。

ここではhmatsu47testという名前で申請しています。

IP ホワイトリストには、EC2 のパブリック IP アドレスを追加指定します。

アカウントを作成します。

まず管理者用の特権アカウント（ここではpgvadmin）、次にアプリケーションからのアクセスに使う標準アカウント（同pgvuser）を作成します。

続いてアプリケーションから使うデータベースを作成します（ここではpgvtest）。

データベース所有者には先ほどの標準アカウントを指定します。

プロダクト環境で使うときは必ず SSL（TLS）での接続を有効化します。

EC2 のコードで`.env`ファイルを書き換える

PolarDB 側の設定に合わせて内容を書き換えます。

.env

PGVECTOR_DRIVER=psycopg2
PGVECTOR_HOST=hmatsu47test.pg.polardb.japan.rds.aliyuncs.com
PGVECTOR_PORT=1921
PGVECTOR_DATABASE=pgvtest
PGVECTOR_USER=pgvuser
PGVECTOR_PASSWORD=【設定した標準アカウントのパスワード】

`psql`コマンドで PolarDB for PostgreSQL に接続し、`vector`拡張機能を有効化する

特権アカウントで接続して作業します。

vector有効化

$ psql -h hmatsu47test.pg.polardb.japan.rds.aliyuncs.com -p 1921 -U pgvadmin -d pgvtest                        
Password for user pgvadmin: 
psql (14.10, server 14.9)
Type "help" for help.

pgvtest=> CREATE EXTENSION vector;
CREATE EXTENSION
pgvtest=> SELECT extversion FROM pg_extension WHERE extname = 'vector';
 extversion 
------------
 0.5.0
(1 row)

pgvtest=> \q

アプリケーションを実行する

アプリケーション実行

$ streamlit run app.py --server.port 8080

Collecting usage statistics. To deactivate, set browser.gatherUsageStats to False.


  You can now view your Streamlit app in your browser.

  Network URL: http://172.31.XX.XX:8080
  External URL: http://35.XX.XX.XX:8080

ブラウザで対象 URL を開きます。

この時点で DB を確認すると、

$ psql -h hmatsu47test.pg.polardb.japan.rds.aliyuncs.com -p 1921 -U pgvuser -d pgvtest
Password for user pgvuser: 
psql (14.10, server 14.9)
Type "help" for help.

pgvtest=> \d
Did not find any relations.

何もありません。

何か文章を入力してみます。

入力した内容がそのままおうむ返しされてきました。

DB を確認します。

pgvtest=> \d
                 List of relations
 Schema |          Name           | Type  |  Owner  
--------+-------------------------+-------+---------
 public | langchain_pg_collection | table | pgvuser
 public | langchain_pg_embedding  | table | pgvuser
(2 rows)

pgvtest=> SELECT * FROM langchain_pg_collection;
     name     | cmetadata |                 uuid                 
--------------+-----------+--------------------------------------
 munou_goroku | null      | 0aa32fd8-2f00-4d4c-b006-29fe541619c8
(1 row)

pgvtest=> \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

pgvtest=> SELECT collection_id, document, cmetadata, custom_id, uuid FROM langchain_pg_embedding;
            collection_id             |   document   | cmetadata |              custom_id               |                 uuid                 
--------------------------------------+--------------+-----------+--------------------------------------+--------------------------------------
 0aa32fd8-2f00-4d4c-b006-29fe541619c8 | こんばんは！ | {}        | 4499a717-a0c9-11ee-bbd5-0adb715db027 | 64682864-63a0-418c-b758-5078ee293ad9
(1 row)

テーブルが 2 つ作成されました。

langchain_pg_embeddingに、入力した文章が 1 行追加されたのがわかります。

embeddings列にベクトル化されたデータが入っていますが、画面には表示しづらいのでここでは表示を省略しました。

さらに別の文章を入力してみます。

先ほど入力された内容が返ってきましたね。

先ほどのテーブルを見てみます。

pgvtest=> SELECT collection_id, document, cmetadata, custom_id, uuid FROM langchain_pg_embedding;
            collection_id             |        document        | cmetadata |              custom_id               |                 uuid                 
--------------------------------------+------------------------+-----------+--------------------------------------+--------------------------------------
 0aa32fd8-2f00-4d4c-b006-29fe541619c8 | こんばんは！           | {}        | 4499a717-a0c9-11ee-bbd5-0adb715db027 | 64682864-63a0-418c-b758-5078ee293ad9
 0aa32fd8-2f00-4d4c-b006-29fe541619c8 | 今日は良い天気ですね！ | {}        | dc61ff4d-a0c9-11ee-aac9-0adb715db027 | 4c581b75-4109-45c7-8929-75b6318f0348
(2 rows)

2 行になりました。

さらに別の文章を入力してみます。

先に 2 つ入力していた文章のうち、意味的に近そうなほうが返ってきました。

テーブルを見てみます。

pgvtest=> SELECT collection_id, document, cmetadata, custom_id, uuid FROM langchain_pg_embedding;
            collection_id             |        document        | cmetadata |              custom_id               |                 uuid                 
--------------------------------------+------------------------+-----------+--------------------------------------+--------------------------------------
 0aa32fd8-2f00-4d4c-b006-29fe541619c8 | こんばんは！           | {}        | 4499a717-a0c9-11ee-bbd5-0adb715db027 | 64682864-63a0-418c-b758-5078ee293ad9
 0aa32fd8-2f00-4d4c-b006-29fe541619c8 | 今日は良い天気ですね！ | {}        | dc61ff4d-a0c9-11ee-aac9-0adb715db027 | 4c581b75-4109-45c7-8929-75b6318f0348
 0aa32fd8-2f00-4d4c-b006-29fe541619c8 | 明日は晴れますかね？   | {}        | 2d34fc70-a0ca-11ee-95c6-0adb715db027 | 2d540208-218c-48f2-9204-a9a9ae0cd2db
(3 rows)

3 行になりました。

見づらいですが、一応ベクトル化されたデータだけ見てみると…

pgvtest=> SELECT embedding FROM langchain_pg_embedding;
（中略）
[1.234375,-0.032470703,-0.28515625,0.22070312,0.25585938,-0.69921875,…
[1.40625,0.765625,-0.104003906,0.10644531,0.296875,-0.26953125,…
[2.34375,0.73828125,-0.22167969,-0.29101562,0.65234375,0.33789062,…
(3 rows)

というような感じで各行 1536 個の配列がベクトルデータとして入っています。

何をしているのか？

Python のコードを見てみます。

13 行目から 20 行目、

CONNECTION_STRING = PGVector.connection_string_from_db_params(
    driver=os.environ.get("PGVECTOR_DRIVER", "psycopg2"),
    host=os.environ.get("PGVECTOR_HOST", "localhost"),
    port=int(os.environ.get("PGVECTOR_PORT", "5432")),
    database=os.environ.get("PGVECTOR_DATABASE", "postgres"),
    user=os.environ.get("PGVECTOR_USER", "postgres"),
    password=os.environ.get("PGVECTOR_PASSWORD", ""),
)

で DB（PolarDB for PostgreSQL）への接続文字列を定義しています。

.envから変数を取得可能な場合はその値が入ります。
なければ 2 つ目のパラメータとして指定した値を（デフォルト値として）使います。

48 行目から 59 行目、

        # ベクターストアから応答を取得
        bedrock_client = boto3.client('bedrock-runtime', region_name="ap-northeast-1")
        embeddings = BedrockEmbeddings(
            client=bedrock_client,
            model_id="amazon.titan-embed-text-v1"
        )
        store = PGVector(
            collection_name=COLLECTION_NAME,
            connection_string=CONNECTION_STRING,
            embedding_function=embeddings,
        )
        docs = store.similarity_search_with_score(trimed_prompt)

で入力文字列に近いレコードをベクトル検索し、

60 行目から 73 行目、

        count = len(docs)
        add_flag = True
        if count == 0:
            # ベクターストアに文章がなければおうむ返しする
            response = trimed_prompt
        else:
            # 近い文章を返す
            response = docs[0][0].page_content
            if trimed_prompt == response:
                # すでに登録済みの文章と同じならベクターストアに登録しない
                add_flag = False
        if add_flag:
            # ベクターストアに入力を追加
            store.add_documents([Document(page_content=trimed_prompt)])

で

検索結果が 0 件なら入力文字列をそのままおうむ返し
近い文書があればそれを返す
- それが入力文字列と違う場合は DB にレコードを追加
  - このとき入力文字列を Titan Embeddings T1 - Text でベクトル化

しています。

LangChain を使っているので SQL 文は完全に隠蔽されています。
実際にどのような SQL 文でベクトル検索できるのか？については、前掲の記事を参照してください。
なお、pgvector を使う場合、デフォルトで<->、つまりユークリッド距離（L2 距離）による最近傍探索が行われるようです（その際、IVFFlat・HNSW インデックスは使いません）。

RAG として使うには

similarity_search_with_score()で返ってきたものをそのまま表示するのではなく、RetrievalQA.from_chain_type()などのパラメータとしてストアのretrieverを渡し、

AWS の Bedrock であれば Anthropic Claude 2.1
OpenAI であれば ChatGPT（3.5/4 など）

などの LLM のプロンプトとして入力文字列（や履歴など）とあわせて送ります。

2024/3/5 追記：
現状は（というかアドベントカレンダーの時点ですでに）RetrievalQA.from_chain_type()などのレガシーな記法ではなく、LCEL の使用が推奨されています。

LangChain の新記法「LangChain Expression Language (LCEL)」入門（oshima_123 さん）

Alibaba Cloud でベクトル検索（with LangChain）

pgvector とは

作って試したもの

AWS 側の準備