向量数据库核心技术与应用实践

from sentence_transformers import SentenceTransformer
import numpy as np

# 加载一个通用中文语义模型（BGE 系列在中文检索上表现优秀）
model = SentenceTransformer('BAAI/bge-base-zh-v1.5')

sentences = [
    "如何申请退款？",
    "退款流程是什么样的？",
    "今天天气怎么样？"
]

# 生成 embedding，normalize_embeddings=True 会输出单位向量，便于用内积代替余弦
embeddings = model.encode(sentences, normalize_embeddings=True)
print("向量维度:", embeddings.shape)  # (3, 768)

# 前两句语义相近，内积（等价于余弦）应明显高于第三句
print("句1 vs 句2:", float(np.dot(embeddings[0], embeddings[1])))
print("句1 vs 句3:", float(np.dot(embeddings[0], embeddings[2])))

from openai import OpenAI

client = OpenAI(api_key="your-api-key")

def get_embedding(text, model="text-embedding-3-small"):
    """调用 OpenAI 生成 embedding，text-3-small 维度 1536，3-large 维度 3072"""
    resp = client.embeddings.create(input=text, model=model)
    return resp.data[0].embedding

vec = get_embedding("向量数据库是 AI 应用的基础设施")
print("维度:", len(vec))  # 1536

# 批量生成更省钱省时：一次请求传入多条文本
texts = ["文档A", "文档B", "文档C"]
resp = client.embeddings.create(input=texts, model="text-embedding-3-small")
vectors = [d.embedding for d in resp.data]

import numpy as np

def l2_normalize(v):
    """L2 归一化，把向量变成单位向量"""
    norm = np.linalg.norm(v)
    if norm == 0:
        return v
    return v / norm

v = np.array([3.0, 4.0])
vn = l2_normalize(v)
print("归一化前模长:", np.linalg.norm(v))   # 5.0
print("归一化后模长:", np.linalg.norm(vn))  # 1.0

import numpy as np

def cosine_similarity(a, b):
    """计算两个向量的余弦相似度"""
    dot_product = np.dot(a, b)
    norm_a = np.linalg.norm(a)
    norm_b = np.linalg.norm(b)
    return dot_product / (norm_a * norm_b)

# 示例
vec1 = np.array([1, 2, 3])
vec2 = np.array([2, 3, 4])
similarity = cosine_similarity(vec1, vec2)
print(f"余弦相似度: {similarity}")

def euclidean_distance(a, b):
    """计算两个向量的欧几里得距离"""
    return np.sqrt(np.sum((a - b) ** 2))

distance = euclidean_distance(vec1, vec2)
print(f"欧几里得距离: {distance}")

def dot_product_similarity(a, b):
    """计算点积相似度"""
    return np.dot(a, b)

dot_sim = dot_product_similarity(vec1, vec2)
print(f"点积相似度: {dot_sim}")

def manhattan_distance(a, b):
    """计算曼哈顿距离（L1 距离）"""
    return np.sum(np.abs(a - b))

print(f"曼哈顿距离: {manhattan_distance(vec1, vec2)}")

import numpy as np

class FlatIndex:
    """暴力精确检索，召回率 100%，作为召回评估的基准"""
    def __init__(self, metric="cosine"):
        self.vectors = None
        self.metric = metric

    def add(self, vectors):
        self.vectors = np.asarray(vectors, dtype=np.float32)

    def search(self, query, k=10):
        if self.metric == "cosine":
            # 归一化后用内积
            vn = self.vectors / np.linalg.norm(self.vectors, axis=1, keepdims=True)
            qn = query / np.linalg.norm(query)
            scores = vn @ qn
            idx = np.argsort(-scores)[:k]
            return [(int(i), float(scores[i])) for i in idx]
        else:  # L2
            dists = np.linalg.norm(self.vectors - query, axis=1)
            idx = np.argsort(dists)[:k]
            return [(int(i), float(dists[i])) for i in idx]

from sklearn.cluster import KMeans
import numpy as np

class IVFIndex:
    def __init__(self, n_clusters=100):
        self.n_clusters = n_clusters
        self.kmeans = KMeans(n_clusters=n_clusters)
        self.clusters = {}
        self.vectors = {}

    def fit(self, vectors):
        """训练聚类模型"""
        self.centroids = self.kmeans.fit(vectors).cluster_centers_

        # 将向量分配到对应聚类
        labels = self.kmeans.predict(vectors)
        for i, label in enumerate(labels):
            if label not in self.clusters:
                self.clusters[label] = []
            self.clusters[label].append(i)
            self.vectors[i] = vectors[i]

    def search(self, query_vector, k=10, nprobe=1):
        """搜索最相似的k个向量，nprobe 控制检查多少个聚类"""
        # 找到查询向量最接近的 nprobe 个聚类
        distances = np.linalg.norm(self.centroids - query_vector, axis=1)
        closest_clusters = np.argsort(distances)[:nprobe]

        # 在这些聚类中搜索
        candidates = []
        for c in closest_clusters:
            candidates.extend(self.clusters.get(c, []))
        similarities = []

        for idx in candidates:
            similarity = cosine_similarity(query_vector, self.vectors[idx])
            similarities.append((idx, similarity))

        # 返回最相似的k个
        similarities.sort(key=lambda x: x[1], reverse=True)
        return similarities[:k]

import hnswlib
import numpy as np

