pgvectorとPythonで作る社内文書RAG：権限フィルタ付き検索の設計と実装

[取り込み] 文書 -> チャンク分割 -> /api/embed (bge-m3) -> documents / chunks テーブル
[応答]     質問 -> /api/embed -> SELECT ... WHERE allowed_groups && :groups ORDER BY embedding <=> :q LIMIT 8
                -> プロンプト組み立て -> /api/chat (gemma4:12b) -> 回答 + 出典 [1][2]

services:
  db:
    image: pgvector/pgvector:pg17        # PostgreSQL 17 に pgvector を同梱した公式イメージ
    environment:
      POSTGRES_DB: rag
      POSTGRES_USER: rag
      POSTGRES_PASSWORD: rag-password    # 本番では .env と secrets に分離する
    ports:
      - "127.0.0.1:5432:5432"
    volumes:
      - pgdata:/var/lib/postgresql/data
volumes:
  pgdata:

docker compose up -d
docker compose exec db psql -U rag -d rag -c "CREATE EXTENSION IF NOT EXISTS vector;"
docker compose exec db psql -U rag -d rag -c "SELECT extversion FROM pg_extension WHERE extname = 'vector';"
# 期待: 0.8.0 以上（iterative_scan はこの版で追加。執筆時点の最新はさらに上）

python3 -m venv .venv && source .venv/bin/activate
pip install "psycopg[binary]" ollama
ollama pull bge-m3
ollama show bge-m3 | grep -i "embedding length"      # 期待: embedding length 1024

CREATE EXTENSION IF NOT EXISTS vector;

CREATE TABLE documents (
    id             bigserial   PRIMARY KEY,
    title          text        NOT NULL,
    source_uri     text        NOT NULL UNIQUE,     -- 再取り込みの突き合わせキー
    owner          text        NOT NULL,
    allowed_groups text[]      NOT NULL DEFAULT '{}',  -- 閲覧を許すグループ（IdP のグループ名）
    content_hash   text        NOT NULL,            -- 変更検知用 SHA-256
    updated_at     timestamptz NOT NULL DEFAULT now()
);

CREATE TABLE chunks (
    id          bigserial PRIMARY KEY,
    document_id bigint    NOT NULL REFERENCES documents(id) ON DELETE CASCADE,
    chunk_index int       NOT NULL,
    content     text      NOT NULL,
    embedding   vector(1024) NOT NULL,              -- bge-m3 の次元。モデルを変えたら列ごと作り直す
    UNIQUE (document_id, chunk_index)
);

CREATE INDEX chunks_embedding_hnsw      ON chunks    USING hnsw (embedding vector_cosine_ops);
CREATE INDEX documents_groups_gin       ON documents USING gin  (allowed_groups);
CREATE INDEX chunks_document_id_idx     ON chunks    (document_id);

import hashlib
import sys
from pathlib import Path

import ollama
import psycopg

DSN = "postgresql://rag:[email protected]:5432/rag"
EMBED_MODEL = "bge-m3"        # 1024 次元。変えたら全チャンクの再生成が必要
CHUNK_SIZE = 600
CHUNK_OVERLAP = 100


def split_text(text: str) -> list[str]:
    """段落境界を優先し、上限を超える段落だけ固定長で刻む。"""
    chunks: list[str] = []
    buf = ""
    for para in (p.strip() for p in text.split("\n\n") if p.strip()):
        if len(buf) + len(para) <= CHUNK_SIZE:
            buf = f"{buf}\n{para}" if buf else para
            continue
        if buf:
            chunks.append(buf)
        while len(para) > CHUNK_SIZE:
            chunks.append(para[:CHUNK_SIZE])
            para = para[CHUNK_SIZE - CHUNK_OVERLAP:]
        buf = para
    if buf:
        chunks.append(buf)
    return chunks


def to_vector_literal(vec: list[float]) -> str:
    """pgvector のテキスト表現。ドライバ側のアダプタ登録なしで渡せる。"""
    return "[" + ",".join(f"{x:.7f}" for x in vec) + "]"


def ingest(path: Path, title: str, owner: str, groups: list[str]) -> None:
    text = path.read_text(encoding="utf-8")
    chunks = [f"{title}\n{c}" for c in split_text(text)]   # 文書名をパンくずとして先頭に付ける
    vectors = ollama.embed(model=EMBED_MODEL, input=chunks)["embeddings"]

    with psycopg.connect(DSN) as conn, conn.cursor() as cur:
        cur.execute(
            """
            INSERT INTO documents (title, source_uri, owner, allowed_groups, content_hash)
            VALUES (%s, %s, %s, %s, %s)
            ON CONFLICT (source_uri) DO UPDATE
               SET title = EXCLUDED.title, allowed_groups = EXCLUDED.allowed_groups,
                   content_hash = EXCLUDED.content_hash, updated_at = now()
            RETURNING id
            """,
            (title, str(path), owner, groups, hashlib.sha256(text.encode()).hexdigest()),
        )
        doc_id = cur.fetchone()[0]
        cur.execute("DELETE FROM chunks WHERE document_id = %s", (doc_id,))
        for i, (chunk, vec) in enumerate(zip(chunks, vectors)):
            cur.execute(
                "INSERT INTO chunks (document_id, chunk_index, content, embedding) "
                "VALUES (%s, %s, %s, %s::vector)",
                (doc_id, i, chunk, to_vector_literal(vec)),
            )
    print(f"{path.name}: {len(chunks)} chunks")


