用 pgvector 做 RAG：别急着上 Pinecone，你的 PostgreSQL 就够了

@startuml
skinparam style strictuml
skinparam backgroundColor #FEFEFE
skinparam sequenceArrowColor #555555

actor User
participant "Application" as App
database "pgvector\n(PostgreSQL)" as PG
participant "Embedding\nModel" as Emb
participant "LLM" as LLM

User -> App: 提问
App -> Emb: 把问题转成向量
Emb --> App: question_embedding
App -> PG: SELECT * FROM docs\nORDER BY embedding <=> ?\nLIMIT 5
PG --> App: 最相关的 5 条文档
App -> LLM: 问题 + 相关文档\n（构建 prompt）
LLM --> App: 生成回答
App --> User: 回答

@enduml

services:
  pgvector:
    image: ankane/pgvector
    container_name: pgvector
    environment:
      POSTGRES_USER: ${DB_USER}
      POSTGRES_PASSWORD: ${DB_PASS}
      POSTGRES_DB: ${DB_NAME}
    ports:
      - "5432:5432"
    volumes:
      - pgvector-data:/var/lib/postgresql/data
    networks:
      - pgvector-net
    restart: unless-stopped

  pgweb:
    image: sosedoff/pgweb
    container_name: pgweb
    ports:
      - "8081:8081"
    environment:
      - PGWEB_DATABASE_URL=${DB_URL}
    depends_on:
      - pgvector
    networks:
      - pgvector-net
    restart: unless-stopped

volumes:
  pgvector-data:

networks:
  pgvector-net:
    driver: bridge

DB_USER=rag_user
DB_PASS=rag_pass_2026
DB_NAME=rag_db
DB_URL=postgres://rag_user:rag_pass_2026@pgvector:5432/rag_db?sslmode=disable

docker compose up -d

pip install psycopg2-binary openai httpx python-dotenv

# 索引一个目录下的所有 Markdown 文件
python pgvector_rag_demo.py index ./docs/

# 单次 RAG 查询
python pgvector_rag_demo.py query "pgvector 怎么建索引？"

# 交互式聊天模式
python pgvector_rag_demo.py chat

# 查看知识库统计
python pgvector_rag_demo.py stats

# 清空所有索引数据（换模型后需要重建）
python pgvector_rag_demo.py reset

# 数据库（二选一）
DATABASE_URL=postgresql://rag_user:rag_pass_2026@localhost:5432/rag_db
# 或者分开写：
# DB_HOST=localhost
# DB_PORT=5432
# DB_NAME=rag_db
# DB_USER=rag_user
# DB_PASS=rag_pass_2026

# Embedding 模型（兼容 OpenAI 接口的都行）
EMBEDDING_BASE_URL=https://api.openai.com/v1
EMBEDDING_API_KEY=sk-xxx
EMBEDDING_MODEL=text-embedding-3-small

# LLM（同样兼容 OpenAI 接口）
LLM_BASE_URL=https://api.openai.com/v1
LLM_API_KEY=sk-xxx
LLM_MODEL=gpt-4o-mini

#!/usr/bin/env python3
"""
pgvector RAG Demo — Read markdown files, embed, store in PostgreSQL, then query.

Usage:
  python pgvector_rag_demo.py index ./path/to/markdown/dir
  python pgvector_rag_demo.py query "How does the RAG engine work?"
  python pgvector_rag_demo.py chat
  python pgvector_rag_demo.py stats
  python pgvector_rag_demo.py reset

Prerequisites:
  pip install psycopg2-binary openai httpx python-dotenv
"""

import os, sys, glob, textwrap
from pathlib import Path
from dotenv import load_dotenv

PROJECT_ROOT = Path(__file__).resolve().parent.parent
load_dotenv(PROJECT_ROOT / ".env")
load_dotenv(override=True)


# ============================================================
# Configuration
# ============================================================

def get_db_config() -> dict:
    db_url = os.getenv("DATABASE_URL", "")
    if db_url.startswith("postgresql"):
        from urllib.parse import urlparse
        p = urlparse(db_url)
        return {
            "host": p.hostname or "localhost",
            "port": p.port or 5432,
            "dbname": p.path.lstrip("/") or "rag_db",
            "user": p.username or "postgres",
            "password": p.password or "postgres",
        }
    return {
        "host": os.getenv("DB_HOST", "localhost"),
        "port": int(os.getenv("DB_PORT", "5432")),
        "dbname": os.getenv("DB_NAME", "rag_db"),
        "user": os.getenv("DB_USER", "postgres"),
        "password": os.getenv("DB_PASS", "postgres"),
    }