dim = 128
num_elements = 100000

# 造一批随机向量作为演示数据
data = np.random.rand(num_elements, dim).astype(np.float32)

# 创建索引，space 可选 'l2' / 'ip'（内积）/ 'cosine'
index = hnswlib.Index(space='cosine', dim=dim)

# 初始化：max_elements 上限，M 与 ef_construction 决定图质量
index.init_index(max_elements=num_elements, ef_construction=200, M=16)

# 批量插入，ids 可自定义
ids = np.arange(num_elements)
index.add_items(data, ids)

# 设置查询时的 ef：越大越准越慢
index.set_ef(64)

# 查询：返回最近的 k 个 id 及距离
query = np.random.rand(1, dim).astype(np.float32)
labels, distances = index.knn_query(query, k=10)
print("命中 id:", labels[0])
print("距离:", distances[0])

# 索引可持久化到磁盘
index.save_index("hnsw_index.bin")

# 之后重新加载
new_index = hnswlib.Index(space='cosine', dim=dim)
new_index.load_index("hnsw_index.bin", max_elements=num_elements)

import faiss
import numpy as np

dim = 128
data = np.random.rand(50000, dim).astype(np.float32)

# 32 是每个节点的邻居数 M，IndexHNSWFlat 底层存原始向量
index = faiss.IndexHNSWFlat(dim, 32)
index.hnsw.efConstruction = 200
index.add(data)

# 查询时设置 efSearch
index.hnsw.efSearch = 64
query = np.random.rand(5, dim).astype(np.float32)
distances, indices = index.search(query, k=10)
print(indices)   # (5, 10)，每行是一个查询的 top-10 结果 id

import numpy as np

class RandomProjectionLSH:
    """随机投影 LSH：用 n_bits 个随机超平面把向量映射成二进制哈希码"""
    def __init__(self, dim, n_bits=16):
        # 每一行是一个随机超平面的法向量
        self.planes = np.random.randn(n_bits, dim)
        self.buckets = {}

    def _hash(self, v):
        # 向量在每个超平面的哪一侧，> 0 记 1，否则记 0
        bits = (self.planes @ v) > 0
        return ''.join('1' if b else '0' for b in bits)

    def add(self, idx, v):
        key = self._hash(v)
        self.buckets.setdefault(key, []).append(idx)

    def query(self, v):
        # 返回落在同一个桶里的候选 id
        return self.buckets.get(self._hash(v), [])

import numpy as np

class ScalarQuantizer:
    """把 float32 向量线性映射到 [0, 255] 的 uint8，压缩 4 倍"""
    def fit(self, vectors):
        self.vmin = vectors.min(axis=0)
        self.vmax = vectors.max(axis=0)
        self.scale = (self.vmax - self.vmin) / 255.0
        self.scale[self.scale == 0] = 1e-8
        return self

    def encode(self, vectors):
        q = np.round((vectors - self.vmin) / self.scale)
        return np.clip(q, 0, 255).astype(np.uint8)

    def decode(self, codes):
        return codes.astype(np.float32) * self.scale + self.vmin

class ProductQuantizer:
    def __init__(self, m=8, k=256):
        self.m = m  # 子向量数量
        self.k = k  # 每个子向量的码本大小
        self.codebooks = None

    def fit(self, vectors):
        """训练乘积量化器"""
        n, d = vectors.shape
        self.d_sub = d // self.m  # 每个子向量的维度

        self.codebooks = np.zeros((self.m, self.k, self.d_sub))

        for i in range(self.m):
            start_idx = i * self.d_sub
            end_idx = (i + 1) * self.d_sub
            sub_vectors = vectors[:, start_idx:end_idx]

            # 对每个子向量聚类
            kmeans = KMeans(n_clusters=self.k)
            kmeans.fit(sub_vectors)
            self.codebooks[i] = kmeans.cluster_centers_

    def encode(self, vectors):
        """编码向量"""
        n, d = vectors.shape
        codes = np.zeros((n, self.m), dtype=np.int32)

        for i in range(self.m):
            start_idx = i * self.d_sub
            end_idx = (i + 1) * self.d_sub
            sub_vectors = vectors[:, start_idx:end_idx]

            for j in range(n):
                distances = np.linalg.norm(
                    self.codebooks[i] - sub_vectors[j], axis=1
                )
                codes[j, i] = np.argmin(distances)

        return codes

    def decode(self, codes):
        """解码向量"""
        n = codes.shape[0]
        reconstructed = np.zeros((n, self.d_sub * self.m))

        for i in range(self.m):
            for j in range(n):
                reconstructed[j, i*self.d_sub:(i+1)*self.d_sub] = self.codebooks[i, codes[j, i]]

        return reconstructed

import faiss
import numpy as np

dim = 128
nlist = 1024        # IVF 聚类数
m = 16              # PQ 子段数（需能整除 dim）
nbits = 8           # 每段 8 bit -> 256 个中心

# 训练数据（真实场景用你自己的向量样本，至少几万条）
train = np.random.rand(100000, dim).astype(np.float32)

# 量化器负责 IVF 的粗聚类
quantizer = faiss.IndexFlatL2(dim)
index = faiss.IndexIVFPQ(quantizer, dim, nlist, m, nbits)

# IVF_PQ 必须先 train 再 add
index.train(train)
index.add(train)

# nprobe 控制检查多少个桶，是召回/速度的核心旋钮
index.nprobe = 16

