向量数据库及应用场景

城市级别 (Y)
   ^
1.0|              
   |             
   |            
   |            
   |          
   |         
   |        
   |       
0.0+-------------------------------------> 差旅报销 (X)
   0                                    1.0

城市级别 (Y)
   ^
1.0|               * 查询 B [0.90, 0.95] (去北京住酒店能报多少)
   |              /.
   |             / . 
   |            /  * 文本 A [0.95, 0.85] (一二线城市不超过500)
   |           /  /
   |          /  / 
   |         /  /
   |        /  /
0.0+-------------------------------------> 差旅报销 (X)
   0                                    1.0

def setup_schema(conn) -> None:
    with conn.cursor() as cur:
        # pgvector 不是一个独立服务，而是 PostgreSQL 扩展。
        # 开启扩展后，PostgreSQL 就多了 vector 类型、向量距离运算符和向量索引能力。
        cur.execute("CREATE EXTENSION IF NOT EXISTS vector")
        cur.execute(
            """
            CREATE TABLE IF NOT EXISTS documents (
                id bigserial PRIMARY KEY,
                title text NOT NULL,
                content text NOT NULL,
                -- embedding 保存 bge-small-zh-v1.5 生成的 512 维语义向量。
                -- 这列相当于“可被相似度搜索的语义特征”。
                embedding vector(512) NOT NULL,
                created_at timestamptz NOT NULL DEFAULT now(),
                updated_at timestamptz NOT NULL DEFAULT now()
            )
            """
        )
        cur.execute(
            """
            -- HNSW 是近似最近邻索引，用来加速 topK 向量检索。
            -- vector_cosine_ops 表示按余弦距离进行向量相似度搜索。
            CREATE INDEX IF NOT EXISTS documents_embedding_hnsw_idx
            ON documents USING hnsw (embedding vector_cosine_ops)
            """
        )
    conn.commit()
 

def insert_documents(
    conn,
    model: SentenceTransformer,
    documents: Iterable[DemoDocument],
) -> list[int]:
    ids: list[int] = []
    with conn.cursor() as cur:
        for doc in documents:
            # 这里用 title + content 生成 embedding，所以检索时二者都会影响语义匹配。
            # 如果只想检索正文，可以改成 encode(model, doc.content)。
            embedding = encode(model, f"{doc.title}\n{doc.content}")
            cur.execute(
                """
                INSERT INTO documents (title, content, embedding)
                VALUES (%s, %s, %s)
                RETURNING id
                """,
                (doc.title, doc.content, embedding),
            )
            ids.append(cur.fetchone()[0])
    conn.commit()
    return ids
 
 
def search(conn, model: SentenceTransformer, query: str, limit: int = 3) -> list[tuple]:
    # 用户输入的查询句子也会先转换成向量，再和数据库里的 embedding 做距离计算。
    query_embedding = encode(model, query)
    with conn.cursor() as cur:
        cur.execute(
            """
            SELECT
                id,
                title,
                content,
                -- <=> 是 pgvector 的余弦距离运算符，距离越小越相似。
                -- 因为前面做了 normalize_embeddings=True，这里用 1 - distance 近似展示相似度。
                round((1 - (embedding <=> %s::vector))::numeric, 4) AS cosine_similarity
            FROM documents
            -- 这就是向量检索的核心：按“文档向量”和“查询向量”的距离排序。
            -- 写法仍然是 SQL，因为 pgvector 把向量能力集成进了 PostgreSQL。
            ORDER BY embedding <=> %s::vector
            LIMIT %s
            """,
            (query_embedding, query_embedding, limit),
        )
        return cur.fetchall()

def main() -> None:
    model_name = os.getenv("MODEL_NAME", "BAAI/bge-small-zh-v1.5")
    print(f"Loading embedding model: {model_name}")
    model = SentenceTransformer(model_name)
 
    with connect(**db_config()) as conn:
        register_vector(conn)
        setup_schema(conn)
 
        docs = [
            DemoDocument("iPhone 15 Pro 性能", "A17 Pro 芯片适合运行大型手游、视频剪辑和高负载应用。"),
            DemoDocument("iPhone 15 拍照", "4800 万像素主摄适合日常拍照，能够保留更多画面细节。"),
            DemoDocument("iPhone 电池续航", "想要长续航可以关注 Plus 或 Pro Max 机型，它们通常电池容量更大。"),
            DemoDocument("苹果手机系统", "iOS 系统更新周期长，应用生态稳定，适合重视长期使用体验的用户。"),
            DemoDocument("Face ID 解锁", "苹果手机使用 Face ID 进行面部识别，暗光环境下也能快速解锁。"),
            DemoDocument("MagSafe 配件", "MagSafe 可以吸附无线充电器、磁吸支架、卡包和移动电源。"),
            DemoDocument("iPhone 存储容量", "经常拍视频或下载大型应用时，256GB 比 128GB 更从容。"),
            DemoDocument("二手 iPhone 选购", "购买二手苹果手机时要检查电池健康、屏幕状态、序列号和是否有隐藏锁。"),
            DemoDocument("iPhone 与 Mac 协同", "苹果手机可以和 Mac 使用隔空投送、接力、短信同步和通用剪贴板。"),
            DemoDocument("USB-C 接口", "新款 iPhone 使用 USB-C 接口，充电线和部分数码设备更容易通用。"),
        ]
 
        if count_documents(conn) == 0:
            ids = insert_documents(conn, model, docs)
            print(f"\n数据库为空，已插入样例 document ids: {ids}")
 
        print_documents(list_documents(conn))
        interactive_search(conn, model)