12-向量数据库¶

from pymilvus import (
    connections, Collection, FieldSchema,
    CollectionSchema, DataType, utility
)

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

# 定义字段
fields = [
    FieldSchema(name="id", dtype=DataType.INT64,
                is_primary=True, auto_id=True),
    FieldSchema(name="title", dtype=DataType.VARCHAR, max_length=512),
    FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=768),
]

# 创建集合
schema = CollectionSchema(fields, description="文档向量集合")
collection = Collection("documents", schema)
print(f"集合创建成功: {collection.name}")

# 创建 HNSW 索引
index_params = {
    "index_type": "HNSW",
    "metric_type": "COSINE",  # 余弦相似度
    "params": {
        "M": 32,               # 每个节点的边数
        "efConstruction": 256  # 构建质量
    }
}
collection.create_index("embedding", index_params)
print("索引创建完成")

# 加载到内存
collection.load()

import numpy as np
from sentence_transformers import SentenceTransformer

# 加载 Embedding 模型
model = SentenceTransformer("BAAI/bge-base-zh-v1.5")

# 准备数据
documents = [
    "向量数据库是 AI 时代的关键基础设施",
    "Milvus 支持十亿级向量的高效检索",
    "HNSW 算法在召回率和速度之间取得了良好平衡",
    "RAG 系统将检索与生成相结合",
    "Product Quantization 可以大幅压缩向量存储",
]

# 生成 Embedding
embeddings = model.encode(documents).tolist()

# 插入数据
data = [documents, embeddings]
result = collection.insert(data)
print(f"插入 {result.insert_count} 条记录")

# 刷新确保数据持久化
collection.flush()

# 查询向量
query = "如何优化大规模向量检索？"
query_embedding = model.encode([query]).tolist()

# 执行搜索
search_params = {
    "metric_type": "COSINE",
    "params": {"ef": 128}  # HNSW 搜索参数
}

results = collection.search(
    data=query_embedding,
    anns_field="embedding",
    param=search_params,
    limit=3,
    output_fields=["title"]
)

# 输出结果
for hits in results:
    for hit in hits:
        print(f"ID: {hit.id}, 距离: {hit.distance:.4f}, "
              f"标题: {hit.entity.get('title')}")

# 带过滤条件的搜索
results = collection.search(
    data=query_embedding,
    anns_field="embedding",
    param=search_params,
    limit=5,
    expr='category == "tech" and publish_year >= 2024',
    output_fields=["title", "category", "publish_year"]
)

# 按条件删除
collection.delete('id in [1, 2, 3]')

# 删除集合
utility.drop_collection("documents")

# 断开连接
connections.disconnect("default")

from langchain_community.vectorstores import Milvus
from langchain_community.embeddings import HuggingFaceBgeEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import TextLoader
from langchain_openai import ChatOpenAI
from langchain.chains import create_retrieval_chain
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain_core.prompts import ChatPromptTemplate

# 1. 加载文档
loader = TextLoader("knowledge_base.txt", encoding="utf-8")
documents = loader.load()

# 2. 文档分块
splitter = RecursiveCharacterTextSplitter(
    chunk_size=500,
    chunk_overlap=50,
    separators=["\n\n", "\n", "。", "！", "？", ".", " "]
)
chunks = splitter.split_documents(documents)
print(f"文档分为 {len(chunks)} 个块")

# 3. Embedding + 存入 Milvus
embeddings = HuggingFaceBgeEmbeddings(
    model_name="BAAI/bge-base-zh-v1.5"
)

vectorstore = Milvus.from_documents(
    documents=chunks,
    embedding=embeddings,
    connection_args={"host": "localhost", "port": "19530"},
    collection_name="rag_knowledge",
    drop_old=True
)

# 4. 构建 RAG 链（LCEL，替代已废弃的 RetrievalQA）
llm = ChatOpenAI(model="gpt-4o", temperature=0)
retriever = vectorstore.as_retriever(search_kwargs={"k": 3})

system_prompt = """基于以下上下文回答问题。如果上下文中没有相关信息，请说"我不知道"。

{context}"""

prompt = ChatPromptTemplate.from_messages([
    ("system", system_prompt),
    ("human", "{input}")
])

question_answer_chain = create_stuff_documents_chain(llm, prompt)
rag_chain = create_retrieval_chain(retriever, question_answer_chain)

# 5. 查询
response = rag_chain.invoke({"input": "向量数据库的核心优势是什么？"})
print(response["answer"])
for doc in response["context"]:
    print(f"  来源: {doc.metadata.get('source', 'N/A')}")

# 场景1：高召回率优先（推荐系统）
index_params_high_recall = {
    "index_type": "HNSW",
    "metric_type": "COSINE",
    "params": {"M": 64, "efConstruction": 512}
}
search_params_high_recall = {"ef": 256}

# 场景2：低延迟优先（实时搜索）
index_params_low_latency = {
    "index_type": "IVF_FLAT",
    "metric_type": "L2",
    "params": {"nlist": 2048}
}
search_params_low_latency = {"nprobe": 16}

# 场景3：大规模低成本（十亿级）
index_params_large_scale = {
    "index_type": "IVF_PQ",
    "metric_type": "L2",
    "params": {"nlist": 4096, "m": 16, "nbits": 8}
}
search_params_large_scale = {"nprobe": 64}

from pymilvus import AnnSearchRequest, RRFRanker

# 向量搜索请求
vector_req = AnnSearchRequest(
    data=query_embedding,
    anns_field="dense_embedding",
    param={"metric_type": "COSINE", "params": {"ef": 128}},
    limit=20
)

# 稀疏向量搜索（BM25 / SPLADE）
sparse_req = AnnSearchRequest(
    data=sparse_query_embedding,
    anns_field="sparse_embedding",
    param={"metric_type": "IP"},
    limit=20
)

# RRF 融合排序
results = collection.hybrid_search(
    reqs=[vector_req, sparse_req],
    ranker=RRFRanker(k=60),
    limit=10,
    output_fields=["title", "content"]
)