pgvector向量库学习教程

单条向量存储大小 ≈ 4 × n + 8 字节（PostgreSQL 的 varlena 头部开销）

示例：
- 768 维向量：4 × 768 + 8 = 3,080 字节 ≈ 3 KB
- 1024 维向量：4 × 1024 + 8 = 4,104 字节 ≈ 4 KB
- 1536 维向量：4 × 1536 + 8 = 6,152 字节 ≈ 6 KB
- 3072 维向量：4 × 3072 + 8 = 12,296 字节 ≈ 12 KB

100 万条 1536 维向量的纯数据存储：约 6 GB
1000 万条 1536 维向量的纯数据存储：约 60 GB

d(A, B) = √(Σᵢ (aᵢ - bᵢ)²)

-- 计算两个向量的欧氏距离
SELECT '[1, 2, 3]'::vector <-> '[4, 5, 6]'::vector AS l2_distance;
-- 结果：5.196152（即 √27 = √((4-1)² + (5-2)² + (6-3)²)）

cosine_similarity(A, B) = (A · B) / (||A|| × ||B||) = Σ(aᵢ × bᵢ) / (√Σaᵢ² × √Σbᵢ²)

cosine_distance(A, B) = 1 - cosine_similarity(A, B)

-- 计算两个向量的余弦距离
SELECT '[1, 2, 3]'::vector <=> '[2, 4, 6]'::vector AS cosine_distance;
-- 结果：0（因为两个向量方向完全相同，[2,4,6] = 2 × [1,2,3]）

inner_product(A, B) = A · B = Σ(aᵢ × bᵢ)

-- 计算两个向量的负内积
SELECT '[1, 2, 3]'::vector <#> '[4, 5, 6]'::vector AS negative_inner_product;
-- 结果：-32（即 -(1×4 + 2×5 + 3×6) = -(4+10+18) = -32）
-- 值为 -32 意味着内积为 32，负内积为 -32
-- 值越小（越负）表示内积越大，向量越相似

┌──────────────────── 离线索引阶段 ────────────────────┐
│                                                       │
│  原始文档 → 文档切片 → Embedding向量化 → 写入pgvector   │
│                                                       │
└───────────────────────────────────────────────────────┘

┌────────────── 在线检索生成阶段 ───────────────────────┐
│                                                       │
│  用户提问 → Query向量化 → pgvector语义检索 →            │
│  Top-K文档片段 → 拼接Prompt → LLM生成回答              │
│                                                       │
└───────────────────────────────────────────────────────┘

# 使用 LangChain 的递归文本切分器
from langchain.text_splitter import RecursiveCharacterTextSplitter

splitter = RecursiveCharacterTextSplitter(
    chunk_size=512,        # 每个 chunk 最大 512 个字符
    chunk_overlap=50,      # 相邻 chunk 重叠 50 个字符
    separators=["\n\n", "\n", "。", ".", " ", ""]  # 按优先级尝试切分
)

chunks = splitter.split_text(document_content)
# 输出示例：
# chunk[0]: "第一章 引言\n本文介绍了基于pgvector的向量检索方案..."
# chunk[1]: "...pgvector作为PostgreSQL扩展，提供了原生的向量数据类型支持..."
# chunk[2]: "第二章 系统架构\n本系统采用pgvector作为向量存储引擎..."

from openai import OpenAI

client = OpenAI()

def get_embedding(text: str, model: str = "text-embedding-3-small") -> list[float]:
    """将文本转换为向量嵌入"""
    response = client.embeddings.create(
        model=model,
        input=text
    )
    return response.data[0].embedding  # 返回 1536 维的浮点数列表

# 对每个 chunk 生成向量
chunk_embeddings = []
for chunk in chunks:
    embedding = get_embedding(chunk)
    chunk_embeddings.append({
        'content': chunk,
        'embedding': embedding,
        'metadata': {'source': 'document.pdf', 'page': 1}
    })

import psycopg
from psycopg.adapt import Dumper

conn = psycopg.connect("postgresql://user:pass@localhost:5432/mydb")
cur = conn.cursor()

# 批量写入
for item in chunk_embeddings:
    cur.execute("""
        INSERT INTO documents (content, metadata, embedding)
        VALUES (%s, %s, %s::vector)
    """, (item['content'], json.dumps(item['metadata']), item['embedding']))

conn.commit()

user_query = "pgvector支持哪些索引类型？"
query_embedding = get_embedding(user_query)

# 语义检索 Top-5 最相关的文档片段
cur.execute("""
    SELECT 
        content,
        metadata,
        1 - (embedding <=> %s::vector) AS similarity_score
    FROM documents
    ORDER BY embedding <=> %s::vector
    LIMIT 5
""", (query_embedding, query_embedding))

results = cur.fetchall()

# 构建 RAG Prompt
context = "\n\n---\n\n".join([row[0] for row in results])

prompt = f"""你是一个知识库助手。请基于以下参考资料回答用户的问题。
如果参考资料中没有相关信息，请明确告知。

参考资料：
{context}

用户问题：{user_query}

请用中文回答："""

# 调用 LLM
response = client.chat.completions.create(
    model="gpt-4o",
    messages=[{"role": "user", "content": prompt}]
)

print(response.choices[0].message.content)

用户输入: "pgvector支持哪些索引类型？"
    ↓
Query Embedding: [0.023, -0.015, ..., 0.089] (1536维)
    ↓
pgvector 语义检索 (ORDER BY embedding <=> query_vec LIMIT 5)
    ↓
检索结果:
  [0.95] "pgvector支持两种索引类型：HNSW和IVFFlat..."
  [0.89] "HNSW索引是一种基于图的近似最近邻搜索算法..."
  [0.85] "IVFFlat索引通过将向量空间划分为多个聚类..."
  [0.82] "创建HNSW索引的SQL语法为：CREATE INDEX..."
  [0.78] "索引选择建议：对于百万级以下的数据量..."
    ↓
拼接 Prompt → LLM 生成回答:
  "pgvector 支持两种主要的索引类型：HNSW（Hierarchical Navigable Small World）
   和 IVFFlat（Inverted File with Flat storage）。HNSW 是一种基于图的..."

# 使用官方 pgvector 镜像（基于 PostgreSQL 17）
docker run -d \
  --name pgvector-demo \
  -e POSTGRES_USER=vector_user \
  -e POSTGRES_PASSWORD=vector_pass \
  -e POSTGRES_DB=vector_db \
  -p 5432:5432 \
  -v pgvector_data:/var/lib/postgresql/data \
  pgvector/pgvector:pg17

# 参数说明：
# -d             : 后台运行
# --name         : 容器名称
# -e POSTGRES_USER     : 数据库用户名
# -e POSTGRES_PASSWORD : 数据库密码
# -e POSTGRES_DB       : 默认数据库名
# -p 5432:5432   : 端口映射（宿主机:容器）
# -v pgvector_data     : 数据持久化卷（避免容器删除后数据丢失）
# pgvector/pgvector:pg17 : 基于 PostgreSQL 17 的 pgvector 镜像

# 等待容器启动完成（约 5-10 秒）
sleep 5

# 连接并验证
docker exec -it pgvector-demo psql -U vector_user -d vector_db -c "CREATE EXTENSION vector;"

version: '3.8'

services:
  postgres:
    image: pgvector/pgvector:pg17
    container_name: pgvector-dev
    environment:
      POSTGRES_USER: vector_user
      POSTGRES_PASSWORD: vector_pass
      POSTGRES_DB: vector_db
      # 性能调优环境变量
      POSTGRES_INITDB_ARGS: "--encoding=UTF-8 --locale=C"
    ports:
      - "5432:5432"
    volumes:
      # 数据持久化
      - pgvector_data:/var/lib/postgresql/data
      # 初始化脚本（可选）
      - ./init.sql:/docker-entrypoint-initdb.d/init.sql
    # 资源配置
    deploy:
      resources:
        limits:
          memory: 4G
          cpus: '2'
    # 健康检查
    healthcheck:
      test: ["CMD-SHELL", "pg_isready -U vector_user -d vector_db"]
      interval: 10s
      timeout: 5s
      retries: 5
    # 重启策略
    restart: unless-stopped

volumes:
  pgvector_data:
    driver: local

-- 启用 pgvector 扩展
CREATE EXTENSION IF NOT EXISTS vector;

-- 创建一个示例表用于验证
CREATE TABLE IF NOT EXISTS test_vectors (
    id bigserial PRIMARY KEY,
    content text,
    embedding vector(3)
);

-- 插入测试数据
INSERT INTO test_vectors (content, embedding) VALUES
    ('hello world', '[1, 2, 3]'),
    ('goodbye world', '[4, 5, 6]'),
    ('pgvector tutorial', '[7, 8, 9]');

# 启动服务
docker-compose up -d

# 查看日志确认启动成功
docker-compose logs -f postgres

# 验证
docker exec -it pgvector-dev psql -U vector_user -d vector_db -c "
    SELECT id, content, embedding, 
           embedding <-> '[3, 1, 2]'::vector AS distance
    FROM test_vectors 
    ORDER BY embedding <-> '[3, 1, 2]'::vector;
"

# 停止服务
docker-compose down

# 停止并清除数据（慎用）
docker-compose down -v

# 1. 安装 PostgreSQL 和开发依赖
sudo apt update
sudo apt install -y postgresql-17 postgresql-server-dev-17 build-essential git

# 如果还没安装 PostgreSQL，先安装
# sudo apt install -y postgresql-17

# 2. 确认 pg_config 可用
pg_config --version
# 输出应类似：PostgreSQL 17.x

# 3. 克隆 pgvector 源码
cd /tmp
git clone --branch v0.8.2 https://github.com/pgvector/pgvector.git
cd pgvector

# 4. 编译
make

# 5. 安装（需要 root 权限或 sudo）
sudo make install

# 6. 验证安装文件
ls -la $(pg_config --pkglibdir)/vector.so
# 应该看到 vector.so 文件存在

# 1. 安装 PostgreSQL 仓库（以 PG 17 为例）
sudo dnf install -y https://download.postgresql.org/pub/repos/yum/reporpms/EL-9-x86_64/pgdg-redhat-repo-latest.noarch.rpm

# 2. 安装 PostgreSQL 和开发依赖
sudo dnf install -y postgresql17-server postgresql17-devel
sudo dnf install -y gcc make git

# 3. 初始化数据库（如果是新安装）
sudo /usr/pgsql-17/bin/postgresql-17-setup initdb
sudo systemctl start postgresql-17
sudo systemctl enable postgresql-17

# 4. 确保 pg_config 在 PATH 中
export PATH=/usr/pgsql-17/bin:$PATH
pg_config --version

# 5. 克隆并编译 pgvector
cd /tmp
git clone --branch v0.8.2 https://github.com/pgvector/pgvector.git
cd pgvector
make
sudo make install

# 使用 Homebrew 安装
brew install pgvector

# 或者从源码编译
brew install postgresql@17
git clone --branch v0.8.2 https://github.com/pgvector/pgvector.git
cd pgvector
make
make install

# 添加 PostgreSQL 官方仓库
sudo sh -c 'echo "deb http://apt.postgresql.org/pub/repos/apt $(lsb_release -cs)-pgdg main" > /etc/apt/sources.list.d/pgdg.list'
wget --quiet -O - https://www.postgresql.org/media/keys/ACCC4CF8.asc | sudo apt-key add -
sudo apt update

# 安装 pgvector 包
sudo apt install -y postgresql-17-pgvector
# 其他 PG 版本：postgresql-16-pgvector, postgresql-15-pgvector, postgresql-14-pgvector

# 安装 PGDG 仓库（如果还没安装）
sudo dnf install -y https://download.postgresql.org/pub/repos/yum/reporpms/EL-9-x86_64/pgdg-redhat-repo-latest.noarch.rpm

# 安装 pgvector
sudo dnf install -y pgvector_17
# 其他版本：pgvector_16, pgvector_15, pgvector_14

# 切换到 postgres 用户（Linux 默认超级用户）
sudo -u postgres psql

# 或者使用你的自定义用户连接
psql -U your_user -d your_database

-- 启用 pgvector 扩展
CREATE EXTENSION vector;

-- 如果需要同时启用半精度向量等功能（v0.7.0+自动包含）
-- 无需额外操作，CREATE EXTENSION vector 已包含所有类型

-- 可选：启用相关扩展以获得更好的性能
CREATE EXTENSION IF NOT EXISTS pg_stat_statements;  -- 查询性能分析

-- 验证
SELECT extname, extversion FROM pg_extension WHERE extname = 'vector';
-- 应输出：vector | 0.8.2

3. **RDS 控制台配置**（可选优化）：
   - 在"参数设置"中调整 `shared_preload_libraries`，确保包含 pgvector 相关配置。
   - 调整 `shared_buffers` 为实例内存的 25%（如 8GB 实例设为 2GB），这对向量索引的缓存性能至关重要。
   - 调整 `effective_cache_size` 为实例内存的 75%。

**腾讯云数据库 PostgreSQL**：

腾讯云 TDSQL-C PostgreSQL 和云数据库 PostgreSQL 均已支持 pgvector。

操作步骤：

1. **确认实例版本**：需要 PostgreSQL 14 及以上版本。