query = np.random.rand(10, dim).astype(np.float32)
distances, indices = index.search(query, k=10)
print("每条查询的 top-10:", indices.shape)  # (10, 10)

# 保存与加载
faiss.write_index(index, "ivfpq.index")
loaded = faiss.read_index("ivfpq.index")

import faiss
# "OPQ16,IVF1024,PQ16" = OPQ 预处理 + 1024 桶 IVF + 16 段 PQ
index = faiss.index_factory(128, "OPQ16,IVF1024,PQ16")
# "HNSW32" = 纯 HNSW；"IVF4096,Flat" = IVF 不压缩

from qdrant_client import QdrantClient
from qdrant_client.models import Filter, FieldCondition, MatchValue, Range

client = QdrantClient(url="http://localhost:6333")

results = client.search(
    collection_name="products",
    query_vector=query_vec,          # 查询向量
    query_filter=Filter(
        must=[
            FieldCondition(key="category", match=MatchValue(value="运动鞋")),
            FieldCondition(key="in_stock", match=MatchValue(value=True)),
            FieldCondition(key="price", range=Range(gte=100, lte=500)),
        ]
    ),
    limit=20,
    with_payload=True,
)
for r in results:
    print(r.id, r.score, r.payload.get("title"))

def reciprocal_rank_fusion(rank_lists, k=60):
    """
    rank_lists: 多路召回结果，每个是按相关性排序的 doc_id 列表
    k: 平滑常数，通常取 60
    返回: 融合后按分数降序的 (doc_id, score) 列表
    """
    scores = {}
    for ranking in rank_lists:
        for rank, doc_id in enumerate(ranking):
            scores[doc_id] = scores.get(doc_id, 0.0) + 1.0 / (k + rank + 1)
    return sorted(scores.items(), key=lambda x: x[1], reverse=True)

# 假设向量检索和 BM25 各返回一个有序 id 列表
dense_rank = ["d3", "d1", "d7", "d2"]
sparse_rank = ["d1", "d5", "d3", "d9"]
fused = reciprocal_rank_fusion([dense_rank, sparse_rank])
print(fused)  # d1、d3 因两路都靠前，融合分最高

from sentence_transformers import CrossEncoder

# 一个中文 / 多语 rerank 模型（BGE reranker 系列常用）
reranker = CrossEncoder('BAAI/bge-reranker-base')

query = "如何重置账户密码？"
# 第一阶段向量召回得到的候选文档
candidates = [
    "点击登录页的忘记密码，通过邮箱验证即可重置。",
    "账户余额可在个人中心查看。",
    "密码需包含大小写字母和数字，长度至少 8 位。",
    "如需注销账户请联系客服。",
]

# cross-encoder 对每个 (query, doc) 打分
pairs = [[query, doc] for doc in candidates]
scores = reranker.predict(pairs)

# 按分数排序，取 top-2 作为最终结果
ranked = sorted(zip(candidates, scores), key=lambda x: x[1], reverse=True)
for doc, score in ranked[:2]:
    print(f"{score:.3f}  {doc}")

from pinecone import Pinecone, ServerlessSpec

pc = Pinecone(api_key="your-api-key")

# 创建索引（Serverless）
pc.create_index(
    name="my-index",
    dimension=1536,
    metric="cosine",
    spec=ServerlessSpec(cloud="aws", region="us-east-1"),
)

index = pc.Index("my-index")

# 添加向量（带 metadata，可用于过滤）
index.upsert(vectors=[
    {"id": "vec1", "values": [0.1, 0.2, 0.3], "metadata": {"genre": "news"}},
    {"id": "vec2", "values": [0.4, 0.5, 0.6], "metadata": {"genre": "blog"}},
])

# 带过滤的查询
result = index.query(
    vector=[0.15, 0.25, 0.35],
    top_k=10,
    include_metadata=True,
    filter={"genre": {"$eq": "news"}},
)

import weaviate

# 连接到Weaviate
client = weaviate.Client("http://localhost:8080")

# 创建类
class_obj = {
    "class": "Document",
    "properties": [
        {
            "name": "content",
            "dataType": ["string"]
        }
    ],
    "vectorizer": "none"  # 手动提供向量
}

client.schema.create_class(class_obj)

# 添加对象
client.data_object.create(
    data_object={
        "content": "这是一个文档内容"
    },
    class_name="Document",
    vector=[0.1, 0.2, 0.3]  # 提供向量
)

# 相似性搜索
near_vector = {"vector": [0.15, 0.25, 0.35]}
result = client.query.get("Document", ["content"]).with_near_vector(
    near_vector
).with_limit(10).do()

from pymilvus import connections, FieldSchema, CollectionSchema, DataType, Collection

# 连接Milvus
connections.connect("default", host="localhost", port="19530")

# 定义schema
fields = [
    FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
    FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=128),
    FieldSchema(name="text", dtype=DataType.VARCHAR, max_length=65535)
]
schema = CollectionSchema(fields, "文档集合")

# 创建collection
collection = Collection("docs", schema)

# 创建索引
index_params = {
    "index_type": "IVF_FLAT",
    "metric_type": "L2",
    "params": {"nlist": 100}
}
collection.create_index(field_name="embedding", index_params=index_params)

# 插入数据
entities = [
    [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]],  # embeddings
    ["文本1", "文本2"]  # texts
]
insert_result = collection.insert(entities)

