データアナリティクス事業本部 機械学習チームの鈴木です。
先日Amazon Aurora PostgreSQLでpgvectorがサポートされたアナウンスがありました。
Amazon RDS for PostgreSQLでは一足先にサポートされていましたが、Auroraでも使えるようになったということで、セットアップしたAurora Serverless v2とLangChainで接続を試してみました。
本記事の内容
以下のLangChainのPGVectorモジュールのガイドを参考に、VPC内に作成したAurora Serverless v2のサーバーレスDBインスタンスへ接続できるかを試してみました。
サーバーレスDBインスタンスは、検証用のVPCを作成し、パブリックIPを有効化して起動しておきました。詳細は後ほど記載します。
pgvectorについては、Amazon RDS for PostgreSQLでの検証を通して以下の記事が非常に分かりやすく記載頂いているのであわせてご確認ください。
Aurora Serverless v2の準備
0. ポイント
今回は、以下の記事を参考にパブリックIPを付与したインスタンスを作成しました。
VPC外からのアクセスで特に注意が必要だった設定は以下になります。
- データベース作成時に
パブリックアクセス
をあり
にする。 - VPCにインターネットゲートウェイを作成しておく。
- Auroraが使うセキュリティグループで、クライアント側からポート番号
5432
でアクセスできるようにIngressの設定をしておく。
1. ネットワーク作成
まず以下のCloudFormationテンプレートでVPCなどのネットワーク用のリソースを作成しました。
特に、ポイントに挙げたように、AuroraSGIngressCidrIp
でクライアントのIPアドレスを指定し、データベースにアクセスできるようにしました。
AWSTemplateFormatVersion: "2010-09-09"
Parameters:
EnvironmentName:
Type: String
VPCCIDR:
Type: String
Default: 10.192.0.0/16
PublicSubnetCIDR1:
Type: String
Default: 10.192.0.0/24
PublicSubnetCIDR2:
Type: String
Default: 10.192.1.0/24
PrivateSubnetCIDR1:
Type: String
Default: 10.192.3.0/24
PrivateSubnetCIDR2:
Type: String
Default: 10.192.4.0/24
AuroraSGIngressCidrIp:
Type: String
Resources:
VPC:
Type: AWS::EC2::VPC
Properties:
CidrBlock: !Ref VPCCIDR
EnableDnsSupport: true
EnableDnsHostnames: true
Tags:
- Key: Name
Value: !Sub ${EnvironmentName}-VPC
# InternetGateway
InternetGateway:
Type: AWS::EC2::InternetGateway
Properties:
Tags:
- Key: Name
Value: !Sub ${EnvironmentName}-igw
AttachGateway:
Type: AWS::EC2::VPCGatewayAttachment
Properties:
VpcId: !Ref VPC
InternetGatewayId: !Ref InternetGateway
# Public Subnetのネットワーク設定
PublicSubnet1:
Type: AWS::EC2::Subnet
Properties:
AvailabilityZone: !Select [0, !GetAZs ""]
CidrBlock: !Ref PublicSubnetCIDR1
VpcId: !Ref VPC
Tags:
- Key: Name
Value: !Sub ${EnvironmentName}-PublicSubnet
PublicRouteTable1:
Type: AWS::EC2::RouteTable
Properties:
VpcId: !Ref VPC
Tags:
- Key: Name
Value: !Sub ${EnvironmentName}-PublicRouteTable1
PublicRoute1:
Type: AWS::EC2::Route
Properties:
RouteTableId: !Ref PublicRouteTable1
DestinationCidrBlock: 0.0.0.0/0
GatewayId: !Ref InternetGateway
PublicSubnetRouteTableAssociation1:
Type: AWS::EC2::SubnetRouteTableAssociation
Properties:
SubnetId: !Ref PublicSubnet1
RouteTableId: !Ref PublicRouteTable1
PublicSubnet2:
Type: AWS::EC2::Subnet
Properties:
AvailabilityZone: !Select [1, !GetAZs ""]
CidrBlock: !Ref PublicSubnetCIDR2
VpcId: !Ref VPC
Tags:
- Key: Name
Value: !Sub ${EnvironmentName}-PublicSubnet
PublicRouteTable2:
Type: AWS::EC2::RouteTable
Properties:
VpcId: !Ref VPC
Tags:
- Key: Name
Value: !Sub ${EnvironmentName}-PublicRouteTable2
PublicRoute2:
Type: AWS::EC2::Route
Properties:
RouteTableId: !Ref PublicRouteTable2
DestinationCidrBlock: 0.0.0.0/0
GatewayId: !Ref InternetGateway
PublicSubnetRouteTableAssociation2:
Type: AWS::EC2::SubnetRouteTableAssociation
Properties:
SubnetId: !Ref PublicSubnet2
RouteTableId: !Ref PublicRouteTable2
# Private Subnetのネットワーク設定
PrivateSubnet1:
Type: AWS::EC2::Subnet
Properties:
AvailabilityZone: !Select [0, !GetAZs ""]
CidrBlock: !Ref PrivateSubnetCIDR1
VpcId: !Ref VPC
Tags:
- Key: Name
Value: !Sub ${EnvironmentName}-PrivateSubnet1
PrivateRouteTable1:
Type: AWS::EC2::RouteTable
Properties:
VpcId: !Ref VPC
Tags:
- Key: Name
Value: !Sub ${EnvironmentName}-PrivateRouteTable1
PrivateSubnetRouteTableAssociation1:
Type: AWS::EC2::SubnetRouteTableAssociation
Properties:
SubnetId: !Ref PrivateSubnet1
RouteTableId: !Ref PrivateRouteTable1
PrivateSubnet2:
Type: AWS::EC2::Subnet
Properties:
AvailabilityZone: !Select [1, !GetAZs ""]
CidrBlock: !Ref PrivateSubnetCIDR2
VpcId: !Ref VPC
Tags:
- Key: Name
Value: !Sub ${EnvironmentName}-PrivateSubnet2
PrivateRouteTable2:
Type: AWS::EC2::RouteTable
Properties:
VpcId: !Ref VPC
Tags:
- Key: Name
Value: !Sub ${EnvironmentName}-PrivateRouteTable2
PrivateSubnetRouteTableAssociation2:
Type: AWS::EC2::SubnetRouteTableAssociation
Properties:
SubnetId: !Ref PrivateSubnet2
RouteTableId: !Ref PrivateRouteTable2
AuroraSecurityGroup:
Type: AWS::EC2::SecurityGroup
Properties:
GroupDescription: AuroraSecurityGroup
VpcId: !Ref VPC
Tags:
- Key: Name
Value: !Sub ${EnvironmentName}-AuroraSecurityGroup
SecurityGroupIngress:
- IpProtocol: tcp
FromPort: 5432
ToPort: 5432
CidrIp: !Ref AuroraSGIngressCidrIp
2. サーバーレスDBインスタンス作成
続いて、VPCにサーバーレスDBインスタンスを作成しました。
RDSのコンソールからデータベースの作成
で、Aurora(postgreSQL Compatible)
のAurora PostgreSQL(Compatible with PostgreSQL 15.3)
を選択しました。このバージョンは冒頭のアナウンスでpgvectorがサポートされているバージョンとなります。
同じ画面で以下のように設定しました。
VPCは先に作成したものを選択しました。パブリックアクセスはここであり
を選択しました。セキュリティグループも作成済みのものを選択しました。
モニタリングのみ、Performaice Insightsが不要だったのでチェックを外し、ほかはデフォルトでデータベースの作成
を押してインスタンスが利用可能になるのを待ちました。
インスタンス作成時に接続情報が得られるのでメモしておきました。
3. SQLクライアントからの接続確認
接続情報を使って、SQLクライアントからインスタンスに接続できることを確認しました。
今回はDBeaverを使いました。
このとき、DBeaverのある端末からのアクセスをセキュリティグループで許可している必要があるのでご注意ください。
4. pgvectorの有効化
SQLクライアントから以下の処理を実行して、pgvectorを有効化しました。
CREATE EXTENSION vector;
LangChainからのアクセス
0. クライアント実行環境
実行環境はGoogle Colaboratoryを使いました。ハードウェアアクセラレータ無し、ランタイム仕様は標準としました。
Pythonのバージョンは以下でした。
!python --version
# Python 3.10.6
また、ライブラリは以下のようにインストールしました。
# pgvectorとの接続に利用する
!pip install psycopg2-binary
!pip install pgvector
!pip install tiktoken
# LangChainで利用する
!pip install python-dotenv
!pip install langchain
!pip install openai
インストールされたライブラリのバージョンは以下でした。
!pip freeze | grep -e "openai" -e "langchain" -e "tiktoken" -e "SQLAlchemy" -e "psycopg2" -e "psycopg2-binary" -e "pgvector"
# langchain==0.0.238
# openai==0.27.8
# pgvector==0.1.8
# psycopg2==2.9.6
# psycopg2-binary==2.9.6
# SQLAlchemy==2.0.18
# tiktoken==0.4.0
1. 環境変数の設定
まず、必要な値を.envファイルに書き込みます。
!echo 'OPENAI_API_KEY="<トークン>"' >.env
!echo 'PGVECTOR_DRIVER="psycopg2"' >>.env
!echo 'PGVECTOR_HOST="<ホストのURL>"' >>.env
!echo 'PGVECTOR_PORT="5432"' >>.env
!echo 'PGVECTOR_DATABASE="<データベース名>"' >>.env
!echo 'PGVECTOR_USER="<データベースユーザー名>"' >>.env
!echo 'PGVECTOR_PASSWORD="<データベースパスワード>"' >>.env
load_dotenv()
で環境変数を読み込みました。
from dotenv import load_dotenv
load_dotenv()
2. ライブラリの読み込み
ガイドにしたがって、ライブラリを読み込みました。
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.text_splitter import CharacterTextSplitter
from langchain.vectorstores.pgvector import PGVector
from langchain.document_loaders import TextLoader
from langchain.docstore.document import Document
3. テキストの準備
特になにも思いつかなかったので、元気がでそうな文面を用意してみました。
!mkdir ./data
!echo "元気な時はいえーいといいましょう。" > ./data/sample_doc.txt
以下のように読み込んでおきました。
loader = TextLoader("./data/sample_doc.txt")
documents = loader.load()
text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
docs = text_splitter.split_documents(documents)
4. 接続情報の作成
CONNECTION_STRING
とCOLLECTION_NAME
を作成しました。前者は、connection_string_from_db_params
メソッドを使って作成しました。
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", "postgres"),
)
COLLECTION_NAME = "state_of_the_union_test"
connection_string_from_db_params
メソッドで作成されるURLについては、例えば、LangChainのレポジトリより実装を確認すると、SQLAlchemyのcreate_engine
メソッドでこの情報を使用していました。SQLAlchemyドキュメントによると、RFC-1738に沿ったURLを接続情報として使っている仕様とのことです。
5. データベースへの接続作成とクエリの実行
ガイドを参考に、データベースへ接続し、クエリを実行してみました。
まず接続です。ここはPythonを実行しているクライアントからインスタンスへの接続をセキュリティグループで許可しないと処理が完了しなかったためご注意ください。
embeddings = OpenAIEmbeddings()
db = PGVector.from_documents(
embedding=embeddings,
documents=docs,
collection_name=COLLECTION_NAME,
connection_string=CONNECTION_STRING,
)
続いてクエリを実行してみました。
query = "元気な時はなんと言ったらいいですか"
docs_with_score = db.similarity_search_with_score(query)
for doc, score in docs_with_score:
print("-" * 80)
print("Score: ", score)
print(doc.page_content)
print("-" * 80)
以下のような結果となりました。
--------------------------------------------------------------------------------
Score: 0.054825832050864776
元気な時はいえーいといいましょう。
--------------------------------------------------------------------------------
期待どおり接続できていそうです。
6. Auroraのテーブルの確認
ここまでで、どのようなテーブルが作成されているのか、SQLクライアントからAuroraに接続して確認してみました。
テーブルを確認すると、langchain_pg_collection
とlangchain_pg_embedding
が作成されていました。
記事執筆時点では、langchain/vectorstores/pgembedding.pyに定義されたクラスにてそれぞれのテーブルの作成が実装されていたので、気になる方はご確認ください。
langchain_pg_collection
テーブルは、名前の通りcollectionに関する情報を保存していました。
langchain_pg_embedding
テーブルは、テキストやそのベクトルなどを保存していました。ベクトルはcollection_id
カラムでどのcollectionに結びついているかを保持しているようです。
最後に
pgvectorを有効化したAurora Serverless v2のサーバーレスDBインスタンスを作成し、AWS外のリソースからLangChainで接続するところを試してみました。今回はAWS外からの利用でしたが、Auroraを使うことによって、AWS内で構成を閉じつつ、とても気軽にベクターストアを利用することができます。Aurora Serverless v2を使うことで、サーバーレスのメリットを受けることもできます。
特にAWSにてベクトルデータベースを用意したい方には、Amazon RDS for PostgreSQLと合わせて有力な選択肢になると思いますので、参考になりましたら幸いです。