pgvectorを有効化したAurora Serverless v2のサーバーレスDBインスタンスにLangChainから接続してみた

データアナリティクス事業本部 機械学習チームの鈴木です。

先日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と合わせて有力な選択肢になると思いますので、参考になりましたら幸いです。

そのた参考資料

• Leverage pgvector and Amazon Aurora PostgreSQL for Natural Language Processing, Chatbots and Sentiment Analysis | AWS Database Blog
• pgvector/pgvector: Open-source vector similarity search for Postgres
• 1分でわかるAmazon Aurora Serverless v2入門 | DevelopersIO
• ベクトルデータベースとは ベクトルデータベースの説明 - AWS