# 搜索
collection.load()
search_params = {"metric_type": "L2", "params": {"nprobe": 10}}
results = collection.search(
    data=[[0.15, 0.25, 0.35]],
    anns_field="embedding",
    param=search_params,
    limit=10,
    output_fields=["text"]
)

import chromadb
from chromadb.utils import embedding_functions

# 创建客户端
client = chromadb.Client()

# 创建集合
collection = client.create_collection("docs")

# 添加文档
collection.add(
    documents=[
        "这是一个文档",
        "这是另一个文档"
    ],
    metadatas=[
        {"source": "web"},
        {"source": "book"}
    ],
    ids=["doc1", "doc2"]
)

# 查询相似文档
results = collection.query(
    query_texts=["查询文本"],
    n_results=2
)

from qdrant_client import QdrantClient
from qdrant_client.models import Distance, VectorParams, PointStruct

client = QdrantClient(url="http://localhost:6333")

# 创建集合
client.recreate_collection(
    collection_name="docs",
    vectors_config=VectorParams(size=768, distance=Distance.COSINE),
)

# 插入点（向量 + payload 元数据）
client.upsert(
    collection_name="docs",
    points=[
        PointStruct(id=1, vector=vec1, payload={"title": "退款政策", "lang": "zh"}),
        PointStruct(id=2, vector=vec2, payload={"title": "配送说明", "lang": "zh"}),
    ],
)

# 检索
hits = client.search(collection_name="docs", query_vector=query_vec, limit=5)

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

-- 建表，embedding 为 768 维
CREATE TABLE documents (
    id bigserial PRIMARY KEY,
    content text,
    category text,
    embedding vector(768)
);

-- 建 HNSW 索引，vector_cosine_ops 表示用余弦距离
CREATE INDEX ON documents
USING hnsw (embedding vector_cosine_ops)
WITH (m = 16, ef_construction = 200);

-- 相似性搜索：<=> 是余弦距离运算符，越小越相似
-- 同时用普通 WHERE 做元数据过滤
SELECT id, content, 1 - (embedding <=> '[0.1, 0.2, ...]') AS similarity
FROM documents
WHERE category = 'faq'
ORDER BY embedding <=> '[0.1, 0.2, ...]'
LIMIT 10;

import redis
from redis.commands.search.field import VectorField, TextField
from redis.commands.search.indexDefinition import IndexDefinition, IndexType

r = redis.Redis(host="localhost", port=6379)

# 创建 HNSW 向量索引
schema = (
    TextField("content"),
    VectorField("embedding", "HNSW", {
        "TYPE": "FLOAT32", "DIM": 768,
        "DISTANCE_METRIC": "COSINE", "M": 16, "EF_CONSTRUCTION": 200,
    }),
)
r.ft("idx:docs").create_index(
    schema,
    definition=IndexDefinition(prefix=["doc:"], index_type=IndexType.HASH),
)

class RAGSystem:
    def __init__(self, vector_db, llm):
        self.vector_db = vector_db
        self.llm = llm
        self.embedding_model = embedding_functions.SentenceTransformerEmbeddingFunction()

    def retrieve(self, query, k=5):
        """从向量数据库检索相关文档"""
        query_embedding = self.embedding_model([query])
        results = self.vector_db.query(
            query_embeddings=query_embedding,
            n_results=k
        )
        return results

    def generate(self, query, context):
        """基于上下文生成答案"""
        prompt = f"""
        基于以下上下文信息回答问题：

        上下文：{context}

        问题：{query}

        请提供准确的回答。如果上下文中没有相关信息，请明确说明"根据已有资料无法回答"，不要编造。
        """
        return self.llm.generate(prompt)

    def query(self, question):
        """完整RAG流程"""
        # 1. 检索相关文档
        retrieved_docs = self.retrieve(question)

        # 2. 构建上下文
        context = " ".join(retrieved_docs['documents'][0])

        # 3. 生成答案
        answer = self.generate(question, context)

        return {
            'answer': answer,
            'sources': retrieved_docs['metadatas'][0]
        }

from langchain.text_splitter import RecursiveCharacterTextSplitter

class DocumentProcessor:
    def __init__(self):
        self.text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=1000,
            chunk_overlap=200,
            separators=["\n\n", "\n", " ", ""]
        )

    def process_document(self, text, doc_id):
        """处理文档并准备插入向量数据库"""
        chunks = self.text_splitter.split_text(text)

        # 为每个chunk生成向量
        embeddings = self.embedding_model(chunks)

        # 准备数据
        data = {
            'ids': [f"{doc_id}_{i}" for i in range(len(chunks))],
            'documents': chunks,
            'embeddings': embeddings,
            'metadatas': [{'doc_id': doc_id, 'chunk_id': i}
                         for i in range(len(chunks))]
        }

        return data

# 一个最小语义搜索：构建 -> 查询
model = SentenceTransformer('BAAI/bge-base-zh-v1.5')

corpus = [
    "退换货需在签收后 7 天内发起申请。",
    "会员积分每消费 1 元累计 1 分。",
    "订单发货后可在物流页面实时追踪。",
]
corpus_emb = model.encode(corpus, normalize_embeddings=True)

def semantic_search(query, top_k=2):
    q = model.encode([query], normalize_embeddings=True)[0]
    scores = corpus_emb @ q          # 归一化后内积即余弦
    top = np.argsort(-scores)[:top_k]
    return [(corpus[i], float(scores[i])) for i in top]