2. **启用扩展**：
```sql
-- 通过控制台或命令行连接实例
psql -h your-instance.postgres.tencentcloudtdsql.com -U tdsqliam_user -d your_database

-- 启用扩展
CREATE EXTENSION IF NOT EXISTS vector;

-- 验证
SELECT extname, extversion FROM pg_extension;

-- 操作步骤与阿里云类似
CREATE EXTENSION IF NOT EXISTS vector;

-- 通过 psql 或其他客户端连接
-- 启用扩展
CREATE EXTENSION IF NOT EXISTS vector;

-- 如果使用 AWS CDK/Terraform 创建实例，可以在参数组中预配置：
-- rds:extension-name: vector

-- ============================================================
-- 环境校验脚本：验证 pgvector 安装是否成功
-- ============================================================

-- 1. 检查扩展是否已启用
SELECT extname, extversion, extnamespace::regnamespace 
FROM pg_extension 
WHERE extname = 'vector';
-- 预期输出：vector | 0.8.2 | public

-- 2. 检查可用的向量相关数据类型
SELECT typname, typcategory
FROM pg_type
WHERE typname IN ('vector', 'halfvec', 'sparsevec', 'bit');
-- 预期输出：vector, halfvec, sparsevec 均存在

-- 3. 检查可用的距离运算符
SELECT oprname, oprleft::regtype, oprright::regtype
FROM pg_operator
WHERE oprname IN ('<->', '<=>', '<#>')
  AND oprleft::regtype::text = 'vector';
-- 预期输出三行：<->, <=>, <#>

-- 4. 创建测试表并执行基本操作
CREATE TABLE IF NOT EXISTS _vector_test (
    id serial PRIMARY KEY,
    embedding vector(3)
);

-- 5. 插入测试数据
INSERT INTO _vector_test (embedding) VALUES 
    ('[1, 0, 0]'),
    ('[0, 1, 0]'),
    ('[0, 0, 1]'),
    ('[1, 1, 1]');

-- 6. 测试三种距离计算
SELECT 
    id,
    embedding,
    embedding <-> '[1, 0, 0]'::vector AS l2_dist,
    embedding <=> '[1, 0, 0]'::vector AS cosine_dist,
    embedding <#> '[1, 0, 0]'::vector AS neg_inner_prod
FROM _vector_test
ORDER BY embedding <-> '[1, 0, 0]'::vector;

-- 预期结果：
-- id=1: l2_dist=0, cosine_dist=0, neg_inner_prod=-1（完全匹配）
-- id=4: l2_dist=1.414, cosine_dist=0.423, neg_inner_prod=-1（方向相似但模长不同）

-- 7. 测试向量维度函数
SELECT vector_dims('[1, 2, 3, 4, 5]'::vector);
-- 预期输出：5

-- 8. 测试 halfvec 类型（v0.7.0+）
SELECT '[1.5, 2.5, 3.5]'::halfvec(3);
-- 预期输出：[1.5, 2.5, 3.5]

-- 9. 测试 sparsevec 类型（v0.7.0+）
SELECT '{1:1.5, 3:2.5}/5'::sparsevec;
-- 预期输出：{1:1.5, 3:2.5}/5（5维稀疏向量，第1和第3位有值）

-- 10. 清理测试表
DROP TABLE IF EXISTS _vector_test;

-- 11. 查看 pgvector 版本详情
SELECT extversion FROM pg_extension WHERE extname = 'vector';

-- 如果所有测试通过，说明 pgvector 环境搭建成功！

# 连接到带 pgvector 的数据库
psql -h localhost -p 5432 -U vector_user -d vector_db

# 常用 psql 命令
\dt          -- 列出所有表
\di          -- 列出所有索引
\d table_name  -- 查看表结构
\x           -- 切换扩展显示模式（适合查看长向量）
\timing      -- 开启查询计时（用于性能测试）
\q           -- 退出

# 安装
pip install "psycopg[binary]"   # 二进制版本，无需编译
# 或
pip install psycopg              # 纯 Python 版本

import psycopg
import numpy as np

# 连接数据库
conn = psycopg.connect(
    host="localhost",
    port=5432,
    dbname="vector_db",
    user="vector_user",
    password="vector_pass"
)

cur = conn.cursor()

# 创建表
cur.execute("""
    CREATE TABLE IF NOT EXISTS docs (
        id bigserial PRIMARY KEY,
        content text,
        embedding vector(3)
    )
""")

# 插入向量（Python 列表直接转为字符串格式）
embedding = [0.1, 0.2, 0.3]
cur.execute(
    "INSERT INTO docs (content, embedding) VALUES (%s, %s::vector)",
    ("测试文档", embedding)
)

# 查询
cur.execute("""
    SELECT content, embedding 
    FROM docs 
    ORDER BY embedding <-> %s::vector 
    LIMIT 5
""", ([0.1, 0.2, 0.3],))

for row in cur.fetchall():
    print(f"内容: {row[0]}, 向量: {row[1]}")

conn.commit()
conn.close()

import numpy as np
import psycopg

# numpy float32 数组可以直接作为向量参数
embedding = np.array([0.1, 0.2, 0.3], dtype=np.float32)

cur.execute(
    "INSERT INTO docs (content, embedding) VALUES (%s, %s::vector)",
    ("测试文档", embedding.tolist())  # 转为 Python 列表后传入
)

pip install asyncpg

import asyncio
import asyncpg

async def main():
    conn = await asyncpg.connect(
        host='localhost', port=5432,
        database='vector_db',
        user='vector_user', password='vector_pass'
    )

    # 需要注册向量类型的编码器/解码器
    await conn.set_type_codec(
        'vector',
        encoder=lambda v: str(v),
        decoder=lambda v: [float(x) for x in v.strip('[]').split(',')]
    )

    # 插入
    embedding = [0.1, 0.2, 0.3]
    await conn.execute(
        "INSERT INTO docs (content, embedding) VALUES ($1, $2::vector)",
        "测试文档", embedding
    )

    # 查询
    rows = await conn.fetch("""
        SELECT content, embedding 
        FROM docs 
        ORDER BY embedding <-> $1::vector 
        LIMIT 5
    """, [0.1, 0.2, 0.3])

    for row in rows:
        print(row)

    await conn.close()

asyncio.run(main())

pip install pgvector sqlalchemy

from sqlalchemy import create_engine, Column, BigInteger, Text
from sqlalchemy.orm import declarative_base, Session
from pgvector.sqlalchemy import Vector

Base = declarative_base()

class Document(Base):
    __tablename__ = 'documents'

    id = Column(BigInteger, primary_key=True, autoincrement=True)
    content = Column(Text)
    embedding = Column(Vector(1024))  # pgvector 的 SQLAlchemy 类型

engine = create_engine("postgresql://vector_user:vector_pass@localhost:5432/vector_db")
Base.metadata.create_all(engine)

# 使用 ORM 插入
with Session(engine) as session:
    doc = Document(
        content="这是测试文档",
        embedding=[0.1, 0.2, 0.3] + [0.0] * 1021  # 需要是 1024 维
    )
    session.add(doc)
    session.commit()

-- OpenAI text-embedding-3-small：1536 维
CREATE TABLE openai_docs (
    id bigserial PRIMARY KEY,
    content text,
    embedding vector(1536)
);

-- 通义千问 text-embedding-v3：1024 维
CREATE TABLE qwen_docs (
    id bigserial PRIMARY KEY,
    content text,
    embedding vector(1024)
);

-- BGE-M3：1024 维
CREATE TABLE bge_docs (
    id bigserial PRIMARY KEY,
    content text,
    embedding vector(1024)
);

-- OpenAI text-embedding-3-large：3072 维
CREATE TABLE openai_large_docs (
    id bigserial PRIMARY KEY,
    content text,
    embedding vector(3072)
);

假设 100 万条文档，使用 text-embedding-3-small（1536 维）：

每条向量存储：4 × 1536 + 8 = 6,152 字节 ≈ 6 KB
100 万条向量总存储：6,152 × 1,000,000 = 6.15 GB（纯数据）
加上 HNSW 索引（通常为数据量的 1.5-2 倍）：约 15-18 GB 总磁盘占用

-- 不推荐：不指定维度
CREATE TABLE bad_design (
    embedding vector  -- 可以存储任意维度，但无法建索引
);

-- 推荐：明确指定维度
CREATE TABLE good_design (
    embedding vector(1536)  -- 固定 1536 维，可以建索引
);

1536 维向量：
- vector(1536)：4 × 1536 + 8 = 6,152 字节
- halfvec(1536)：2 × 1536 + 8 = 3,080 字节（节省 50%）

100 万条 1536 维向量：
- vector：约 6.15 GB
- halfvec：约 3.08 GB（节省约 3 GB 磁盘 + 内存）

-- 创建使用半精度向量的表
CREATE TABLE docs_halfvec (
    id bigserial PRIMARY KEY,
    content text,
    embedding halfvec(1536)
);

-- 插入数据（语法与 vector 完全相同）
INSERT INTO docs_halfvec (content, embedding) 
VALUES ('测试文档', '[0.1234, 0.5678, -0.9012]');

-- 距离查询（运算符相同）
SELECT content, 
       embedding <=> '[0.1, 0.5, -0.9]'::halfvec AS cosine_dist
FROM docs_halfvec
ORDER BY embedding <=> '[0.1, 0.5, -0.9]'::halfvec
LIMIT 5;

-- vector 与 halfvec 之间可以相互转换
SELECT '[1.5, 2.5, 3.5]'::vector(3)::halfvec(3);
-- 将 vector 转为 halfvec

-- 创建稀疏向量表
CREATE TABLE docs_sparse (
    id bigserial PRIMARY KEY,
    content text,
    -- 假设使用 BM25 特征，词汇表 10000 维，但每个文档只有约 100 个非零项
    sparse_embedding sparsevec(100, 10000)  
    -- sparsevec(最大非零元素数, 总维度)
);

-- 插入稀疏向量：10000维向量，只有第2、第5、第100位有值
INSERT INTO docs_sparse (content, sparse_embedding) 
VALUES (
    '这是一篇关于 PostgreSQL 的文档',
    '{2:0.85, 5:1.23, 100:0.67}/10000'::sparsevec
);

-- 稀疏向量的距离查询
SELECT content,
       sparse_embedding <-> '{2:0.9, 5:1.0, 50:0.5}/10000'::sparsevec AS l2_dist
FROM docs_sparse
ORDER BY sparse_embedding <-> '{2:0.9, 5:1.0, 50:0.5}/10000'::sparsevec
LIMIT 5;

使用 vector(10000) 稠密存储：4 × 10000 + 8 = 40,008 字节 ≈ 40 KB
使用 sparsevec(100, 10000) 稀疏存储：约 (4 + 4) × 100 + 头部 ≈ 800 + 头部 字节 ≈ 1 KB
存储节省约 97.5%！

1024 维向量：
- vector(1024)：4 × 1024 + 8 = 4,104 字节
- bit(1024)：1024 / 8 + 8 = 136 字节（节省约 97%）

1536 维向量：
- vector(1536)：6,152 字节
- bit(1536)：1536 / 8 + 8 = 200 字节（节省约 97%）

-- 创建二进制向量表
CREATE TABLE image_fingerprints (
    id bigserial PRIMARY KEY,
    image_url text,
    fingerprint bit(256)  -- 256位图片指纹
);

-- 插入二进制向量（使用 B'...' 格式或十六进制格式）
INSERT INTO image_fingerprints (image_url, fingerprint) VALUES
    ('https://example.com/photo1.jpg', B'101100101011001010110010101100101011001010110010101100101011001010110010101100101011001010110010101100101011001010110010101100101011001010110010101100101011001010110010101100101011001010110010101100101011001010110010101100101011001010110010101100101011001010110010'),
    ('https://example.com/photo2.jpg', B'1100110011001100110011001100110011001100110011001100110011001100110011001100110011001100110011001100110011001100110011001100110011001100110011001100110011001100110011001100110011001100110011001100110011001100110011001100110011001100110011001100110011001100');

-- 使用汉明距离检索相似图片
-- 汉明距离运算符：<~>（pgvector 0.7.0+）
SELECT image_url,
       fingerprint <~> B'1011001010110010101100101011001010110010101100101011001010110010101100101011001010110010101100101011001010110010101100101011001010110010101100101011001010110010101100101011001010110010101100101011001010110010101100101011001010110010'::bit(256) AS hamming_dist
FROM image_fingerprints
ORDER BY fingerprint <~> B'1011001010110010101100101011001010110010101100101011001010110010101100101011001010110010101100101011001010110010101100101011001010110010101100101011001010110010101100101011001010110010101100101011001010110010101100101011001010110010'::bit(256)
LIMIT 5;

-- 也可以使用 Jaccard 距离：<%>
SELECT image_url,
       fingerprint <%> B'10110...'::bit(256) AS jaccard_dist
FROM image_fingerprints
ORDER BY fingerprint <%> B'10110...'::bit(256)
LIMIT 5;

# 方式 1：Python 列表 → 字符串格式（最通用）
embedding_list = [0.1, 0.2, 0.3, 0.4, 0.5]
vector_str = str(embedding_list)  # "[0.1, 0.2, 0.3, 0.4, 0.5]"
# 在 SQL 中使用：INSERT ... VALUES (... , vector_str::vector)

# 方式 2：Numpy 数组 → 列表 → 字符串
import numpy as np
embedding_np = np.array([0.1, 0.2, 0.3, 0.4, 0.5], dtype=np.float32)
vector_str = str(embedding_np.tolist())  # "[0.1, 0.2, 0.3, 0.4, 0.5]"

# 方式 3：使用 pgvector-python 库的自动转换（推荐）
from pgvector.psycopg import register_vector
import psycopg

conn = psycopg.connect("postgresql://...")
register_vector(conn)  # 注册后，可以直接传入列表/numpy数组

cur = conn.cursor()
cur.execute(
    "INSERT INTO docs (embedding) VALUES (%s)",
    (embedding_list,)  # 直接传入 Python 列表，无需手动转换
)

-- 精度测试
SELECT '[0.123456789012345]'::vector(1);
-- 输出：[0.12345679]  （只保留了约 7 位有效数字）

-- 如果你的模型输出需要更高精度，考虑使用 double precision 数组辅助存储
-- 但注意：距离运算和索引仍然基于 float32

-- pgvector 不支持空向量（维度为 0）
SELECT '[]'::vector;  -- ERROR: vector must have at least 1 dimension

-- 如果需要表示"无向量"的状态，使用 NULL
CREATE TABLE docs (
    id bigserial PRIMARY KEY,
    content text,
    embedding vector(1536) DEFAULT NULL  -- 允许为空
);

-- 插入时不提供向量
INSERT INTO docs (content) VALUES ('待向量化的文档');
-- embedding 字段为 NULL

-- 查询时过滤掉 NULL 向量
SELECT * FROM docs WHERE embedding IS NOT NULL;

-- ============================================================
-- RAG 标准文档向量表设计
-- 适用于知识库、文档问答等场景
-- ============================================================

CREATE TABLE knowledge_base (
    -- 主键：自增长整型，用于唯一标识每条记录
    id bigserial PRIMARY KEY,

    -- 文档级字段
    doc_id text NOT NULL,                    -- 原始文档的唯一标识（如文件名、URL、数据库ID）
    title text,                              -- 文档标题

    -- 切片级字段
    chunk_id text NOT NULL,                  -- 切片的唯一标识（如 doc_id + chunk_index）
    chunk_index integer NOT NULL DEFAULT 0,  -- 切片在原文档中的序号
    content text NOT NULL,                   -- 切片文本内容
    content_length integer,                  -- 文本字符长度（冗余字段，便于查询过滤）

    -- 元数据字段（JSONB 提供灵活的键值存储）
    metadata jsonb DEFAULT '{}'::jsonb,      -- 灵活元数据，如：
    -- {
    --   "source": "company_handbook.pdf",
    --   "page": 42,
    --   "category": "HR",
    --   "language": "zh",
    --   "author": "张三",
    --   "tags": ["入职", "规章"]
    -- }

    -- 向量字段：存储 Embedding 向量
    embedding vector(1024),                  -- 1024 维向量（适配 BGE-M3/通义千问等模型）

    -- 审计字段
    created_at timestamptz NOT NULL DEFAULT now(),   -- 创建时间（带时区）
    updated_at timestamptz NOT NULL DEFAULT now()    -- 更新时间
);

-- 为常用查询字段创建索引
CREATE UNIQUE INDEX idx_knowledge_base_chunk_id ON knowledge_base (chunk_id);
CREATE INDEX idx_knowledge_base_doc_id ON knowledge_base (doc_id);
CREATE INDEX idx_knowledge_base_metadata ON knowledge_base USING GIN (metadata);
CREATE INDEX idx_knowledge_base_created_at ON knowledge_base (created_at);
CREATE INDEX idx_knowledge_base_category ON knowledge_base ((metadata->>'category'));

-- 添加表注释
COMMENT ON TABLE knowledge_base IS 'RAG 知识库文档向量表';
COMMENT ON COLUMN knowledge_base.embedding IS '文档切片的 Embedding 向量，1024 维';
COMMENT ON COLUMN knowledge_base.metadata IS 'JSONB 格式的灵活元数据';

-- 创建一个自动更新 updated_at 的触发器
CREATE OR REPLACE FUNCTION update_updated_at_column()
RETURNS TRIGGER AS $$
BEGIN
    NEW.updated_at = now();
    RETURN NEW;
END;
$$ LANGUAGE plpgsql;

CREATE TRIGGER tr_knowledge_base_updated_at
    BEFORE UPDATE ON knowledge_base
    FOR EACH ROW
    EXECUTE FUNCTION update_updated_at_column();

-- 基本单行插入
INSERT INTO knowledge_base (doc_id, chunk_id, chunk_index, title, content, metadata, embedding)
VALUES (
    'doc_001',
    'doc_001_chunk_0',
    0,
    'PostgreSQL 入门指南',
    'PostgreSQL 是一个强大的开源关系型数据库管理系统，以其可靠性、功能丰富和性能著称。',
    '{"source": "pg_guide.pdf", "page": 1, "category": "database", "language": "zh"}'::jsonb,
    '[0.023, -0.015, 0.087, 0.042, -0.033]'::vector  -- 实际应为 1024 维，此处简化
);

-- 插入后获取生成的 id
INSERT INTO knowledge_base (doc_id, chunk_id, content, embedding)
VALUES ('doc_002', 'doc_002_chunk_0', 'pgvector 支持多种向量数据类型', '[0.1, 0.2, 0.3]'::vector)
RETURNING id, chunk_id;

-- 批量插入多个切片（同一文档的连续切片）
INSERT INTO knowledge_base (doc_id, chunk_id, chunk_index, title, content, metadata, embedding)
VALUES 
    ('doc_001', 'doc_001_chunk_1', 1, 'PostgreSQL 入门指南',
     'PostgreSQL 支持丰富的数据类型，包括整数、浮点数、字符串、JSON、数组等。',
     '{"source": "pg_guide.pdf", "page": 2, "category": "database"}'::jsonb,
     '[0.031, -0.022, 0.095, 0.051, -0.041]'::vector),

    ('doc_001', 'doc_001_chunk_2', 2, 'PostgreSQL 入门指南',
     'PostgreSQL 的索引机制非常强大，支持 B-Tree、Hash、GiST、GIN 等多种索引类型。',
     '{"source": "pg_guide.pdf", "page": 3, "category": "database"}'::jsonb,
     '[0.045, -0.018, 0.073, 0.063, -0.029]'::vector),

    ('doc_001', 'doc_001_chunk_3', 3, 'PostgreSQL 入门指南',
     '事务是 PostgreSQL 的核心特性之一，支持完整的 ACID 保证。',
     '{"source": "pg_guide.pdf", "page": 4, "category": "database"}'::jsonb,
     '[0.052, -0.011, 0.081, 0.038, -0.045]'::vector);

-- 准备 CSV 文件：vectors_data.csv
-- 格式：doc_id,chunk_id,chunk_index,content,embedding
-- 内容示例：
-- doc_003,doc_003_chunk_0,0,"第一条文档内容","[0.1, 0.2, 0.3, ...]"
-- doc_003,doc_003_chunk_1,1,"第二条文档内容","[0.4, 0.5, 0.6, ...]"

-- 使用 COPY 命令导入（在 psql 中执行）
\COPY knowledge_base (doc_id, chunk_id, chunk_index, content, embedding) 
FROM '/path/to/vectors_data.csv' 
WITH (FORMAT csv, HEADER true);

-- 或使用 PostgreSQL 的 COPY 命令（需要超级用户权限）
COPY knowledge_base (doc_id, chunk_id, chunk_index, content, embedding) 
FROM '/path/to/vectors_data.csv' 
WITH (FORMAT csv, HEADER true);

import psycopg
from pgvector.psycopg import register_vector
import numpy as np

conn = psycopg.connect("postgresql://vector_user:vector_pass@localhost:5432/vector_db")
register_vector(conn)
cur = conn.cursor()

# 模拟 1000 条文档数据
documents = []
for i in range(1000):
    embedding = np.random.rand(1024).astype(np.float32)  # 模拟 1024 维向量
    documents.append((
        f'doc_{i//10:04d}',           # doc_id
        f'doc_{i//10:04d}_chunk_{i%10}',  # chunk_id
        i % 10,                        # chunk_index
        f'这是第 {i} 个文档切片的文本内容',  # content
        embedding                      # embedding
    ))

# 方式 1：executemany 批量插入
cur.executemany("""
    INSERT INTO knowledge_base (doc_id, chunk_id, chunk_index, content, embedding)
    VALUES (%s, %s, %s, %s, %s)
""", documents)
conn.commit()

# 方式 2：使用 COPY（更快，适合 10 万级以上数据）
from psycopg import sql
with cur.copy("COPY knowledge_base (doc_id, chunk_id, chunk_index, content, embedding) FROM STDIN") as copy:
    for doc in documents:
        # COPY 需要将向量转为字符串格式
        vec_str = '[' + ','.join(str(x) for x in doc[4]) + ']'
        copy.write_row((doc[0], doc[1], doc[2], doc[3], vec_str))
conn.commit()

print(f"成功插入 {len(documents)} 条记录")
cur.close()
conn.close()

-- 1. 查看前 5 条记录（包括向量字段）
SELECT id, chunk_id, title, content, embedding
FROM knowledge_base
LIMIT 5;

-- 输出示例：
-- id | chunk_id           | title              | content              | embedding
--  1 | doc_001_chunk_0    | PostgreSQL 入门指南  | PostgreSQL 是一个...  | [0.023, -0.015, ...]

-- 2. 查看向量的维度
SELECT id, chunk_id, vector_dims(embedding) AS dims
FROM knowledge_base
WHERE embedding IS NOT NULL
LIMIT 5;

-- 输出：
-- id | chunk_id           | dims
--  1 | doc_001_chunk_0    | 1024

-- 3. 格式化查看向量的部分维度（只看前 5 个元素）
-- 注意：pgvector 没有内置的切片语法，需要通过转换为数组实现
SELECT id, chunk_id, 
       embedding::text AS embedding_text
FROM knowledge_base
WHERE id = 1;

-- 4. 统计各文档的切片数量
SELECT doc_id, title, COUNT(*) AS chunk_count
FROM knowledge_base
GROUP BY doc_id, title
ORDER BY chunk_count DESC;

-- 5. 查看向量的 L2 范数（模长）
SELECT id, chunk_id, 
       embedding <-> '0'::vector AS l2_norm  -- 与零向量的距离即模长
       -- 注意：需要维度匹配，'0' 不正确，应该用全零的 1024 维向量
FROM knowledge_base
WHERE id = 1;

-- 正确计算向量模长（使用 vector 的算术运算）
-- pgvector 没有直接的 norm 函数，但可以通过 l2_distance(zero_vector) 计算
-- 或者使用 sqrt(embedding <#> embedding 的负值)，不过最直接的方式是：
SELECT id, chunk_id,
       sqrt((embedding <#> embedding) * -1) AS l2_norm
FROM knowledge_base
WHERE id = 1;
-- 原理：向量与自身的内积 = ||v||²，负内积 = -||v||²，取负后开方 = ||v||

-- 6. 查看向量字段的统计信息
SELECT 
    COUNT(*) AS total_rows,
    COUNT(embedding) AS non_null_vectors,
    COUNT(*) - COUNT(embedding) AS null_vectors
FROM knowledge_base;

-- ============================================================
-- 更新操作
-- ============================================================

-- 1. 单独更新向量字段（如重新生成 Embedding 后回填）
UPDATE knowledge_base
SET embedding = '[0.099, -0.088, 0.077, 0.066, -0.055]'::vector,
    updated_at = now()
WHERE chunk_id = 'doc_001_chunk_0';

-- 2. 批量更新某个文档的所有切片向量（如更换 Embedding 模型后）
UPDATE knowledge_base
SET embedding = NULL,        -- 先清空旧向量
    updated_at = now()
WHERE doc_id = 'doc_001';
-- 然后异步重新生成向量并回填

-- 3. 更新元数据
UPDATE knowledge_base
SET metadata = metadata || '{"reviewed": true, "reviewer": "李四"}'::jsonb
WHERE doc_id = 'doc_001';

-- ============================================================
-- 删除操作
-- ============================================================

-- 4. 按结构化条件删除（如删除某个文档的所有切片）
DELETE FROM knowledge_base
WHERE doc_id = 'doc_001';
-- 这个操作会删除 doc_001 的所有切片记录，包括它们的向量数据

-- 5. 按元数据条件删除
DELETE FROM knowledge_base
WHERE metadata->>'category' = 'deprecated';

-- 6. 按时间条件删除（如清理过期数据）
DELETE FROM knowledge_base
WHERE created_at < '2025-01-01'::timestamptz;

-- ============================================================
-- 事务一致性演示
-- ============================================================

-- 场景：更新文档元数据的同时更新向量，保证原子性
BEGIN;

-- 步骤 1：更新文档元数据
UPDATE knowledge_base
SET metadata = jsonb_set(metadata, '{version}', '"2.0"')
WHERE doc_id = 'doc_002';

-- 步骤 2：更新对应的向量
UPDATE knowledge_base
SET embedding = '[0.111, 0.222, 0.333, 0.444, 0.555]'::vector
WHERE chunk_id = 'doc_002_chunk_0';

-- 如果任何一步出错，整个事务回滚
COMMIT;
-- 如果出错：ROLLBACK;

-- 演示回滚
BEGIN;
    -- 尝试插入一条重复 chunk_id 的记录（会触发唯一索引冲突）
    INSERT INTO knowledge_base (doc_id, chunk_id, content, embedding)
    VALUES ('doc_999', 'doc_001_chunk_0', '冲突测试', '[0,0,0]'::vector);
    -- ERROR: duplicate key value violates unique constraint
    -- 此时事务中的所有操作都不会被提交
ROLLBACK;
-- 数据库状态与 BEGIN 之前完全一致

-- ============================================================
-- 准备测试数据
-- ============================================================
CREATE TABLE demo_vectors (
    id serial PRIMARY KEY,
    label text,
    embedding vector(5)
);

INSERT INTO demo_vectors (label, embedding) VALUES
    ('A: 原点附近',   '[0.0, 0.0, 0.0, 0.0, 0.0]'),
    ('B: 正方向',     '[1.0, 1.0, 1.0, 1.0, 1.0]'),
    ('C: 负方向',     '[-1.0, -1.0, -1.0, -1.0, -1.0]'),
    ('D: 同向但模长不同', '[2.0, 2.0, 2.0, 2.0, 2.0]'),
    ('E: 正交方向',   '[1.0, 0.0, 0.0, 0.0, 0.0]');

-- 查询向量：使用 B 向量作为查询基准
-- ============================================================

-- 1. 欧氏距离（L2 Distance）<->
-- 计算的是空间中的直线距离，值越小越相似
SELECT 
    label,
    embedding,
    embedding <-> '[1.0, 1.0, 1.0, 1.0, 1.0]'::vector AS l2_distance
FROM demo_vectors
ORDER BY embedding <-> '[1.0, 1.0, 1.0, 1.0, 1.0]'::vector;

-- 预期结果：
-- B: 正方向          | [1,1,1,1,1]     | 0.000  （完全匹配，距离为0）
-- D: 同向但模长不同    | [2,2,2,2,2]     | 2.236  （方向相同但距离较远）
-- E: 正交方向        | [1,0,0,0,0]     | 2.000  （只有第一个维度匹配）
-- A: 原点附近        | [0,0,0,0,0]     | 2.236  （对角线距离）
-- C: 负方向          | [-1,-1,-1,-1,-1] | 4.472  （最远）

-- 2. 余弦距离 <=>
-- 计算的是方向差异，忽略模长，值越小方向越一致
SELECT 
    label,
    embedding,
    embedding <=> '[1.0, 1.0, 1.0, 1.0, 1.0]'::vector AS cosine_distance
FROM demo_vectors
WHERE id != 1  -- 排除零向量（零向量无法计算余弦距离）
ORDER BY embedding <=> '[1.0, 1.0, 1.0, 1.0, 1.0]'::vector;

-- 预期结果：
-- B: 正方向          | 0.000  （方向完全一致）
-- D: 同向但模长不同    | 0.000  （方向也完全一致！余弦距离忽略模长）
-- E: 正交方向        | 0.553  （部分维度同向）
-- C: 负方向          | 2.000  （方向完全相反）

-- 注意：B 和 D 的余弦距离都是 0，因为 [2,2,2,2,2] 与 [1,1,1,1,1] 方向完全相同

-- 3. 负内积 <#>
-- 内积同时考虑方向和模长，pgvector 返回的是负内积，值越小（越负）表示越相似
SELECT 
    label,
    embedding,
    embedding <#> '[1.0, 1.0, 1.0, 1.0, 1.0]'::vector AS neg_inner_product
FROM demo_vectors
ORDER BY embedding <#> '[1.0, 1.0, 1.0, 1.0, 1.0]'::vector;

-- 预期结果：
-- D: 同向但模长不同    | -10.000  （内积=10，最相似，因为模长最大且方向一致）
-- B: 正方向          | -5.000   （内积=5）
-- A: 原点附近        | 0.000    （内积=0）
-- E: 正交方向        | -1.000   （内积=1）
-- C: 负方向          | 5.000    （内积=-5，方向相反）

-- 注意排序：负内积值越小（越负）排越前面，表示内积越大越相似

查询向量 Q = [1, 1, 1, 1, 1]，对比向量 B = [1, 1, 1, 1, 1] 和 D = [2, 2, 2, 2, 2]

- 欧氏距离：B(0.0) < D(2.236)   → B 更近（因为空间距离更近）
- 余弦距离：B(0.0) = D(0.0)      → 一样近（因为方向相同）
- 负内积：  D(-10) < B(-5)       → D 更相似（因为方向相同且模长更大）

-- ============================================================
-- Top-K 语义相似度检索
-- ============================================================

-- 基本 Top-K 查询：找出与查询向量最相似的 5 条记录
-- 假设查询向量为 [0.03, -0.02, 0.08, ...]（1024 维，此处简化为 5 维演示）
SELECT 
    id,
    doc_id,
    chunk_id,
    title,
    content,
    embedding <=> '[0.03, -0.02, 0.08, 0.04, -0.03]'::vector AS cosine_distance
FROM knowledge_base
WHERE embedding IS NOT NULL
ORDER BY embedding <=> '[0.03, -0.02, 0.08, 0.04, -0.03]'::vector
LIMIT 5;

-- ============================================================
-- 带相似度阈值过滤的 Top-K
-- ============================================================

-- 只返回余弦距离小于 0.3 的结果（即相似度大于 70%）
SELECT 
    id,
    chunk_id,
    content,
    embedding <=> '[0.03, -0.02, 0.08, 0.04, -0.03]'::vector AS cosine_distance,
    1 - (embedding <=> '[0.03, -0.02, 0.08, 0.04, -0.03]'::vector) AS similarity_score
FROM knowledge_base
WHERE embedding IS NOT NULL
  AND embedding <=> '[0.03, -0.02, 0.08, 0.04, -0.03]'::vector < 0.3  -- 阈值过滤
ORDER BY embedding <=> '[0.03, -0.02, 0.08, 0.04, -0.03]'::vector
LIMIT 10;

-- ============================================================
-- 使用 CTE（Common Table Expression）使查询更清晰
-- ============================================================

WITH query AS (
    -- 定义查询向量（可以来自函数参数或子查询）
    SELECT '[0.03, -0.02, 0.08, 0.04, -0.03]'::vector AS q_vec
),
ranked AS (
    SELECT 
        kb.id,
        kb.chunk_id,
        kb.content,
        kb.metadata,
        kb.embedding <=> query.q_vec AS cosine_dist,
        1 - (kb.embedding <=> query.q_vec) AS similarity
    FROM knowledge_base kb, query
    WHERE kb.embedding IS NOT NULL
    ORDER BY kb.embedding <=> query.q_vec
    LIMIT 10
)
SELECT * FROM ranked
WHERE similarity > 0.7  -- 在 Top-10 中进一步筛选相似度 > 70% 的结果
ORDER BY similarity DESC;

-- ============================================================
-- 场景 1：在特定分类中进行语义检索
-- ============================================================

-- 只在"数据库"分类的文档中搜索与"索引优化"语义最相关的 5 条记录
SELECT 
    id,
    chunk_id,
    title,
    content,
    metadata->>'category' AS category,
    1 - (embedding <=> '[0.05, 0.03, -0.02, 0.08, 0.01]'::vector) AS similarity
FROM knowledge_base
WHERE embedding IS NOT NULL
  AND metadata->>'category' = 'database'     -- 结构化过滤：分类为 database
ORDER BY embedding <=> '[0.05, 0.03, -0.02, 0.08, 0.01]'::vector
LIMIT 5;

-- ============================================================
-- 场景 2：时间范围 + 语义检索
-- ============================================================

-- 在最近 30 天新增的文档中搜索
SELECT 
    id,
    chunk_id,
    title,
    content,
    created_at,
    1 - (embedding <=> '[0.05, 0.03, -0.02, 0.08, 0.01]'::vector) AS similarity
FROM knowledge_base
WHERE embedding IS NOT NULL
  AND created_at >= now() - interval '30 days'  -- 时间范围过滤
ORDER BY embedding <=> '[0.05, 0.03, -0.02, 0.08, 0.01]'::vector
LIMIT 5;

-- ============================================================
-- 场景 3：多条件联合过滤 + 语义检索
-- ============================================================

-- 在中文文档中，搜索 HR 分类下最近审核过的内容
SELECT 
    id,
    chunk_id,
    title,
    content,
    metadata->>'category' AS category,
    metadata->>'language' AS language,
    1 - (embedding <=> '[0.02, 0.06, -0.01, 0.04, 0.03]'::vector) AS similarity
FROM knowledge_base
WHERE embedding IS NOT NULL
  AND metadata->>'language' = 'zh'               -- 语言为中文
  AND metadata->>'category' = 'HR'               -- 分类为 HR
  AND metadata->>'reviewed' = 'true'             -- 已审核
ORDER BY embedding <=> '[0.02, 0.06, -0.01, 0.04, 0.03]'::vector
LIMIT 5;

-- ============================================================
-- 场景 4：先向量检索后结构化过滤（Pre-filtering vs Post-filtering）
-- ============================================================

-- 方式 A：Post-filtering（先向量检索取 Top-50，再从中过滤）
-- 适用于：过滤条件选择性较高（大部分记录满足条件）的情况
WITH vector_results AS (
    SELECT 
        id, chunk_id, content, metadata,
        1 - (embedding <=> '[0.05, 0.03, -0.02, 0.08, 0.01]'::vector) AS similarity
    FROM knowledge_base
    WHERE embedding IS NOT NULL
    ORDER BY embedding <=> '[0.05, 0.03, -0.02, 0.08, 0.01]'::vector
    LIMIT 50  -- 先取较大的 Top-K
)
SELECT * FROM vector_results
WHERE metadata->>'category' = 'database'  -- 再过滤
ORDER BY similarity DESC
LIMIT 5;

-- 方式 B：Pre-filtering（WHERE 子句同时包含结构化和向量条件）
-- 适用于：PostgreSQL 查询优化器会自动选择最优执行计划
SELECT 
    id, chunk_id, content, metadata,
    1 - (embedding <=> '[0.05, 0.03, -0.02, 0.08, 0.01]'::vector) AS similarity
FROM knowledge_base
WHERE embedding IS NOT NULL
  AND metadata->>'category' = 'database'
ORDER BY embedding <=> '[0.05, 0.03, -0.02, 0.08, 0.01]'::vector
LIMIT 5;

-- 建议：对于大多数场景，直接使用方式 B（让 PostgreSQL 优化器决定执行计划）
-- 如果发现过滤后结果太少，可以增大向量检索的候选集（在子查询中取更大的 LIMIT）

-- ============================================================
-- 余弦相似度分数 = 1 - 余弦距离
-- ============================================================

-- 余弦距离范围：[0, 2]
-- 余弦相似度 = 1 - 余弦距离，范围：[-1, 1]
-- 对于归一化向量，余弦距离范围：[0, 1]（不会出现负值）

-- 基本用法
SELECT 
    id,
    chunk_id,
    content,
    embedding <=> '[0.03, -0.02, 0.08, 0.04, -0.03]'::vector AS cosine_distance,
    1 - (embedding <=> '[0.03, -0.02, 0.08, 0.04, -0.03]'::vector) AS cosine_similarity
FROM knowledge_base
WHERE embedding IS NOT NULL
ORDER BY cosine_similarity DESC
LIMIT 5;

-- 输出示例：
-- cosine_distance | cosine_similarity | 含义
-- 0.05            | 0.95             | 非常相似
-- 0.15            | 0.85             | 高度相似
-- 0.35            | 0.65             | 中度相似
-- 0.60            | 0.40             | 较低相似
-- 0.85            | 0.15             | 不太相似

-- ============================================================
-- 将相似度映射到百分制（更直观）
-- ============================================================

SELECT 
    id,
    chunk_id,
    content,
    ROUND(
        (1 - (embedding <=> '[0.03, -0.02, 0.08, 0.04, -0.03]'::vector)) * 100, 
        2
    ) AS similarity_percent
FROM knowledge_base
WHERE embedding IS NOT NULL
ORDER BY similarity_percent DESC
LIMIT 5;

-- 输出示例：
-- similarity_percent | 含义
-- 95.23             | 95.23% 相似
-- 85.67             | 85.67% 相似

-- ============================================================
-- 欧氏距离转相似度（需要归一化）
-- ============================================================

-- 欧氏距离没有固定的上界，无法直接转为 0-1 相似度
-- 常用的转换方式：similarity = 1 / (1 + distance)
SELECT 
    id,
    chunk_id,
    embedding <-> '[0.03, -0.02, 0.08, 0.04, -0.03]'::vector AS l2_distance,
    1.0 / (1.0 + (embedding <-> '[0.03, -0.02, 0.08, 0.04, -0.03]'::vector)) AS l2_similarity
FROM knowledge_base
WHERE embedding IS NOT NULL
ORDER BY l2_similarity DESC
LIMIT 5;

-- ============================================================
-- 内积转相似度
-- ============================================================

-- 对于归一化向量，负内积 = -cosine_similarity
-- 所以 cosine_similarity = -(neg_inner_product) = inner_product
SELECT 
    id,
    chunk_id,
    embedding <#> '[0.03, -0.02, 0.08, 0.04, -0.03]'::vector AS neg_inner_product,
    -(embedding <#> '[0.03, -0.02, 0.08, 0.04, -0.03]'::vector) AS inner_product_score
FROM knowledge_base
WHERE embedding IS NOT NULL
ORDER BY inner_product_score DESC
LIMIT 5;

-- ============================================================
-- 生产级 RAG 检索模板
-- 参数：$1 = 查询向量, $2 = 分类过滤(可选), $3 = Top-K
-- ============================================================

WITH search_results AS (
    SELECT 
        kb.id,
        kb.doc_id,
        kb.chunk_id,
        kb.title,
        kb.content,
        kb.metadata,
        1 - (kb.embedding <=> $1::vector) AS similarity_score,
        kb.created_at
    FROM knowledge_base kb
    WHERE kb.embedding IS NOT NULL
      AND ($2::text IS NULL OR kb.metadata->>'category' = $2::text)  -- 可选分类过滤
    ORDER BY kb.embedding <=> $1::vector
    LIMIT $3::integer
)
SELECT 
    id,
    doc_id,
    chunk_id,
    title,
    content,
    metadata,
    ROUND(similarity_score::numeric, 4) AS similarity_score,
    created_at
FROM search_results
WHERE similarity_score >= 0.5  -- 最低相似度阈值
ORDER BY similarity_score DESC;

ERROR:  different vector dimensions, in [1024] and [1536]

-- 表定义为 1024 维
CREATE TABLE docs (embedding vector(1024));

-- 错误：插入 1536 维向量
INSERT INTO docs (embedding) VALUES ('[0.1, 0.2, ...（共1536个值）]'::vector);
-- ERROR: different vector dimensions, in [1024] and [1536]

-- 错误：查询向量维度不匹配
SELECT * FROM docs ORDER BY embedding <=> '[0.1, 0.2, 0.3]'::vector LIMIT 5;
-- ERROR: different vector dimensions, in [1024] and [3]

-- 1. 确认表的向量维度
SELECT column_name, udt_name, character_maximum_length
FROM information_schema.columns
WHERE table_name = 'docs' AND column_name = 'embedding';

-- 或者使用 pgvector 的 vector_dims 函数
SELECT vector_dims(embedding) FROM docs LIMIT 1;

-- 2. 确保插入和查询的向量维度与定义一致
-- 如果确实需要改变维度（如更换了 Embedding 模型），需要重建表：
-- 步骤 A：新增列
ALTER TABLE docs ADD COLUMN embedding_new vector(1536);
-- 步骤 B：重新生成所有向量并填入新列
-- 步骤 C：删除旧列，重命名新列
ALTER TABLE docs DROP COLUMN embedding;
ALTER TABLE docs RENAME COLUMN embedding_new TO embedding;
-- 步骤 D：重建索引

ERROR:  type "vector" does not exist
LINE 1: CREATE TABLE docs (embedding vector(1024));
                                     ^

-- 1. 首先启用扩展（需要数据库 CREATE 权限）
CREATE EXTENSION IF NOT EXISTS vector;

-- 2. 如果提示 "permission denied"，使用超级用户：
-- sudo -u postgres psql -d your_database -c "CREATE EXTENSION vector;"

-- 3. 如果提示 "could not open extension control file"，说明 pgvector 没有安装
-- 需要先按照 1.2 节的步骤安装 pgvector

-- 4. 验证扩展是否启用
SELECT extname, extversion FROM pg_extension WHERE extname = 'vector';

-- 注意：CREATE EXTENSION 是针对单个数据库的
-- 如果你在 database_A 中启用了，database_B 中仍然需要单独启用
-- 如果希望所有新数据库自动启用，可以在 template1 中启用：
-- sudo -u postgres psql -d template1 -c "CREATE EXTENSION vector;"

ERROR:  invalid input syntax for type vector: "1, 2, 3"

-- 错误：缺少方括号
SELECT '1, 2, 3'::vector;
-- ERROR: invalid input syntax

-- 正确：必须有方括号
SELECT '[1, 2, 3]'::vector;

-- 错误：使用了花括号（这是 PostgreSQL 数组的格式）
SELECT '{1, 2, 3}'::vector;
-- ERROR: invalid input syntax

-- 错误：使用了圆括号
SELECT '(1, 2, 3)'::vector;
-- ERROR: invalid input syntax

-- 错误：包含非数字字符
SELECT '[1, 2, abc]'::vector;
-- ERROR: invalid input syntax

-- 错误：维度为 0（空向量）
SELECT '[]'::vector;
-- ERROR: vector must have at least 1 dimension

-- 正确：标准格式
SELECT '[1.0, 2.0, 3.0]'::vector;          -- 浮点数
SELECT '[1, 2, 3]'::vector;                -- 整数（自动转为 float32）
SELECT '[1e-5, 2.5e3, -0.001]'::vector;    -- 科学计数法

-- 正确：指定维度
SELECT '[1, 2, 3]'::vector(3);             -- 明确指定 3 维

-- 错误：指定维度与实际不匹配
SELECT '[1, 2, 3]'::vector(5);
-- ERROR: different vector dimensions, in [5] and [3]

# 错误：直接传入 numpy 数组（某些驱动不支持）
import numpy as np
embedding = np.array([0.1, 0.2, 0.3])
cur.execute("INSERT INTO docs (embedding) VALUES (%s)", (embedding,))
# 可能报错

# 正确：转为 Python 列表
cur.execute("INSERT INTO docs (embedding) VALUES (%s::vector)", (embedding.tolist(),))

# 正确：转为字符串格式
vector_str = '[' + ','.join(str(x) for x in embedding) + ']'
cur.execute("INSERT INTO docs (embedding) VALUES (%s::vector)", (vector_str,))

# 最推荐：使用 pgvector-python 注册自动转换
from pgvector.psycopg import register_vector
register_vector(conn)
cur.execute("INSERT INTO docs (embedding) VALUES (%s)", (embedding.tolist(),))

-- 现象：插入的向量值与查询出来的不完全一致
INSERT INTO test_vec (embedding) VALUES ('[0.123456789012345678]'::vector(1));
SELECT embedding FROM test_vec;
-- 结果：[0.12345679]  （精度被截断到约 7 位有效数字）

-- 这不是 bug！vector 类型使用 float32（单精度浮点），精度约为 7 位十进制数字
-- float32 能精确表示的范围：±1.17549435e-38 到 ±3.40282347e+38

-- 如果确实需要更高精度（不推荐，会增加存储和计算成本），可以：
-- 1. 将高精度值存储为额外的 text 列（仅用于审计/追溯）
-- 2. 接受 float32 精度限制（绝大多数 AI 应用场景足够）

-- 精度对检索的影响测试
SELECT 
    '[0.1234567]'::vector(1) <-> '[0.1234568]'::vector(1) AS diff_7th_digit;
-- 结果：约 1e-7，对检索排序几乎没有影响

-- 结论：float32 的精度对语义检索完全够用，不必纠结精度损失

ERROR:  invalid byte sequence for encoding "UTF8": 0xXX

-- 1. 检查数据库编码
SHOW server_encoding;
-- 应为 UTF8

-- 2. 检查客户端编码
SHOW client_encoding;
-- 应为 UTF8

-- 3. 如果数据库不是 UTF-8，需要重建数据库（创建时指定）
CREATE DATABASE vector_db 
    ENCODING 'UTF8' 
    LC_COLLATE 'en_US.UTF-8' 
    LC_CTYPE 'en_US.UTF-8'
    TEMPLATE template0;

-- 4. 如果客户端编码不匹配，在连接时指定
-- psql "host=localhost dbname=vector_db user=vector_user client_encoding=UTF8"

-- 5. 在 Python 中确保使用 UTF-8 连接
import psycopg
conn = psycopg.connect(
    "host=localhost dbname=vector_db user=vector_user",
    client_encoding="UTF8"
)

-- 6. 处理特殊字符
-- 确保 Embedding 模型返回的文本内容使用 UTF-8 编码
-- 在写入前清理文本中的非法字符
import re
def clean_text(text: str) -> str:
    # 移除非法 UTF-8 字符
    return text.encode('utf-8', errors='ignore').decode('utf-8')

ERROR:  column "embedding" does not have 64-bit aligned values

ERROR:  index row size XXXX exceeds maximum 2712 for index "xxx"

-- HNSW 索引对于高维向量的限制：
-- vector(n) 的 HNSW 索引：每个维度需要约 4 字节存储，加上图结构的指针开销
-- 通常 vector(2000) 以上的维度需要注意索引行大小

-- 解决方案 1：使用 halfvec 类型减少存储（2 字节/维）
ALTER TABLE docs ALTER COLUMN embedding TYPE halfvec(2000);

-- 解决方案 2：对于超高维向量，使用 IVFFlat 索引（对行大小限制更宽松）
CREATE INDEX ON docs USING ivfflat (embedding vector_cosine_ops) WITH (lists = 100);

-- 解决方案 3：降维（如果 Embedding 模型支持）
-- OpenAI text-embedding-3-* 支持维度截断，可以直接请求 1024 维而非 3072 维

-- 解决方案 4：拆分向量列（不推荐，仅作为最后手段）
-- 将 3072 维拆分为两个 1536 维的列

-- 1. pgvector 扩展是否已启用？
SELECT extname, extversion FROM pg_extension WHERE extname = 'vector';

-- 2. 表的向量维度是多少？
SELECT attname, atttypid::regtype 
FROM pg_attribute 
WHERE attrelid = 'knowledge_base'::regclass 
  AND atttypid::regtype::text LIKE 'vector%';

-- 3. 查询向量维度是否匹配？
SELECT vector_dims('[0.1, 0.2, 0.3]'::vector);  -- 应为表的定义维度

-- 4. 是否有 NULL 向量？
SELECT COUNT(*) AS null_count 
FROM knowledge_base 
WHERE embedding IS NULL;

-- 5. 数据库中是否有不同维度的向量混存？
SELECT vector_dims(embedding), COUNT(*) 
FROM knowledge_base 
WHERE embedding IS NOT NULL 
GROUP BY vector_dims(embedding);
-- 如果输出多行，说明存在维度不一致的问题

-- 6. 索引状态检查
SELECT indexname, indexdef 
FROM pg_indexes 
WHERE tablename = 'knowledge_base';

-- 7. 表和索引大小
SELECT 
    pg_size_pretty(pg_total_relation_size('knowledge_base')) AS total_size,
    pg_size_pretty(pg_relation_size('knowledge_base')) AS table_size,
    pg_size_pretty(pg_indexes_size('knowledge_base')) AS index_size;

对于查询向量 q，逐一计算它与表中每一个向量 v_i 的距离 d(q, v_i)，
然后将所有距离排序，返回距离最小的 Top-K 个结果。

EXPLAIN ANALYZE
SELECT *, embedding <=> '[0.1, 0.2, ...]' AS distance
FROM items
ORDER BY distance
LIMIT 10;

-- 典型的执行计划输出：
-- Limit  (cost=... rows=10)
--   ->  Sort  (cost=...)
--         Sort Key: (embedding <=> '[...]'::vector)
--         ->  Seq Scan on items  (cost=...)

-- 生成测试数据
CREATE TABLE perf_test (
    id serial PRIMARY KEY,
    embedding vector(768)
);

-- 插入10万条随机向量
INSERT INTO perf_test (embedding)
SELECT random_vector(768)
FROM generate_series(1, 100000);

-- 辅助函数：生成随机向量
CREATE OR REPLACE FUNCTION random_vector(dim int)
RETURNS vector
LANGUAGE plpgsql AS $$
DECLARE
    arr float8[];
    i int;
BEGIN
    arr := ARRAY[]::float8[];
    FOR i IN 1..dim LOOP
        arr := arr || random();
    END LOOP;
    RETURN arr::vector;
END;
$$;

-- 无索引测试
EXPLAIN ANALYZE
SELECT id, embedding <=> (SELECT embedding FROM perf_test LIMIT 1) AS dist
FROM perf_test
ORDER BY dist
LIMIT 10;

输入：查询向量 q，参数 ef_search
1. 从最高层的入口节点开始
2. 对于每一层（从最高层到第1层）：
   - 贪心搜索，找到该层中距离 q 最近的节点
3. 在第0层执行 beam search：
   - 维护一个大小为 ef_search 的候选优先队列
   - 维护一个大小为 K 的结果列表
   - 不断从候选队列中取出最近节点，探索其邻居
   - 如果邻居比当前结果列表中最远的更近，则加入结果
4. 返回结果列表中的 Top-K

-- 调整 HNSW 搜索时的候选集大小（会话级别）
SET hnsw.ef_search = 100;  -- 默认值为 40

-- 使用余弦距离创建 HNSW 索引（最常用）
CREATE INDEX ON items USING hnsw (embedding vector_cosine_ops)
WITH (m = 16, ef_construction = 64);

-- 使用 L2（欧氏）距离创建 HNSW 索引
CREATE INDEX ON items USING hnsw (embedding vector_l2_ops)
WITH (m = 16, ef_construction = 64);

-- 使用内积距离创建 HNSW 索引
CREATE INDEX ON items USING hnsw (embedding vector_ip_ops)
WITH (m = 16, ef_construction = 64);

-- 高质量索引配置（适合生产环境）
CREATE INDEX ON items USING hnsw (embedding vector_cosine_ops)
WITH (m = 32, ef_construction = 128);

-- 索引创建后查询时自动生效
SELECT *, embedding <=> '[0.1, 0.2, ...]' AS distance
FROM items
ORDER BY distance
LIMIT 10;

输入：查询向量 q，参数 nprobe
1. 计算 q 到所有聚类中心的距离
2. 选出距离最近的 nprobe 个聚类中心
3. 在这 nprobe 个聚类对应的倒排列表中，逐一计算 q 与每个向量的距离
4. 合并所有结果，返回 Top-K

-- 设置查询时探测的聚类数（会话级别）
SET ivfflat.probes = 10;  -- 默认值为 1

-- 使用余弦距离创建 IVFFlat 索引
-- 注意：必须确保表中已有足够数据
CREATE INDEX ON items USING ivfflat (embedding vector_cosine_ops)
WITH (lists = 100);

-- 使用 L2 距离创建 IVFFlat 索引
CREATE INDEX ON items USING ivfflat (embedding vector_l2_ops)
WITH (lists = 100);

-- 百万级数据的推荐配置
CREATE INDEX ON items USING ivfflat (embedding vector_cosine_ops)
WITH (lists = 1000);

-- 设置查询参数
SET ivfflat.probes = 10;

-- 查询
SELECT *, embedding <=> '[0.1, 0.2, ...]' AS distance
FROM items
ORDER BY distance
LIMIT 10;

-- 无索引查询（暴力检索），无需任何额外配置
SELECT *, embedding <=> '[0.1, 0.2, ...]' AS distance
FROM items
ORDER BY distance
LIMIT 10;

-- 场景一：通用平衡配置（推荐大多数场景）
CREATE INDEX idx_embedding_hnsw ON items
USING hnsw (embedding vector_cosine_ops)
WITH (m = 16, ef_construction = 64);

-- 场景二：高召回率配置（推荐对精度要求高的场景）
-- 更高的 m 值带来更好的图连通性，但内存占用增加约一倍
CREATE INDEX idx_embedding_hnsw ON items
USING hnsw (embedding vector_cosine_ops)
WITH (m = 32, ef_construction = 128);

-- 场景三：低内存配置（资源受限时使用）
-- m=8 可显著降低内存，但在高维数据上可能损失 2-5% 的召回率
CREATE INDEX idx_embedding_hnsw ON items
USING hnsw (embedding vector_cosine_ops)
WITH (m = 8, ef_construction = 40);

-- 在会话级别调整搜索精度（不影响索引构建）
-- 值越大，结果越精确，但查询越慢

-- 低延迟优先（在线服务，允许轻微精度损失）
SET hnsw.ef_search = 40;   -- 默认值

-- 平衡模式（推荐大多数场景）
SET hnsw.ef_search = 80;

-- 高精度优先（离线分析、对召回率要求高的场景）
SET hnsw.ef_search = 200;

-- 也可以在单条查询中通过子查询设置
SET LOCAL hnsw.ef_search = 150;  -- 仅在当前事务中生效

-- 首先确认表中向量数量
SELECT COUNT(*) FROM items;

-- 假设结果为 500,000 行
-- 推荐 lists = 500000 / 1000 = 500
CREATE INDEX idx_embedding_ivfflat ON items
USING ivfflat (embedding vector_cosine_ops)
WITH (lists = 500);

-- 假设结果为 5,000,000 行（百万级以上，用平方根法则）
-- 推荐 lists = sqrt(5000000) ≈ 2236
CREATE INDEX idx_embedding_ivfflat ON items
USING ivfflat (embedding vector_cosine_ops)
WITH (lists = 2236);

-- 设置会话级别的 probes 参数
-- 推荐起始值：probes = sqrt(lists)

-- lists=100 时
SET ivfflat.probes = 10;

-- lists=500 时
SET ivfflat.probes = 22;

-- lists=2236 时
SET ivfflat.probes = 47;

-- 逐步调优：从 sqrt(lists) 开始，增加直到召回率满足要求
SET ivfflat.probes = 50;  -- 提高召回率
SET ivfflat.probes = 5;   -- 降低延迟

-- 方案一：先导入数据，再创建索引（推荐）
-- 步骤1：先批量导入所有数据
INSERT INTO items (embedding) VALUES (...), (...), ...;
-- 步骤2：数据导入完成后创建索引
CREATE INDEX ON items USING ivfflat (embedding vector_cosine_ops) WITH (lists = 100);

-- 方案二：如果数据是持续写入的，分两步走
-- 步骤1：数据量达到预期 50% 以上时先创建初始索引
CREATE INDEX ON items USING ivfflat (embedding vector_cosine_ops) WITH (lists = 100);
-- 步骤2：数据量翻倍后重建索引
REINDEX INDEX CONCURRENTLY idx_items_ivfflat;

-- 普通重建（会锁表，阻塞写入）
REINDEX INDEX idx_embedding_hnsw;

-- 并发重建（不阻塞读写，推荐生产环境使用）
-- PostgreSQL 12+ 支持
REINDEX INDEX CONCURRENTLY idx_embedding_hnsw;

-- 重建整个表的所有索引
REINDEX TABLE items;

-- 重建整个数据库的索引（谨慎使用）
REINDEX DATABASE mydb;

-- 普通删除
DROP INDEX idx_embedding_hnsw;

-- 如果存在则删除（避免报错）
DROP INDEX IF EXISTS idx_embedding_hnsw;

-- 并发删除（PostgreSQL 不提供 DROP INDEX CONCURRENTLY，
-- 但 DROP INDEX 本身持有排他锁的时间很短，通常不是问题）

-- 安装 pgstattuple 扩展（需要超级用户权限）
CREATE EXTENSION IF NOT EXISTS pgstattuple;

-- 检测索引的内部统计信息
SELECT * FROM pgstatindex('idx_embedding_hnsw');

-- 关键字段：
-- avg_leaf_density: 叶页面利用率（低于 70% 可能需要重建）
-- internal_pages / leaf_pages: 内部页和叶页数量
-- empty_pages: 空页面数量（过多表示膨胀）

-- 对于 HNSW 索引，使用 pgstatindex 查看：
SELECT * FROM pgstatindex('idx_embedding_hnsw');

-- 简单的索引大小监控
SELECT
    indexname,
    pg_size_pretty(pg_relation_size(indexrelid)) AS index_size
FROM pg_stat_user_indexes
WHERE relname = 'items';

-- 设置并行工作进程数（默认通常为 2）
SET max_parallel_maintenance_workers = 4;

-- 确保有足够的共享内存
SET maintenance_work_mem = '2GB';  -- 根据可用内存调整

-- 然后创建索引（会自动使用并行）
CREATE INDEX idx_embedding_hnsw ON items
USING hnsw (embedding vector_cosine_ops)
WITH (m = 16, ef_construction = 64);

-- 检查并行是否生效
SET client_min_messages = DEBUG1;
CREATE INDEX idx_embedding_hnsw ON items
USING hnsw (embedding vector_cosine_ops)
WITH (m = 16, ef_construction = 64);
-- 输出中会显示 "using N parallel workers"

-- 余弦距离（最常用，适合归一化向量，衡量方向相似性）
CREATE INDEX ON items USING hnsw (embedding vector_cosine_ops);
-- 对应查询：embedding <=> query_vector

-- L2 欧氏距离（适合非归一化向量，同时考虑方向和大小）
CREATE INDEX ON items USING hnsw (embedding vector_l2_ops);
-- 对应查询：embedding <-> query_vector

-- 负内积（适合已经归一化的向量，等价于余弦相似度）
CREATE INDEX ON items USING hnsw (embedding vector_ip_ops);
-- 对应查询：embedding <#> query_vector

-- 正确：用 vector_cosine_ops 建索引，用 <=> 查询
CREATE INDEX ON items USING hnsw (embedding vector_cosine_ops);
SELECT * FROM items ORDER BY embedding <=> '[...]' LIMIT 5;  -- ✅ 走索引

-- 错误：用 vector_cosine_ops 建索引，用 <-> 查询
CREATE INDEX ON items USING hnsw (embedding vector_cosine_ops);
SELECT * FROM items ORDER BY embedding <-> '[...]' LIMIT 5;  -- ❌ 不走索引，退化为全表扫描

EXPLAIN ANALYZE
SELECT * FROM items ORDER BY embedding <=> '[...]' LIMIT 5;

-- 如果走索引，会看到 "Index Scan using idx_xxx"
-- 如果没走索引，会看到 "Seq Scan" + "Sort"

-- halfvec 使用 16 位浮点数，内存占用减半
-- 适合对精度要求不是极高但需要节省内存的场景
CREATE TABLE items_half (
    id serial PRIMARY KEY,
    embedding halfvec(768)
);

CREATE INDEX ON items_half USING hnsw (embedding halfvec_cosine_ops)
WITH (m = 16, ef_construction = 64);

SELECT * FROM items_half ORDER BY embedding <=> '[0.1, 0.2, ...]'::halfvec LIMIT 5;

-- sparsevec 适合高维稀疏向量（如 BM25 特征）
CREATE TABLE items_sparse (
    id serial PRIMARY KEY,
    embedding sparsevec(30000)
);

CREATE INDEX ON items_sparse USING hnsw (embedding sparsevec_l2_ops);

-- bit 类型用于二进制嵌入，极度压缩存储空间
CREATE TABLE items_bit (
    id serial PRIMARY KEY,
    embedding bit(768)
);

-- 汉明距离索引
CREATE INDEX ON items_bit USING hnsw (embedding bit_hamming_ops);
SELECT * FROM items_bit ORDER BY embedding <~> B'101...' LIMIT 5;

-- 杰卡德距离索引
CREATE INDEX ON items_bit USING hnsw (embedding bit_jaccard_ops);
SELECT * FROM items_bit ORDER BY embedding <%> B'101...' LIMIT 5;

                    ┌──────────────────────────────────────┐
                    │          向量数据量是多少？            │
                    └──────────────┬───────────────────────┘
                                   │
              ┌────────────────────┼────────────────────┐
              ▼                    ▼                     ▼
         < 1 万              1万 - 100万            > 100万
              │                    │                     │
              ▼                    ▼                     ▼
      Brute Force          内存充足吗？            内存充足吗？
      （无索引）            │          │            │          │
                       Yes ▼     No ▼         Yes ▼     No ▼
                       HNSW      HNSW         HNSW    IVFFlat
                    m=16,ef=64  m=8,ef=40   m=16,ef=128  lists=sqrt(n)

pip install psycopg[binary] pgvector numpy tenacity

"""
pgvector + psycopg3 基础对接示例
演示：连接数据库、注册向量类型、插入与查询
"""
import psycopg
from pgvector.psycopg import register_vector
import numpy as np


def get_connection():
    """建立数据库连接并注册向量类型"""
    conn = psycopg.connect(
        host="localhost",
        port=5432,
        dbname="vector_db",
        user="postgres",
        password="your_password",
        autocommit=True
    )
    # 注册向量类型，使 psycopg 能自动转换 Python 列表/numpy 数组 <-> pgvector
    register_vector(conn)
    return conn


def create_table(conn):
    """创建向量表"""
    with conn.cursor() as cur:
        cur.execute("""
            CREATE TABLE IF NOT EXISTS documents (
                id serial PRIMARY KEY,
                content text NOT NULL,
                embedding vector(768),
                metadata jsonb DEFAULT '{}'::jsonb,
                created_at timestamptz DEFAULT now()
            )
        """)


def insert_document(conn, content: str, embedding: np.ndarray, metadata: dict = None):
    """插入单条文档"""
    with conn.cursor() as cur:
        cur.execute(
            """
            INSERT INTO documents (content, embedding, metadata)
            VALUES (%s, %s, %s)
            RETURNING id
            """,
            (content, embedding, metadata or {})
        )
        return cur.fetchone()[0]


def search_similar(conn, query_embedding: np.ndarray, top_k: int = 5):
    """余弦相似度搜索"""
    with conn.cursor() as cur:
        cur.execute(
            """
            SELECT id, content, embedding <=> %s AS distance
            FROM documents
            ORDER BY distance
            LIMIT %s
            """,
            (query_embedding, top_k)
        )
        results = []
        for row in cur.fetchall():
            results.append({
                "id": row[0],
                "content": row[1],
                "distance": float(row[2])
            })
        return results


# 主流程
if __name__ == "__main__":
    conn = get_connection()
    create_table(conn)

    # 创建 HNSW 索引
    with conn.cursor() as cur:
        cur.execute("""
            CREATE INDEX IF NOT EXISTS idx_documents_embedding
            ON documents USING hnsw (embedding vector_cosine_ops)
            WITH (m = 16, ef_construction = 64)
        """)

    # 插入测试数据
    sample_embedding = np.random.rand(768).astype(np.float32)
    doc_id = insert_document(conn, "这是一条测试文档", sample_embedding)
    print(f"插入文档 ID: {doc_id}")

    # 查询
    query_vec = np.random.rand(768).astype(np.float32)
    results = search_similar(conn, query_vec, top_k=5)
    for r in results:
        print(f"ID: {r['id']}, 距离: {r['distance']:.4f}, 内容: {r['content'][:50]}")

    conn.close()

"""
高效批量写入向量数据
方式一：executemany（简单可靠）
方式二：COPY（极致性能，适合百万级数据）
"""
import psycopg
from pgvector.psycopg import register_vector
import numpy as np
from io import BytesIO


def batch_insert_executemany(conn, documents: list[dict]):
    """
    使用 executemany 批量写入
    documents: [{"content": "...", "embedding": np.array, "metadata": {}}, ...]
    """
    with conn.cursor() as cur:
        cur.executemany(
            """
            INSERT INTO documents (content, embedding, metadata)
            VALUES (%s, %s, %s)
            """,
            [
                (doc["content"], doc["embedding"], doc.get("metadata", {}))
                for doc in documents
            ]
        )
    print(f"通过 executemany 插入了 {len(documents)} 条记录")


def batch_insert_copy(conn, documents: list[dict]):
    """
    使用 COPY 协议批量写入（性能最优）
    适合百万级数据导入，速度比 executemany 快 5-10 倍
    """
    with conn.cursor() as cur:
        # 使用 copy 接口
        with cur.copy(
            "COPY documents (content, embedding, metadata) FROM STDIN"
        ) as copy:
            for doc in documents:
                # 将 embedding 转换为 pgvector 的文本格式
                emb_str = "[" + ",".join(str(x) for x in doc["embedding"]) + "]"
                meta_str = str(doc.get("metadata", {})).replace("'", '"')
                row = f"{doc['content']}\t{emb_str}\t{meta_str}\n"
                copy.write(row.encode("utf-8"))
    print(f"通过 COPY 插入了 {len(documents)} 条记录")


# 使用示例
if __name__ == "__main__":
    conn = psycopg.connect(
        "postgresql://postgres:your_password@localhost:5432/vector_db",
        autocommit=True
    )
    register_vector(conn)

    # 准备批量数据
    batch_data = []
    for i in range(10000):
        batch_data.append({
            "content": f"这是第 {i} 条测试文档的内容",
            "embedding": np.random.rand(768).astype(np.float32),
            "metadata": {"source": "test", "index": i}
        })

    # 方式一：executemany（推荐大多数场景）
    batch_insert_executemany(conn, batch_data[:1000])

    # 方式二：COPY（大数据量场景）
    batch_insert_copy(conn, batch_data[1000:2000])

    conn.close()

"""
向量检索分页封装
支持：基于偏移量的分页 + 基于游标的分页（推荐）
"""
import psycopg
from pgvector.psycopg import register_vector
import numpy as np
from dataclasses import dataclass


@dataclass
class SearchResult:
    id: int
    content: str
    distance: float
    metadata: dict


@dataclass
class PaginatedResults:
    results: list[SearchResult]
    total_estimated: int
    has_next: bool


def search_paginated(
    conn,
    query_embedding: np.ndarray,
    top_k: int = 10,
    offset: int = 0,
    max_distance: float = None
) -> PaginatedResults:
    """
    基于偏移量的分页向量检索

    参数:
        conn: 数据库连接
        query_embedding: 查询向量
        top_k: 每页返回数量
        offset: 偏移量
        max_distance: 最大距离阈值（可选，过滤掉不相关的结果）
    """
    with conn.cursor() as cur:
        # 构建查询
        query = """
            SELECT id, content, embedding <=> %s AS distance, metadata
            FROM documents
        """
        params = [query_embedding]

        if max_distance is not None:
            query += " WHERE embedding <=> %s <= %s "
            params.extend([query_embedding, max_distance])

        query += " ORDER BY distance LIMIT %s OFFSET %s"
        params.extend([top_k + 1, offset])  # 多取一条判断是否有下一页

        cur.execute(query, params)
        rows = cur.fetchall()

        has_next = len(rows) > top_k
        results = []
        for row in rows[:top_k]:
            results.append(SearchResult(
                id=row[0],
                content=row[1],
                distance=float(row[2]),
                metadata=row[3] if row[3] else {}
            ))

        return PaginatedResults(
            results=results,
            total_estimated=offset + len(results),
            has_next=has_next
        )


def search_cursor_paginated(
    conn,
    query_embedding: np.ndarray,
    top_k: int = 10,
    last_distance: float = None,
    last_id: int = None
) -> PaginatedResults:
    """
    基于游标的分页（推荐，性能更好，结果更稳定）
    使用上一页最后一条记录的距离和 ID 作为游标
    """
    with conn.cursor() as cur:
        if last_distance is not None and last_id is not None:
            cur.execute(
                """
                SELECT id, content, embedding <=> %s AS distance, metadata
                FROM documents
                WHERE (embedding <=> %s, id) > (%s, %s)
                ORDER BY distance, id
                LIMIT %s
                """,
                (query_embedding, query_embedding, last_distance, last_id, top_k + 1)
            )
        else:
            cur.execute(
                """
                SELECT id, content, embedding <=> %s AS distance, metadata
                FROM documents
                ORDER BY distance, id
                LIMIT %s
                """,
                (query_embedding, top_k + 1)
            )

        rows = cur.fetchall()
        has_next = len(rows) > top_k
        results = []
        for row in rows[:top_k]:
            results.append(SearchResult(
                id=row[0],
                content=row[1],
                distance=float(row[2]),
                metadata=row[3] if row[3] else {}
            ))

        return PaginatedResults(
            results=results,
            total_estimated=0,  # 游标分页不便计算总数
            has_next=has_next
        )

"""
使用 tenacity 库实现数据库操作的重试机制
处理：连接超时、死锁、临时网络故障等瞬时错误
"""
import psycopg
from pgvector.psycopg import register_vector
from tenacity import (
    retry,
    stop_after_attempt,
    wait_exponential,
    retry_if_exception_type,
    before_sleep_log
)
import logging
import numpy as np

logger = logging.getLogger(__name__)


# 定义可重试的异常类型
RETRYABLE_EXCEPTIONS = (
    psycopg.OperationalError,   # 连接失败、超时
    psycopg.errors.DeadlockDetected,  # 死锁
    psycopg.errors.SerializationFailure,  # 序列化失败
    ConnectionError,
    TimeoutError,
)


@retry(
    stop=stop_after_attempt(5),           # 最多重试 5 次
    wait=wait_exponential(                # 指数退避：1s, 2s, 4s, 8s
        multiplier=1,
        min=1,
        max=30
    ),
    retry=retry_if_exception_type(RETRYABLE_EXCEPTIONS),
    before_sleep=before_sleep_log(logger, logging.WARNING),
    reraise=True
)
def robust_search(
    conn_str: str,
    query_embedding: np.ndarray,
    top_k: int = 5
) -> list[dict]:
    """
    带重试机制的向量检索

    使用独立连接，确保重试时能重新建立连接
    """
    conn = psycopg.connect(conn_str, autocommit=True)
    try:
        register_vector(conn)
        with conn.cursor() as cur:
            cur.execute(
                """
                SELECT id, content, embedding <=> %s AS distance
                FROM documents
                ORDER BY distance
                LIMIT %s
                """,
                (query_embedding, top_k)
            )
            return [
                {"id": r[0], "content": r[1], "distance": float(r[2])}
                for r in cur.fetchall()
            ]
    finally:
        conn.close()


@retry(
    stop=stop_after_attempt(3),
    wait=wait_exponential(multiplier=1, min=2, max=20),
    retry=retry_if_exception_type(RETRYABLE_EXCEPTIONS),
    reraise=True
)
def robust_insert(
    conn_str: str,
    content: str,
    embedding: np.ndarray,
    metadata: dict = None
) -> int:
    """带重试机制的文档插入"""
    conn = psycopg.connect(conn_str, autocommit=True)
    try:
        register_vector(conn)
        with conn.cursor() as cur:
            cur.execute(
                """
                INSERT INTO documents (content, embedding, metadata)
                VALUES (%s, %s, %s) RETURNING id
                """,
                (content, embedding, metadata or {})
            )
            return cur.fetchone()[0]
    finally:
        conn.close()


# 使用示例
if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    conn_str = "postgresql://postgres:your_password@localhost:5432/vector_db"

    query = np.random.rand(768).astype(np.float32)
    results = robust_search(conn_str, query, top_k=5)
    print(f"检索到 {len(results)} 条结果")

"""
psycopg 连接池配置
适用于高并发 Web 应用场景
"""
import psycopg
from psycopg_pool import ConnectionPool
from pgvector.psycopg import register_vector
import numpy as np


def create_production_pool(
    conninfo: str,
    min_size: int = 5,
    max_size: int = 20,
    **kwargs
) -> ConnectionPool:
    """
    创建生产级连接池

    参数:
        conninfo: 连接字符串，如 "postgresql://user:pass@host:port/db"
        min_size: 最小连接数（始终保持的空闲连接）
        max_size: 最大连接数（超过后请求会等待）
    """
    pool = ConnectionPool(
        conninfo=conninfo,
        min_size=min_size,
        max_size=max_size,
        # 连接超时（等待可用连接的最大时间，秒）
        timeout=30,
        # 每个连接的最大存活时间（秒），防止长时间持有旧连接
        max_lifetime=3600,
        # 空闲连接的最大存活时间（秒）
        max_idle=300,
        # 连接创建后的初始化回调
        open=True,
        configure=lambda conn: _configure_connection(conn),
        **kwargs
    )
    return pool


def _configure_connection(conn):
    """连接初始化配置"""
    register_vector(conn)
    with conn.cursor() as cur:
        # 设置 HNSW 搜索参数
        cur.execute("SET hnsw.ef_search = 80")
        # 设置 IVFFlat 探测参数
        cur.execute("SET ivfflat.probes = 10")
        # 设置语句超时（防止慢查询拖垮系统）
        cur.execute("SET statement_timeout = '5s'")
    conn.commit()


# 全局连接池（在应用启动时创建）
_pool: ConnectionPool = None


def get_pool() -> ConnectionPool:
    """获取全局连接池（懒加载单例）"""
    global _pool
    if _pool is None:
        _pool = create_production_pool(
            conninfo="postgresql://postgres:your_password@localhost:5432/vector_db",
            min_size=5,    # 启动时创建 5 个连接
            max_size=20,   # 最多 20 个连接
        )
    return _pool


def search_with_pool(query_embedding: np.ndarray, top_k: int = 5) -> list[dict]:
    """使用连接池执行检索"""
    pool = get_pool()
    with pool.connection() as conn:
        with conn.cursor() as cur:
            cur.execute(
                """
                SELECT id, content, embedding <=> %s AS distance
                FROM documents
                ORDER BY distance
                LIMIT %s
                """,
                (query_embedding, top_k)
            )
            return [
                {"id": r[0], "content": r[1], "distance": float(r[2])}
                for r in cur.fetchall()
            ]


def close_pool():
    """关闭连接池（应用关闭时调用）"""
    global _pool
    if _pool is not None:
        _pool.close()
        _pool = None


# ---- 高并发场景的连接池参数建议 ----
#
# 小型应用（< 50 并发用户）:
#   min_size=3, max_size=10
#
# 中型应用（50-200 并发用户）:
#   min_size=5, max_size=20
#
# 大型应用（200-500 并发用户）:
#   min_size=10, max_size=50
#   注意：PostgreSQL 默认最大连接数为 100，
#   需要相应调整 postgresql.conf 中的 max_connections
#
# 超大型应用（> 500 并发用户）:
#   使用 PgBouncer 作为外部连接池，
#   应用连接到 PgBouncer，PgBouncer 再连接到 PostgreSQL
#   min_size=20, max_size=100（PgBouncer 层面）


if __name__ == "__main__":
    # 测试连接池
    results = search_with_pool(np.random.rand(768).astype(np.float32), top_k=3)
    print(f"检索结果: {results}")
    close_pool()

pip install langchain-postgresql psycopg[binary] langchain-core langchain-text-splitters

"""
LangChain + PGVector 完整集成示例
演示：初始化、文档入库、相似度搜索、MMR 检索
"""
from langchain_postgresql import PGVector
from langchain_core.documents import Document
from langchain_core.embeddings import Embeddings
import numpy as np


# ---- 自定义 Embedding 模型（适配本地模型） ----
class LocalEmbeddings(Embeddings):
    """
    适配本地 sentence-transformers 模型的 Embeddings 实现
    如果使用 OpenAI 等 API，可直接使用 langchain-openai 的 OpenAIEmbeddings
    """

    def __init__(self, model_name: str = "BAAI/bge-m3"):
        from sentence_transformers import SentenceTransformer
        self.model = SentenceTransformer(model_name)

    def embed_documents(self, texts: list[str]) -> list[list[float]]:
        """批量生成文档向量"""
        embeddings = self.model.encode(
            texts,
            normalize_embeddings=True,  # 归一化，使余弦距离等价于内积
            batch_size=64,
            show_progress_bar=True
        )
        return embeddings.tolist()

    def embed_query(self, text: str) -> list[float]:
        """生成查询向量"""
        embedding = self.model.encode(
            text,
            normalize_embeddings=True
        )
        return embedding.tolist()


def create_vector_store(
    connection_string: str,
    collection_name: str = "my_documents",
    embeddings: Embeddings = None
) -> PGVector:
    """
    初始化 PGVector 向量存储

    参数:
        connection_string: PostgreSQL 连接字符串
        collection_name: 集合名称（类似表名，LangChain 会自动创建）
        embeddings: Embedding 模型实例
    """
    if embeddings is None:
        embeddings = LocalEmbeddings()

    vector_store = PGVector(
        embeddings=embeddings,
        collection_name=collection_name,
        connection=connection_string,
        use_jsonb=True,  # 使用 JSONB 存储元数据
    )
    return vector_store


def ingest_documents(
    vector_store: PGVector,
    documents: list[Document]
) -> list[str]:
    """
    文档入库：将文档向量化并存入 pgvector

    返回文档 ID 列表
    """
    ids = vector_store.add_documents(documents)
    print(f"成功入库 {len(ids)} 条文档")
    return ids


def similarity_search(
    vector_store: PGVector,
    query: str,
    top_k: int = 5,
    filter: dict = None
) -> list[Document]:
    """
    基础相似度搜索

    参数:
        query: 查询文本
        top_k: 返回结果数
        filter: 元数据过滤条件，如 {"source": "pdf"}
    """
    results = vector_store.similarity_search(
        query=query,
        k=top_k,
        filter=filter
    )
    return results


def similarity_search_with_score(
    vector_store: PGVector,
    query: str,
    top_k: int = 5
) -> list[tuple[Document, float]]:
    """
    带距离分数的相似度搜索
    返回 (文档, 距离) 元组列表
    """
    results = vector_store.similarity_search_with_score(
        query=query,
        k=top_k
    )
    return results


def mmr_search(
    vector_store: PGVector,
    query: str,
    top_k: int = 5,
    fetch_k: int = 20,
    lambda_mult: float = 0.5
) -> list[Document]:
    """
    MMR（最大边际相关性）去重检索

    MMR 在相关性和多样性之间取得平衡：
    - 先检索 fetch_k 个最相关的文档
    - 然后从中选出 top_k 个既相关又多样的文档

    参数:
        query: 查询文本
        top_k: 最终返回的结果数
        fetch_k: 预检索的候选数（越大，多样性选择空间越大）
        lambda_mult: 多样性参数（0=最大多样性，1=最大相关性，默认0.5）
    """
    results = vector_store.max_marginal_relevance_search(
        query=query,
        k=top_k,
        fetch_k=fetch_k,
        lambda_mult=lambda_mult
    )
    return results


# ---- 主流程 ----
if __name__ == "__main__":
    CONNECTION_STRING = "postgresql://postgres:your_password@localhost:5432/vector_db"

    # 1. 初始化
    embeddings = LocalEmbeddings("BAAI/bge-m3")
    store = create_vector_store(CONNECTION_STRING, "demo_docs", embeddings)

    # 2. 准备文档
    docs = [
        Document(
            page_content="PostgreSQL 是一个强大的开源关系数据库管理系统",
            metadata={"source": "database_guide", "page": 1}
        ),
        Document(
            page_content="pgvector 是 PostgreSQL 的向量搜索扩展，支持 ANN 检索",
            metadata={"source": "pgvector_docs", "page": 2}
        ),
        Document(
            page_content="HNSW 是一种基于图的近似最近邻搜索算法",
            metadata={"source": "algorithm_paper", "page": 3}
        ),
        Document(
            page_content="RAG 通过检索增强生成，结合了检索系统和大型语言模型",
            metadata={"source": "rag_tutorial", "page": 4}
        ),
        Document(
            page_content="向量嵌入将文本映射到高维空间中的稠密向量表示",
            metadata={"source": "nlp_basics", "page": 5}
        ),
    ]

    # 3. 文档入库
    ids = ingest_documents(store, docs)

    # 4. 相似度搜索
    print("\n--- 相似度搜索 ---")
    results = similarity_search(store, "什么是向量搜索？", top_k=3)
    for doc in results:
        print(f"  [{doc.metadata['source']}] {doc.page_content}")

    # 5. 带分数的搜索
    print("\n--- 带分数的搜索 ---")
    scored_results = similarity_search_with_score(store, "HNSW 算法原理", top_k=3)
    for doc, score in scored_results:
        print(f"  距离: {score:.4f} | {doc.page_content}")

    # 6. MMR 去重搜索
    print("\n--- MMR 去重搜索 ---")
    mmr_results = mmr_search(
        store, "数据库技术", top_k=3, fetch_k=10, lambda_mult=0.7
    )
    for doc in mmr_results:
        print(f"  [{doc.metadata['source']}] {doc.page_content}")

    # 7. 带过滤条件的搜索
    print("\n--- 带过滤条件的搜索 ---")
    filtered = similarity_search(
        store, "向量相关内容", top_k=2,
        filter={"source": "nlp_basics"}
    )
    for doc in filtered:
        print(f"  [{doc.metadata['source']}] {doc.page_content}")

pip install llama-index llama-index-vector-stores-postgres psycopg[binary] sentence-transformers

"""
LlamaIndex + pgvector 完整集成示例
演示：PGVectorStore 配置、本地 Embedding 模型适配、文档索引与检索
"""
from llama_index.core import VectorStoreIndex, StorageContext, Document
from llama_index.core.embeddings import BaseEmbedding
from llama_index.vector_stores.postgres import PGVectorStore
from sentence_transformers import SentenceTransformer
import numpy as np
from typing import List


# ---- 自定义本地 Embedding 模型 ----
class BGEEmbedding(BaseEmbedding):
    """
    适配 BGE-M3 模型的 LlamaIndex Embedding 实现
    """

    def __init__(self, model_name: str = "BAAI/bge-m3", **kwargs):
        super().__init__(**kwargs)
        self._model = SentenceTransformer(model_name)

    def _get_query_embedding(self, query: str) -> List[float]:
        """生成查询向量"""
        embedding = self._model.encode(query, normalize_embeddings=True)
        return embedding.tolist()

    def _get_text_embedding(self, text: str) -> List[float]:
        """生成文本向量"""
        embedding = self._model.encode(text, normalize_embeddings=True)
        return embedding.tolist()

    def _get_text_embeddings(self, texts: List[str]) -> List[List[float]]:
        """批量生成文本向量"""
        embeddings = self._model.encode(
            texts,
            normalize_embeddings=True,
            batch_size=64,
            show_progress_bar=True
        )
        return embeddings.tolist()

    async def _aget_query_embedding(self, query: str) -> List[float]:
        return self._get_query_embedding(query)

    async def _aget_text_embedding(self, text: str) -> List[float]:
        return self._get_text_embedding(text)


def create_llamaindex_vector_store(
    host: str = "localhost",
    port: int = 5432,
    database: str = "vector_db",
    user: str = "postgres",
    password: str = "your_password",
    table_name: str = "llamaindex_docs",
    embed_dim: int = 1024  # BGE-M3 的向量维度
) -> PGVectorStore:
    """
    创建 LlamaIndex PGVectorStore 实例
    """
    vector_store = PGVectorStore.from_params(
        host=host,
        port=str(port),
        database=database,
        user=user,
        password=password,
        table_name=table_name,
        embed_dim=embed_dim,
        # 使用 HNSW 索引
        hnsw_kwargs={
            "m": 16,
            "ef_construction": 64,
        }
    )
    return vector_store


def build_index(
    vector_store: PGVectorStore,
    documents: List[Document],
    embed_model: BaseEmbedding
) -> VectorStoreIndex:
    """
    构建向量索引

    将文档向量化并存入 pgvector，同时创建 HNSW 索引
    """
    storage_context = StorageContext.from_defaults(vector_store=vector_store)

    index = VectorStoreIndex.from_documents(
        documents,
        storage_context=storage_context,
        embed_model=embed_model,
        show_progress=True
    )
    return index


def query_index(
    index: VectorStoreIndex,
    query: str,
    top_k: int = 5
) -> list:
    """
    执行语义检索
    """
    retriever = index.as_retriever(similarity_top_k=top_k)
    nodes = retriever.retrieve(query)
    return nodes


# ---- 主流程 ----
if __name__ == "__main__":
    # 1. 初始化 Embedding 模型
    embed_model = BGEEmbedding("BAAI/bge-m3")

    # 2. 创建向量存储
    vector_store = create_llamaindex_vector_store(
        host="localhost",
        database="vector_db",
        table_name="llamaindex_docs",
        embed_dim=1024  # BGE-M3 输出维度
    )

    # 3. 准备文档
    documents = [
        Document(
            text="PostgreSQL 是世界上最先进的开源关系数据库。",
            metadata={"source": "postgres_intro", "category": "database"}
        ),
        Document(
            text="pgvector 扩展让 PostgreSQL 具备了向量搜索能力。",
            metadata={"source": "pgvector_intro", "category": "extension"}
        ),
        Document(
            text="HNSW 索引通过构建分层图结构实现高效的近似最近邻搜索。",
            metadata={"source": "hnsw_paper", "category": "algorithm"}
        ),
        Document(
            text="IVFFlat 索引使用 K-Means 聚类来加速向量检索。",
            metadata={"source": "ivfflat_paper", "category": "algorithm"}
        ),
        Document(
            text="RAG 系统将检索到的相关文档作为上下文注入到 LLM 的提示中。",
            metadata={"source": "rag_paper", "category": "application"}
        ),
    ]

    # 4. 构建索引
    print("正在构建向量索引...")
    index = build_index(vector_store, documents, embed_model)
    print("索引构建完成")

    # 5. 执行查询
    print("\n--- 语义检索 ---")
    query = "pgvector 支持哪些索引类型？"
    nodes = query_index(index, query, top_k=3)

    for i, node in enumerate(nodes):
        print(f"\n结果 {i + 1}:")
        print(f"  文本: {node.get_content()}")
        print(f"  分数: {node.get_score():.4f}")
        print(f"  来源: {node.metadata.get('source')}")

用户提问 → 问题向量化 → pgvector 语义检索 Top-K → 召回原文片段
                                                       ↓
                        LLM 生成回答 ← 构建 Prompt（问题 + 上下文）

"""
端到端 RAG（检索增强生成）系统完整实现

功能：
1. PDF/TXT 文档切片
2. 向量化（sentence-transformers + BGE-M3）
3. 存入 pgvector + HNSW 索引
4. 语义检索 + LLM 生成回答

依赖安装：
    pip install psycopg[binary] pgvector numpy sentence-transformers
    pip install tenacity openai PyPDF2 langchain-text-splitters

使用方法：
    python rag_system.py
"""
import os
import re
import numpy as np
import psycopg
from pgvector.psycopg import register_vector
from sentence_transformers import SentenceTransformer
from tenacity import retry, stop_after_attempt, wait_exponential
from dataclasses import dataclass, field
from typing import List, Optional
import logging

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s"
)
logger = logging.getLogger(__name__)


# ============================================================
# 配置
# ============================================================
@dataclass
class RAGConfig:
    """RAG 系统配置"""
    # 数据库配置
    db_host: str = "localhost"
    db_port: int = 5432
    db_name: str = "vector_db"
    db_user: str = "postgres"
    db_password: str = "your_password"
    table_name: str = "rag_documents"

    # Embedding 配置
    embedding_model: str = "BAAI/bge-m3"
    embedding_dim: int = 1024  # BGE-M3 输出维度

    # 切片配置
    chunk_size: int = 500       # 每个切片的最大字符数
    chunk_overlap: int = 50     # 切片之间的重叠字符数

    # 检索配置
    top_k: int = 5              # 召回的文档数量
    ef_search: int = 80         # HNSW 搜索精度

    # LLM 配置（以 OpenAI 兼容 API 为例）
    llm_api_base: str = "https://api.openai.com/v1"  # 可替换为其他兼容 API
    llm_api_key: str = os.environ.get("OPENAI_API_KEY", "your-api-key")
    llm_model: str = "gpt-4o-mini"  # 或其他模型

    @property
    def connection_string(self) -> str:
        return f"postgresql://{self.db_user}:{self.db_password}@{self.db_host}:{self.db_port}/{self.db_name}"


# ============================================================
# 第一步：数据源准备 - 文档切片
# ============================================================
@dataclass
class TextChunk:
    """文本切片"""
    content: str
    source: str
    chunk_index: int
    metadata: dict = field(default_factory=dict)


def load_text_file(file_path: str) -> str:
    """加载 TXT 文件"""
    with open(file_path, "r", encoding="utf-8") as f:
        return f.read()


def load_pdf_file(file_path: str) -> str:
    """加载 PDF 文件（使用 PyPDF2）"""
    try:
        from PyPDF2 import PdfReader
    except ImportError:
        raise ImportError("请安装 PyPDF2: pip install PyPDF2")

    reader = PdfReader(file_path)
    text = ""
    for i, page in enumerate(reader.pages):
        page_text = page.extract_text() or ""
        text += f"\n--- 第 {i + 1} 页 ---\n" + page_text
    return text


def split_text(
    text: str,
    source: str,
    chunk_size: int = 500,
    chunk_overlap: int = 50
) -> List[TextChunk]:
    """
    将长文本分割为重叠的切片

    使用 RecursiveCharacterTextSplitter 的思想：
    1. 先按大段落分割（\n\n）
    2. 再按段落分割（\n）
    3. 再按句子分割（。！？.!?）
    4. 最后按固定长度分割

    切片之间保持 chunk_overlap 个字符的重叠，确保上下文不丢失。
    """
    # 分隔符列表，从大到小
    separators = ["\n\n", "\n", "。", "！", "？", ".", "!", "?", " ", ""]

    chunks = []
    # 使用递归分割
    parts = _recursive_split(text, separators, chunk_size)

    # 合并过小的片段，并添加重叠
    current_chunk = ""
    chunk_index = 0

    for part in parts:
        if len(current_chunk) + len(part) <= chunk_size:
            current_chunk += part
        else:
            if current_chunk:
                chunks.append(TextChunk(
                    content=current_chunk.strip(),
                    source=source,
                    chunk_index=chunk_index
                ))
                chunk_index += 1
                # 保留重叠部分
                if len(current_chunk) > chunk_overlap:
                    current_chunk = current_chunk[-chunk_overlap:] + part
                else:
                    current_chunk = part
            else:
                current_chunk = part

    # 处理最后一块
    if current_chunk.strip():
        chunks.append(TextChunk(
            content=current_chunk.strip(),
            source=source,
            chunk_index=chunk_index
        ))

    logger.info(f"文档 '{source}' 被分割为 {len(chunks)} 个切片")
    return chunks


def _recursive_split(text: str, separators: list, chunk_size: int) -> list:
    """递归分割文本"""
    if len(text) <= chunk_size:
        return [text]

    if not separators:
        # 强制按长度切割
        return [text[i:i + chunk_size] for i in range(0, len(text), chunk_size)]

    sep = separators[0]
    remaining_seps = separators[1:]

    if sep == "":
        return [text[i:i + chunk_size] for i in range(0, len(text), chunk_size)]

    parts = text.split(sep)
    # 把分隔符加回去
    result_parts = []
    for i, part in enumerate(parts):
        if i < len(parts) - 1:
            result_parts.append(part + sep)
        else:
            result_parts.append(part)

    # 递归处理过长的部分
    final_parts = []
    for part in result_parts:
        if len(part) > chunk_size:
            final_parts.extend(_recursive_split(part, remaining_seps, chunk_size))
        else:
            final_parts.append(part)

    return final_parts


# ============================================================
# 第二步：向量化
# ============================================================
class EmbeddingService:
    """Embedding 向量化服务"""

    def __init__(self, model_name: str = "BAAI/bge-m3"):
        logger.info(f"正在加载 Embedding 模型: {model_name}")
        self.model = SentenceTransformer(model_name)
        logger.info("Embedding 模型加载完成")

    def embed_texts(self, texts: List[str], batch_size: int = 64) -> np.ndarray:
        """
        批量生成文本向量

        参数:
            texts: 文本列表
            batch_size: 批处理大小

        返回:
            numpy 数组，形状 (len(texts), embedding_dim)
        """
        embeddings = self.model.encode(
            texts,
            normalize_embeddings=True,  # L2 归一化
            batch_size=batch_size,
            show_progress_bar=True,
            convert_to_numpy=True
        )
        return embeddings

    def embed_query(self, text: str) -> np.ndarray:
        """生成查询向量"""
        embedding = self.model.encode(
            text,
            normalize_embeddings=True,
            convert_to_numpy=True
        )
        return embedding


# ============================================================
# 第三步：存储 - pgvector 操作
# ============================================================
class VectorStore:
    """pgvector 向量存储层"""

    def __init__(self, config: RAGConfig):
        self.config = config
        self.conn = psycopg.connect(config.connection_string, autocommit=True)
        register_vector(self.conn)
        self._init_table()
        self._ensure_index()

    def _init_table(self):
        """创建向量表"""
        with self.conn.cursor() as cur:
            cur.execute(f"""
                CREATE TABLE IF NOT EXISTS {self.config.table_name} (
                    id serial PRIMARY KEY,
                    content text NOT NULL,
                    embedding vector({self.config.embedding_dim}),
                    source text,
                    chunk_index integer,
                    metadata jsonb DEFAULT '{{}}'::jsonb,
                    created_at timestamptz DEFAULT now()
                )
            """)
        logger.info(f"表 {self.config.table_name} 已就绪")

    def _ensure_index(self):
        """确保 HNSW 索引存在"""
        index_name = f"idx_{self.config.table_name}_embedding"
        with self.conn.cursor() as cur:
            cur.execute(f"""
                CREATE INDEX IF NOT EXISTS {index_name}
                ON {self.config.table_name}
                USING hnsw (embedding vector_cosine_ops)
                WITH (m = 16, ef_construction = 64)
            """)
        logger.info("HNSW 索引已就绪")

    def upsert_chunks(self, chunks: List[TextChunk], embeddings: np.ndarray):
        """
        批量写入切片及其向量

        使用 source + chunk_index 作为唯一键，支持幂等写入
        """
        assert len(chunks) == len(embeddings), "切片数量与向量数量不匹配"

        with self.conn.cursor() as cur:
            for chunk, emb in zip(chunks, embeddings):
                cur.execute(f"""
                    INSERT INTO {self.config.table_name}
                        (content, embedding, source, chunk_index, metadata)
                    VALUES (%s, %s, %s, %s, %s)
                    ON CONFLICT DO NOTHING
                """, (
                    chunk.content,
                    emb.astype(np.float32),
                    chunk.source,
                    chunk.chunk_index,
                    chunk.metadata
                ))

        logger.info(f"已写入 {len(chunks)} 个切片到 {self.config.table_name}")

    def search(
        self,
        query_embedding: np.ndarray,
        top_k: int = 5,
        filter_source: str = None
    ) -> List[dict]:
        """
        向量相似度检索

        参数:
            query_embedding: 查询向量
            top_k: 返回结果数
            filter_source: 按来源过滤（可选）
        """
        with self.conn.cursor() as cur:
            # 设置搜索精度
            cur.execute(f"SET LOCAL hnsw.ef_search = {self.config.ef_search}")

            if filter_source:
                cur.execute(f"""
                    SELECT id, content, source, chunk_index,
                           embedding <=> %s AS distance
                    FROM {self.config.table_name}
                    WHERE source = %s
                    ORDER BY distance
                    LIMIT %s
                """, (query_embedding, filter_source, top_k))
            else:
                cur.execute(f"""
                    SELECT id, content, source, chunk_index,
                           embedding <=> %s AS distance
                    FROM {self.config.table_name}
                    ORDER BY distance
                    LIMIT %s
                """, (query_embedding, top_k))

            results = []
            for row in cur.fetchall():
                results.append({
                    "id": row[0],
                    "content": row[1],
                    "source": row[2],
                    "chunk_index": row[3],
                    "distance": float(row[4])
                })
            return results

    def count(self) -> int:
        """获取表中文档数量"""
        with self.conn.cursor() as cur:
            cur.execute(f"SELECT COUNT(*) FROM {self.config.table_name}")
            return cur.fetchone()[0]

    def close(self):
        """关闭连接"""
        self.conn.close()


# ============================================================
# 第四步：检索增强生成（RAG）
# ============================================================
class RAGGenerator:
    """RAG 生成器：检索 + LLM 生成"""

    def __init__(self, config: RAGConfig):
        self.config = config

    def build_prompt(self, query: str, contexts: List[dict]) -> str:
        """
        构建 RAG 提示词

        将检索到的上下文和用户问题组装成结构化的 Prompt
        """
        # 拼接上下文
        context_text = ""
        for i, ctx in enumerate(contexts):
            source_info = f"[来源: {ctx['source']}"
            if ctx.get("chunk_index") is not None:
                source_info += f", 片段 {ctx['chunk_index']}"
            source_info += "]"
            context_text += f"\n\n段落 {i + 1} {source_info}:\n{ctx['content']}"

        prompt = f"""你是一个智能问答助手。请根据以下提供的参考文档回答用户的问题。
如果参考文档中没有足够的信息来回答问题，请明确说明"根据已有文档无法完全回答该问题"。
不要编造文档中没有的信息。

---参考文档---
{context_text}
---参考文档结束---

用户问题：{query}

请给出详细、准确的回答："""

        return prompt

    def generate(self, prompt: str) -> str:
        """
        调用 LLM 生成回答

        使用 OpenAI 兼容 API（可替换为通义千问等任何兼容接口）
        """
        try:
            from openai import OpenAI
        except ImportError:
            raise ImportError("请安装 openai: pip install openai")

        client = OpenAI(
            api_key=self.config.llm_api_key,
            base_url=self.config.llm_api_base
        )

        response = client.chat.completions.create(
            model=self.config.llm_model,
            messages=[
                {
                    "role": "system",
                    "content": "你是一个专业的知识问答助手，基于提供的文档内容回答问题。"
                },
                {
                    "role": "user",
                    "content": prompt
                }
            ],
            temperature=0.1,  # 低温度，提高回答的确定性
            max_tokens=1024
        )

        return response.choices[0].message.content


# ============================================================
# 主流程：完整的 RAG Pipeline
# ============================================================
class RAGPipeline:
    """RAG 系统主入口"""

    def __init__(self, config: RAGConfig = None):
        self.config = config or RAGConfig()
        self.embedding_service = EmbeddingService(self.config.embedding_model)
        self.vector_store = VectorStore(self.config)
        self.generator = RAGGenerator(self.config)

    def ingest_file(self, file_path: str):
        """导入单个文件"""
        logger.info(f"正在处理文件: {file_path}")

        # 1. 加载文件
        if file_path.endswith(".pdf"):
            text = load_pdf_file(file_path)
        elif file_path.endswith(".txt"):
            text = load_text_file(file_path)
        else:
            raise ValueError(f"不支持的文件格式: {file_path}")

        # 2. 文本切片
        chunks = split_text(
            text,
            source=os.path.basename(file_path),
            chunk_size=self.config.chunk_size,
            chunk_overlap=self.config.chunk_overlap
        )

        if not chunks:
            logger.warning(f"文件 {file_path} 切片结果为空，跳过")
            return

        # 3. 向量化
        texts = [c.content for c in chunks]
        embeddings = self.embedding_service.embed_texts(texts)

        # 4. 存入数据库
        self.vector_store.upsert_chunks(chunks, embeddings)
        logger.info(f"文件 {file_path} 处理完成，共 {len(chunks)} 个切片")

    def ingest_directory(self, dir_path: str, extensions: tuple = (".txt", ".pdf")):
        """批量导入目录下的所有文件"""
        for root, dirs, files in os.walk(dir_path):
            for file_name in files:
                if file_name.endswith(extensions):
                    file_path = os.path.join(root, file_name)
                    self.ingest_file(file_path)

    def ask(self, question: str, top_k: int = None) -> dict:
        """
        RAG 问答：检索 + 生成

        参数:
            question: 用户问题
            top_k: 召回的文档数量

        返回:
            {
                "answer": LLM 生成的回答,
                "sources": 引用的文档片段列表,
                "question": 原始问题
            }
        """
        top_k = top_k or self.config.top_k
        logger.info(f"正在处理问题: {question}")

        # 1. 问题向量化
        query_embedding = self.embedding_service.embed_query(question)

        # 2. 语义检索
        contexts = self.vector_store.search(query_embedding, top_k=top_k)
        logger.info(f"检索到 {len(contexts)} 个相关片段")

        if not contexts:
            return {
                "answer": "抱歉，知识库中没有找到与该问题相关的内容。",
                "sources": [],
                "question": question
            }

        # 3. 构建 Prompt
        prompt = self.generator.build_prompt(question, contexts)

        # 4. LLM 生成回答
        answer = self.generator.generate(prompt)

        return {
            "answer": answer,
            "sources": [
                {
                    "source": ctx["source"],
                    "chunk_index": ctx["chunk_index"],
                    "distance": ctx["distance"],
                    "preview": ctx["content"][:100] + "..."
                }
                for ctx in contexts
            ],
            "question": question
        }

    def close(self):
        """关闭资源"""
        self.vector_store.close()


# ============================================================
# 运行示例
# ============================================================
if __name__ == "__main__":
    # 配置（请根据实际环境修改）
    config = RAGConfig(
        db_host="localhost",
        db_port=5432,
        db_name="vector_db",
        db_user="postgres",
        db_password="your_password",
        table_name="rag_documents",
        embedding_model="BAAI/bge-m3",
        embedding_dim=1024,
        chunk_size=500,
        chunk_overlap=50,
        top_k=5,
        ef_search=80,
        # 可替换为通义千问 DashScope 兼容 API
        # llm_api_base="https://dashscope.aliyuncs.com/compatible-mode/v1",
        # llm_api_key="your-dashscope-key",
        # llm_model="qwen-plus",
    )

    # 初始化 Pipeline
    pipeline = RAGPipeline(config)

    # ---- 方式一：导入本地文件 ----
    # pipeline.ingest_file("./my_document.pdf")
    # pipeline.ingest_file("./notes.txt")

    # ---- 方式二：批量导入目录 ----
    # pipeline.ingest_directory("./documents/")

    # ---- 方式三：直接导入文本（演示用） ----
    sample_texts = [
        "pgvector 是 PostgreSQL 的开源向量搜索扩展，支持 HNSW、IVFFlat 等索引类型。",
        "HNSW（分层可导航小世界图）是目前最高效的近似最近邻搜索算法之一。",
        "RAG（检索增强生成）是一种将信息检索与大型语言模型结合的技术范式。",
        "向量嵌入是将文本、图像等非结构化数据映射为高维稠密向量的过程。",
        "PostgreSQL 支持通过 CREATE INDEX 语句为向量列创建 HNSW 或 IVFFlat 索引。",
        "在 RAG 系统中，文档首先被切分为小的文本片段，然后每个片段被转换为向量嵌入。",
        "余弦相似度通过计算两个向量之间的夹角来衡量它们的语义相似程度。",
        "BGE-M3 是北京智源研究院发布的多语言、多功能嵌入模型，支持 100+ 种语言。",
    ]

    from pathlib import Path

    # 写入临时文件用于演示
    demo_file = Path("demo_knowledge_base.txt")
    demo_file.write_text("\n\n".join(sample_texts), encoding="utf-8")
    pipeline.ingest_file(str(demo_file))
    demo_file.unlink()  # 清理临时文件

    # ---- 执行问答 ----
    print("\n" + "=" * 60)
    print("RAG 系统已就绪，开始问答")
    print("=" * 60)

    questions = [
        "pgvector 支持哪些索引类型？",
        "什么是 RAG？它是怎么工作的？",
        "BGE-M3 模型有什么特点？",
    ]

    for q in questions:
        print(f"\n{'─' * 40}")
        print(f"问题: {q}")
        result = pipeline.ask(q)
        print(f"回答: {result['answer']}")
        print(f"引用来源:")
        for src in result["sources"]:
            print(f"  - [{src['source']}] 片段{src['chunk_index']} "
                  f"(距离: {src['distance']:.4f})")

    pipeline.close()

┌──────────────────────────────────────────────────────────────┐
│                    RAG 系统完整流程                           │
├──────────────────────────────────────────────────────────────┤
│                                                              │
│  离线阶段（文档入库）:                                        │
│                                                              │
│  PDF/TXT ──→ 文本提取 ──→ 切片 ──→ 向量化 ──→ pgvector 存储  │
│             load_pdf()   split()  encode()   upsert_chunks() │
│                         500字/片  BGE-M3     HNSW 索引       │
│                         50字重叠  1024维                     │
│                                                              │
│  在线阶段（问答检索）:                                        │
│                                                              │
│  用户问题 ──→ 问题向量化 ──→ pgvector 检索 ──→ Top-K 召回    │
│              embed_query() <=> 余弦距离     距离排序          │
│                                                              │
│  Top-K 上下文 + 用户问题 ──→ 构建 Prompt ──→ LLM ──→ 回答    │
│                              build_prompt()  GPT/Qwen        │
│                                                              │
└──────────────────────────────────────────────────────────────┘

"""
向量聚合计算示例
"""
import psycopg
from pgvector.psycopg import register_vector
import numpy as np


def compute_average_embedding(conn, table: str = "documents") -> np.ndarray:
    """
    计算所有文档的平均向量（即"中心向量"）

    用途：
    - 了解文档集合的整体语义方向
    - 作为聚类初始中心
    - 检测异常文档（距离平均向量过远的文档）
    """
    with conn.cursor() as cur:
        # 使用 AVG 聚合计算每个维度的平均值
        # 注意：pgvector 0.7.0+ 支持 avg() 聚合函数
        cur.execute(f"""
            SELECT avg(embedding) FROM {table}
        """)
        result = cur.fetchone()
        if result and result[0] is not None:
            # avg() 返回的是 vector 类型，psycopg 会自动转为 numpy 数组
            avg_vec = np.array(result[0])
            logger.info(f"平均向量维度: {avg_vec.shape}")
            return avg_vec
        return None


def compute_centroid_per_source(conn, table: str = "documents") -> dict:
    """
    按来源分组计算每个来源的中心向量

    用途：
    - 比较不同文档/来源的语义差异
    - 构建来源级别的语义索引
    """
    with conn.cursor() as cur:
        cur.execute(f"""
            SELECT source, avg(embedding) AS centroid
            FROM {table}
            GROUP BY source
            ORDER BY source
        """)
        centroids = {}
        for row in cur.fetchall():
            centroids[row[0]] = np.array(row[1])
            logger.info(f"来源 '{row[0]}' 的中心向量已计算")
        return centroids


def vector_arithmetic(conn):
    """
    向量加减法运算

    应用场景：
    - 语义类比：king - man + woman ≈ queen
    - 文档差异：doc_new - doc_old = 变化方向
    - 混合查询：0.7 * query_vec + 0.3 * category_vec
    """
    with conn.cursor() as cur:
        # 获取两个文档的向量
        cur.execute("""
            SELECT id, content, embedding
            FROM documents
            LIMIT 2
        """)
        rows = cur.fetchall()
        if len(rows) < 2:
            print("需要至少 2 条文档")
            return

        vec_a = np.array(rows[0][2])
        vec_b = np.array(rows[1][2])

        # 向量加法
        vec_sum = vec_a + vec_b
        print(f"向量 A + B: 前5维 = {vec_sum[:5]}")

        # 向量减法
        vec_diff = vec_a - vec_b
        print(f"向量 A - B: 前5维 = {vec_diff[:5]}")

        # 加权混合（常用于混合查询）
        alpha = 0.7
        vec_mixed = alpha * vec_a + (1 - alpha) * vec_b
        print(f"0.7*A + 0.3*B: 前5维 = {vec_mixed[:5]}")

        # 使用混合向量进行搜索
        cur.execute("""
            SELECT id, content, embedding <=> %s AS distance
            FROM documents
            ORDER BY distance
            LIMIT 5
        """, (vec_mixed.astype(np.float32),))

        print("\n混合向量搜索结果:")
        for row in cur.fetchall():
            print(f"  ID: {row[0]}, 距离: {float(row[2]):.4f}, 内容: {row[1][:60]}")


# 使用示例
if __name__ == "__main__":
    conn = psycopg.connect(
        "postgresql://postgres:your_password@localhost:5432/vector_db",
        autocommit=True
    )
    register_vector(conn)

    # 计算平均向量
    avg = compute_average_embedding(conn)
    if avg is not None:
        print(f"平均向量 (前5维): {avg[:5]}")

    # 按来源计算中心向量
    centroids = compute_centroid_per_source(conn)
    for source, centroid in centroids.items():
        print(f"来源 '{source}' 中心向量 (前5维): {centroid[:5]}")

    # 向量运算
    vector_arithmetic(conn)

    conn.close()

MMR = argmax_{d_i ∈ R\S} [ λ * Sim(d_i, Q) - (1-λ) * max_{d_j ∈ S} Sim(d_i, d_j) ]

其中：
- R: 全部候选文档集合
- S: 已选中的文档集合
- Q: 查询向量
- λ: 相关性-多样性权衡参数（0=最大多样性，1=最大相关性）
- Sim: 相似度函数

"""
MMR（最大边际相关性）多样性检索的完整实现
不依赖框架，纯 Python + numpy 实现
"""
import numpy as np
import psycopg
from pgvector.psycopg import register_vector
from typing import List


def cosine_similarity_matrix(
    vectors_a: np.ndarray,
    vectors_b: np.ndarray
) -> np.ndarray:
    """
    计算两组向量之间的余弦相似度矩阵

    参数:
        vectors_a: 形状 (m, dim)
        vectors_b: 形状 (n, dim)

    返回:
        相似度矩阵，形状 (m, n)
    """
    # 假设向量已经归一化
    return vectors_a @ vectors_b.T


def mmr_search(
    conn,
    query_embedding: np.ndarray,
    table: str = "documents",
    top_k: int = 5,
    fetch_k: int = 20,
    lambda_mult: float = 0.5
) -> List[dict]:
    """
    MMR 多样性检索

    参数:
        conn: 数据库连接
        query_embedding: 查询向量（已归一化）
        table: 表名
        top_k: 最终返回的结果数
        fetch_k: 预检索的候选数量（fetch_k >= top_k）
        lambda_mult: 多样性参数
            - 1.0: 纯相关性排序（等价于普通 Top-K）
            - 0.0: 纯多样性排序
            - 0.5: 平衡相关性与多样性（推荐默认值）
            - 0.7: 偏向相关性，适度去重（推荐用于 RAG）

    返回:
        MMR 选中的文档列表
    """
    assert fetch_k >= top_k, "fetch_k 必须大于等于 top_k"

    # 第一步：预检索 fetch_k 个最相关的候选
    with conn.cursor() as cur:
        cur.execute(f"""
            SELECT id, content, source, embedding,
                   embedding <=> %s AS distance
            FROM {table}
            ORDER BY distance
            LIMIT %s
        """, (query_embedding, fetch_k))

        candidates = []
        candidate_embeddings = []
        for row in cur.fetchall():
            candidates.append({
                "id": row[0],
                "content": row[1],
                "source": row[2],
                "distance": float(row[4])
            })
            candidate_embeddings.append(np.array(row[3]))

    if len(candidates) <= top_k:
        return candidates

    candidate_embeddings = np.array(candidate_embeddings)

    # 第二步：计算查询与所有候选的相似度
    # 余弦相似度 = 1 - 余弦距离
    query_sims = 1.0 - np.array([c["distance"] for c in candidates])

    # 第三步：计算候选之间的相似度矩阵
    sim_matrix = cosine_similarity_matrix(candidate_embeddings, candidate_embeddings)

    # 第四步：MMR 迭代选择
    selected_indices = []
    remaining_indices = list(range(len(candidates)))

    for _ in range(min(top_k, len(candidates))):
        best_score = -float("inf")
        best_idx = None

        for idx in remaining_indices:
            # 与查询的相关性
            relevance = query_sims[idx]

            # 与已选文档的最大相似度（冗余度）
            if selected_indices:
                max_sim_to_selected = max(
                    sim_matrix[idx][sel_idx] for sel_idx in selected_indices
                )
            else:
                max_sim_to_selected = 0.0

            # MMR 分数
            mmr_score = (
                lambda_mult * relevance
                - (1 - lambda_mult) * max_sim_to_selected
            )

            if mmr_score > best_score:
                best_score = mmr_score
                best_idx = idx

        selected_indices.append(best_idx)
        remaining_indices.remove(best_idx)

    # 返回选中的文档（按 MMR 选择顺序排列）
    return [candidates[i] for i in selected_indices]


# ---- 使用示例 ----
if __name__ == "__main__":
    conn = psycopg.connect(
        "postgresql://postgres:your_password@localhost:5432/vector_db",
        autocommit=True
    )
    register_vector(conn)

    # 模拟查询向量
    query_vec = np.random.rand(1024).astype(np.float32)
    query_vec = query_vec / np.linalg.norm(query_vec)  # 归一化

    # 普通 Top-K 检索（可能有重复内容）
    print("=== 普通 Top-K 检索 ===")
    with conn.cursor() as cur:
        cur.execute("""
            SELECT id, content, embedding <=> %s AS distance
            FROM documents
            ORDER BY distance
            LIMIT 5
        """, (query_vec,))
        for row in cur.fetchall():
            print(f"  ID={row[0]}, 距离={float(row[2]):.4f}, "
                  f"内容={row[1][:60]}...")

    # MMR 多样性检索（去重后的结果）
    print("\n=== MMR 多样性检索 (λ=0.7) ===")
    mmr_results = mmr_search(
        conn, query_vec,
        table="documents",
        top_k=5,
        fetch_k=20,
        lambda_mult=0.7
    )
    for r in mmr_results:
        print(f"  ID={r['id']}, 距离={r['distance']:.4f}, "
              f"内容={r['content'][:60]}...")

    conn.close()

"""
距离归一化：将向量距离转换为 0-1 的相似度分数
"""
import numpy as np


def normalize_distance(distance: float, distance_type: str = "cosine") -> float:
    """
    将距离转换为 0-1 的相似度分数
    1.0 = 完全相似，0.0 = 完全不相似

    参数:
        distance: 原始距离值
        distance_type: 距离类型
            - "cosine": 余弦距离，范围 [0, 2]
            - "l2": L2 欧氏距离，范围 [0, +inf)
            - "ip": 负内积，范围 [-1, +inf)（归一化向量时为 [-1, 1]）
    """
    if distance_type == "cosine":
        # 余弦距离范围 [0, 2]
        # 距离 0 → 相似度 1.0（完全相同方向）
        # 距离 1 → 相似度 0.5（正交）
        # 距离 2 → 相似度 0.0（完全相反方向）
        similarity = 1.0 - (distance / 2.0)
        return max(0.0, min(1.0, similarity))

    elif distance_type == "l2":
        # L2 距离范围 [0, +inf)
        # 使用指数衰减函数映射到 (0, 1]
        # distance=0 → 1.0, distance=1 → ~0.37, distance=5 → ~0.007
        similarity = np.exp(-distance)
        return float(similarity)

    elif distance_type == "ip":
        # 负内积（pgvector 的 <#> 返回的是负内积）
        # 对于归一化向量：内积范围 [-1, 1]，负内积范围 [-1, 1]
        # 实际内积 = -distance
        # 相似度 = (实际内积 + 1) / 2
        inner_product = -distance
        similarity = (inner_product + 1.0) / 2.0
        return max(0.0, min(1.0, similarity))

    else:
        raise ValueError(f"不支持的距离类型: {distance_type}")


def batch_normalize_distances(
    distances: list[float],
    distance_type: str = "cosine"
) -> list[float]:
    """批量归一化距离"""
    return [normalize_distance(d, distance_type) for d in distances]


# ---- 使用示例 ----
if __name__ == "__main__":
    # 余弦距离示例
    cosine_distances = [0.0, 0.1, 0.3, 0.5, 1.0, 1.5, 2.0]
    print("余弦距离 → 相似度:")
    for d in cosine_distances:
        s = normalize_distance(d, "cosine")
        print(f"  距离 {d:.2f} → 相似度 {s:.4f}")

    # L2 距离示例
    l2_distances = [0.0, 0.5, 1.0, 2.0, 5.0]
    print("\nL2 距离 → 相似度:")
    for d in l2_distances:
        s = normalize_distance(d, "l2")
        print(f"  距离 {d:.2f} → 相似度 {s:.4f}")

    # 在实际检索中使用
    import psycopg
    from pgvector.psycopg import register_vector

    conn = psycopg.connect(
        "postgresql://postgres:your_password@localhost:5432/vector_db",
        autocommit=True
    )
    register_vector(conn)

    query_vec = np.random.rand(1024).astype(np.float32)
    query_vec = query_vec / np.linalg.norm(query_vec)

    with conn.cursor() as cur:
        cur.execute("""
            SELECT id, content, embedding <=> %s AS distance
            FROM documents
            ORDER BY distance
            LIMIT 5
        """, (query_vec,))

        print("\n带相似度分数的检索结果:")
        for row in cur.fetchall():
            raw_distance = float(row[2])
            similarity = normalize_distance(raw_distance, "cosine")
            print(f"  ID={row[0]}, "
                  f"原始距离={raw_distance:.4f}, "
                  f"相似度={similarity:.4f}, "
                  f"内容={row[1][:50]}...")

    conn.close()

"""
将相似度分数转换为置信度评级
基于距离分布的阈值设定策略
"""
import numpy as np
from dataclasses import dataclass
from enum import Enum


class Confidence(Enum):
    """置信度等级"""
    HIGH = "高置信度"      # 非常确定答案准确
    MEDIUM = "中等置信度"  # 答案可能准确
    LOW = "低置信度"       # 答案可能不准确
    VERY_LOW = "极低置信度"  # 不建议依赖此答案


@dataclass
class ConfidenceResult:
    """置信度评估结果"""
    confidence: Confidence
    score: float           # 0-1 的置信度分数
    recommendation: str    # 对应用的建议


def compute_confidence(
    similarity: float,
    thresholds: dict = None
) -> ConfidenceResult:
    """
    基于相似度分数计算置信度

    参数:
        similarity: 归一化后的相似度分数 (0-1)
        thresholds: 自定义阈值，默认使用经验值

    默认阈值（基于余弦相似度）：
        - ≥ 0.85: 高置信度（文档高度匹配）
        - ≥ 0.70: 中等置信度（文档相关但可能不完全匹配）
        - ≥ 0.50: 低置信度（文档部分相关）
        - < 0.50: 极低置信度（文档不太相关）
    """
    if thresholds is None:
        thresholds = {
            "high": 0.85,
            "medium": 0.70,
            "low": 0.50,
        }

    if similarity >= thresholds["high"]:
        return ConfidenceResult(
            confidence=Confidence.HIGH,
            score=similarity,
            recommendation="可以直接使用检索结果回答用户问题"
        )
    elif similarity >= thresholds["medium"]:
        return ConfidenceResult(
            confidence=Confidence.MEDIUM,
            score=similarity,
            recommendation="可以参考检索结果，但建议提示用户验证"
        )
    elif similarity >= thresholds["low"]:
        return ConfidenceResult(
            confidence=Confidence.LOW,
            score=similarity,
            recommendation="检索结果相关性一般，建议告知用户可能无法准确回答"
        )
    else:
        return ConfidenceResult(
            confidence=Confidence.VERY_LOW,
            score=similarity,
            recommendation="检索结果不相关，建议直接告知用户知识库中没有相关信息"
        )


def adaptive_threshold_calibration(
    conn,
    table: str = "documents",
    sample_queries: list[str] = None,
    embedding_fn=None
) -> dict:
    """
    自适应阈值校准

    通过分析查询向量与库中向量的距离分布，自动确定合理的阈值

    原理：
    - 随机采样一组查询，计算它们与库中所有向量的距离
    - 分析距离分布的统计特征（均值、标准差、分位数）
    - 基于分布特征自动设定阈值
    """
    if sample_queries is None or embedding_fn is None:
        # 使用库中已有向量作为模拟查询
        with conn.cursor() as cur:
            cur.execute(f"""
                SELECT embedding FROM {table}
                ORDER BY RANDOM()
                LIMIT 20
            """)
            sample_embeddings = [np.array(row[0]) for row in cur.fetchall()]
    else:
        sample_embeddings = [embedding_fn(q) for q in sample_queries]

    all_distances = []
    with conn.cursor() as cur:
        for emb in sample_embeddings:
            cur.execute(f"""
                SELECT embedding <=> %s AS distance
                FROM {table}
                ORDER BY distance
                LIMIT 10
            """, (emb,))
            for row in cur.fetchall():
                all_distances.append(float(row[0]))

    if not all_distances:
        return {"high": 0.85, "medium": 0.70, "low": 0.50}

    all_distances = np.array(all_distances)

    # 将距离转换为相似度
    similarities = 1.0 - (all_distances / 2.0)

    # 基于统计分布设定阈值
    mean_sim = np.mean(similarities)
    std_sim = np.std(similarities)

    # 高置信度：均值 + 1 个标准差（高于平均水平的匹配）
    high_threshold = float(np.clip(mean_sim + std_sim, 0.7, 0.95))
    # 中等置信度：均值附近
    medium_threshold = float(np.clip(mean_sim, 0.5, 0.8))
    # 低置信度：均值 - 1 个标准差
    low_threshold = float(np.clip(mean_sim - std_sim, 0.3, 0.6))

    calibrated = {
        "high": round(high_threshold, 3),
        "medium": round(medium_threshold, 3),
        "low": round(low_threshold, 3),
    }

    print(f"距离分布统计:")
    print(f"  相似度均值: {mean_sim:.4f}")
    print(f"  相似度标准差: {std_sim:.4f}")
    print(f"  校准后阈值: {calibrated}")

    return calibrated


# ---- 使用示例 ----
if __name__ == "__main__":
    # 基本置信度评估
    test_similarities = [0.95, 0.80, 0.65, 0.45, 0.30]
    for sim in test_similarities:
        result = compute_confidence(sim)
        print(f"相似度 {sim:.2f} → {result.confidence.value} "
              f"(分数: {result.score:.2f}) | 建议: {result.recommendation}")

    # 自适应校准（需要有数据库连接）
    # conn = psycopg.connect(...)
    # thresholds = adaptive_threshold_calibration(conn)
    # 然后使用校准后的阈值:
    # result = compute_confidence(similarity, thresholds=thresholds)

"""
多路召回与结果融合排序

场景：
- 同一个问题用不同的 Embedding 模型编码，分别检索后融合
- 混合检索：向量检索 + 全文检索 + 元数据过滤，融合排序
- 多查询扩展：将一个问题扩展为多个子查询，分别检索后合并
"""
import psycopg
from pgvector.psycopg import register_vector
import numpy as np
from typing import List
from collections import defaultdict


def reciprocal_rank_fusion(
    result_lists: List[List[dict]],
    k: int = 60
) -> List[dict]:
    """
    RRF（Reciprocal Rank Fusion）倒数排名融合

    这是信息检索中最经典的结果融合算法。
    不依赖于原始分数，只依赖于排名位置。

    RRF 分数公式：
        score(d) = Σ 1 / (k + rank_i(d))

    其中 k 是平滑常数（默认 60），rank_i(d) 是文档 d 在第 i 个列表中的排名（从 1 开始）。

    参数:
        result_lists: 多路检索结果列表，每路是一个按相关性排序的文档列表
        k: RRF 平滑参数（越大，排名差异的影响越小）

    返回:
        融合排序后的文档列表
    """
    # 用文档 ID 作为唯一键
    doc_scores = defaultdict(float)
    doc_info = {}

    for results in result_lists:
        for rank, doc in enumerate(results):
            doc_id = doc["id"]
            # RRF 分数
            doc_scores[doc_id] += 1.0 / (k + rank + 1)  # rank 从 0 开始，+1 转为从 1 开始
            # 保留最新的文档信息
            if doc_id not in doc_info:
                doc_info[doc_id] = doc

    # 按 RRF 分数降序排序
    sorted_ids = sorted(doc_scores.keys(), key=lambda x: doc_scores[x], reverse=True)

    fused_results = []
    for doc_id in sorted_ids:
        doc = doc_info[doc_id]
        doc["rrf_score"] = doc_scores[doc_id]
        fused_results.append(doc)

    return fused_results


def weighted_score_fusion(
    result_lists: List[List[dict]],
    weights: List[float] = None,
    score_key: str = "distance"
) -> List[dict]:
    """
    加权分数融合

    将多路检索的分数按权重加权后排序。
    注意：需要将距离转换为"越大越好"的分数。

    参数:
        result_lists: 多路检索结果列表
        weights: 每路的权重（默认等权重）
        score_key: 分数字段名
    """
    if weights is None:
        weights = [1.0 / len(result_lists)] * len(result_lists)

    assert len(result_lists) == len(weights), "结果列表数量与权重数量不匹配"

    doc_scores = defaultdict(float)
    doc_info = {}

    for results, weight in zip(result_lists, weights):
        for doc in results:
            doc_id = doc["id"]
            # 将距离转换为分数（距离越小，分数越高）
            # 使用 1 / (1 + distance) 作为分数
            score = 1.0 / (1.0 + doc.get(score_key, 0.0))
            doc_scores[doc_id] += weight * score
            if doc_id not in doc_info:
                doc_info[doc_id] = doc

    sorted_ids = sorted(doc_scores.keys(), key=lambda x: doc_scores[x], reverse=True)

    fused_results = []
    for doc_id in sorted_ids:
        doc = doc_info[doc_id]
        doc["fused_score"] = doc_scores[doc_id]
        fused_results.append(doc)

    return fused_results


def multi_vector_search(
    conn,
    query_embeddings: List[np.ndarray],
    table: str = "documents",
    top_k: int = 5,
    fetch_k: int = 10,
    fusion_method: str = "rrf"
) -> List[dict]:
    """
    多向量融合检索

    使用多个查询向量分别检索，然后融合结果。

    适用场景：
    - 多查询扩展：一个问题生成多个变体查询
    - 多模型检索：使用不同的 Embedding 模型
    - 多语言检索：不同语言的查询向量

    参数:
        conn: 数据库连接
        query_embeddings: 查询向量列表
        table: 表名
        top_k: 最终返回的结果数
        fetch_k: 每路检索的候选数
        fusion_method: 融合方法 ("rrf" 或 "weighted")
    """
    result_lists = []

    with conn.cursor() as cur:
        for emb in query_embeddings:
            cur.execute(f"""
                SELECT id, content, source,
                       embedding <=> %s AS distance
                FROM {table}
                ORDER BY distance
                LIMIT %s
            """, (emb, fetch_k))

            results = []
            for row in cur.fetchall():
                results.append({
                    "id": row[0],
                    "content": row[1],
                    "source": row[2],
                    "distance": float(row[3])
                })
            result_lists.append(results)

    # 融合
    if fusion_method == "rrf":
        fused = reciprocal_rank_fusion(result_lists)
    else:
        fused = weighted_score_fusion(result_lists)

    return fused[:top_k]


def hybrid_search(
    conn,
    query_embedding: np.ndarray,
    query_text: str,
    table: str = "documents",
    top_k: int = 5,
    vector_weight: float = 0.7,
    text_weight: float = 0.3
) -> List[dict]:
    """
    混合检索：向量语义检索 + 全文关键词检索

    将两路检索结果通过 RRF 或加权融合。

    参数:
        query_embedding: 查询向量
        query_text: 查询文本（用于全文检索）
        table: 表名
        top_k: 返回结果数
        vector_weight: 向量检索权重
        text_weight: 全文检索权重
    """
    vector_results = []
    text_results = []

    with conn.cursor() as cur:
        # 向量检索
        cur.execute(f"""
            SELECT id, content, source,
                   embedding <=> %s AS distance
            FROM {table}
            ORDER BY distance
            LIMIT %s
        """, (query_embedding, top_k * 3))

        for row in cur.fetchall():
            vector_results.append({
                "id": row[0],
                "content": row[1],
                "source": row[2],
                "distance": float(row[3])
            })

        # 全文检索（使用 PostgreSQL 的全文搜索）
        cur.execute(f"""
            SELECT id, content, source,
                   ts_rank(
                       to_tsvector('simple', content),
                       plainto_tsquery('simple', %s)
                   ) AS rank
            FROM {table}
            WHERE to_tsvector('simple', content) @@ plainto_tsquery('simple', %s)
            ORDER BY rank DESC
            LIMIT %s
        """, (query_text, query_text, top_k * 3))

        for row in cur.fetchall():
            text_results.append({
                "id": row[0],
                "content": row[1],
                "source": row[2],
                # 将 rank 转换为距离形式（越大越好 → 越小越好）
                "distance": 1.0 / (1.0 + float(row[3])) if row[3] else 1.0
            })

    # 加权融合
    fused = weighted_score_fusion(
        [vector_results, text_results],
        weights=[vector_weight, text_weight],
        score_key="distance"
    )

    return fused[:top_k]


# ---- 使用示例 ----
if __name__ == "__main__":
    conn = psycopg.connect(
        "postgresql://postgres:your_password@localhost:5432/vector_db",
        autocommit=True
    )
    register_vector(conn)

    # 多向量融合检索（模拟查询扩展）
    # 实际场景中，这3个向量可能来自：原始问题、改写后的问题、翻译后的问题
    query_vecs = [
        np.random.rand(1024).astype(np.float32) for _ in range(3)
    ]
    for v in query_vecs:
        v /= np.linalg.norm(v)

    print("=== 多向量 RRF 融合检索 ===")
    fused = multi_vector_search(
        conn, query_vecs,
        table="documents",
        top_k=5,
        fusion_method="rrf"
    )
    for r in fused:
        print(f"  ID={r['id']}, RRF分数={r['rrf_score']:.4f}, "
              f"内容={r['content'][:60]}...")

    # 混合检索（向量 + 全文）
    print("\n=== 混合检索（向量 70% + 全文 30%）===")
    hybrid = hybrid_search(
        conn,
        query_embedding=query_vecs[0],
        query_text="pgvector 向量搜索",
        table="documents",
        top_k=5,
        vector_weight=0.7,
        text_weight=0.3
    )
    for r in hybrid:
        print(f"  ID={r['id']}, 融合分数={r['fused_score']:.4f}, "
              f"内容={r['content'][:60]}...")

    conn.close()

# postgresql.conf

# 假设服务器有 64GB 内存
shared_buffers = 16GB          # 系统内存的 25%

# 如果是专门的向量检索服务器（128GB 内存）
# shared_buffers = 40GB        # 可适当增大，但不超过 40%

# postgresql.conf

# 通用场景
maintenance_work_mem = 512MB

# 创建 HNSW 索引时（推荐临时调大）
# 方法一：全局设置（索引构建完成后建议调回）
maintenance_work_mem = 2GB

# 方法二：会话级设置（推荐，只对当前会话的索引构建生效）
SET maintenance_work_mem = '2GB';
CREATE INDEX ON documents USING hnsw (embedding vector_cosine_ops)
  WITH (m = 16, ef_construction = 64);
-- 索引创建完成后，该会话设置自动失效

# postgresql.conf

# 通用场景
work_mem = 16MB

# 复杂向量检索场景
work_mem = 64MB

# 会话级动态调整（对复杂查询临时调大）
SET work_mem = '128MB';
SELECT id, embedding <=> query_vector AS distance
FROM documents
ORDER BY distance
LIMIT 100;

# postgresql.conf

# 假设服务器有 64GB 内存
effective_cache_size = 48GB    # 系统内存的 75%

# postgresql.conf

# 推荐设为 CPU 核心数的一半，但不超过 8
max_parallel_maintenance_workers = 4   # 4核以上服务器

# 创建索引时可以会话级临时调大
SET max_parallel_maintenance_workers = 8;
SET maintenance_work_mem = '2GB';
CREATE INDEX ON items USING hnsw (embedding vector_cosine_ops);

# postgresql.conf
max_parallel_workers_per_gather = 4
max_parallel_workers = 8        # 系统级并行工作进程上限

# postgresql.conf

wal_buffers = 64MB              # WAL 缓冲区，默认 -1（自动计算）通常够用
                                 # 向量批量写入场景可显式设为 64MB

max_wal_size = 4GB              # WAL 日志最大总量，默认 1GB
                                 # 向量批量写入时建议 4GB~8GB

min_wal_size = 1GB              # WAL 日志最小保留量

checkpoint_completion_target = 0.9  # 检查点完成目标，默认 0.5
                                     # 设为 0.9 可以平滑写入压力
                                     # 让检查点更慢地完成，减少IO突发

# postgresql.conf

# SSD 环境（推荐）
random_page_cost = 1.1

# NVMe SSD 环境
random_page_cost = 1.0

# 机械硬盘（保持默认）
# random_page_cost = 4.0

-- 会话级设置（仅当前会话有效）
SET hnsw.ef_search = 100;      -- 推荐范围：40~200

-- 在事务内设置（事务结束自动恢复）
BEGIN;
SET LOCAL hnsw.ef_search = 200;  -- 高精度模式
SELECT id, title
FROM documents
ORDER BY embedding <=> '[0.1, 0.2, ...]'::vector
LIMIT 10;
COMMIT;
-- COMMIT 后 ef_search 自动恢复为会话级值

-- 全局设置（需要超级用户权限，写入 postgresql.conf）
-- postgresql.conf 中添加：
-- hnsw.ef_search = 100
-- 然后执行：
SELECT pg_reload_conf();

-- 会话级设置
SET ivfflat.probes = 10;       -- 推荐范围：5~50

-- 经验法则：probes ≈ sqrt(lists)
-- 如果创建索引时 lists = 1000，则 probes ≈ 32
-- 如果创建索引时 lists = 100，则 probes ≈ 10

-- 全局设置方式同上
-- postgresql.conf 中添加：
-- ivfflat.probes = 10

# ============================================================
# pgvector 生产环境调优配置模板
# 服务器配置：64GB RAM, 8核 CPU, NVMe SSD
# ============================================================

# --- 内存参数 ---
shared_buffers = 16GB
maintenance_work_mem = 2GB
work_mem = 32MB
effective_cache_size = 48GB

# --- IO 与并发 ---
max_parallel_maintenance_workers = 4
max_parallel_workers_per_gather = 4
max_parallel_workers = 8
random_page_cost = 1.1

# --- WAL 参数 ---
wal_buffers = 64MB
max_wal_size = 4GB
min_wal_size = 1GB
checkpoint_completion_target = 0.9

# --- pgvector 专属参数 ---
hnsw.ef_search = 100
ivfflat.probes = 10

# --- 连接与超时 ---
max_connections = 200
statement_timeout = 5000           # 5秒查询超时保护

-- 手动分表：按月创建独立表
CREATE TABLE documents_2025_01 (
    id BIGSERIAL PRIMARY KEY,
    title TEXT,
    embedding vector(1536),
    created_at TIMESTAMPTZ DEFAULT now()
);

CREATE TABLE documents_2025_02 (LIKE documents_2025_01 INCLUDING ALL);
CREATE TABLE documents_2025_03 (LIKE documents_2025_01 INCLUDING ALL);
-- ... 每月一张表

-- 为每张表独立创建向量索引
CREATE INDEX ON documents_2025_01 USING hnsw (embedding vector_cosine_ops);
CREATE INDEX ON documents_2025_02 USING hnsw (embedding vector_cosine_ops);
CREATE INDEX ON documents_2025_03 USING hnsw (embedding vector_cosine_ops);

-- 查询时通过应用层路由到具体表
SELECT id, title, embedding <=> $1 AS distance
FROM documents_2025_06
ORDER BY distance
LIMIT 10;

-- 按业务类型分表
CREATE TABLE product_embeddings (
    product_id BIGINT PRIMARY KEY,
    embedding vector(1536),
    category TEXT,
    updated_at TIMESTAMPTZ DEFAULT now()
);

CREATE TABLE article_embeddings (
    article_id BIGINT PRIMARY KEY,
    embedding vector(1536),
    source TEXT,
    updated_at TIMESTAMPTZ DEFAULT now()
);

CREATE TABLE user_embeddings (
    user_id BIGINT PRIMARY KEY,
    embedding vector(768),           -- 用户向量维度可能与商品不同
    segment TEXT,
    updated_at TIMESTAMPTZ DEFAULT now()
);

-- 每张表独立建索引，参数可以按业务需求调整
CREATE INDEX ON product_embeddings USING hnsw (embedding vector_cosine_ops)
  WITH (m = 32, ef_construction = 128);   -- 商品向量精度高，参数调大

CREATE INDEX ON user_embeddings USING hnsw (embedding vector_cosine_ops)
  WITH (m = 16, ef_construction = 64);    -- 用户向量精度要求低，参数适中

-- 创建分区主表
CREATE TABLE vectors (
    id BIGSERIAL,
    content TEXT,
    embedding vector(1536),
    created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
    category TEXT
) PARTITION BY RANGE (created_at);

-- 创建月度分区
CREATE TABLE vectors_2025_q1 PARTITION OF vectors
    FOR VALUES FROM ('2025-01-01') TO ('2025-04-01');

CREATE TABLE vectors_2025_q2 PARTITION OF vectors
    FOR VALUES FROM ('2025-04-01') TO ('2025-07-01');

CREATE TABLE vectors_2025_q3 PARTITION OF vectors
    FOR VALUES FROM ('2025-07-01') TO ('2025-10-01');

CREATE TABLE vectors_2025_q4 PARTITION OF vectors
    FOR VALUES FROM ('2025-10-01') TO ('2026-01-01');

-- 为每个分区独立创建向量索引（分区索引比全局索引更小、更快）
CREATE INDEX ON vectors_2025_q1 USING hnsw (embedding vector_cosine_ops);
CREATE INDEX ON vectors_2025_q2 USING hnsw (embedding vector_cosine_ops);
CREATE INDEX ON vectors_2025_q3 USING hnsw (embedding vector_cosine_ops);
CREATE INDEX ON vectors_2025_q4 USING hnsw (embedding vector_cosine_ops);

-- 查询时自动分区裁剪（Partition Pruning）
-- 如果查询包含分区键条件，PostgreSQL 只扫描相关分区
SELECT id, content, embedding <=> $1 AS distance
FROM vectors
WHERE created_at >= '2025-07-01'    -- 只会扫描 Q3、Q4 分区
ORDER BY distance
LIMIT 10;

-- 使用 pg_partman 自动管理分区
CREATE EXTENSION pg_partman;

SELECT partman.create_parent(
    p_parent_table   => 'public.vectors',
    p_control        => 'created_at',
    p_type           => 'range',
    p_interval       => '3 months',
    p_premake        => 2        -- 提前创建未来 2 个分区
);

-- 或者用简单的定时 SQL 手动创建下季度分区
-- 建议在 pg_cron 或外部调度中执行
CREATE TABLE IF NOT EXISTS vectors_2026_q1
    PARTITION OF vectors
    FOR VALUES FROM ('2026-01-01') TO ('2026-04-01');

CREATE INDEX IF NOT EXISTS vectors_2026_q1_hnsw_idx
    ON vectors_2026_q1 USING hnsw (embedding vector_cosine_ops);

-- 安装 pg_cron 扩展
CREATE EXTENSION IF NOT EXISTS pg_cron;

-- 每天凌晨 2 点清理 6 个月前的向量数据
SELECT cron.schedule(
    'cleanup-old-vectors',          -- 任务名称
    '0 2 * * *',                     -- cron 表达式：每天 02:00
    $$
    DELETE FROM vectors
    WHERE created_at < now() - INTERVAL '6 months'
    $$
);

-- 清理后回收空间（每周日凌晨 3 点）
SELECT cron.schedule(
    'vacuum-vectors',
    '0 3 * * 0',
    $$VACUUM ANALYZE vectors$$
);

-- 查看已注册的定时任务
SELECT * FROM cron.job;

-- 取消定时任务
SELECT cron.unschedule('cleanup-old-vectors');

-- 快速移除过期分区（O(1) 操作，不产生大量 WAL）
ALTER TABLE vectors DETACH PARTITION vectors_2024_q1;
DROP TABLE vectors_2024_q1;

-- 如果需要先归档再删除
ALTER TABLE vectors DETACH PARTITION vectors_2024_q1;
-- 使用 pg_dump 备份该分区
-- pg_dump -t vectors_2024_q1 mydb > vectors_2024_q1_backup.sql
DROP TABLE vectors_2024_q1;

#!/usr/bin/env python3
"""vector_cleanup.py - 过期向量清理脚本"""
import psycopg2
from datetime import datetime, timedelta

DB_CONFIG = {
    "host": "localhost",
    "port": 5432,
    "dbname": "vectordb",
    "user": "admin",
    "password": "secure_password"
}

RETENTION_DAYS = 180
BATCH_SIZE = 10000  # 分批删除，避免长事务

def cleanup_expired_vectors():
    conn = psycopg2.connect(**DB_CONFIG)
    conn.autocommit = False
    cur = conn.cursor()

    cutoff = datetime.now() - timedelta(days=RETENTION_DAYS)
    total_deleted = 0

    while True:
        cur.execute("""
            DELETE FROM documents
            WHERE id IN (
                SELECT id FROM documents
                WHERE created_at < %s
                LIMIT %s
            )
        """, (cutoff, BATCH_SIZE))
        deleted = cur.rowcount
        conn.commit()
        total_deleted += deleted

        if deleted < BATCH_SIZE:
            break
        print(f"已删除 {total_deleted} 条过期向量...")

    # 更新统计信息
    cur.execute("ANALYZE documents")
    conn.commit()

    print(f"清理完成，共删除 {total_deleted} 条记录")
    cur.close()
    conn.close()

if __name__ == "__main__":
    cleanup_expired_vectors()

# crontab 配置：每天凌晨 2 点执行
0 2 * * * /usr/bin/python3 /opt/scripts/vector_cleanup.py >> /var/log/vector_cleanup.log 2>&1

-- 热数据表：最近 3 个月的向量，SSD 存储，HNSW 索引
CREATE TABLE vectors_hot (
    id BIGSERIAL PRIMARY KEY,
    embedding vector(1536) NOT NULL,
    content TEXT,
    created_at TIMESTAMPTZ DEFAULT now(),
    access_count INT DEFAULT 0      -- 访问计数器
);

-- HNSW 索引，高精度参数
CREATE INDEX ON vectors_hot USING hnsw (embedding vector_cosine_ops)
  WITH (m = 32, ef_construction = 128);

-- 放在 SSD 表空间
ALTER TABLE vectors_hot SET TABLESPACE ssd_tablespace;
ALTER INDEX vectors_hot_embedding_idx SET TABLESPACE ssd_tablespace;

-- 冷数据表：3 个月前的向量，HDD 存储，IVFFlat 索引或无索引
CREATE TABLE vectors_cold (
    id BIGINT PRIMARY KEY,
    embedding vector(1536) NOT NULL,
    content TEXT,
    created_at TIMESTAMPTZ,
    archived_at TIMESTAMPTZ DEFAULT now()
);

-- IVFFlat 索引，节省存储和内存
CREATE INDEX ON vectors_cold USING ivfflat (embedding vector_cosine_ops)
  WITH (lists = 500);

-- 放在 HDD 表空间
ALTER TABLE vectors_cold SET TABLESPACE hdd_tablespace;

-- 定期将热数据中的过期记录迁移到冷表
BEGIN;

-- 1. 将 3 个月前的热数据插入冷表
INSERT INTO vectors_cold (id, embedding, content, created_at)
SELECT id, embedding, content, created_at
FROM vectors_hot
WHERE created_at < now() - INTERVAL '3 months';

-- 2. 从热表中删除
DELETE FROM vectors_hot
WHERE created_at < now() - INTERVAL '3 months';

-- 3. 回收热表空间
VACUUM ANALYZE vectors_hot;

COMMIT;

-- 创建 SSD 表空间（需要操作系统先挂载 SSD）
CREATE TABLESPACE ssd_tablespace LOCATION '/mnt/ssd/pg_data';

-- 创建 HDD 表空间
CREATE TABLESPACE hdd_tablespace LOCATION '/mnt/hdd/pg_data';

-- 安装扩展
CREATE EXTENSION IF NOT EXISTS pg_stat_statements;

-- 查看向量检索相关查询的性能统计
SELECT
    query,
    calls AS total_calls,
    round(mean_exec_time::numeric, 2) AS avg_ms,
    round(max_exec_time::numeric, 2) AS max_ms,
    round(min_exec_time::numeric, 2) AS min_ms,
    round(stddev_exec_time::numeric, 2) AS stddev_ms,
    rows AS total_rows
FROM pg_stat_statements
WHERE query LIKE '%<=>%'
   OR query LIKE '%<#>%'
   OR query LIKE '%<+>%'
   OR query LIKE '%vector%'
ORDER BY mean_exec_time DESC
LIMIT 20;

-- 重置统计信息（定期重置，避免历史数据累积）
SELECT pg_stat_statements_reset();

-- 检查向量索引的使用情况
SELECT
    schemaname,
    relname AS table_name,
    indexrelname AS index_name,
    idx_scan AS index_scans,         -- 索引扫描次数
    idx_tup_read AS tuples_read,     -- 通过索引读取的行数
    idx_tup_fetch AS tuples_fetched  -- 通过索引实际获取的行数
FROM pg_stat_user_indexes
WHERE indexrelname LIKE '%hnsw%'
   OR indexrelname LIKE '%ivfflat%'
   OR indexrelname LIKE '%vector%'
ORDER BY idx_scan DESC;

-- 检查索引是否被有效使用
-- 如果 idx_scan 长期为 0，说明索引未被使用（可能是查询无法命中索引）

-- 检查表的索引扫描 vs 全表扫描比例
SELECT
    schemaname,
    relname AS table_name,
    seq_scan AS full_table_scans,
    idx_scan AS index_scans,
    CASE WHEN seq_scan + idx_scan > 0
         THEN round(100.0 * idx_scan / (seq_scan + idx_scan), 1)
         ELSE 0
    END AS index_hit_rate_pct
FROM pg_stat_user_tables
WHERE relname LIKE '%vector%'
   OR relname LIKE '%embedding%'
   OR relname LIKE '%document%'
ORDER BY index_hit_rate_pct ASC;

-- 安装 pgstattuple 扩展
CREATE EXTENSION IF NOT EXISTS pgstattuple;

-- 检查向量表的膨胀情况
SELECT
    table_len,
    tuple_count,
    tuple_len,
    round(100.0 * tuple_len / NULLIF(table_len, 0), 1) AS live_pct,
    dead_tuple_count,
    dead_tuple_len,
    round(100.0 * dead_tuple_len / NULLIF(table_len, 0), 1) AS dead_pct,
    free_space,
    round(100.0 * free_space / NULLIF(table_len, 0), 1) AS free_pct
FROM pgstattuple('documents');

-- 如果 dead_pct > 20%，说明需要执行 VACUUM FULL 或等待 autovacuum 清理
-- 如果 free_pct 很高但 dead_pct 低，说明空间已被标记为可复用，等待 autovacuum

-- 监控向量表的写入活动
SELECT
    schemaname,
    relname AS table_name,
    n_tup_ins AS inserts,      -- 插入行数
    n_tup_upd AS updates,      -- 更新行数
    n_tup_del AS deletes,      -- 删除行数
    n_live_tup AS live_tuples, -- 当前活跃行数
    n_dead_tup AS dead_tuples, -- 当前死元组数
    last_vacuum,
    last_autovacuum,
    last_analyze,
    last_autoanalyze
FROM pg_stat_user_tables
WHERE relname LIKE '%vector%'
   OR relname LIKE '%embedding%'
ORDER BY n_tup_ins DESC;

-- 当前活跃连接及状态
SELECT
    pid,
    usename,
    application_name,
    client_addr,
    state,
    wait_event_type,
    wait_event,
    query_start,
    now() - query_start AS duration,
    left(query, 100) AS query_preview
FROM pg_stat_activity
WHERE state != 'idle'
ORDER BY duration DESC;

-- 连接数统计
SELECT
    state,
    count(*) AS count,
    count(*) FILTER (WHERE query LIKE '%vector%' OR query LIKE '%<=>%')
        AS vector_queries
FROM pg_stat_activity
GROUP BY state;

-- 检查是否有锁等待
SELECT
    blocked_locks.pid AS blocked_pid,
    blocked_activity.usename AS blocked_user,
    blocking_locks.pid AS blocking_pid,
    blocking_activity.usename AS blocking_user,
    blocked_activity.query AS blocked_query
FROM pg_catalog.pg_locks blocked_locks
JOIN pg_catalog.pg_stat_activity blocked_activity ON blocked_activity.pid = blocked_locks.pid
JOIN pg_catalog.pg_locks blocking_locks ON blocking_locks.locktype = blocked_locks.locktype
    AND blocking_locks.database IS NOT DISTINCT FROM blocked_locks.database
    AND blocking_locks.relation IS NOT DISTINCT FROM blocked_locks.relation
    AND blocking_locks.pid != blocked_locks.pid
JOIN pg_catalog.pg_stat_activity blocking_activity ON blocking_activity.pid = blocking_locks.pid
WHERE NOT blocked_locks.granted;

# 备份整个数据库（包含向量数据和索引定义）
pg_dump -h localhost -U postgres -Fc vectordb > vectordb_backup.dump

# 只备份特定向量表
pg_dump -h localhost -U postgres -t documents -Fc vectordb > documents_backup.dump

# 备份表结构（不包含数据，用于查看索引定义）
pg_dump -h localhost -U postgres -s vectordb > schema_backup.sql

# 恢复逻辑备份
pg_restore -h localhost -U postgres -d vectordb_restored vectordb_backup.dump

# 注意事项：
# 1. 恢复时会自动重建向量索引，大数据量下可能需要较长时间
# 2. 如果只恢复数据不恢复索引（加快恢复速度）：
pg_restore --no-indexes -h localhost -U postgres -d vectordb_restored vectordb_backup.dump
# 然后手动创建索引（可以调大 maintenance_work_mem 加速）

# 执行基础物理备份
pg_basebackup -h localhost -U replicator -D /backup/pg_base_20250612 \
    -Ft -z -P -Xstream

# 参数说明：
# -Ft     : tar 格式
# -z      : gzip 压缩
# -P      : 显示进度
# -Xstream: 流式传输 WAL 日志，保证备份一致性

# 物理备份包含所有向量数据和索引的完整二进制文件
# 恢复时不需要重建索引（索引文件直接复制）

# 步骤 1：在源实例导出数据（不含索引，加快导出速度）
pg_dump -h source_host -U postgres -t documents --no-indexes -Fc vectordb > docs_data.dump

# 步骤 2：在目标实例创建数据库和扩展
psql -h target_host -U postgres -c "CREATE DATABASE vectordb_new;"
psql -h target_host -U postgres -d vectordb_new -c "CREATE EXTENSION vector;"

# 步骤 3：恢复数据
pg_restore -h target_host -U postgres -d vectordb_new docs_data.dump

# 步骤 4：在目标实例创建索引（利用目标服务器的计算资源）
psql -h target_host -U postgres -d vectordb_new << 'EOF'
SET maintenance_work_mem = '2GB';
SET max_parallel_maintenance_workers = 4;
CREATE INDEX ON documents USING hnsw (embedding vector_cosine_ops)
  WITH (m = 16, ef_construction = 64);
ANALYZE documents;
EOF

# 步骤 5：验证数据完整性
psql -h target_host -U postgres -d vectordb_new -c "
SELECT count(*) AS total_rows FROM documents;
SELECT avg(vector_dims(embedding)) AS avg_dims FROM documents;
"

-- 重建单个向量索引
REINDEX INDEX CONCURRENTLY documents_embedding_idx;

-- CONCURRENTLY 关键字很重要：
-- 它允许在重建索引期间继续处理查询（不锁定表）
-- 但重建时间会更长

-- 重建整个表的所有索引
REINDEX TABLE CONCURRENTLY documents;

-- 建议：每月或每季度执行一次 REINDEX，或者在大批量数据导入后执行

-- 日常维护：VACUUM（由 autovacuum 自动执行）
VACUUM ANALYZE documents;

-- 深度清理：VACUUM FULL（需要维护窗口）
-- 注意：VACUUM FULL 会锁定表，期间无法读写！
-- 建议在业务低峰期执行
VACUUM FULL documents;

-- 替代方案：使用 pg_repack 在线重组（无需排他锁）
-- 安装：CREATE EXTENSION pg_repack;
-- 执行：pg_repack -t documents -d vectordb

-- 为特定向量表设置更激进的 autovacuum 参数
ALTER TABLE documents SET (
    autovacuum_vacuum_threshold = 1000,        -- 基础阈值：1000 行
    autovacuum_vacuum_scale_factor = 0.01,     -- 比例因子：1%（默认 20%）
    autovacuum_analyze_threshold = 500,         -- 分析阈值
    autovacuum_analyze_scale_factor = 0.005,   -- 分析比例因子
    autovacuum_vacuum_cost_delay = 2,          -- 降低延迟，加快清理速度
    autovacuum_vacuum_cost_limit = 1000        -- 提高成本限制
);

-- 含义：当 1000 + 1% * 总行数 的行发生变更时，触发 autovacuum
-- 例如 100 万行的表，当 11000 行变更时就会触发（而非默认的 201000 行）

-- 检查索引的碎片情况（通过 pgstatindex 函数）
CREATE EXTENSION IF NOT EXISTS pgstattuple;

-- 对于 HNSW 索引，主要关注索引大小变化
SELECT
    indexrelname,
    pg_size_pretty(pg_relation_size(indexrelid)) AS index_size,
    idx_scan,
    idx_tup_read,
    idx_tup_fetch
FROM pg_stat_user_indexes
WHERE relname = 'documents';

-- 如果索引大小持续增长但数据量未显著增加，说明存在碎片
-- 此时应执行 REINDEX CONCURRENTLY 重建索引

-- 方案 1：调大 maintenance_work_mem（治标）
SET maintenance_work_mem = '2GB';

-- 方案 2：降低 ef_construction（治本，但会降低索引质量）
-- 重建索引时使用较小的 ef_construction
CREATE INDEX ON documents USING hnsw (embedding vector_cosine_ops)
  WITH (m = 16, ef_construction = 40);  -- 从 64 降到 40

-- 方案 3：分批写入（推荐的生产实践）
-- 不要在已有 HNSW 索引的情况下大批量 INSERT
-- 而是先删除索引，批量导入，再重建索引

-- 步骤 1：删除索引
DROP INDEX documents_embedding_idx;

-- 步骤 2：批量导入（无索引维护，速度快）
\COPY documents (title, embedding) FROM '/data/embeddings.csv' CSV;

-- 步骤 3：重建索引
SET maintenance_work_mem = '2GB';
SET max_parallel_maintenance_workers = 4;
CREATE INDEX ON documents USING hnsw (embedding vector_cosine_ops)
  WITH (m = 16, ef_construction = 64);

-- 方案 4：增量写入时控制批次大小
-- 每次 INSERT 不超过 1000 行，间隔 100ms
DO $$
DECLARE
    batch_size INT := 1000;
    offset_val INT := 0;
    total INT;
BEGIN
    SELECT count(*) INTO total FROM staging_vectors;
    WHILE offset_val < total LOOP
        INSERT INTO documents (title, embedding)
        SELECT title, embedding
        FROM staging_vectors
        ORDER BY id
        LIMIT batch_size OFFSET offset_val;

        offset_val := offset_val + batch_size;
        PERFORM pg_sleep(0.1);  -- 间隔 100ms，让 autovacuum 有时间工作
    END LOOP;
END $$;

-- 错误示例：创建了 L2 距离索引，但查询用余弦距离
CREATE INDEX ON docs USING hnsw (embedding vector_l2_ops);

-- 这条查询不会使用上面的索引！
SELECT id FROM docs ORDER BY embedding <=> '[0.1, 0.2]'::vector LIMIT 5;
-- <=> 是余弦距离运算符，但索引是 L2 距离的

-- 正确做法：确保操作符类和运算符匹配
-- 余弦距离：vector_cosine_ops + <=>
CREATE INDEX ON docs USING hnsw (embedding vector_cosine_ops);
SELECT id FROM docs ORDER BY embedding <=> '[0.1, 0.2]'::vector LIMIT 5;

-- L2 距离：vector_l2_ops + <->
CREATE INDEX ON docs USING hnsw (embedding vector_l2_ops);
SELECT id FROM docs ORDER BY embedding <-> '[0.1, 0.2]'::vector LIMIT 5;

-- 内积距离：vector_ip_ops + <#>（注意返回负内积）
CREATE INDEX ON docs USING hnsw (embedding vector_ip_ops);
SELECT id FROM docs ORDER BY embedding <#> '[0.1, 0.2]'::vector LIMIT 5;

-- L1 距离：vector_l1_ops + <+>
CREATE INDEX ON docs USING hnsw (embedding vector_l1_ops);
SELECT id FROM docs ORDER BY embedding <+> '[0.1, 0.2]'::vector LIMIT 5;

-- 错误：对向量列使用了函数转换
CREATE INDEX ON docs USING hnsw (embedding vector_cosine_ops);

-- 不会走索引！因为 embedding 被 l2_normalize() 转换了
SELECT id FROM docs
ORDER BY l2_normalize(embedding) <=> '[0.1, 0.2]'::vector
LIMIT 5;

-- 正确做法 A：存储时就做好归一化，查询时直接用原始列
-- 写入时：
INSERT INTO docs (embedding) VALUES (l2_normalize('[1, 2, 3]'::vector));
-- 查询时：
SELECT id FROM docs ORDER BY embedding <=> '[0.1, 0.2]'::vector LIMIT 5;

-- 正确做法 B：创建表达式索引（如果确实需要动态转换）
CREATE INDEX ON docs USING hnsw (l2_normalize(embedding) vector_cosine_ops);
SELECT id FROM docs
ORDER BY l2_normalize(embedding) <=> '[0.1, 0.2]'::vector
LIMIT 5;

-- IVFFlat 索引，probes 设置过小导致无法命中正确的聚类
SET ivfflat.probes = 1;  -- 只搜索 1 个聚类，极可能错过最近邻

-- HNSW 索引，ef_search 过小导致搜索宽度不够
SET hnsw.ef_search = 1;  -- 几乎等于贪婪搜索，召回率极低

-- 诊断方法：使用 EXPLAIN ANALYZE 检查执行计划
EXPLAIN ANALYZE
SELECT id FROM docs ORDER BY embedding <=> '[0.1, 0.2]'::vector LIMIT 5;

-- 如果看到 "Index Scan using ... (never executed)" 或者执行计划中没有
-- 使用索引，检查上述参数设置

-- 错误做法：先创建 IVFFlat 索引，再导入数据
CREATE TABLE docs (id SERIAL, embedding vector(3));
CREATE INDEX ON docs USING ivfflat (embedding vector_l2_ops) WITH (lists = 100);
-- 此时表为空！K-means 聚类产生空聚类中心，索引质量极差

INSERT INTO docs (embedding) SELECT ... FROM generate_series(1, 100000);
-- 数据导入后，索引基于空数据的聚类中心，效果很差

-- 正确做法：先导入数据，再创建 IVFFlat 索引
CREATE TABLE docs (id SERIAL, embedding vector(3));
INSERT INTO docs (embedding) SELECT ... FROM generate_series(1, 100000);
-- 数据导入完成后再创建索引
CREATE INDEX ON docs USING ivfflat (embedding vector_l2_ops) WITH (lists = 100);
-- K-means 基于实际数据聚类，索引质量好

-- 如果已经犯了空表建索引的错误，重建索引
REINDEX INDEX docs_embedding_idx;

-- 当表中只有几百条数据时，PostgreSQL 优化器可能认为
-- "全表扫描比走索引更快"（因为索引扫描有额外的随机 IO 开销）

-- 诊断：
EXPLAIN ANALYZE
SELECT id FROM docs ORDER BY embedding <=> '[0.1, 0.2]'::vector LIMIT 5;

-- 如果结果是 Seq Scan 而非 Index Scan，可能是数据量太小
-- 对于小表（< 几千行），全表扫描可能确实更快，这不是问题

-- 如果想强制使用索引（仅用于测试）：
SET enable_seqscan = off;
EXPLAIN ANALYZE
SELECT id FROM docs ORDER BY embedding <=> '[0.1, 0.2]'::vector LIMIT 5;
-- 注意：测试完成后一定要恢复
SET enable_seqscan = on;

-- 当怀疑索引不生效时，按以下步骤排查：

-- 第 1 步：确认索引存在
SELECT indexname, indexdef
FROM pg_indexes
WHERE tablename = 'documents';

-- 第 2 步：检查执行计划
EXPLAIN (ANALYZE, BUFFERS, FORMAT TEXT)
SELECT id, title
FROM documents
ORDER BY embedding <=> '[0.1, 0.2, ...]'::vector
LIMIT 10;

-- 第 3 步：检查操作符类是否匹配
-- 查看索引定义中的 ops 类型，对比查询中使用的运算符

-- 第 4 步：检查运行时参数
SHOW hnsw.ef_search;
SHOW ivfflat.probes;

-- 第 5 步：检查索引统计
SELECT idx_scan, idx_tup_read, idx_tup_fetch
FROM pg_stat_user_indexes
WHERE relname = 'documents';

; pgbouncer.ini

[databases]
vectordb = host=127.0.0.1 port=5432 dbname=vectordb

[pgbouncer]
listen_port = 6432
listen_addr = 0.0.0.0
auth_type = md5
auth_file = /etc/pgbouncer/userlist.txt

; 连接池模式
pool_mode = transaction       ; 事务级复用（推荐）
; pool_mode = session         ; 会话级（兼容 SET 命令，但复用率低）

; 连接数控制
max_client_conn = 1000        ; 允许的最大客户端连接数
default_pool_size = 50        ; 每个数据库的默认连接池大小
min_pool_size = 10            ; 最小保留连接数
reserve_pool_size = 10        ; 预留连接数（突发流量）
reserve_pool_timeout = 3      ; 预留连接等待时间（秒）

; 超时设置
server_idle_timeout = 300     ; 服务端空闲连接超时（秒）
client_idle_timeout = 600     ; 客户端空闲连接超时
query_timeout = 5             ; 查询超时（秒）
query_wait_timeout = 120      ; 等待可用连接的超时

; 对于向量检索场景，如果应用使用 SET hnsw.ef_search，
; 需要使用 session 模式或在每个事务开始时重新 SET

-- 全局查询超时保护（防止单个慢查询占满连接）
ALTER DATABASE vectordb SET statement_timeout = '5000';  -- 5 秒

-- 特定查询超时
SET LOCAL statement_timeout = '2000';  -- 2 秒（仅当前事务）
SELECT id, title
FROM documents
ORDER BY embedding <=> $1::vector
LIMIT 10;

-- 锁等待超时
SET LOCAL lock_timeout = '1000';  -- 1 秒

-- 查看当前索引的构建参数
SELECT indexname, indexdef
FROM pg_indexes
WHERE tablename = 'documents'
  AND indexdef LIKE '%hnsw%' OR indexdef LIKE '%ivfflat%';

-- HNSW 索引：m 和 ef_construction 越大，索引质量越高
-- 如果 m=8, ef_construction=32，对于高维向量（>512维）可能不够
-- 建议：m >= 16, ef_construction >= 64

-- IVFFlat 索引：lists 参数决定聚类数量
-- 如果 lists 过小（如 10），每个聚类包含太多向量，精度低
-- 经验公式：lists = sqrt(N)，N 为数据量
-- 100 万条数据：lists ≈ 1000
-- 1000 万条数据：lists ≈ 3162

-- 如果原始向量未归一化，使用余弦距离可能产生不准确的结果
-- 检查向量的 L2 范数
SELECT
    avg(l2_norm(embedding)) AS avg_norm,
    min(l2_norm(embedding)) AS min_norm,
    max(l2_norm(embedding)) AS max_norm
FROM documents
LIMIT 10000;

-- 如果 avg_norm 不接近 1.0，说明向量未归一化
-- 解决方案：存储时归一化
UPDATE documents
SET embedding = l2_normalize(embedding)
WHERE l2_norm(embedding) BETWEEN 0.99 AND 1.01 = false;

-- 或者改用 L2 距离（对未归一化向量更鲁棒）
-- 但需要注意 L2 距离和余弦距离的语义差异

-- 使用"暴力搜索"（无索引）作为基准，对比索引搜索结果
-- 先关闭索引扫描，强制全表搜索得到"真实"的 Top-K
SET enable_indexscan = off;
SET enable_indexonlyscan = off;

SELECT id, title, embedding <=> query_vec AS distance
FROM documents, (SELECT '[0.1, 0.2, ...]'::vector AS query_vec) q
ORDER BY distance
LIMIT 10;
-- 记录这些结果作为"黄金标准"

-- 重新开启索引扫描
SET enable_indexscan = on;
SET enable_indexonlyscan = on;

-- 使用索引搜索
SET hnsw.ef_search = 40;   -- 先用小值
SELECT id, title, embedding <=> query_vec AS distance
FROM documents, (SELECT '[0.1, 0.2, ...]'::vector AS query_vec) q
ORDER BY distance
LIMIT 10;

-- 逐步增大 ef_search，直到结果与"黄金标准"一致
SET hnsw.ef_search = 100;
SET hnsw.ef_search = 200;
-- 当 ef_search = 200 时结果与暴力搜索一致，说明 200 是足够的

-- 数据分布漂移可能导致 IVFFlat 的聚类中心不再最优
-- 检查最近新增数据与历史数据的向量分布差异

-- 随机采样检查向量的各维度统计
SELECT
    d AS dim_index,
    avg(val) AS avg_val,
    stddev(val) AS stddev_val,
    min(val) AS min_val,
    max(val) AS max_val
FROM documents,
     LATERAL unnest(embedding::real[]) WITH ORDINALITY AS t(val, d)
WHERE d IN (1, 2, 3, 100, 500)   -- 抽样检查几个维度
GROUP BY d
ORDER BY d;

-- 如果发现分布发生显著变化（如某个维度的均值从 0 变为 1），
-- 说明 IVFFlat 的聚类中心可能已过时，需要 REINDEX 重建

REINDEX INDEX CONCURRENTLY documents_embedding_idx;

-- tsvector：将文本分解为词元（lexeme）的有序列表
SELECT to_tsvector('english', 'The quick brown fox jumps over the lazy dog');
-- 结果：'brown':3 'dog':9 'fox':4 'jump':5 'lazi':8 'quick':2

-- tsquery：构造查询表达式
SELECT to_tsquery('english', 'quick & fox');
-- 结果：'quick' & 'fox'

-- 匹配操作 @@
SELECT to_tsvector('english', 'The quick brown fox')
    @@ to_tsquery('english', 'quick & fox') AS match;
-- 结果：true

-- 相关度排序
SELECT
    title,
    ts_rank(
        to_tsvector('english', title || ' ' || content),
        to_tsquery('english', 'database & performance')
    ) AS rank
FROM articles
WHERE to_tsvector('english', title || ' ' || content)
    @@ to_tsquery('english', 'database & performance')
ORDER BY rank DESC
LIMIT 10;

-- 创建 GIN 索引加速全文检索
CREATE INDEX ON articles USING gin (to_tsvector('english', title || ' ' || content));

-- 安装 pg_trgm
CREATE EXTENSION IF NOT EXISTS pg_trgm;

-- 创建 GIN 三元组索引
CREATE INDEX ON products USING gin (name gin_trgm_ops);

-- 模糊搜索（使用索引）
SELECT name, similarity(name, 'iphnoe 15') AS sim  -- 注意拼写错误
FROM products
WHERE name % 'iphnoe 15'   -- % 运算符：相似度超过阈值则匹配
ORDER BY sim DESC
LIMIT 5;
-- 结果可能包含 'iPhone 15'（自动纠正拼写错误）

-- 结合 LIKE 查询
SELECT name
FROM products
WHERE name ILIKE '%iphon%'    -- 使用 GIN 索引加速
ORDER BY similarity(name, 'iphone') DESC;

RRF_score(d) = SUM(1 / (k + rank_i(d)))

-- 混合检索：关键词搜索 + 语义向量搜索 + RRF 融合

-- 假设我们有一个文档表，同时包含全文索引和向量索引
CREATE TABLE documents (
    id BIGSERIAL PRIMARY KEY,
    title TEXT NOT NULL,
    content TEXT NOT NULL,
    embedding vector(1536),
    created_at TIMESTAMPTZ DEFAULT now()
);

-- 全文检索索引
ALTER TABLE documents ADD COLUMN search_vector tsvector
    GENERATED ALWAYS AS (
        to_tsvector('english', coalesce(title, '') || ' ' || coalesce(content, ''))
    ) STORED;
CREATE INDEX ON documents USING gin (search_vector);

-- 向量索引
CREATE INDEX ON documents USING hnsw (embedding vector_cosine_ops)
    WITH (m = 16, ef_construction = 64);

-- 混合检索查询
WITH
-- 通道 1：关键词检索（BM25 / ts_rank 排序）
keyword_results AS (
    SELECT
        id,
        title,
        ts_rank(search_vector, query) AS keyword_score,
        row_number() OVER (
            ORDER BY ts_rank(search_vector, query) DESC
        ) AS keyword_rank
    FROM documents, to_tsquery('english', 'database & optimization & performance') AS query
    WHERE search_vector @@ query
    ORDER BY keyword_score DESC
    LIMIT 50     -- 每个通道取 Top-50 候选
),
-- 通道 2：语义向量检索
vector_results AS (
    SELECT
        id,
        title,
        1 - (embedding <=> $1::vector) AS semantic_score,  -- 转换为相似度
        row_number() OVER (
            ORDER BY embedding <=> $1::vector
        ) AS vector_rank
    FROM documents
    ORDER BY embedding <=> $1::vector
    LIMIT 50     -- 每个通道取 Top-50 候选
),
-- RRF 融合
rrf_fusion AS (
    SELECT
        COALESCE(k.id, v.id) AS id,
        COALESCE(k.title, v.title) AS title,
        COALESCE(k.keyword_rank, 999) AS keyword_rank,   -- 未出现则排名靠后
        COALESCE(v.vector_rank, 999) AS vector_rank,
        k.keyword_score,
        v.semantic_score,
        -- RRF 分数计算（k=60）
        1.0 / (60 + COALESCE(k.keyword_rank, 999))
        + 1.0 / (60 + COALESCE(v.vector_rank, 999))
        AS rrf_score
    FROM keyword_results k
    FULL OUTER JOIN vector_results v ON k.id = v.id
)
SELECT
    id,
    title,
    keyword_rank,
    vector_rank,
    round(keyword_score::numeric, 4) AS keyword_score,
    round(semantic_score::numeric, 4) AS semantic_score,
    round(rrf_score::numeric, 6) AS rrf_score
FROM rrf_fusion
ORDER BY rrf_score DESC
LIMIT 10;

-- 加权融合：关键词权重 0.3，语义权重 0.7
-- 适用于语义搜索更重要的场景（如问答系统）
WITH
keyword_results AS (
    SELECT id, row_number() OVER (
        ORDER BY ts_rank(search_vector, query) DESC
    ) AS keyword_rank
    FROM documents, to_tsquery('english', 'machine & learning') AS query
    WHERE search_vector @@ query
    LIMIT 50
),
vector_results AS (
    SELECT id, row_number() OVER (
        ORDER BY embedding <=> $1::vector
    ) AS vector_rank
    FROM documents
    ORDER BY embedding <=> $1::vector
    LIMIT 50
),
weighted_fusion AS (
    SELECT
        COALESCE(k.id, v.id) AS id,
        -- 加权 RRF
        0.3 * (1.0 / (60 + COALESCE(k.keyword_rank, 999)))
        + 0.7 * (1.0 / (60 + COALESCE(v.vector_rank, 999)))
        AS weighted_score
    FROM keyword_results k
    FULL OUTER JOIN vector_results v ON k.id = v.id
)
SELECT * FROM weighted_fusion
ORDER BY weighted_score DESC
LIMIT 10;

-- 综合查询示例：
-- 在"技术类"文档中，搜索 2025 年发布的、包含"性能优化"关键词的、
-- 语义上与"数据库查询加速"最相似的 Top 10 文档

SELECT
    d.id,
    d.title,
    d.category,
    d.created_at,
    d.embedding <=> $1::vector AS semantic_distance,
    ts_rank(d.search_vector, query) AS keyword_rank
FROM documents d,
     to_tsquery('english', 'performance & optimization') AS query
WHERE
    d.category = 'technology'                           -- 标签过滤
    AND d.created_at >= '2025-01-01'                    -- 时间范围
    AND d.search_vector @@ query                        -- 全文匹配
ORDER BY
    d.embedding <=> $1::vector                          -- 向量距离排序
LIMIT 10;

-- 为过滤条件创建 B-tree 索引
CREATE INDEX ON documents (category);
CREATE INDEX ON documents (created_at);

-- 为向量检索创建向量索引
CREATE INDEX ON documents USING hnsw (embedding vector_cosine_ops);

-- 为全文检索创建 GIN 索引
CREATE INDEX ON documents USING gin (search_vector);

-- PostgreSQL 优化器会自动选择最优的索引组合
-- 通常的策略是：先用 B-tree/GIN 索引过滤，再在过滤后的结果上做向量检索

-- 只对"活跃"文档建立向量索引（假设只有 20% 的文档是活跃的）
CREATE INDEX ON documents USING hnsw (embedding vector_cosine_ops)
    WHERE status = 'active';

-- 只对最近 6 个月的文档建立向量索引
CREATE INDEX ON documents USING hnsw (embedding vector_cosine_ops)
    WHERE created_at >= '2025-06-01';

-- 查询时如果 WHERE 条件匹配部分索引的谓词，优化器会使用部分索引
SELECT id, title
FROM documents
WHERE status = 'active'                          -- 匹配部分索引谓词
ORDER BY embedding <=> $1::vector
LIMIT 10;
-- 使用部分索引，索引大小只有全量索引的 20%，查询更快

                    ┌─────────────────┐
                    │  Coordinator    │  ← 接收查询，分发到 Worker 节点
                    │  (Citus 协调器) │
                    └───────┬─────────┘
                            │
              ┌─────────────┼─────────────┐
              │             │             │
        ┌─────┴─────┐ ┌────┴─────┐ ┌─────┴─────┐
        │ Worker 1  │ │ Worker 2 │ │ Worker 3  │
        │ 分片 1,4  │ │ 分片 2,5 │ │ 分片 3,6  │
        │ HNSW 索引 │ │ HNSW索引 │ │ HNSW 索引 │
        └───────────┘ └──────────┘ └───────────┘

-- 在每个节点（Coordinator + 所有 Worker）上安装扩展
CREATE EXTENSION citus;
CREATE EXTENSION vector;

-- 在 Coordinator 上注册 Worker 节点
SELECT citus_add_node('worker1.example.com', 5432);
SELECT citus_add_node('worker2.example.com', 5432);
SELECT citus_add_node('worker3.example.com', 5432);

-- 创建分布式向量表
CREATE TABLE documents (
    id BIGSERIAL,
    tenant_id BIGINT NOT NULL,
    title TEXT,
    embedding vector(1536),
    created_at TIMESTAMPTZ DEFAULT now()
);

-- 按 tenant_id 哈希分片（默认 32 个分片）
SELECT create_distributed_table('documents', 'tenant_id');

-- 为每个分片创建 HNSW 索引（自动传播到所有 Worker）
CREATE INDEX ON documents USING hnsw (embedding vector_cosine_ops)
    WITH (m = 16, ef_construction = 64);

-- 方案 A：按 tenant_id 哈希分片（多租户场景）
-- 优点：同一租户的数据在同一分片，支持租户内的向量检索
-- 缺点：如果某些租户数据量远大于其他租户，可能导致热点
SELECT create_distributed_table('documents', 'tenant_id');

-- 方案 B：按 id 哈希分片（通用场景）
-- 优点：数据均匀分布
-- 缺点：跨分片的向量检索需要合并多个分片的结果
SELECT create_distributed_table('documents', 'id');

-- 方案 C：引用表（小维度表）
-- 对于元数据表等小表，使用引用表（复制到每个 Worker）
SELECT create_reference_table('categories');

-- 查询：在 tenant_id = 42 的文档中搜索最相似的 10 个
SELECT id, title, embedding <=> $1::vector AS distance
FROM documents
WHERE tenant_id = 42
ORDER BY distance
LIMIT 10;

-- 执行过程：
-- 1. Coordinator 根据 tenant_id = 42 的哈希值定位到具体分片（如 Worker 2）
-- 2. 查询被下推到 Worker 2，Worker 2 在自己的分片上执行 HNSW 索引扫描
-- 3. Worker 2 返回 Top-10 结果给 Coordinator
-- 4. Coordinator 直接返回给客户端
-- 注意：因为 tenant_id 过滤后只涉及单个分片，效率接近单节点查询

-- 跨分片查询（无 tenant_id 过滤）
SELECT id, title, embedding <=> $1::vector AS distance
FROM documents
ORDER BY distance
LIMIT 10;

-- 执行过程：
-- 1. Coordinator 将查询下推到所有 Worker
-- 2. 每个 Worker 在自己的分片上执行 HNSW 搜索，返回 Top-10
-- 3. Coordinator 收集所有 Worker 的 Top-10（共 30 条），全局排序取 Top-10
-- 注意：跨分片查询的延迟 = 最慢 Worker 的延迟 + 合并开销
-- 如果分片数量为 N，每个 Worker 返回 Top-K，Coordinator 需要合并 N*K 条记录

# Python 应用层分片示例
import hashlib
import psycopg2

SHARD_CONFIGS = [
    {"host": "shard1.db.example.com", "port": 5432, "dbname": "vectors_1"},
    {"host": "shard2.db.example.com", "port": 5432, "dbname": "vectors_2"},
    {"host": "shard3.db.example.com", "port": 5432, "dbname": "vectors_3"},
    {"host": "shard4.db.example.com", "port": 5432, "dbname": "vectors_4"},
]

def get_shard(doc_id: int) -> dict:
    """根据文档 ID 确定所在分片"""
    shard_index = doc_id % len(SHARD_CONFIGS)
    return SHARD_CONFIGS[shard_index]

def vector_search(query_vector: list, top_k: int = 10) -> list:
    """跨分片向量检索"""
    all_results = []

    for shard in SHARD_CONFIGS:
        conn = psycopg2.connect(**shard)
        cur = conn.cursor()
        cur.execute("""
            SELECT id, title, embedding <=> %s::vector AS distance
            FROM documents
            ORDER BY distance
            LIMIT %s
        """, (str(query_vector), top_k))
        all_results.extend(cur.fetchall())
        cur.close()
        conn.close()

    # 合并所有分片的结果，全局排序
    all_results.sort(key=lambda x: x[2])  # 按 distance 排序
    return all_results[:top_k]

                    ┌─────────────────┐
  写入（INSERT） ──→│  Primary 主节点  │
                    └────────┬────────┘
                             │ WAL 流复制
              ┌──────────────┼──────────────┐
              │              │              │
        ┌─────┴─────┐ ┌─────┴─────┐ ┌─────┴─────┐
        │ Replica 1 │ │ Replica 2 │ │ Replica 3 │
        │ 只读查询   │ │ 只读查询   │ │ 只读查询   │
        │ HNSW 索引  │ │ HNSW 索引  │ │ HNSW 索引  │
        └───────────┘ └───────────┘ └───────────┘
              ↑              ↑              ↑
              └──────────────┴──────────────┘
                    向量检索查询（负载均衡）

# postgresql.conf（主节点）
wal_level = replica
max_wal_senders = 10
wal_keep_size = 1GB        # 保留足够的 WAL 供副本追赶

# 在副本节点上，可以独立调整向量检索参数
# 例如副本可以设置更大的 hnsw.ef_search（因为不承担写入压力）
hnsw.ef_search = 200

BEGIN;
  -- 插入新文档
  INSERT INTO documents (title, content) VALUES ('New Doc', '...') RETURNING id;
  -- 调用外部服务生成 embedding 并存储
  UPDATE documents SET embedding = $1 WHERE id = lastval();
  -- 同时更新关联的元数据表
  INSERT INTO document_metadata (doc_id, category, tags) VALUES (lastval(), 'tech', '{ai,ml}');
COMMIT;  -- 所有操作要么全部成功，要么全部回滚

你的向量数据量有多大？
│
├── < 100 万条
│   ├── 已有 PostgreSQL？ ───────────────→ pgvector（零额外成本）
│   ├── 快速原型/实验？ ─────────────────→ Chroma（最简单）
│   └── 需要托管服务？ ──────────────────→ Pinecone（免运维）
│
├── 100 万 ~ 1000 万条
│   ├── 需要复杂混合查询/事务？ ─────────→ pgvector + 调优
│   ├── 纯向量检索 + 低延迟？ ──────────→ Qdrant 或 Milvus
│   └── 不想运维？ ──────────────────────→ Pinecone / Zilliz Cloud
│
├── 1000 万 ~ 1 亿条
│   ├── 已有 PG 运维能力 + 混合查询？ ──→ pgvector + Citus
│   ├── 纯向量检索为主？ ───────────────→ Milvus（原生分布式）
│   └── 低延迟优先？ ───────────────────→ Qdrant（分布式模式）
│
└── > 1 亿条
    ├── 预算充足 + 不想运维？ ──────────→ Pinecone / Zilliz Cloud
    └── 自建集群？ ──────────────────────→ Milvus（经过大规模验证）

-- 启用迭代扫描（通过增加 ef_search 间接控制）
-- 迭代扫描会在 ef_search 的搜索预算内自动进行多轮搜索
SET hnsw.ef_search = 200;

-- 迭代扫描对以下场景特别有效：
-- 1. 图结构因频繁 INSERT/DELETE 而退化（尚未 REINDEX）
-- 2. 高维向量（>1024 维），图结构中"空隙"更多
-- 3. 需要极高召回率（>99%）的关键业务场景

-- 概念预览：未来的 GPU 加速向量查询
-- （目前尚未可用，仅为概念展示）

-- 假设未来有 pgvector_gpu 扩展
CREATE EXTENSION pgvector_gpu;

-- 使用 GPU 加速的索引类型
CREATE INDEX ON documents USING hnsw_gpu (embedding vector_cosine_ops)
    WITH (device = 'cuda:0');  -- 指定 GPU 设备

-- 查询语法不变，自动利用 GPU
SELECT id, title
FROM documents
ORDER BY embedding <=> $1::vector
LIMIT 100;
-- GPU 可以同时计算数千个向量的距离，对大批量查询加速显著

-- VectorChord 使用示例（需要安装 vchord 扩展）
CREATE EXTENSION vchord;

-- 使用 ScaNN 索引（比 HNSW 更省内存）
CREATE INDEX ON documents USING scann (embedding vector_cosine_ops)
    WITH (n_lists = 1000, n_probes = 10);

-- ScaNN 索引的优势：
-- 1. 内存占用比 HNSW 低 50%~70%（通过量化压缩）
-- 2. 构建速度比 HNSW 快 3~5 倍
-- 3. 查询延迟与 HNSW 相当
-- 4. 适合 1000 万+ 向量且内存预算有限的场景

-- 概念示例：使用 pgai 在数据库内直接生成 embedding（未来方向）
-- 目前需要通过应用层调用 embedding API

-- 目标：
SELECT ai_embedding('text-embedding-3-small', content) AS embedding
FROM documents
WHERE id = 42;

-- 结合触发器实现自动向量化
CREATE OR REPLACE FUNCTION auto_embed()
RETURNS TRIGGER AS $$
BEGIN
    NEW.embedding := ai_embedding('text-embedding-3-small', NEW.content);
    RETURN NEW;
END;
$$ LANGUAGE plpgsql;

CREATE TRIGGER embed_trigger
    BEFORE INSERT OR UPDATE OF content ON documents
    FOR EACH ROW EXECUTE FUNCTION auto_embed();