PostgreSQL+pgvector：大模型 RAG 的「结构化 + 向量」混合检索

version: '3.8'
services:
  pgvector:
    image: pgvector/pgvector:pg15
    container_name: pgvector-llm
    ports:
      - "5432:5432"
    environment:
      POSTGRES_USER: llm_user     # 数据库用户（后续代码用）
      POSTGRES_PASSWORD: llm_pass # 数据库密码（后续代码用）
      POSTGRES_DB: llm_db         # 数据库名（后续代码用）
    volumes:
      - pgvector-data:/var/lib/postgresql/data # 数据持久化
volumes:
  pgvector-data:

docker-compose up -d

# 连接命令（替换为你的密码/端口）
psql -h localhost -p 5432 -U llm_user -d llm_db

-- 1. 启用pgvector扩展（PostgreSQL 15+自动安装，仅需启用）
CREATE EXTENSION IF NOT EXISTS vector;

-- 2. 创建RAG文档表【严格复用系列第2篇的SQL语法】
CREATE TABLE rag_documents (
    id SERIAL PRIMARY KEY,                -- 自增ID（系列第2篇：SERIAL语法）
    doc_title VARCHAR(255) NOT NULL,      -- 文档标题（VARCHAR：短文本）
    doc_content TEXT NOT NULL,            -- 文档内容（TEXT：长文本，系列第2篇知识点）
    doc_meta JSONB NOT NULL,              -- 动态元数据（JSONB：系列第2篇核心知识点）
    doc_embedding VECTOR(1536) NOT NULL   -- 向量字段：1536维（OpenAI/智谱等主流模型）
);

-- 3. 创建索引【严格复用系列第2篇的索引优化知识点】
-- ① JSONB元数据索引：用于结构化过滤
CREATE INDEX idx_doc_meta ON rag_documents USING GIN(doc_meta);
-- ② 向量索引：用于语义检索（IVFFlat：pgvector推荐的近似索引，召回率≥95%）
CREATE INDEX idx_doc_embedding ON rag_documents USING ivfflat(doc_embedding vector_cosine_ops) WITH (nlist=100);
-- ③ 复合索引：加速混合检索的前过滤
CREATE INDEX idx_doc_meta_year ON rag_documents ((doc_meta ->> 'year')::INT);

-- 插入测试数据：模拟3篇关于大模型的文档
INSERT INTO rag_documents (doc_title, doc_content, doc_meta, doc_embedding)
VALUES 
-- 文档1：2024年发布的大模型推理加速文档
('大模型推理加速指南', 'vLLM是一款高性能大模型推理框架...', '{"year": 2024, "type": "技术文档", "visible": "all"}', '[0.123, 0.456, 0.789, ...]'::VECTOR(1536)),
-- 文档2：2023年发布的大模型训练文档
('大模型训练调参手册', 'LoRA是一种高效的模型微调技术...', '{"year": 2023, "type": "技术文档", "visible": "all"}', '[0.234, 0.567, 0.890, ...]'::VECTOR(1536)),
-- 文档3：2024年发布的VIP专属文档
('大模型RAG企业级落地方案', 'PostgreSQL+pgvector是企业级RAG的最佳选择...', '{"year": 2024, "type": "解决方案", "visible": "vip"}', '[0.345, 0.678, 0.901, ...]'::VECTOR(1536));

-- 【混合检索核心SQL】
-- 1. 结构化过滤：先筛选2024年发布的文档（前过滤，减少向量检索的数据量）
-- 2. 向量检索：再查找与「大模型推理加速」语义相似的文档
-- 3. 排序：按相似度降序，取前3条
SELECT 
    doc_title, 
    doc_content,
    doc_meta ->> 'year' AS publish_year,
    1 - (doc_embedding <=> '[0.123, 0.456, ...]'::VECTOR(1536)) AS similarity  -- 余弦相似度（1-距离=相似度）
FROM rag_documents
WHERE 
    doc_meta ->> 'year' = '2024'  -- 结构化过滤：2024年发布
    AND doc_meta ->> 'visible' = 'all'  -- 结构化过滤：所有人可见
ORDER BY similarity DESC  -- 按语义相似度排序
LIMIT 3;

fastapi==0.104.1
uvicorn[standard]==0.24.0
sqlalchemy==2.0.23
psycopg2-binary==2.9.9  # PostgreSQL的Python驱动
pydantic==2.5.0

pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple

from fastapi import FastAPI, Depends, HTTPException
from sqlalchemy import create_engine, JSON, Text, String, Integer, Column, Index, VECTOR
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
from pydantic import BaseModel
from typing import List

# --------------------------
# 1. 数据库连接（复用系列第3篇的架构）
# --------------------------
DATABASE_URL = "postgresql://llm_user:llm_pass@localhost:5432/llm_db"
engine = create_engine(DATABASE_URL)
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
Base = declarative_base()

# --------------------------
# 2. 定义RAG文档表ORM模型（复用系列第2篇的结构）
# --------------------------
class RagDocument(Base):
    __tablename__ = "rag_documents"
    id = Column(Integer, primary_key=True, index=True)
    doc_title = Column(String(255), index=True)
    doc_content = Column(Text)
    doc_meta = Column(JSON)  # 对应SQL的JSONB类型
    doc_embedding = Column(VECTOR(length=1536))  # 向量字段，1536维

# --------------------------
# 3. FastAPI初始化+依赖
# --------------------------
app = FastAPI(title="大模型RAG混合检索API", version="1.0")

def get_db():
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()

# --------------------------
# 4. 请求/响应模型
# --------------------------
class RAGRequest(BaseModel):
    prompt: str  # 用户提问
    year: int  # 结构化过滤：发布年份
    visible: str = "all"  # 结构化过滤：可见权限

class RAGResponse(BaseModel):
    doc_title: str
    publish_year: str
    content: str
    similarity: float

# --------------------------
# 5. 混合检索接口（核心功能）
# --------------------------
@app.post(
    "/rag/hybrid",
    response_model=List[RAGResponse],
    summary="大模型RAG混合检索接口",
    description="结构化元数据过滤+向量语义检索"
)
async def rag_hybrid_search(request: RAGRequest, db: Session = Depends(get_db)):
    try:
        # 模拟生成用户提问的向量（真实场景用OpenAI/智谱API生成）
        user_embedding = [0.123, 0.456, ...]  # 替换为真实向量
        
        # 【混合检索】执行SQL（ORM自动生成对应SQL语句）
        # 对应前面的核心SQL：WHERE year=2024 AND visible=all ORDER BY similarity DESC
        docs = db.query(RagDocument)\
                .filter(RagDocument.doc_meta["year"] == request.year)\
                .filter(RagDocument.doc_meta["visible"] == request.visible)\
                .order_by(RagDocument.doc_embedding <=> user_embedding)\
                .limit(3)\
                .all()
        
        # 转换为响应格式
        return [
            RAGResponse(
                doc_title=doc.doc_title,
                publish_year=str(doc.doc_meta["year"]),
                content=doc.doc_content,
                similarity=float(1 - (doc.doc_embedding <=> user_embedding))
            ) for doc in docs
        ]
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"检索失败：{str(e)}")

# --------------------------
# 启动命令：uvicorn main:app --reload
# --------------------------

{
  "prompt": "大模型推理加速",
  "year": 2024,
  "visible": "all"
}