Azure PostgreSQL の pgvector と LangChain を用いたベクトル検索

ベクトル検索は、AI アプリケーションにおける類似文書の検索やレコメンデーションシステムなど、さまざまな場面で活用されています。

本記事では、PostgreSQL の拡張機能である pgvector を利用したベクトル検索の手順を解説します。さらに、Azure Database for PostgreSQL フレキシブル サーバーで pgvector を利用する方法についても説明します。この記事ではタイマーによる定期実行を想定し、Azure Functions を利用していますが、もちろん他のアプリケーションからでも問題なく利用できます。

また、LangChain を利用してベクトル検索を行う方法についても触れていきます。

まず pgvector の概要と基本的な使い方を説明し、次に Azure Database for PostgreSQL フレキシブル サーバーでの設定方法、最後に LangChain を使った pgvector を利用したベクトル検索の実装例を紹介します。

開発環境の devcontainer の作成

今回ローカルでの開発に、Visual Studio Code の devcontainer を利用します。

新規ディレクトリを作成し、移動します。

mkdir langchain-pgvector
cd langchain-pgvector

そのディレクトリを Visual Studio Code を開きます。

code .

開発環境として Visual Studio Code の devcontainer を利用します。

まずは、開発環境の構築を行います。

次のようなディレクトリ構成とします。

langchain-pgvector/
├── .devcontainer/       # 開発環境用の devcontainer の設定
│   ├── app/
|   |   └── Dockerfile   # 開発環境用 Dockerfile
│   ├── db/
|   |   └── initdb/      # データベースの初期化用 SQL ファイル
|   |   └── Dockerfile   # 開発環境データベース用 Dockerfile
│   └── devcontainer.json
├── app/                 # アプリケーションコード
└── docker-compose.yml

• .devcontainer/ に開発環境用の devcontainer の設定ファイルを配置します。
• app ディレクトリに Azure Functions のアプリケーションコードを配置します。devcontainer での作業ディレクトリとなります。

開発環境の Dockerfile を作成します。

Python で Azure Functions を開発するので、Microsoft が用意してくれている Azure Functions Core Tools がインストール済みの開発用のイメージを利用します。

開発中に PostgreSQL に接続するため、postgresql-client-17 もインストールしています。

.devcontainer/app/Dockerfile

FROM mcr.microsoft.com/azure-functions/python:4-python3.11-core-tools

WORKDIR /app