print(semantic_search("怎么退货"))   # 命中"退换货"那条，尽管没有"退货"二字

import numpy as np

def build_user_vector(clicked_item_vectors, weights=None):
    """用用户点击过的商品向量加权平均，作为用户兴趣向量"""
    vecs = np.asarray(clicked_item_vectors, dtype=np.float32)
    if weights is None:
        user_vec = vecs.mean(axis=0)
    else:
        w = np.asarray(weights, dtype=np.float32).reshape(-1, 1)
        user_vec = (vecs * w).sum(axis=0) / w.sum()
    # 归一化后用于向量库检索
    return user_vec / np.linalg.norm(user_vec)

# 用户最近浏览的 3 个商品向量（越近的权重越高）
user_vec = build_user_vector([item_a, item_b, item_c], weights=[1.0, 2.0, 3.0])

# 在商品向量库中召回 50 个候选，过滤掉已购买和已下架商品
candidates = client.search(
    collection_name="items",
    query_vector=user_vec.tolist(),
    query_filter=Filter(must_not=[FieldCondition(key="status", match=MatchValue(value="offline"))]),
    limit=50,
)

import torch
import clip
from PIL import Image

device = "cuda" if torch.cuda.is_available() else "cpu"
model, preprocess = clip.load("ViT-B/32", device=device)

def encode_image(path):
    img = preprocess(Image.open(path)).unsqueeze(0).to(device)
    with torch.no_grad():
        feat = model.encode_image(img)
    feat = feat / feat.norm(dim=-1, keepdim=True)   # 归一化
    return feat.cpu().numpy()[0]

def encode_text(text):
    tokens = clip.tokenize([text]).to(device)
    with torch.no_grad():
        feat = model.encode_text(tokens)
    feat = feat / feat.norm(dim=-1, keepdim=True)
    return feat.cpu().numpy()[0]

# 以图搜图：encode_image(用户图) 去检索
# 以文搜图：encode_text("红色连衣裙") 去同一个向量库检索
query_vec = encode_text("红色连衣裙")

def is_duplicate(new_vec, index, threshold=0.95):
    """若库中已存在相似度 >= 阈值的向量，则判为重复"""
    hits = index.search(new_vec, k=1)
    if hits and hits[0][1] >= threshold:
        return True, hits[0][0]   # 返回重复对象的 id
    return False, None

# 流式去重：先查重，非重复才入库
dup, dup_id = is_duplicate(vec, index, threshold=0.92)
if not dup:
    index.add([vec])

def anomaly_score(vec, index, k=10):
    """用到 k 近邻的平均距离衡量异常程度，越大越异常"""
    neighbors = index.search(vec, k=k)   # 返回 (id, distance)
    dists = [d for _, d in neighbors]
    return float(np.mean(dists))

# 超过经验阈值即告警
score = anomaly_score(new_event_vec, index, k=20)
if score > ANOMALY_THRESHOLD:
    raise_alert(new_event_vec)

def batch_upsert(vectors, batch_size=1000):
    """批量插入向量"""
    for i in range(0, len(vectors), batch_size):
        batch = vectors[i:i+batch_size]
        index.upsert(vectors=batch)
        print(f"已插入 {i+len(batch)} / {len(vectors)} 个向量")

def embed_in_batches(texts, model, batch_size=64):
    """分批生成 embedding，避免单次请求过大或过多小请求"""
    all_vecs = []
    for i in range(0, len(texts), batch_size):
        batch = texts[i:i + batch_size]
        vecs = model.encode(batch, normalize_embeddings=True)
        all_vecs.extend(vecs)
    return all_vecs

def merge_shard_results(shard_results, k=10):
    """归并多个分片各自返回的 top-k，取全局 top-k"""
    merged = []
    for res in shard_results:      # 每个是 [(id, score), ...]
        merged.extend(res)
    merged.sort(key=lambda x: x[1], reverse=True)
    return merged[:k]

import numpy as np

def recall_at_k(ann_index, flat_index, queries, k=10):
    """
    ann_index: 待评估的近似索引
    flat_index: 暴力精确索引（真值）
    返回平均 recall@k
    """
    total = 0.0
    for q in queries:
        truth = set(i for i, _ in flat_index.search(q, k))
        approx = set(i for i, _ in ann_index.search(q, k))
        total += len(truth & approx) / k
    return total / len(queries)

# 用一批查询评估，逐步调大 nprobe/efSearch，观察 recall 变化
r = recall_at_k(my_hnsw, my_flat, sample_queries, k=10)
print(f"recall@10 = {r:.3f}")

# 方案一：共享集合 + tenant_id 过滤（务必强制注入，绝不可遗漏）
def tenant_search(client, tenant_id, query_vec, k=10):
    return client.search(
        collection_name="shared",
        query_vector=query_vec,
        query_filter=Filter(must=[
            FieldCondition(key="tenant_id", match=MatchValue(value=tenant_id))
        ]),
        limit=k,
    )

# 方案二：Pinecone 用 namespace 做隔离
index.upsert(vectors=vectors, namespace=f"tenant_{tenant_id}")
index.query(vector=query_vec, top_k=10, namespace=f"tenant_{tenant_id}")

from annoy import AnnoyIndex
import numpy as np

dim = 128
t = AnnoyIndex(dim, 'angular')   # angular 近似余弦
for i in range(10000):
    t.add_item(i, np.random.rand(dim))

