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 を利用します。

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

```bash
mkdir langchain-pgvector
cd langchain-pgvector
```

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

```bash
code .
```

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

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

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

```txt
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` もインストールしています。

```Dockerfile:.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` を作成します。

```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 が有効化されます。

```sql:.devcontainer/db/initdb/01-create-extension.sql
CREATE EXTENSION IF NOT EXISTS vector;
```

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

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

```yml:.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 の設定ファイルを作成します。

```json:.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"
        }
    }
}

```

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

![langchain-pgvecotr-devcon-tree.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/231348/3ef89b4d-5504-b883-0927-45b820abf3f4.png)


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

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

![reopne_in_container.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/231348/93fcc106-a44b-3710-0ced-09fe89414f22.png)


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

```txt
Development container started
```

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


```python:app/app.py
print("Hello, World!")
```

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

```bash
python app.py
```

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

```bash
Hello, World!
```

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

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

```bash
psql -h db -U postgres -d db
```

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

```sql
\dx
```

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

```txt
                             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...` を選択します。

![create-funciton.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/231348/fc5895fd-e3b0-3e0b-2bd5-a0f2a82781e2.png)


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

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

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

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

```bash
rm -rf .git/
rm .gitignore
```

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

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

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

```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 を起動します。

```bash
func start --port 8080
```

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

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

```txt
Hello, test. This HTTP triggered function executed successfully.
```

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

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

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

```bash
pip install python-dotenv
```

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

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


```bash
pip install psycopg2-binary langchain langchain_postgres openai langchain-openai
```

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

```bash
pip freeze > requirements.txt
```

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

```diff_txt:.funcignore
+ .env
```


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

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

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


```py: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` を呼び出し、検索結果を返します。


```py: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 を起動します。

```bash
func start --port 8080
```

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

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

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


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

```bash
psql -h db -U postgres -d db
```

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

```bash
\dt
```

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

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


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

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

```sql
SELECT * FROM langchain_pg_collection;
```

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


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

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

```sql
\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` を選択し、開き直します。

![reopen-folder-locally.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/231348/6dc215ba-d433-fcce-019b-0750d79083e8.png)



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

:::note info
開発環境用の 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 環境ホスティングを利用します。


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

https://qiita.com/rysh25/items/0d969eb00327c8b0f196

### PostgreSQL サーバーの作成

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

https://learn.microsoft.com/ja-jp/azure/postgresql/flexible-server/how-to-use-pgvector


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


### pgvector 拡張機能の有効化

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

![azuredatabase-postgresql-extension-pgvector.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/231348/3426474a-df06-4158-e984-8cd9c811a08b.png)



## Azure 環境にデプロイ

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

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

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

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

https://qiita.com/rysh25/items/0d969eb00327c8b0f196


## Azure Functions の実行

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


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