pgvectorとハイブリッド検索でプロダクションRAGを構築する

WITH vector_search AS (
  SELECT id, content, metadata,
         RANK() OVER (ORDER BY embedding <=> $1) AS rank
  FROM documents
  ORDER BY embedding <=> $1
  LIMIT 50
),
keyword_search AS (
  SELECT id, content, metadata,
         RANK() OVER (
           ORDER BY ts_rank_cd(fts, plainto_tsquery('english', $2)) DESC
         ) AS rank
  FROM documents
  WHERE fts @@ plainto_tsquery('english', $2)
  ORDER BY ts_rank_cd(fts, plainto_tsquery('english', $2)) DESC
  LIMIT 50
)
SELECT
  COALESCE(v.id, k.id)               AS id,
  COALESCE(v.content, k.content)     AS content,
  COALESCE(v.metadata, k.metadata)   AS metadata,
  (
    COALESCE(1.0 / (60 + v.rank), 0.0) +
    COALESCE(1.0 / (60 + k.rank), 0.0)
  )                                   AS rrf_score
FROM vector_search v
FULL OUTER JOIN keyword_search k ON v.id = k.id
ORDER BY rrf_score DESC
LIMIT 10;
-- $1: embedding vector, $2: query text string

EXPLAIN (ANALYZE, BUFFERS, FORMAT TEXT)
SELECT id, content, embedding <=> '[0.1, 0.2, ...]' AS dist
FROM documents
ORDER BY embedding <=> '[0.1, 0.2, ...]'
LIMIT 10;

import asyncpg
import openai
from typing import Any

async def embed(text: str) -> list[float]:
    resp = await openai.AsyncOpenAI().embeddings.create(
        model="text-embedding-3-small",
        input=text,
    )
    return resp.data[0].embedding

async def hybrid_search(
    pool: asyncpg.Pool,
    query: str,
    limit: int = 10,
) -> list[dict[str, Any]]:
    embedding = await embed(query)
    # asyncpg expects vectors as lists of floats
    vector_str = "[" + ",".join(str(x) for x in embedding) + "]"

    rows = await pool.fetch(
        """
        WITH vector_search AS (
          SELECT id, content, metadata,
                 RANK() OVER (ORDER BY embedding <=> $1::vector) AS rank
          FROM documents ORDER BY embedding <=> $1::vector LIMIT 50
        ),
        keyword_search AS (
          SELECT id, content, metadata,
                 RANK() OVER (
                   ORDER BY ts_rank_cd(fts, plainto_tsquery('english', $2)) DESC
                 ) AS rank
          FROM documents
          WHERE fts @@ plainto_tsquery('english', $2)
          LIMIT 50
        )
        SELECT
          COALESCE(v.id, k.id)             AS id,
          COALESCE(v.content, k.content)   AS content,
          COALESCE(v.metadata, k.metadata) AS metadata,
          (COALESCE(1.0/(60+v.rank),0) + COALESCE(1.0/(60+k.rank),0)) AS score
        FROM vector_search v
        FULL OUTER JOIN keyword_search k ON v.id = k.id
        ORDER BY score DESC LIMIT $3
        """,
        vector_str,
        query,
        limit,
    )
    return [dict(r) for r in rows]

from sentence_transformers import CrossEncoder

reranker = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')

def rerank(query: str, docs: list[dict]) -> list[dict]:
    pairs = [(query, d['content']) for d in docs]
    scores = reranker.predict(pairs)
    ranked = sorted(zip(docs, scores), key=lambda x: x[1], reverse=True)
    return [d for d, _ in ranked]