t.build(50)                       # 50 棵树，越多越准越慢
t.save('annoy.ann')

# 查询：返回最近的 10 个 id（可带距离）
ids, dists = t.get_nns_by_vector(np.random.rand(dim), 10, include_distances=True)
print(ids)

import numpy as np

def binary_quantize(vectors):
    """按符号二值化：>0 -> 1，否则 0，打包成 uint8 比特"""
    bits = (vectors > 0).astype(np.uint8)
    # 每 8 位打包成一个字节，节省存储
    return np.packbits(bits, axis=1)

def hamming_distance(a_packed, b_packed):
    """汉明距离：异或后统计 1 的个数"""
    xor = np.bitwise_xor(a_packed, b_packed)
    return int(np.unpackbits(xor).sum())

vecs = np.random.randn(1000, 1536).astype(np.float32)
packed = binary_quantize(vecs)
print("原始字节:", vecs.nbytes, " 二值后字节:", packed.nbytes)  # 约 32 倍压缩

def bq_search_with_rescore(query, packed_index, raw_store, k=10, rescore_n=200):
    """
    packed_index: 二值索引，快速返回候选
    raw_store: 原始向量存储（可在磁盘），用于精排
    """
    # 1. 二值粗筛
    q_packed = binary_quantize(query.reshape(1, -1))[0]
    candidates = packed_index.hamming_topk(q_packed, rescore_n)  # 返回候选 id
    # 2. 用原始向量精排
    q = query / np.linalg.norm(query)
    scored = []
    for cid in candidates:
        v = raw_store.get(cid)
        v = v / np.linalg.norm(v)
        scored.append((cid, float(np.dot(q, v))))
    scored.sort(key=lambda x: x[1], reverse=True)
    return scored[:k]

from qdrant_client.models import ScalarQuantization, ScalarQuantizationConfig, ScalarType

client.create_collection(
    collection_name="docs",
    vectors_config=VectorParams(size=1536, distance=Distance.COSINE),
    quantization_config=ScalarQuantization(
        scalar=ScalarQuantizationConfig(
            type=ScalarType.INT8,
            quantile=0.99,      # 用分位数裁剪极端值，减小量化误差
            always_ram=True,    # 量化向量常驻内存，原始向量可落盘
        )
    ),
)

import numpy as np

def colbert_maxsim(query_token_vecs, doc_token_vecs):
    """
    query_token_vecs: (Lq, d) 查询每个 token 的向量（已归一化）
    doc_token_vecs:   (Ld, d) 文档每个 token 的向量（已归一化）
    返回 MaxSim 分数：每个查询 token 取与文档所有 token 的最大相似度再求和
    """
    sim = query_token_vecs @ doc_token_vecs.T   # (Lq, Ld) 相似度矩阵
    return float(sim.max(axis=1).sum())         # 每行取最大，再求和

import math, time

def freshness_rerank(hits, half_life_days=7.0):
    """
    hits: [(id, score, timestamp), ...]  timestamp 为发布时间(秒)
    用指数衰减：距今越久权重越低，half_life_days 天后权重减半
    """
    now = time.time()
    lam = math.log(2) / (half_life_days * 86400)
    rescored = []
    for _id, score, ts in hits:
        decay = math.exp(-lam * (now - ts))
        rescored.append((_id, score * decay))
    rescored.sort(key=lambda x: x[1], reverse=True)
    return rescored

# Upsert：内容更新时重新 embed 并覆盖同 id
def upsert_document(client, doc_id, new_text, model):
    new_vec = model.encode([new_text], normalize_embeddings=True)[0]
    client.upsert(
        collection_name="docs",
        points=[PointStruct(
            id=doc_id,
            vector=new_vec.tolist(),
            payload={"text": new_text, "updated_at": int(time.time())},
        )],
    )

# 删除：按 id 或按过滤条件批量删
def delete_expired(client, before_ts):
    client.delete(
        collection_name="docs",
        points_selector=Filter(must=[
            FieldCondition(key="updated_at", range=Range(lt=before_ts))
        ]),
    )

def reindex_in_batches(source, target, model, batch_size=500, rate_limit_s=0.1):
    """把 source 集合的原始文本用新模型重嵌入写入 target 集合"""
    offset = None
    while True:
        records, offset = source.scroll(
            collection_name="docs_old", limit=batch_size, offset=offset,
            with_payload=True, with_vectors=False,
        )
        if not records:
            break
        texts = [r.payload["text"] for r in records]
        vecs = model.encode(texts, normalize_embeddings=True)
        points = [
            PointStruct(id=r.id, vector=v.tolist(), payload=r.payload)
            for r, v in zip(records, vecs)
        ]
        target.upsert(collection_name="docs_new", points=points)
        time.sleep(rate_limit_s)   # 限速，保护线上
        if offset is None:
            break

import time
import numpy as np

