一、安装python依赖包

pip install psycopg2-binary sentence-transformers numpy

• psycopg2-binary：连接PostgreSQL数据库并进行数据库读写操作，支持 PostgreSQL 特有类型如向量的存入和相似度搜索。
• sentence-transformers：生成文本嵌入Text Embedding，加载bge-small-zh-v1.5预训练模型，将自然语言文本句子、段落等转换为语义向量。

二、编写程序

程序逻辑：

1）初始化PostgreSQL向量表。

2）加载本地bge-small-zh-v1.5模型，将实例知识库文本数据向量化后加载到documents表中。

3）针对查询问题进行相似度搜索。

import os
import psycopg2
from psycopg2.extras import RealDictCursor
from sentence_transformers import SentenceTransformer
import numpy as np

# ======================
# 配置部分
# ======================
# PostgreSQL 连接配置
DB_CONFIG = {
    'host': 'localhost',
    'port': 5432,
    'database': 'postgres',      # 替换为你的数据库名
    'user': 'postgres',         # 替换为你的用户名
    'password': 'postgres'      # 替换为你的密码
}

# 向量维度（bge-small-zh-v1.5 输出 512 维）
VECTOR_DIM = 512

# 示例知识库数据（可替换为你的文档）
KNOWLEDGE_DOCS = [
    "自然语言处理是人工智能和语言学领域的分支，致力于让计算机能理解人类语言。",
    "深度学习是一种基于神经网络的机器学习方法，近年来在图像识别、语音识别等领域取得突破。",
    "大语言模型（LLM）如 ChatGPT 能够生成高质量文本，广泛应用于问答、创作等场景。",
    "向量数据库用于高效存储和检索高维向量，常用于语义搜索和推荐系统。",
    "RAG（检索增强生成）结合了信息检索与文本生成，能提升大模型回答的准确性和时效性。"
]

# ======================
# 初始化数据库表
# ======================
def init_db():
    conn = psycopg2.connect(**DB_CONFIG)
    cur = conn.cursor()
    
    # 创建表（如果不存在）
    cur.execute(f"""
        CREATE TABLE IF NOT EXISTS documents (
            id SERIAL PRIMARY KEY,
            content TEXT NOT NULL,
            embedding VECTOR({VECTOR_DIM})
        );
    """)
    
    # 创建向量索引（HNSW，需 pgvector >= 0.5.0）
    cur.execute("""
        CREATE INDEX IF NOT EXISTS idx_embedding ON documents 
        USING hnsw (embedding vector_cosine_ops) 
        WITH (m = 16, ef_construction = 64);
    """)
    
    conn.commit()
    cur.close()
    conn.close()
    print("数据库表和索引初始化完成")

# ======================
# 加载嵌入模型
# ======================
def load_model():
    print("正在加载 bge-small-zh-v1.5 模型...")
    model = SentenceTransformer('/opt/knowly/bge-small-zh-v1.5')
    print("模型加载成功")
    return model

# ======================
# 插入文档到数据库
# ======================
def insert_documents(model, docs):
    conn = psycopg2.connect(**DB_CONFIG)
    cur = conn.cursor()
    
    print(f"正在将 {len(docs)} 条文档向量化并插入数据库...")
    for doc in docs:
        # 生成嵌入向量
        embedding = model.encode(doc, normalize_embeddings=True)  # 归一化便于余弦相似度
        # 转为 PG 向量格式：[x1,x2,...]
        embedding_str = '[' + ','.join(map(str, embedding.tolist())) + ']'
        
        cur.execute(
            "INSERT INTO documents (content, embedding) VALUES (%s, %s::vector)",
            (doc, embedding_str)
        )
    
    conn.commit()
    cur.close()
    conn.close()
    print("文档插入完成")

# ======================
# 语义搜索函数
# ======================
def semantic_search(model, query: str, top_k: int = 3):
    conn = psycopg2.connect(**DB_CONFIG)
    cur = conn.cursor(cursor_factory=RealDictCursor)
    
    # 对查询进行向量化
    query_embedding = model.encode(query, normalize_embeddings=True)
    query_str = '[' + ','.join(map(str, query_embedding.tolist())) + ']'
    
    # 使用余弦相似度（归一化后等价于内积）
    cur.execute("""
        SELECT id, content, 
               1 - (embedding <=> %s::vector) AS similarity
        FROM documents
        ORDER BY embedding <=> %s::vector
        LIMIT %s;
    """, (query_str, query_str, top_k))
    
    results = cur.fetchall()
    cur.close()
    conn.close()
    return results

# ======================
# 主程序
# ======================
def main():
    # 初始化
    init_db()
    model = load_model()
    
    # 清空旧数据（可选）
    conn = psycopg2.connect(**DB_CONFIG)
    conn.cursor().execute("DELETE FROM documents;")
    conn.commit()
    conn.close()
    
    # 插入知识库
    insert_documents(model, KNOWLEDGE_DOCS)
    
    # 示例查询
    print("\n测试查询：")
    test_queries = [
        "什么是自然语言处理？",
        "如何用向量数据库做语义搜索？",
        "大模型和深度学习有什么关系？"
    ]
    
    for query in test_queries:
        print(f"\n查询: {query}")
        results = semantic_search(model, query, top_k=2)
        for i, res in enumerate(results, 1):
            print(f"  {i}. 相似度: {res['similarity']:.4f} | 内容: {res['content']}")

if __name__ == "__main__":
    main()

三、测试

(env) root@zsscqmate:/opt/knowly# python3 rag_with_pgvector.py 
数据库表和索引初始化完成
正在加载 bge-small-zh-v1.5 模型...
模型加载成功
正在将 5 条文档向量化并插入数据库...
文档插入完成

测试查询：

查询: 什么是自然语言处理？
  1. 相似度: 0.8198 | 内容: 自然语言处理是人工智能和语言学领域的分支，致力于让计算机能理解人类语言。
  2. 相似度: 0.4880 | 内容: 大语言模型（LLM）如 ChatGPT 能够生成高质量文本，广泛应用于问答、创作等场景。

查询: 如何用向量数据库做语义搜索？
  1. 相似度: 0.7674 | 内容: 向量数据库用于高效存储和检索高维向量，常用于语义搜索和推荐系统。
  2. 相似度: 0.4757 | 内容: 大语言模型（LLM）如 ChatGPT 能够生成高质量文本，广泛应用于问答、创作等场景。

查询: 大模型和深度学习有什么关系？
  1. 相似度: 0.6210 | 内容: 深度学习是一种基于神经网络的机器学习方法，近年来在图像识别、语音识别等领域取得突破。
  2. 相似度: 0.5992 | 内容: 大语言模型（LLM）如 ChatGPT 能够生成高质量文本，广泛应用于问答、创作等场景。