EMBEDDING_BASE_URL = (os.getenv("EMBEDDING_BASE_URL")
                      or os.getenv("LLM_BASE_URL", "https://api.openai.com/v1"))
EMBEDDING_API_KEY  = (os.getenv("EMBEDDING_API_KEY")
                      or os.getenv("LLM_API_KEY", ""))
EMBEDDING_MODEL    = (os.getenv("EMBEDDING_MODEL")
                      or os.getenv("LLM_EMBEDDING_MODEL", "text-embedding-3-small"))
LLM_BASE_URL   = os.getenv("LLM_BASE_URL", "https://api.openai.com/v1")
LLM_API_KEY    = os.getenv("LLM_API_KEY", "")
LLM_MODEL      = os.getenv("LLM_MODEL", "gpt-4o-mini")
LLM_VERIFY_SSL = os.getenv("LLM_VERIFY_SSL", "true").lower() not in ("false", "0", "no")

TOP_K = 5
CHUNK_SIZE = 512
CHUNK_OVERLAP = 50


# ============================================================
# Embedding (OpenAI-compatible API)
# ============================================================

def get_embedding_client():
    import httpx
    from openai import OpenAI
    http_client = httpx.Client(
        headers={"X-Api-Key": EMBEDDING_API_KEY},
        verify=LLM_VERIFY_SSL,
    )
    return OpenAI(
        api_key=EMBEDDING_API_KEY,
        base_url=EMBEDDING_BASE_URL,
        http_client=http_client,
    )

def get_embeddings(texts: list[str]) -> list[list[float]]:
    client = get_embedding_client()
    response = client.embeddings.create(input=texts, model=EMBEDDING_MODEL)
    return [item.embedding for item in response.data]

def get_embedding(text: str) -> list[float]:
    return get_embeddings([text])[0]


# ============================================================
# Text chunking (sentence-aware)
# ============================================================