if __name__ == "__main__":
    # 例: python ingest.py docs/keihi.md 経費精算規程 keiri-admin grp-all,grp-keiri
    file, title, owner, groups = sys.argv[1:5]
    ingest(Path(file), title, owner, groups.split(","))

SET hnsw.iterative_scan = relaxed_order;   -- pgvector 0.8.0 以降。フィルタで足りない分を索引から読み足す

SELECT c.id, d.title, d.source_uri, c.content,
       1 - (c.embedding <=> $1::vector) AS score      -- コサイン類似度（1 が同一）
FROM chunks c
JOIN documents d ON d.id = c.document_id
WHERE d.allowed_groups && $2::text[]                 -- 利用者の所属グループのどれかを含む文書だけ
ORDER BY c.embedding <=> $1::vector                  -- 索引と同じ演算子でないと索引が使われない
LIMIT 8;

import ollama
import psycopg

DSN = "postgresql://rag:[email protected]:5432/rag"
EMBED_MODEL = "bge-m3"
CHAT_MODEL = "gemma4:12b"

SEARCH_SQL = """
SELECT c.id, d.title, d.source_uri, c.content,
       1 - (c.embedding <=> %(q)s::vector) AS score
FROM chunks c
JOIN documents d ON d.id = c.document_id
WHERE d.allowed_groups && %(groups)s::text[]
ORDER BY c.embedding <=> %(q)s::vector
LIMIT %(k)s
"""

SYSTEM_PROMPT = """あなたは社内文書に基づいて回答するアシスタントです。
- 回答は「参考資料」に書かれている内容だけを根拠にしてください。
- 参考資料に根拠がない場合は「該当する記載が見つかりませんでした」とだけ答えてください。
- 根拠にした資料の番号を文末に [1] の形式で示してください。"""


def search(question: str, groups: list[str], k: int = 8) -> list[dict]:
    qvec = ollama.embed(model=EMBED_MODEL, input=question)["embeddings"][0]
    q = "[" + ",".join(f"{x:.7f}" for x in qvec) + "]"
    with psycopg.connect(DSN) as conn, conn.cursor() as cur:
        cur.execute("SET hnsw.iterative_scan = relaxed_order")
        cur.execute(SEARCH_SQL, {"q": q, "groups": groups, "k": k})
        rows = cur.fetchall()
    return [dict(id=r[0], title=r[1], uri=r[2], content=r[3], score=r[4]) for r in rows]


def answer(question: str, groups: list[str]) -> str:
    hits = search(question, groups)
    if not hits:   # 権限内に根拠がなければ LLM を呼ばない
        return "閲覧権限のある文書の中に、該当する記載が見つかりませんでした。"
    context = "\n\n".join(
        f"[{i}] {h['title']}（{h['uri']}）\n{h['content']}" for i, h in enumerate(hits, 1)
    )
    res = ollama.chat(
        model=CHAT_MODEL,
        options={"num_ctx": 8192},   # 8 チャンク分の文脈が既定の 4096 に収まらない
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": f"参考資料:\n{context}\n\n質問: {question}"},
        ],
    )
    return res["message"]["content"]


if __name__ == "__main__":
    print(answer("経費精算の締め日はいつですか", groups=["grp-all", "grp-keiri"]))

$ python ingest.py docs/keihi.md 経費精算規程 keiri-admin grp-all,grp-keiri
keihi.md: 23 chunks

$ docker compose exec db psql -U rag -d rag -c "SELECT count(*) FROM chunks;"
 count
-------
    23

$ docker compose exec db psql -U rag -d rag -c "EXPLAIN SELECT c.id FROM chunks c JOIN documents d ON d.id = c.document_id WHERE d.allowed_groups && ARRAY['grp-keiri'] ORDER BY c.embedding <=> (SELECT embedding FROM chunks LIMIT 1) LIMIT 8;"
 ...
   ->  Index Scan using chunks_embedding_hnsw on chunks c     ← HNSW 索引が使われている
 ...

$ python search.py
経費精算の締め日は毎月5日です。前月分はワークフローで締め日までに申請してください [1]。

$ python -c "from search import search; print(len(search('経費精算の締め日', groups=['grp-soumu'])))"
0     ← 権限のないグループでは 1 件も返らない