AI 工作流 13：向量数据库怎么选与实战

原始文档 → 切片 → embedding 模型 → 向量 → 【向量数据库】存起来
                                                          ↑
                            用户提问 → embedding → 去【向量数据库】找最像的 → 喂给 LLM

"猫"     → [0.21, 0.55, 0.10]
"小猫"   → [0.20, 0.56, 0.09]   ← 和"猫"非常接近
"狗"     → [0.18, 0.50, 0.12]   ← 也比较近（都是动物）
"汽车"   → [0.90, -0.30, 0.77]  ← 和前三个都很远

# similarity_demo.py: 三种相似度度量
import numpy as np

a = np.array([1.0, 2.0, 3.0])
b = np.array([2.0, 4.0, 6.0])  # 和 a 方向一致，长度翻倍

# 1. 余弦相似度（忽略长度，方向一致就为 1）
cosine = np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))
print("余弦相似度:", cosine)  # 1.0（方向完全一致）

# 2. 欧氏距离（看绝对位置）
euclidean = np.linalg.norm(a - b)
print("欧氏距离:", euclidean)  # 3.74（位置差不少）

# 3. 点积（受长度影响）
dot = np.dot(a, b)
print("点积:", dot)  # 28.0

# 启动一个带 pgvector 的 PostgreSQL 16
docker run -d \
  --name pgvector-demo \
  -e POSTGRES_PASSWORD=mysecret \
  -e POSTGRES_DB=vectordb \
  -p 5432:5432 \
  pgvector/pgvector:pg16

# 检查是否起来了
docker ps | grep pgvector

-- 1. 启用 vector 扩展（每个数据库只需执行一次）
CREATE EXTENSION IF NOT EXISTS vector;

-- 2. 建一张表：既存业务字段，又存向量
CREATE TABLE IF NOT EXISTS docs (
    id          BIGSERIAL PRIMARY KEY,
    title       TEXT,                       -- 文档标题
    content     TEXT,                       -- 文档正文
    department  TEXT,                       -- 所属部门（用于过滤）
    embedding   vector(1536)                -- 向量列！1536 维（text-embedding-3-small）
);

-- 3. 创建 HNSW 索引（用余弦距离）
CREATE INDEX IF NOT EXISTS idx_docs_embedding
ON docs USING hnsw (embedding vector_cosine_ops)
WITH (m = 16, ef_construction = 64);

-- 插入两条带向量的记录（这里用 3 维向量简化演示，真实是 1536 维）
INSERT INTO docs (title, content, department, embedding)
VALUES (
    'RAG 入门',
    'RAG 是检索增强生成……',
    'tech',
    '[0.1, 0.2, 0.3]'
);

-- 找和给定向量最像的 5 条记录，按相似度升序
SELECT id, title, department,
       1 - (embedding <=> '[0.1, 0.2, 0.3]') AS similarity
FROM docs
WHERE department = 'tech'                   -- 先按业务字段过滤
ORDER BY embedding <=> '[0.1, 0.2, 0.3]'    -- 再按向量距离排序
LIMIT 5;

# pgvector_demo.py: pgvector 完整 CRUD + 相似度检索
import os
import psycopg
from psycopg.rows import dict_row_factory
from openai import OpenAI
from dotenv import load_dotenv

load_dotenv()

# ---------- 1. 初始化客户端 ----------
oai = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))   # 国内可换 base_url
DB_DSN = os.getenv("PG_DSN", "postgresql://postgres:mysecret@localhost:5432/vectordb")


def get_embedding(text: str) -> list[float]:
    """把文字转向量（1536 维）"""
    resp = oai.embeddings.create(
        model="text-embedding-3-small",
        input=text,
    )
    return resp.data[0].embedding


# ---------- 2. 建表（带 HNSW 索引）----------
SCHEMA_SQL = """
CREATE EXTENSION IF NOT EXISTS vector;

CREATE TABLE IF NOT EXISTS docs (
    id          BIGSERIAL PRIMARY KEY,
    title       TEXT,
    content     TEXT,
    department  TEXT,
    embedding   vector(1536)
);

CREATE INDEX IF NOT EXISTS idx_docs_embedding
ON docs USING hnsw (embedding vector_cosine_ops)
WITH (m = 16, ef_construction = 64);
"""


def init_db():
    with psycopg.connect(DB_DSN) as conn:
        with conn.cursor() as cur:
            cur.execute(SCHEMA_SQL)
        conn.commit()
    print("[OK] 表和索引已就绪")


# ---------- 3. 插入（带向量）----------
INSERT_SQL = """
INSERT INTO docs (title, content, department, embedding)
VALUES (%s, %s, %s, %s)
RETURNING id;
"""


def insert_doc(title: str, content: str, department: str):
    vec = get_embedding(content)         # 先算 embedding
    with psycopg.connect(DB_DSN) as conn:
        with conn.cursor() as cur:
            # 注意：向量要转成字符串 '[0.1,0.2,...]' 传给 PG
            cur.execute(INSERT_SQL, (title, content, department, str(vec)))
            doc_id = cur.fetchone()[0]
        conn.commit()
    print(f"[OK] 插入成功 id={doc_id} title={title}")
    return doc_id


# ---------- 4. 相似度检索 ----------
SEARCH_SQL = """
SELECT id, title, department,
       1 - (embedding <=> %s::vector) AS similarity
FROM docs
WHERE department = %s
ORDER BY embedding <=> %s::vector
LIMIT %s;
"""