def benchmark(index, flat_truth, queries, k=10, param_grid=None, set_param=None):
    """
    index: 待测索引对象
    flat_truth: dict[query_id -> set(真值 topk id)]，用暴力搜索预先算好
    queries: [(query_id, query_vec), ...]
    param_grid: 要扫描的参数取值列表，例如 [16, 32, 64, 128]
    set_param: 回调，set_param(index, p) 把参数 p 应用到索引（如 efSearch）
    """
    results = []
    for p in (param_grid or [None]):
        if set_param and p is not None:
            set_param(index, p)

        latencies = []
        recalls = []
        for qid, qv in queries:
            t0 = time.perf_counter()
            hits = index.search(qv, k)
            latencies.append((time.perf_counter() - t0) * 1000)  # ms
            approx = set(i for i, _ in hits)
            truth = flat_truth[qid]
            recalls.append(len(approx & truth) / k)

        latencies.sort()
        n = len(latencies)
        results.append({
            "param": p,
            "recall@k": round(float(np.mean(recalls)), 4),
            "qps": round(n / (sum(latencies) / 1000), 1),
            "p50_ms": round(latencies[int(n * 0.50)], 2),
            "p99_ms": round(latencies[min(int(n * 0.99), n - 1)], 2),
        })
    return results

# 用法：扫描 efSearch，打印召回-延迟曲线
def set_ef(idx, ef):
    idx.hnsw.efSearch = ef   # faiss HNSW

# report = benchmark(hnsw_index, truth, test_queries, k=10,
#                    param_grid=[16, 32, 64, 128, 256], set_param=set_ef)
# for row in report: print(row)

import time
from contextlib import contextmanager

@contextmanager
def observe(metric_name, exporter):
    """一个极简的埋点上下文，记录耗时与成败并上报"""
    start = time.perf_counter()
    ok = True
    try:
        yield
    except Exception:
        ok = False
        raise
    finally:
        elapsed_ms = (time.perf_counter() - start) * 1000
        exporter.emit(metric_name, elapsed_ms, tags={"ok": ok})

# 用法
# with observe("vector_search_latency", metrics):
#     hits = index.search(query_vec, k=10)

import hashlib

def assign_bucket(user_id, experiment="rerank_test", buckets=("control", "treatment")):
    """用一致性哈希把用户稳定分桶，保证同一用户始终在同一组"""
    h = hashlib.md5(f"{experiment}:{user_id}".encode()).hexdigest()
    idx = int(h, 16) % len(buckets)
    return buckets[idx]

def search_with_experiment(user_id, query_vec, k=10):
    bucket = assign_bucket(user_id)
    hits = base_retrieve(query_vec, k=50)
    if bucket == "treatment":
        hits = rerank(hits)[:k]     # 实验组加重排
    else:
        hits = hits[:k]             # 对照组不加
    log_impression(user_id, bucket, [h[0] for h in hits])  # 记录曝光用于后续分析
    return hits

class ProductionRAG:
    def __init__(self, client, dense_model, reranker, llm):
        self.client = client            # 向量数据库客户端
        self.dense_model = dense_model  # 稠密 embedding 模型
        self.reranker = reranker        # cross-encoder 重排模型
        self.llm = llm

    def retrieve(self, query, tenant_id, recall_k=50, final_k=5):
        # 1. 稠密向量召回（带多租户过滤）
        q_vec = self.dense_model.encode([query], normalize_embeddings=True)[0]
        dense_hits = self.client.search(
            collection_name="kb",
            query_vector=q_vec.tolist(),
            query_filter=Filter(must=[
                FieldCondition(key="tenant_id", match=MatchValue(value=tenant_id))
            ]),
            limit=recall_k,
            with_payload=True,
        )
        # 2. cross-encoder 重排，选出最终 top-k
        pairs = [[query, h.payload["text"]] for h in dense_hits]
        scores = self.reranker.predict(pairs)
        ranked = sorted(zip(dense_hits, scores), key=lambda x: x[1], reverse=True)
        return [h for h, _ in ranked[:final_k]]

    def answer(self, query, tenant_id):
        docs = self.retrieve(query, tenant_id)
        # 3. 构建带编号引用的上下文，便于溯源
        context = "\n\n".join(
            f"[{i+1}] {d.payload['text']}" for i, d in enumerate(docs)
        )
        prompt = f"""你是企业知识库助手。仅依据下列资料回答，并在句末用 [编号] 标注来源。
若资料不足以回答，请直接说明无法回答，不要编造。

资料：
{context}

问题：{query}
回答："""
        answer = self.llm.generate(prompt)
        # 4. 返回答案与来源，供前端展示引用
        return {
            "answer": answer,
            "sources": [{"id": d.id, "title": d.payload.get("title")} for d in docs],
        }

import hnswlib
import numpy as np

dim, num = 768, 100000
data = np.random.rand(num, dim).astype("float32")

index = hnswlib.Index(space="cosine", dim=dim)
index.init_index(max_elements=num, ef_construction=200, M=32)
index.add_items(data, np.arange(num))

# 查询期动态调 efSearch，权衡召回与延迟
index.set_ef(120)
labels, distances = index.knn_query(data[:5], k=10)
print(labels.shape)  # (5, 10)

import faiss
import numpy as np

dim = 768
data = np.random.rand(200000, dim).astype("float32")

quantizer = faiss.IndexFlatL2(dim)
# nlist=1024 个簇，PQ 把 768 维切成 m=48 段，每段 8 bit
index = faiss.IndexIVFPQ(quantizer, dim, 1024, 48, 8)
index.train(data)      # PQ/IVF 都需要先训练
index.add(data)

index.nprobe = 16      # 查询时搜 16 个最近簇
D, I = index.search(data[:5], 10)

