pgvector PostgreSQL 向量扩展¶

# 确保已安装 PostgreSQL 15+
sudo apt install postgresql-16-pgvector  # 安装 pgvector 扩展

# 连接到 PostgreSQL
sudo -u postgres psql  # 以 postgres 用户连接

# 在数据库中启用扩展
CREATE EXTENSION vector;  -- 启用 pgvector 扩展

# 使用自带 pgvector 的 PostgreSQL Docker 镜像
docker run -d \
  --name postgres-pgvector \
  -e POSTGRES_PASSWORD=mysecretpassword \  # 设置密码
  -p 5432:5432 \                            # 映射端口
  -v pgdata:/var/lib/postgresql/data \      # 数据持久化
  pgvector/pgvector:pg16                    # pgvector 官方镜像

# 连接并启用扩展
docker exec -it postgres-pgvector psql -U postgres -c "CREATE EXTENSION vector;"

# 下载并编译（需要 PostgreSQL 开发头文件）
git clone https://github.com/pgvector/pgvector.git  # 下载源码
cd pgvector
make                    # 编译
sudo make install       # 安装

# 在 PostgreSQL 中启用
psql -U postgres -c "CREATE EXTENSION vector;"

pip install psycopg2-binary  # PostgreSQL Python 驱动
# 或
pip install pgvector         # pgvector 专用 Python 工具

-- 创建包含向量列的表
CREATE TABLE papers (
    id SERIAL PRIMARY KEY,                -- 自增主键
    title TEXT NOT NULL,                   -- 论文标题
    abstract TEXT,                         -- 摘要
    year INTEGER,                          -- 年份
    topic TEXT,                            -- 主题
    embedding vector(384)                  -- 384维向量列
);

-- 插入数据（向量用方括号表示）
INSERT INTO papers (title, abstract, year, topic, embedding) VALUES
    ('宏基因组分析方法综述',
     '本文综述了宏基因组分析的主要方法...',
     2024,
     'metagenomics',
     '[0.1, 0.2, 0.3, ...]'              -- 384维向量
    ),
    ('16S rRNA测序技术进展',
     '16S rRNA是细菌分类的金标准...',
     2023,
     'amplicon',
     '[0.3, 0.1, 0.2, ...]'
    );

-- 余弦距离搜索（最常用）
SELECT title, year,
       1 - (embedding <=> '[0.15, 0.25, ...]') AS similarity  -- 余弦相似度
FROM papers
ORDER BY embedding <=> '[0.15, 0.25, ...]'  -- <=> 是余弦距离运算符
LIMIT 5;                                     -- 返回最相似的 5 条

-- L2 距离搜索（欧几里得距离）
SELECT title, embedding <-> '[0.15, 0.25, ...]' AS distance  -- <-> 是 L2 距离
FROM papers
ORDER BY embedding <-> '[0.15, 0.25, ...]'
LIMIT 5;

-- 内积搜索
SELECT title, embedding <#> '[0.15, 0.25, ...]' AS neg_inner_product  -- <#> 是负内积
FROM papers
ORDER BY embedding <#> '[0.15, 0.25, ...]'
LIMIT 5;

-- 向量搜索 + WHERE 条件（SQL 的强大之处）
SELECT title, year,
       1 - (embedding <=> '[0.15, 0.25, ...]') AS similarity
FROM papers
WHERE year >= 2024                         -- 只搜 2024 年及以后
  AND topic = 'metagenomics'               -- 只搜宏基因组主题
ORDER BY embedding <=> '[0.15, 0.25, ...]'
LIMIT 5;

-- 还可以 JOIN 其他表
SELECT p.title, a.name AS author
FROM papers p
JOIN authors a ON p.id = a.paper_id        -- 关联作者表
ORDER BY p.embedding <=> '[0.15, 0.25, ...]'
LIMIT 5;

-- 创建 HNSW 索引（Hierarchical Navigable Small World）
-- 这是搜索速度最快的索引类型
CREATE INDEX ON papers
USING hnsw (embedding vector_cosine_ops)  -- vector_cosine_ops = 余弦距离
WITH (m = 16, ef_construction = 64);       -- m: 连接数, ef_construction: 构建精度