def chunk_text(text: str, chunk_size: int = CHUNK_SIZE,
               overlap: int = CHUNK_OVERLAP) -> list[str]:
    """按语义边界切块，带重叠窗口。"""
    text = text.strip()
    if not text:
        return []
    if len(text) <= chunk_size:
        return [text]

    chunks, start = [], 0
    while start < len(text):
        end = start + chunk_size
        if end < len(text):
            for sep in ["\n\n", "\n", ". ", "。", "! ", "? "]:
                last_sep = text.rfind(sep, start + chunk_size // 2, end)
                if last_sep != -1:
                    end = last_sep + len(sep)
                    break
        chunk = text[start:end].strip()
        if chunk:
            chunks.append(chunk)
        new_start = end - overlap
        if new_start <= start:
            new_start = end
        start = new_start
    return chunks


# ============================================================
# Markdown file reading
# ============================================================

def read_markdown_files(directory: str) -> list[dict]:
    docs, seen = [], set()
    for pattern in [os.path.join(directory, "*.md"),
                    os.path.join(directory, "**/*.md")]:
        for filepath in glob.glob(pattern, recursive=True):
            filepath = os.path.abspath(filepath)
            if filepath in seen:
                continue
            seen.add(filepath)
            try:
                with open(filepath, "r", encoding="utf-8") as f:
                    content = f.read()
            except Exception as e:
                print(f"  Skipping {filepath}: {e}")
                continue
            if not content.strip():
                continue
            title = Path(filepath).stem
            for line in content.split("\n"):
                line = line.strip()
                if line.startswith("# "):
                    title = line.lstrip("# ").strip()
                    break
                if line.startswith("Title:"):
                    title = line.split(":", 1)[1].strip()
                    break
            docs.append({"title": title, "content": content, "source": filepath})
    return docs


# ============================================================
# Database
# ============================================================

def get_connection():
    import psycopg2
    cfg = get_db_config()
    conn = psycopg2.connect(**cfg)
    print(f"Connected to PostgreSQL: {cfg['host']}:{cfg['port']}/{cfg['dbname']}")
    return conn

def _detect_embedding_dim() -> int:
    """发一条短文本给 embedding 模型，自动探测输出维度。"""
    try:
        emb = get_embedding("dimension probe")
        dim = len(emb)
        print(f"Detected embedding dimension: {dim} (model: {EMBEDDING_MODEL})")
        return dim
    except Exception as e:
        print(f"Could not probe dimension ({e}), defaulting to 1536")
        return 1536

def init_db(conn):
    dim = _detect_embedding_dim()
    with conn.cursor() as cur:
        cur.execute("CREATE EXTENSION IF NOT EXISTS vector;")
        cur.execute("""
            SELECT EXISTS (
                SELECT 1 FROM information_schema.tables
                WHERE table_name = 'rag_documents'
            );
        """)
        if not cur.fetchone()[0]:
            cur.execute(f"""
                CREATE TABLE rag_documents (
                    id SERIAL PRIMARY KEY,
                    title TEXT NOT NULL,
                    content TEXT NOT NULL,
                    source TEXT DEFAULT '',
                    chunk_index INTEGER DEFAULT 0,
                    embedding vector({dim}),
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                );
            """)
            cur.execute(f"""
                CREATE INDEX idx_rag_docs_embedding
                ON rag_documents
                USING hnsw (embedding vector_cosine_ops)
                WITH (m = 16, ef_construction = 64);
            """)
            print(f"Created table rag_documents with vector({dim}) + HNSW index")
        else:
            cur.execute("""
                SELECT atttypmod FROM pg_attribute
                WHERE attrelid = 'rag_documents'::regclass
                  AND attname = 'embedding';
            """)
            row = cur.fetchone()
            existing_dim = row[0] if row and row[0] > 0 else None
            if existing_dim and existing_dim != dim:
                print(f"Warning: existing column is vector({existing_dim}), "
                      f"but model produces {dim}-dim. Run 'reset' and re-index.")
            print(f"Table rag_documents already exists (dim: {existing_dim or 'unset'})")
        conn.commit()

def index_documents(conn, docs: list[dict]):
    from psycopg2.extras import execute_values
    total_chunks = 0
    for doc in docs:
        chunks = chunk_text(doc["content"])
        if not chunks:
            continue
        print(f"  {doc['title']}: {len(chunks)} chunks", end="")
        batch_size = 64
        all_embeddings = []
        for i in range(0, len(chunks), batch_size):
            all_embeddings.extend(get_embeddings(chunks[i:i + batch_size]))
            print(".", end="", flush=True)
        print()
        rows = [(doc["title"], c, doc.get("source", ""), i, e)
                for i, (c, e) in enumerate(zip(chunks, all_embeddings))]
        with conn.cursor() as cur:
            execute_values(
                cur,
                """INSERT INTO rag_documents
                   (title, content, source, chunk_index, embedding)
                   VALUES %s""",
                rows,
                template="(%s, %s, %s, %s, %s::vector)",
            )
        conn.commit()
        total_chunks += len(rows)
    print(f"Indexed {total_chunks} chunks from {len(docs)} documents")

def search_similar(conn, query: str, top_k: int = TOP_K) -> list[dict]:
    query_emb = get_embedding(query)
    with conn.cursor() as cur:
        cur.execute("""
            SELECT id, title, content, source, chunk_index,
                   1 - (embedding <=> %s::vector) AS similarity
            FROM rag_documents
            ORDER BY embedding <=> %s::vector
            LIMIT %s
        """, (query_emb, query_emb, top_k))
        columns = [desc[0] for desc in cur.description]
        return [dict(zip(columns, row)) for row in cur.fetchall()]

def get_stats(conn) -> dict:
    with conn.cursor() as cur:
        cur.execute("SELECT COUNT(*) FROM rag_documents;")
        total = cur.fetchone()[0]
        cur.execute("SELECT COUNT(DISTINCT title) FROM rag_documents;")
        docs = cur.fetchone()[0]
        cur.execute("SELECT COUNT(DISTINCT source) FROM rag_documents;")
        sources = cur.fetchone()[0]
    return {"total_chunks": total, "documents": docs, "sources": sources}


# ============================================================
# RAG: Retrieve + Generate
# ============================================================

def rag_query(conn, question: str) -> str:
    import httpx
    from openai import OpenAI

    results = search_similar(conn, question)
    print(f"\nQuestion: {question}")
    print(f"Retrieved {len(results)} chunks:")
    for i, r in enumerate(results, 1):
        print(f"  [{i}] {r['title']} (chunk #{r['chunk_index']}, "
              f"similarity: {r['similarity']:.4f})")

    context_parts = []
    for r in results:
        snippet = r["content"][:500]
        context_parts.append(f"### {r['title']}\n{snippet}\n(source: {r['source']})")
    context = "\n\n".join(context_parts)

    http_client = httpx.Client(
        headers={"X-Api-Key": LLM_API_KEY}, verify=LLM_VERIFY_SSL,
    )
    client = OpenAI(
        api_key=LLM_API_KEY, base_url=LLM_BASE_URL, http_client=http_client,
    )
    messages = [
        {"role": "system", "content": (
            "You are a knowledgeable assistant. Answer based on "
            "the provided reference materials. If insufficient, say so. "
            "Cite sources. Answer in the same language as the question."
        )},
        {"role": "user",
         "content": f"Reference materials:\n{context}\n\nQuestion: {question}"},
    ]
    response = client.chat.completions.create(
        model=LLM_MODEL, messages=messages, temperature=0.3, max_tokens=1024,
    )
    answer = response.choices[0].message.content
    print(f"\nAnswer:\n{textwrap.fill(answer, width=80)}")
    return answer


# ============================================================
# CLI
# ============================================================

def main():
    if len(sys.argv) < 2:
        print("Usage: python pgvector_rag_demo.py <index|query|chat|stats|reset>")
        sys.exit(1)

    command = sys.argv[1]
    conn = get_connection()
    try:
        init_db(conn)

        if command == "index":
            directory = sys.argv[2] if len(sys.argv) > 2 else "."
            docs = read_markdown_files(directory)
            if docs:
                index_documents(conn, docs)
            else:
                print("No markdown files found.")

        elif command == "query":
            question = " ".join(sys.argv[2:])
            rag_query(conn, question)

        elif command == "chat":
            print("Interactive RAG chat (type 'quit' to exit)\n")
            while True:
                try:
                    q = input("You: ").strip()
                except (EOFError, KeyboardInterrupt):
                    break
                if not q or q.lower() in ("quit", "exit", "q"):
                    break
                rag_query(conn, q)
                print()

        elif command == "stats":
            s = get_stats(conn)
            print(f"Chunks: {s['total_chunks']}, "
                  f"Documents: {s['documents']}, Sources: {s['sources']}")

        elif command == "reset":
            with conn.cursor() as cur:
                cur.execute("DROP TABLE IF EXISTS rag_documents;")
                conn.commit()
            print("All indexed data cleared. Run 'index' to re-index.")

    finally:
        conn.close()

if __name__ == "__main__":
    main()

def _detect_embedding_dim() -> int:
    try:
        emb = get_embedding("dimension probe")
        return len(emb)
    except Exception:
        return 1536  # fallback

def chunk_text(text, chunk_size=512, overlap=50):
    # ...
    for sep in ["\n\n", "\n", ". ", "。", "! ", "? "]:
        last_sep = text.rfind(sep, start + chunk_size // 2, end)
        if last_sep != -1:
            end = last_sep + len(sep)
            break
    # ...

CREATE INDEX idx_rag_docs_embedding
ON rag_documents
USING hnsw (embedding vector_cosine_ops)
WITH (m = 16, ef_construction = 64);

SET hnsw.ef_search = 100;  -- 默认 40，越大越精确但越慢

SELECT pg_size_pretty(pg_relation_size('idx_rag_docs_embedding'));

@startmindmap
* pgvector 做 RAG
** 为什么用 pgvector
*** 复用 PostgreSQL
*** ACID + JOIN
*** 零额外成本
*** 会 SQL 就行
** 核心步骤
*** 部署 pgvector (Docker)
*** 创建 vector 列（自动探测维度）
*** 建 HNSW 索引
*** 语义分块 + 算 embedding
*** 余弦相似度检索
*** 拼 prompt → LLM
** CLI 工具
*** index: 索引 Markdown 目录
*** query: 单次 RAG 查询
*** chat: 交互式聊天
*** stats: 知识库统计
*** reset: 清空重建
** 性能优化
*** 选对 embedding 维度
*** 批量插入（batch 64）
*** HNSW 参数调优
*** 预热索引
** 适用边界
*** < 500 万条：pgvector
*** > 千万条：专用向量库
*** 需要 JOIN：pgvector
*** 多租户/sharding：专用
** 踩坑清单
*** CREATE EXTENSION 别忘
*** 维度不匹配要 reset
*** IVFFlat 需要先有数据
*** Docker 网络 hostname
@endmindmap