def reciprocal_rank_fusion(result_lists, k=60):
    scores = {}
    for results in result_lists:
        for rank, doc_id in enumerate(results):
            scores[doc_id] = scores.get(doc_id, 0) + 1 / (k + rank + 1)
    return sorted(scores, key=scores.get, reverse=True)

dense_ids = vector_search(query, k=20)     # 向量召回
sparse_ids = bm25_search(query, k=20)      # 关键词召回
fused = reciprocal_rank_fusion([dense_ids, sparse_ids])[:20]

# 再用 Cross-Encoder 对融合结果精排，取前 5
reranked = cross_encoder_rerank(query, fused)[:5]

-- pgvector：建表、建 HNSW 索引、查询
CREATE EXTENSION IF NOT EXISTS vector;
CREATE TABLE docs (id bigserial PRIMARY KEY, content text, embedding vector(768));
CREATE INDEX ON docs USING hnsw (embedding vector_cosine_ops)
  WITH (m = 32, ef_construction = 200);

-- 查询期调 efSearch
SET hnsw.ef_search = 120;
SELECT id, content, 1 - (embedding <=> $1) AS score
FROM docs ORDER BY embedding <=> $1 LIMIT 10;

# 用评估集横向对比 embedding 模型
def eval_embedding(embed_fn, eval_set, index_builder):
    index = index_builder(embed_fn)
    recalls = []
    for case in eval_set:
        q_vec = embed_fn(case["question"])
        got = index.search(q_vec, k=10)
        hit = len(set(got) & set(case["gold_ids"]))
        recalls.append(hit / len(case["gold_ids"]))
    return sum(recalls) / len(recalls)

for name, fn in {"3-small": embed_openai, "bge-m3": embed_bge}.items():
    print(name, eval_embedding(fn, eval_set, build_index))

# Qdrant：带元数据过滤的检索
from qdrant_client import QdrantClient
from qdrant_client.models import Filter, FieldCondition, MatchValue

client = QdrantClient(url="http://localhost:6333")
results = client.search(
    collection_name="docs",
    query_vector=query_vec,
    query_filter=Filter(must=[
        FieldCondition(key="user_id", match=MatchValue(value="u-42")),
        FieldCondition(key="category", match=MatchValue(value="faq")),
    ]),
    limit=10,
)

import torch
from transformers import CLIPModel, CLIPProcessor

model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
proc = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")

# 图片入库：编码成向量
img_inputs = proc(images=images, return_tensors="pt")
with torch.no_grad():
    img_vecs = model.get_image_features(**img_inputs)
    img_vecs = torch.nn.functional.normalize(img_vecs, dim=-1)

# 文本查询：编码到同一空间
txt_inputs = proc(text=["一只在沙滩上的狗"], return_tensors="pt", padding=True)
with torch.no_grad():
    txt_vec = model.get_text_features(**txt_inputs)
    txt_vec = torch.nn.functional.normalize(txt_vec, dim=-1)

# 之后把 img_vecs 存进向量库，用 txt_vec 检索即可跨模态匹配

import numpy as np
from rank_bm25 import BM25Okapi

class HybridStore:
    def __init__(self, embed_fn, index):
        self.embed = embed_fn
        self.index = index          # 如 hnswlib / faiss 实例
        self.docs = []              # 存原文与元数据
        self.bm25 = None

    def add(self, chunks):
        # 1) 向量化并 L2 归一化（cosine 前提）
        vecs = self.embed([c["text"] for c in chunks])
        vecs = vecs / np.linalg.norm(vecs, axis=1, keepdims=True)
        start = len(self.docs)
        self.index.add_items(vecs, np.arange(start, start + len(chunks)))
        self.docs.extend(chunks)
        # 2) 重建 BM25 稀疏索引
        self.bm25 = BM25Okapi([c["text"].split() for c in self.docs])

    def search(self, query, k=20, filters=None):
        q = self.embed([query])[0]
        q = q / np.linalg.norm(q)
        self.index.set_ef(120)                      # 查询期调召回
        dense_ids, _ = self.index.knn_query(q, k=k)
        dense_ids = dense_ids[0].tolist()
        # 3) BM25 稀疏召回
        scores = self.bm25.get_scores(query.split())
        sparse_ids = np.argsort(scores)[::-1][:k].tolist()
        # 4) RRF 融合
        fused = self._rrf([dense_ids, sparse_ids])
        # 5) 元数据后过滤
        if filters:
            fused = [i for i in fused if self._match(i, filters)]
        return [self.docs[i] for i in fused[:k]]

    def _rrf(self, lists, k=60):
        s = {}
        for lst in lists:
            for rank, i in enumerate(lst):
                s[i] = s.get(i, 0) + 1 / (k + rank + 1)
        return sorted(s, key=s.get, reverse=True)

    def _match(self, i, filters):
        meta = self.docs[i].get("meta", {})
        return all(meta.get(kk) == vv for kk, vv in filters.items())

import time

def benchmark(store, queries, gold, ef_list=(50, 100, 200, 400)):
    for ef in ef_list:
        store.index.set_ef(ef)
        t0 = time.time()
        recalls = []
        for q, g in zip(queries, gold):
            got = [d["id"] for d in store.search(q, k=10)]
            recalls.append(len(set(got) & set(g)) / len(g))
        latency = (time.time() - t0) / len(queries) * 1000
        print(f"efSearch={ef}: recall={sum(recalls)/len(recalls):.3f} "
              f"latency={latency:.1f}ms")