Zenn
Open33

fly postgres に pgvector を入れてセマンティックサーチ: まだできず

Prisma
Remix
PostgreSQL
Fly.io
ChatGPT
Coji MizoguchiCoji Mizoguchi

flyio で postgres を中心に apps の Scale-To-Zero ができたので、今度はこの postgres に pgvector を入れて、text-embedding-ada-002 でベクトル化した文書を格納して、SQL でセマンティックサーチをできるようにしたい。

Coji MizoguchiCoji Mizoguchi

Dockerfile。元記事からコピペしてバージョン番号最新にしただけ。

FROM flyio/postgres-flex:15.3

# Install build dependencies
RUN apt-get update && \
    apt-get install -y --no-install-recommends \
        build-essential \
        curl \
        postgresql-server-dev-all

# Set the pgvector version
ARG PGVECTOR_VERSION=0.4.4

# Download and extract the pgvector release, build the extension, and install it
RUN curl -L -o pgvector.tar.gz "https://github.com/ankane/pgvector/archive/refs/tags/v${PGVECTOR_VERSION}.tar.gz" && \
    tar -xzf pgvector.tar.gz && \
    cd "pgvector-${PGVECTOR_VERSION}" && \
    make && \
    make install

# Clean up build dependencies and temporary files
RUN apt-get remove -y build-essential curl postgresql-server-dev-all && \
    apt-get autoremove -y && \
    apt-get clean && \
    rm -rf /var/lib/apt/lists/* && \
    rm -rf /pgvector.tar.gz /pgvector-${PGVECTOR_VERSION}
Coji MizoguchiCoji Mizoguchi

push 長いので待ってる間に作成済みの postgres インスタンスのイメージ更新方法をしらべる

Coji MizoguchiCoji Mizoguchi

fly image のリファレンス によると --image string で独自のイメージ名指定できるみたい。

flyctl image update [flags]

  -a, --app string           Application name
  -c, --config string        Path to application configuration file
      --detach               Return immediately instead of monitoring update progress. (Nomad only)
  -h, --help                 help for update
      --image string         Target a specific image. (Machines only)
      --skip-health-checks   Skip waiting for health checks inbetween VM updates. (Machines only)
      --strategy string      Deployment strategy. (Nomad only)
  -y, --yes                  Accept all confirmations
Coji MizoguchiCoji Mizoguchi

お、いけそう感。

$ fly image update -a coji-db --image coji/fly-pg-postgres
The following changes will be applied to all Postgres machines.
Machines not running the official Postgres image will be skipped.

        ... // 85 identical lines
            }
          },
-         "image": "flyio/postgres-flex:15.3@sha256:c380a6108f9f49609d64e5e83a3117397ca3b5c3202d0bf0996883ec3dbb80c8",
+         "image": "coji/fly-pg-postgres",
          "restart": {
            "policy": "on-failure",
        ... // 8 identical lines
  
? Apply changes? (y/N)
Coji MizoguchiCoji Mizoguchi

おっと

? Apply changes? Yes
Identifying cluster role(s)
  Machine 17811616a52548: primary
Updating machine 17811616a52548
Error: could not update machine 17811616a52548: failed to update VM 17811616a52548: Authentication required to access image "docker.io/coji/fly-pg-postgres:latest"
Coji MizoguchiCoji Mizoguchi

普通にイメージ名まちがえてただけでした。

✗ $ fly image update -a coji-db --image coji/fly-pg-postgres
◯ $ fly image update -a coji-db --image coji/fly-pg-pgvector

Coji MizoguchiCoji Mizoguchi
Identifying cluster role(s)
  Machine 17811616a52548: primary
Updating machine 17811616a52548
  Waiting for 17811616a52548 to become healthy (started, 3/3)
Machine 17811616a52548 updated successfully!
Postgres cluster has been successfully updated!
Coji MizoguchiCoji Mizoguchi

TablePlus で接続して試してみる。

まず proxy でポートフォワード開けておいて

$ fly proxy 5432:5432 -a <postgres-app-name>

TablePlus で接続します。
fly pg attach でアタッチしたときにできた DATABASE_URL のユーザ名、パスワード、データベースを使って localhost:5432 に接続。

Coji MizoguchiCoji Mizoguchi

pgvector の Getting Startedをそのままやる

拡張を有効に。

CREATE EXTENSION vector;

Query 1 OK: CREATE EXTENSION

ベクトル付きのテーブル作成

CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3));

Query 1 OK: CREATE TABLE

ベクトルのレコードを２件INSERT

INSERT INTO items (embedding) VALUES ('[1,2,3]'), ('[4,5,6]');

Query 1 OK: INSERT 0 2, 2 rows affected

近い順に検索

SELECT * FROM items ORDER BY embedding <-> '[3,1,2]' LIMIT 5;

id	embedding
1	[1,2,3]
2	[4,5,6]
Coji MizoguchiCoji Mizoguchi

これでまず select しようとすると

vector-test.tsx
import { Box } from '@chakra-ui/react'
import { json } from '@remix-run/node'
import { useLoaderData } from '@remix-run/react'
import { prisma } from '~/services/database.server'

export const loader = async () => {
  const list = await prisma.$queryRaw`SELECT * FROM items`
  return json({ list })
}

export default function VectorTestPage() {
  const { list } = useLoaderData<typeof loader>()
  return <Box>{JSON.stringify(list)}</Box>
}

エラー

Error: 
Invalid `prisma.$queryRaw()` invocation:


Raw query failed. Code: `N/A`. Message: `Failed to deserialize column of type 'vector'. If you're using $queryRaw and this column is explicitly marked as `Unsupported` in your Prisma schema, try casting this column to any supported Prisma type such as `String`.`
    at An.handleRequestError (/Users/coji/progs/nickname-gpt/node_modules/.pnpm/@prisma+client@4.15.0_prisma@4.15.0/node_modules/@prisma/client/runtime/library.js:174:6929)
    at An.handleAndLogRequestError (/Users/coji/progs/nickname-gpt/node_modules/.pnpm/@prisma+client@4.15.0_prisma@4.15.0/node_modules/@prisma/client/runtime/library.js:174:6358)
    at An.request (/Users/coji/progs/nickname-gpt/node_modules/.pnpm/@prisma+client@4.15.0_prisma@4.15.0/node_modules/@prisma/client/runtime/library.js:174:6237)
    at loader5 (/Users/coji/progs/nickname-gpt/app/routes/vector-test.tsx:7:16)
    at Object.callRouteLoaderRR (/Users/coji/progs/nickname-gpt/node_modules/.pnpm/@remix-run+server-runtime@1.17.0/node_modules/@remix-run/server-runtime/dist/data.js:52:16)
    at callLoaderOrAction (/Users/coji/progs/nickname-gpt/node_modules/.pnpm/@remix-run+router@1.6.3/node_modules/@remix-run/router/router.ts:3568:16)
    at async Promise.all (index 1)
    at loadRouteData (/Users/coji/progs/nickname-gpt/node_modules/.pnpm/@remix-run+router@1.6.3/node_modules/@remix-run/router/router.ts:3001:19)
    at queryImpl (/Users/coji/progs/nickname-gpt/node_modules/.pnpm/@remix-run+router@1.6.3/node_modules/@remix-run/router/router.ts:2780:20)
    at Object.query (/Users/coji/progs/nickname-gpt/node_modules/.pnpm/@remix-run+router@1.6.3/node_modules/@remix-run/router/router.ts:2656:18)
Coji MizoguchiCoji Mizoguchi

prisma の github issues に pgvector サポートのスレがあったのだけど、良き方法が。
https://github.com/prisma/prisma/issues/18442#issuecomment-1518982987

マイグレーション、これでいいみたい。

generator client {
  provider        = "prisma-client-js"
  previewFeatures = ["postgresqlExtensions"]
}

datasource db {
  provider   = "postgres"
  url        = env("DATABASE_URL")
  extensions = [pgvector(map: "vector", schema: "public")]
}

model items {
  id        Int                   @id @default(autoincrement())
  embedding Unsupported("vector")
}

これで prisma migrate dev で作られた SQL の先頭にこれ CREATE EXTENSION 文を入れると prisma migrate deploy したときに自動的に有効になって、テーブル作成できるよと。

-- CreateExtension
CREATE EXTENSION IF NOT EXISTS "vector" WITH SCHEMA "public";

-- CreateTable
CREATE TABLE "items" (
    "id" SERIAL NOT NULL,
    "embedding" vector NOT NULL,

    CONSTRAINT "items_pkey" PRIMARY KEY ("id")
);
Coji MizoguchiCoji Mizoguchi

というわけで queryRaw で cast してやってみる。

vector-test.tsx
import { Box } from '@chakra-ui/react'
import { json } from '@remix-run/node'
import { useLoaderData } from '@remix-run/react'
import { prisma } from '~/services/database.server'

export const loader = async () => {
  const list = await prisma.$queryRaw`
SELECT
  id,
  cast(embedding as text)
FROM
  items
ORDER BY
  embedding <-> '[3,1,2]'
`
  return json({ list })
}

export default function VectorTestPage() {
  const { list } = useLoaderData<typeof loader>()
  return (
    <Box>
      <pre>{JSON.stringify(list, null, 2)}</pre>
    </Box>
  )
}

実行結果

[
  {
    "id": 1,
    "embedding": "[1,2,3]"
  },
  {
    "id": 2,
    "embedding": "[4,5,6]"
  }
]

yay!

Coji MizoguchiCoji Mizoguchi

pgvector の nodejs サンプルに prisma での使い方載ってた。

https://github.com/pgvector/pgvector-node#prisma

Import the library

import pgvector from 'pgvector/utils';

Add the extension to the schema

generator client {
  provider        = "prisma-client-js"
  previewFeatures = ["postgresqlExtensions"]
}

datasource db {
  provider   = "postgresql"
  url        = env("DATABASE_URL")
  extensions = [vector]
}

Add a vector column to the schema

model Item {
  id        Int                       @id @default(autoincrement())
  embedding Unsupported("vector(3)")?
}

Insert a vector

const embedding = pgvector.toSql([1, 1, 1])
await prisma.$executeRaw`INSERT INTO items (embedding) VALUES (${embedding}::vector)`

Get the nearest neighbors to a vector

const embedding = pgvector.toSql([1, 1, 1])
const items = await prisma.$queryRaw`SELECT id, embedding::text FROM items ORDER BY embedding <-> ${embedding}::vector LIMIT 5`
Coji MizoguchiCoji Mizoguchi

見慣れない ${embedding}::vector とか、 embedding::text とかの :: は postgres 独自の型キャスト演算子のようで CAST(embedding as text) とかにするのと同じみたい。つらい。

Coji MizoguchiCoji Mizoguchi

order by のこれ、なんなんだろう。距離っぽいけど、pgvector で定義されてるのかなあ。

ORDER BY embedding <-> ${embedding}::vector

Coji MizoguchiCoji Mizoguchi

pgvector の README に下記記載があるから、やっぱり <-> は pgvector 定義の演算子っぽい

Vector Operators

Operator Description
+ element-wise addition
- element-wise subtraction
<-> Euclidean distance
<#> negative inner product
<=> cosine distance

Vector Functions

Function Description
cosine_distance(vector, vector) → double precision cosine distance
inner_product(vector, vector) → double precision inner product
l2_distance(vector, vector) → double precision Euclidean distance
vector_dims(vector) → integer number of dimensions
vector_norm(vector) → double precision Euclidean norm

Aggregate Functions

Function Description
avg(vector) → vector arithmetic mean
Coji MizoguchiCoji Mizoguchi

7次元までは大丈夫だけど、8次元以上だと select 時にプロセスが落ちちゃう。どういうことなの。。

Coji MizoguchiCoji Mizoguchi

dmesg で以下のログが出てた。 invalid opcode ということで。ぜんぜんわかんない

[203708.729520] traps: postgres[21870] trap invalid opcode ip:7f7a9cc316f2 sp:7ffe482c95c0 error:0 in vector.so[7f7a9cc2b000+8000]