-- 查询时调整搜索精度
SET hnsw.ef_search = 100;  -- 越大越精确，越慢（默认 40）

-- L2 距离的索引
CREATE INDEX ON papers USING hnsw (embedding vector_l2_ops);

-- 内积距离的索引
CREATE INDEX ON papers USING hnsw (embedding vector_ip_ops);

-- v0.8.0 新增：迭代扫描避免过滤后结果太少
SET hnsw.iterative_scan = relaxed_order;  -- 启用迭代扫描

-- 即使过滤条件很严格，也能返回足够的结果
SELECT title
FROM papers
WHERE year = 2025 AND topic = 'rare_topic'  -- 很严格的过滤
ORDER BY embedding <=> '[0.15, 0.25, ...]'
LIMIT 10;

-- 设回默认
SET hnsw.iterative_scan = off;

import psycopg2                              # PostgreSQL 驱动
from sentence_transformers import SentenceTransformer  # 嵌入模型
import openai                                # LLM 客户端

# 初始化嵌入模型
model = SentenceTransformer("all-MiniLM-L6-v2")  # 384 维

# 连接数据库
conn = psycopg2.connect(
    host="localhost",
    port=5432,
    dbname="mydb",
    user="postgres",
    password="password"
)
conn.autocommit = True
cur = conn.cursor()

# 启用 pgvector 扩展
cur.execute("CREATE EXTENSION IF NOT EXISTS vector")

# 创建表
cur.execute("""
    CREATE TABLE IF NOT EXISTS knowledge (
        id SERIAL PRIMARY KEY,
        content TEXT,
        embedding vector(384)
    )
""")

# 插入知识文档
docs = [
    "2型糖尿病患者的肠道菌群Alpha多样性显著降低",
    "Akkermansia muciniphila 与胰岛素敏感性正相关",
    "短链脂肪酸（SCFAs）能增强肠道屏障功能",
]

for doc in docs:
    embedding = model.encode(doc).tolist()     # 生成嵌入向量
    cur.execute(
        "INSERT INTO knowledge (content, embedding) VALUES (%s, %s)",
        (doc, str(embedding))                  # 插入文本和向量
    )

# 创建 HNSW 索引
cur.execute("""
    CREATE INDEX IF NOT EXISTS knowledge_embedding_idx
    ON knowledge USING hnsw (embedding vector_cosine_ops)
""")

# RAG 搜索
def rag_search(question, top_k=3):
    query_embedding = model.encode(question).tolist()
    cur.execute("""
        SELECT content, 1 - (embedding <=> %s::vector) AS similarity
        FROM knowledge
        ORDER BY embedding <=> %s::vector
        LIMIT %s
    """, (str(query_embedding), str(query_embedding), top_k))

    return cur.fetchall()

# 搜索 + LLM 回答
results = rag_search("糖尿病和肠道菌群有什么关系？")
context = "\n".join([r[0] for r in results])

client = openai.Client(base_url="http://localhost:11434/v1", api_key="ollama")
response = client.chat.completions.create(
    model="llama3.1:8b",
    messages=[
        {"role": "system", "content": f"根据以下知识回答：\n{context}"},
        {"role": "user", "content": "糖尿病和肠道菌群有什么关系？"}
    ]
)
print(response.choices[0].message.content)

# 清理
cur.close()
conn.close()

-- === pgvector SQL 速查 ===

-- 启用扩展
CREATE EXTENSION vector;

-- 创建向量列
ALTER TABLE t ADD COLUMN embedding vector(384);

-- 距离运算符
-- <-> L2 距离
-- <=> 余弦距离
-- <#> 负内积

-- 创建索引
CREATE INDEX ON t USING hnsw (embedding vector_cosine_ops);
CREATE INDEX ON t USING ivfflat (embedding vector_cosine_ops) WITH (lists = 100);

-- 搜索
SELECT * FROM t ORDER BY embedding <=> '[...]' LIMIT 5;

-- 设置搜索精度
SET hnsw.ef_search = 100;
SET ivfflat.probes = 10;