def search(query: str, department: str, top_k: int = 5):
    qvec = get_embedding(query)
    vec_str = str(qvec)
    with psycopg.connect(DB_DSN, row_factory=dict_row_factory) as conn:
        with conn.cursor() as cur:
            cur.execute(SEARCH_SQL, (vec_str, department, vec_str, top_k))
            return cur.fetchall()


# ---------- 5. 主流程 ----------
if __name__ == "__main__":
    init_db()

    # 插入几条文档（演示）
    insert_doc("RAG 入门", "RAG 是检索增强生成，先检索再生成。", "tech")
    insert_doc("向量检索", "向量数据库做相似度检索，找最像的。", "tech")
    insert_doc("公司放假通知", "国庆放假 7 天，10 月 1 日至 7 日。", "hr")

    # 检索："怎么找相似的文字" → 应该命中前两条 tech 文档
    results = search("如何让模型先查资料再回答", department="tech", top_k=2)
    print("\n检索结果：")
    for r in results:
        print(f"  [{r['similarity']:.3f}] {r['title']}  (id={r['id']})")

pip install pgvector psycopg

# 用 pgvector 包，向量可直接传 list，无需 str()
from pgvector.psycopg import register_vector
import psycopg

conn = psycopg.connect(DB_DSN)
register_vector(conn)   # 注册 vector 类型适配器

# 现在可以直接传 list，PG 自动识别
cur = conn.cursor()
cur.execute(
    "INSERT INTO docs (title, embedding) VALUES (%s, %s)",
    ("test", [0.1, 0.2, 0.3])   # ← 直接传 list
)

docker run -d \
  --name qdrant \
  -p 6333:6333 \
  -p 6334:6334 \
  -v $(pwd)/qdrant_data:/qdrant/storage \
  qdrant/qdrant

pip install qdrant-client openai python-dotenv

# qdrant_demo.py: Qdrant 完整 CRUD + 相似度检索
import os
from qdrant_client import QdrantClient
from qdrant_client.models import Distance, VectorParams, PointStruct, Filter, FieldCondition, MatchValue
from openai import OpenAI
from dotenv import load_dotenv

load_dotenv()

oai = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
COLLECTION = "docs"
DIM = 1536  # text-embedding-3-small 维度

# ---------- 1. 连接 + 建 collection ----------
client = QdrantClient(url="http://localhost:6333")

# 如果不存在就建（已存在会跳过）
if COLLECTION not in [c.name for c in client.get_collections().collections]:
    client.create_collection(
        collection_name=COLLECTION,
        vectors_config=VectorParams(size=DIM, distance=Distance.COSINE),
    )
    print(f"[OK] 已创建 collection: {COLLECTION}")


# ---------- 2. 工具：算 embedding ----------
def get_embedding(text: str) -> list[float]:
    resp = oai.embeddings.create(model="text-embedding-3-small", input=text)
    return resp.data[0].embedding


# ---------- 3. 插入（upsert）----------
def upsert_docs():
    points = [
        PointStruct(
            id=1,
            vector=get_embedding("RAG 是检索增强生成，先检索再生成。"),
            payload={"title": "RAG 入门", "department": "tech"},
        ),
        PointStruct(
            id=2,
            vector=get_embedding("向量数据库做相似度检索。"),
            payload={"title": "向量检索", "department": "tech"},
        ),
        PointStruct(
            id=3,
            vector=get_embedding("国庆放假 7 天。"),
            payload={"title": "放假通知", "department": "hr"},
        ),
    ]
    client.upsert(collection_name=COLLECTION, points=points)
    print(f"[OK] 已插入 {len(points)} 条 point")


# ---------- 4. 相似度检索（带过滤）----------
def search(query: str, department: str, top_k: int = 5):
    results = client.search(
        collection_name=COLLECTION,
        query_vector=get_embedding(query),
        query_filter=Filter(
            must=[FieldCondition(key="department", match=MatchValue(value=department))]
        ),
        limit=top_k,
    )
    return results


# ---------- 5. 主流程 ----------
if __name__ == "__main__":
    upsert_docs()

    print("\n检索：'怎么找相似的文字'")
    for r in search("如何让模型先查资料再回答", department="tech", top_k=2):
        print(f"  [{r.score:.3f}] {r.payload['title']}  (id={r.id})")

-- 临时调高查询时的搜索宽度（值越大越准但越慢）
SET hnsw.ef_search = 100;

-- 然后再查
SELECT id, title FROM docs ORDER BY embedding <=> '[...]' LIMIT 5;

SELECT id, title
FROM docs
WHERE department = 'tech'           -- 过滤在 ORDER BY 之前
  AND created_at > '2025-01-01'
ORDER BY embedding <=> $1
LIMIT 5;

client.search(
    collection_name="docs",
    query_vector=qvec,
    query_filter=Filter(must=[
        FieldCondition(key="department", match=MatchValue(value="tech")),
    ]),
    limit=5,
)

from psycopg import sql
import psycopg

rows = [("doc1", [0.1, 0.2]), ("doc2", [0.3, 0.4])]  # 你的批量数据

with psycopg.connect(DB_DSN) as conn:
    with conn.cursor() as cur:
        # 一次性插多行
        args = b",".join(
            cur.mogrify("(%s, %s)", (title, str(vec)))
            for title, vec in rows
        )
        cur.execute(b"INSERT INTO docs (title, embedding) VALUES " + args)
    conn.commit()

# Qdrant 的 upsert 本身就支持批量
points = [PointStruct(id=i, vector=vecs[i], payload={"i": i}) for i in range(1000)]
client.upsert(collection_name="docs", points=points)   # 一把插入 1000 条