RUN apt-get update && \
    apt-get install -y --no-install-recommends libpq-dev wget gnupg lsb-release && \
    wget -qO - https://www.postgresql.org/media/keys/ACCC4CF8.asc | apt-key add - && \
    echo "deb http://apt.postgresql.org/pub/repos/apt/ $(lsb_release -cs)-pgdg main" | tee /etc/apt/sources.list.d/pgdg.list && \
    apt update && \
    apt install -y postgresql-client-17 && \
    rm -rf /var/lib/apt/lists/* && \
    pip install --upgrade pip

# Expose the port
EXPOSE 8080

次は開発環境のデータベースです。

PostgreSQL postgres:17-alpine をベースにして拡張機能の pgvector をインストールします。

次のように Dockerfile を作成します。

.devcontainer/db/Dockerfile

FROM postgres:17-alpine

RUN apk update && apk add --no-cache \
    build-base \
    git \
    postgresql-dev \
    clang \
    llvm19

RUN git clone https://github.com/pgvector/pgvector.git /tmp/pgvector \
 && cd /tmp/pgvector \
 && make \
 && make install \
 && rm -rf /tmp/pgvector

次に、データベースの初期化用の SQL ファイルを作成します。

イメージビルド時に、この SQL ファイルが実行され、拡張機能の pgvector が有効化されます。

.devcontainer/db/initdb/01-create-extension.sql

CREATE EXTENSION IF NOT EXISTS vector;

docker-compose.yml にて、開発用のコンテナとデータベース用のコンテナを定義します。

作業ディレクトリを /app に設定します。

.devcontainer/app/docker-compose.yml

version: "3.8"
services:
  db:
    build: ./db
    container_name: postgres
    ports:
      - 5432:5432
    volumes:
      - pgdata:/var/lib/postgresql/data
      - ./db/initdb:/docker-entrypoint-initdb.d
    environment:
      POSTGRES_USER: postgres
      POSTGRES_PASSWORD: password
      POSTGRES_INITDB_ARGS: "--encoding=UTF-8"
      POSTGRES_DB: db
    hostname: postgres
    restart: always
    user: root
    healthcheck:
      test: ["CMD-SHELL", "pg_isready -U postgres -d db"]
      interval: 10s
      timeout: 5s
      retries: 5

  app:
    build:
      context: .
      dockerfile: ./app/Dockerfile # 開発環境用の Dockerfile
    depends_on:
      db:
        condition: service_healthy
    environment:
      AzureWebJobsStorage: "UseDevelopmentStorage=true" # Required for local development
      FUNCTIONS_WORKER_RUNTIME: python
      DATABASE_URL: postgresql://postgres:password@db:5432/db
    working_dir: /app
    volumes:
      - ../app:/app
    # Keep the container running even if the command exits
    tty: true
    stdin_open: true
    ports:
      - "8080:8080" # Expose the port

volumes:
  pgdata:

開発環境用の devcontainer の設定ファイルを作成します。

.devcontainer/app/devcontainer.json

{
    "name": "Langchain pgvector",
    "dockerComposeFile": ["docker-compose.yml"],
    "service": "app",
    "workspaceFolder": "/app",
    "shutdownAction": "stopCompose",
    "customizations": {
        "vscode": {
            "extensions": [
                "ms-python.python",
                "ms-python.vscode-pylance",
                "ms-azuretools.vscode-docker",
                "ms-azuretools.vscode-azurefunctions"
            ]
        }
    },
    "postCreateCommand": "echo 'Development container started'",
    "remoteUser": "root",
    "features": {
        "ghcr.io/devcontainers/features/common-utils:2": {
            "installZsh": true,
            "username": "root",
            "userUid": "0",
            "userGid": "0",
            "upgradePackages": true
        },
        "ghcr.io/devcontainers/features/docker-in-docker:2": {
            "moby": true,
            "azureDnsAutoDetection": true,
            "installDockerBuildx": true,
            "version": "latest",
            "dockerDashComposeVersion": "v2"
        }
    }
}

次のようにファイルが作成されます。

開発環境用の devcontainer の設定ファイルの作成ができましたので、Visual Studio Code で devcontainer を開きます。

Ctrl + Shift + P でコマンドパレットを開き、 Dev Containers: Reopen in Container を選択します。

コンテナーのビルドが完了すると、次のように表示されます。

Development container started

簡単な Python コードを実行し、開発環境が正しく動作していることを確認します。

app/app.py

print("Hello, World!")

Visual Studio Code 内の Terminal で、Python コードを実行します。

python app.py

次のように表示されれば成功です。

Hello, World!

app.py は削除しておきます。

次に、データベースに接続できることと、拡張機能の pgvector がインストールされていることを確認します。

psql -h db -U postgres -d db

インストールされている拡張機能の一覧を表示します。

\dx

次のように、 vector が表示されればインストールされています。

                             List of installed extensions
  Name   | Version |   Schema   |                     Description
---------+---------+------------+------------------------------------------------------
 plpgsql | 1.0     | pg_catalog | PL/pgSQL procedural language
 vector  | 0.8.0   | public     | vector data type and ivfflat and hnsw access methods
(2 rows)

Function の作成

引き続き、devcontainer 内で作業を行います。

Visual Studio Codeにて、コマンドパレットで、Azure Functions: Create Function... を選択します。

指示に従って、ディレクトリに app、言語に Python、プログラミングモデルに Model V2 (Recommended) を選択して作成します。

コンテナー内に作成しているため venv は不要なので Skip virtual environment を選択します。

最初のトリガーを作成するかも聞かれるため、テストのため HTTP トリガーの Function を http_trigger という名前で作成します。

自動的に git リポジトリーが作成されるため、不要な場合 .git/ ディレクトリを削除します。

rm -rf .git/
rm .gitignore

以下のファイルが作成されます。

├── function_app.py
├── host.json
├── local.settings.json
└── requirements.txt

function_app.py は次のようになっています。

function_app.py

import azure.functions as func
import logging

app = func.FunctionApp(http_auth_level=func.AuthLevel.FUNCTION)

@app.route(route="http_trigger")
def http_trigger(req: func.HttpRequest) -> func.HttpResponse:
    logging.info('Python HTTP trigger function processed a request.')

    name = req.params.get('name')
    if not name:
        try:
            req_body = req.get_json()
        except ValueError:
            pass
        else:
            name = req_body.get('name')

    if name:
        return func.HttpResponse(f"Hello, {name}. This HTTP triggered function executed successfully.")
    else:
        return func.HttpResponse(
             "This HTTP triggered function executed successfully. Pass a name in the query string or in the request body for a personalized response.",
             status_code=200
        )

http_trigger Function が実行できるか確認します。

次のコマンドで Function App を起動します。

func start --port 8080

ホストのブラウザーで、http://localhost:8080/api/http_trigger?name=test にアクセスします。

次のように表示されれば成功です。

Hello, test. This HTTP triggered function executed successfully.

pgvector へのドキュメントの追加と検索

devcontainer 環境で、LangChain を利用して、pgvector にドキュメントを追加し、検索するコードを記述します。

PostgreSQL の接続文字列などを .env ファイルから取得できるようにするため、python-dotenv をインストールします。

pip install python-dotenv

.env ファイルを作成し、接続情報を記述します。

.env

DATABASE_URL="postgresql://postgres:password@localhost:5432/db"

pip install psycopg2-binary langchain langchain_postgres openai langchain-openai

インストールした pip パッケージを requirements.txt に保存します。

pip freeze > requirements.txt

また、後述するデプロイのために .funcignore ファイルに .env 追記し、ローカル環境の .env ファイルがデプロイされないようにします。

.funcignore

+ .env

langchain_pgvector.py を作成し、LangChain を使って、pgvector にドキュメントを追加し、検索するコードを記述します。

LangChain の PGVector は、最初に add_documents メソッドでドキュメントを追加した際に、PostgreSQL データベースに、ドキュメントのベクトルを保存する用のテーブルを作成します。

ここでは、Azure OpenAI を利用しています。各自環境に合わせて変更および .env ファイルの設定を行ってください。

langchain_pgvector.py

import datetime
import os
from logging import INFO, basicConfig, getLogger

from dotenv import load_dotenv
from langchain_core.documents import Document

from langchain_openai import AzureOpenAIEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_postgres.vectorstores import PGVector

load_dotenv()
basicConfig(level=INFO)
logger = getLogger(__name__)


CONNECTION_STRING = os.getenv("DATABASE_URL")

def add_and_search_documents():

    # Azure OpenAI Embedding
    embeddings = AzureOpenAIEmbeddings(
        azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"),
        azure_deployment=os.getenv("AZURE_OPENAI_DEPLOYMENT_EMBEDDINGS_NAME"),
        api_version="2024-06-01",
    )

    # PGVector インスタンスを作成
    # ここでは、"my_collection" という名前のコレクションにドキュメントを追加し
    # 検索します。
    vector_store = PGVector(
        embeddings=embeddings,
        collection_name="my_collection",
        connection=CONNECTION_STRING,
        use_jsonb=True,
    )

    # 動作確認用検索対象のドキュメント
    docs = [
        Document(
            page_content="I went to the store and bought some apples.",
            metadata={"source": "doc1.txt", "timestamp": datetime.datetime.now().timestamp()},
        ),
        Document(
            page_content="You want to go to London, but I want to go to Paris.",
            metadata={"author": "doc2.txt", "timestamp": datetime.datetime.now().timestamp()},
        ),
    ]

    # テキストを分割するための splitter を作成
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=500,
        chunk_overlap=50,
        separators=["\n\n", "\n", ".", "。", " "],
    )

    # ドキュメントを分割 500 文字ごとのチャンクに分割
    splited_docs = splitter.split_documents(docs)

    vector_store.add_documents(splited_docs)

    logger.info("Documents added successfully.")

    # 類似性検索
    result_docs = vector_store.similarity_search(
        query="London",
        k=1,
    )

    logger.info(f"Similarity search result length: {len(result_docs)}")

    for doc in result_docs:
        logger.info(f"Similarity search result: {doc}")

    return result_docs

http_trigger Function から add_and_search_documents を呼び出し、検索結果を返します。

function_app.py

import azure.functions as func
import logging

from langchain_pgvector import add_and_search_documents

app = func.FunctionApp(http_auth_level=func.AuthLevel.FUNCTION)

@app.route(route="http_trigger")
def http_trigger(req: func.HttpRequest) -> func.HttpResponse:
    logging.info('Python HTTP trigger function processed a request.')

    result_docs = add_and_search_documents()

    # 検索結果を返す
    return func.HttpResponse(
        f"Search result: {result_docs[0].page_content if result_docs else 'No results found'}",
        status_code=200
    )

次のコマンドで Function App を起動します。

func start --port 8080

ホストのブラウザーで、http://localhost:8080/api/http_trigger にアクセスします。

次のような検索結果に2つめのドキュメントが表示されるはずです。

Search result: You want to go to London, but I want to go to Paris.

テーブルが作成され、データが追加されることを確認します。

psql -h db -U postgres -d db

\dt コマンドでテーブルの一覧を表示します。

\dt

次のように、langchain_pg_collection と langchain_pg_embedding テーブルが表示されています。

langchain_pg_collection テーブルには、コレクションの情報が保存され、langchain_pg_embedding テーブルには、ドキュメントのベクトルが保存されています。

                  List of relations
 Schema |          Name           | Type  |  Owner
--------+-------------------------+-------+----------
 public | langchain_pg_collection | table | postgres
 public | langchain_pg_embedding  | table | postgres

langchain_pg_collection テーブルの内容を確認します。

SELECT * FROM langchain_pg_collection;

次のように、先ほど add_documents を呼び出した際に作成されたコレクション my_collection が表示されます。

                 uuid                 |     name      | cmetadata
--------------------------------------+---------------+-----------
 898b418e-375c-4622-b1a4-bcae75feddc7 | my_collection | null

langchain_pg_embedding テーブルの定義を確認します。

\d langchain_pg_embedding

それぞれのカラムには、次のようにデータが保存されます。

• collection_id: 指定したコレクションの uuid
• embedding: ドキュメントのベクトル
• document: ドキュメントのテキスト
• cmetadata: ドキュメントのメタデータ

               Table "public.langchain_pg_embedding"
    Column     |       Type        | Collation | Nullable | Default
---------------+-------------------+-----------+----------+---------
 id            | character varying |           | not null |
 collection_id | uuid              |           |          |
 embedding     | vector            |           |          |
 document      | character varying |           |          |
 cmetadata     | jsonb             |           |          |
Indexes:
    "langchain_pg_embedding_pkey" PRIMARY KEY, btree (id)
    "ix_cmetadata_gin" gin (cmetadata jsonb_path_ops)
    "ix_langchain_pg_embedding_id" UNIQUE, btree (id)
Foreign-key constraints:
    "langchain_pg_embedding_collection_id_fkey" FOREIGN KEY (collection_id) REFERENCES langchain_pg_collection(uuid) ON DELETE CASCADE

Azure 環境デプロイ用の設定

Azure にデプロイするための設定を行います。devcontainer から抜けて、ローカルファイルで Visual Studio Code を開き直します。

Visual Studio Code で、コマンドパレットで、Dev Containers: Reopen Folder Locally を選択し、開き直します。

今回は、Container Apps 環境ホスティングの Azure Functions を利用するため、そのための Dockerfile を作成します。

開発環境用の Dockerfile には、本番環境では不要なツールが含まれているため、新たに本番環境用の Dockerfile を作成します。

app/Dockerfile

FROM mcr.microsoft.com/azure-functions/python:4-python3.12

ENV AzureWebJobsScriptRoot=/home/site/wwwroot \
    AzureFunctionsJobHost__Logging__Console__IsEnabled=true

COPY requirements.txt /
RUN apt-get update && \
    apt-get install -y libpq-dev && \
    pip install -r /requirements.txt

COPY . /home/site/wwwroot

Azure リソースの作成

Azure Functions、Azure Container Registry の作成

Azure ポータルにログインし、Azure Functions と Azure Container Registry を作成します。

前述の通り、Azure Functions は、Container Apps 環境ホスティングを利用します。

以前 Azure Functions の Container Apps 環境のホスティングの設定について記事にしているので参考にしてください。

PostgreSQL サーバーの作成

Azure Database for PostgreSQL フレキシブル サーバー では、pgvector を利用可能です。

Azure Database for PostgreSQL フレキシブル サーバー を作成します。

pgvector 拡張機能の有効化

Azure Database for PostgreSQL フレキシブル サーバー の Server parameters > Extensions から、VECTOR を選択し、保存します。これで pgvector 拡張機能が有効化されます。

Azure 環境にデプロイ

次のコマンドで Azure Container Registry にイメージを登録します。

az acr build --registry <CONTAINER_REGISTRY_NAME> --image langchain-pgvector:v0.0.1 app/.

そして、Function App でこのイメージを使用するように構成することで、Azure Functions にデプロイできます。

同じく以前の記事にて Azure Functions へのデプロイについて記事にしているので参考にしてください。

Azure Functions の実行

Azure Functions がデプロイされたら、Azure Functions の URL にアクセスし、ローカルでの確認時と同じように、検索結果が表示されることを確認します。

これで、Azure PostgreSQL の pgvector と Azure Functions を利用して、ベクトル検索を行う環境が構築の完成です。