samples/pgvector_1 — awesome-ai-stack

예제

pgvector — Postgres 안의 벡터 유사도 검색

실제로 실행 가능한 미니 프로젝트입니다. 받아서 Docker로 실행해 보세요.

GitHub에서 samples/pgvector_1 보기 ↗

2026년 6월 28일

pgvector — Postgres 안의 벡터 유사도 검색

아주 작은 pgvector 데모입니다. vector 컬럼을 만들고, 3차원 장난감 "임베딩" 몇 개를 저장한 뒤, 코사인 거리 연산자 (<=>)로 최근접 이웃을 조회합니다. 장난감 벡터를 임베딩 모델의 출력으로 바꾸면 SQL은 그대로입니다.

Docker Compose로 실행 (Postgres 포함)

compose 파일이 pgvector/pgvector Postgres를 띄우고 그 위에서 앱을 실행합니다:

cd samples/pgvector_1
docker compose up --build

여기서는 .env가 필요 없습니다 — compose가 DATABASE_URL을 자동으로 설정합니다.

docker-compose.yml의 질의어(command: ["cat"])를 바꿔 보세요 — rocket으로 하면 이웃이 바뀝니다.

Docker Compose로 실행 (DooD를 쓰는 devcontainer에서)

호스트 Docker 데몬과 통신하는 dev container(Docker-outside-of-Docker)에서는 attach된 docker compose up이 앱을 실제로 실행했는데도 출력을 보여주지 않을 수 있습니다. 실시간 attach 스트림이 VM 경계에서 출력을 흘려버리기 때문이니, 로그로 다시 확인하세요:

cd samples/pgvector_1
docker compose up --build -d
docker compose logs -f app
docker compose down -v

로컬에서 실행

vector 확장을 쓸 수 있는 Postgres로 DATABASE_URL을 지정하세요:

cd samples/pgvector_1
cp .env.sample .env          # 필요하면 DATABASE_URL 수정
pip install -r requirements.txt
python app.py cat

.env는 gitignore 처리되어 있고, .env.sample만 커밋됩니다.

실행 결과

LLM 샘플과 달리 이 샘플은 결정적입니다 — toy 벡터가 고정이라 매번 같은 출력이 나옵니다.

nearest to 'cat' by cosine distance:
  cat        0.0000  (query)
  kitten     0.0129
  rocket     0.8901
  airplane   0.9901

.env.sample

# pgvector needs a Postgres with the `vector` extension. The bundled
# docker-compose.yml starts one for you and sets this automatically.
#
# For a local run (`python app.py`), copy this to `.env` and point it at your
# own Postgres:
#   cp .env.sample .env
DATABASE_URL=postgresql://postgres:postgres@localhost:5432/postgres

app.py

"""pgvector: vector similarity search inside Postgres.

Creates a `vector` column, stores a few toy "embeddings", and runs a
nearest-neighbour query with the cosine-distance operator (`<=>`). A real app
swaps the toy vectors for an embedding model's output — the SQL is identical.

Run (with the bundled Postgres):
    docker compose up --build
"""

import os
import sys

import psycopg
from dotenv import load_dotenv

# Load .env when running locally; docker-compose sets DATABASE_URL directly.
load_dotenv()

DSN = os.environ.get("DATABASE_URL", "postgresql://postgres:postgres@localhost:5432/postgres")

# Toy 3-dim "embeddings" — in practice these come from an embedding model.
DOCS = [
    ("cat", "[1.0, 0.1, 0.0]"),
    ("kitten", "[0.9, 0.2, 0.1]"),
    ("airplane", "[0.0, 0.1, 1.0]"),
    ("rocket", "[0.1, 0.0, 0.9]"),
]


def main() -> None:
    target = (sys.argv[1] if len(sys.argv) > 1 else "cat").lower()
    query_vec = dict(DOCS).get(target, DOCS[0][1])

    with psycopg.connect(DSN, autocommit=True) as conn:
        conn.execute("CREATE EXTENSION IF NOT EXISTS vector")
        conn.execute("DROP TABLE IF EXISTS docs")
        conn.execute("CREATE TABLE docs (id serial PRIMARY KEY, label text, embedding vector(3))")
        for label, embedding in DOCS:
            conn.execute("INSERT INTO docs (label, embedding) VALUES (%s, %s)", (label, embedding))

        rows = conn.execute(
            "SELECT label, embedding <=> %s::vector AS distance FROM docs ORDER BY distance LIMIT 4",
            (query_vec,),
        ).fetchall()

    print(f"nearest to '{target}' by cosine distance:")
    for label, distance in rows:
        marker = "  (query)" if label == target else ""
        print(f"  {label:10s} {distance:.4f}{marker}")


if __name__ == "__main__":
    main()

docker-compose.yml

# pgvector needs a Postgres with the `vector` extension. This bundles one
# (the official pgvector image) and runs the sample app against it.
#   docker compose up --build
services:
  db:
    image: pgvector/pgvector:pg16
    environment:
      POSTGRES_PASSWORD: postgres
    healthcheck:
      test: ["CMD-SHELL", "pg_isready -U postgres"]
      interval: 2s
      timeout: 3s
      retries: 15

  app:
    build: .
    environment:
      DATABASE_URL: postgresql://postgres:postgres@db:5432/postgres
    depends_on:
      db:
        condition: service_healthy
    # The query word — try "rocket" to see the neighbours flip.
    command: ["cat"]

Dockerfile

# Minimal image for the pgvector sample app.
# Postgres (with the vector extension) runs as a separate service — see
# docker-compose.yml.
FROM python:3.13-slim

# Stream stdout/stderr unbuffered so the output appears live.
ENV PYTHONUNBUFFERED=1

WORKDIR /app

COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

COPY app.py .

# DATABASE_URL points at your Postgres (the bundled one is set in compose):
#   docker compose up --build
ENTRYPOINT ["python", "app.py"]