Milvus 分布式向量数据库

# Milvus Lite 无需额外安装服务器，随 PyMilvus 一起使用
from pymilvus import MilvusClient

# 使用本地文件作为存储
client = MilvusClient("./milvus_demo.db")

from pymilvus import connections

# 连接到 Milvus 服务器
connections.connect(
    alias="default",
    host="localhost",
    port="19530"
)
print("Milvus 连接成功！")

from pymilvus import connections

connections.connect(
    alias="default",   # 连接别名，后续操作可引用
    host="localhost",  # Milvus 服务器地址
    port="19530",      # gRPC 端口号，默认19530
    user="",           # 用户名（如启用了认证）
    password="",       # 密码
    db_name="default", # 数据库名称
    token="",          # 认证令牌（替代user/password）
    server_name="",    # 服务器名称（TLS）
    server_pem_path="",# TLS 证书路径
)

connections.disconnect(alias="default")

conns = connections.list_connections()
# 返回格式: [("alias", <address>), ...]

from pymilvus import CollectionSchema, FieldSchema, DataType

# 主键字段
pk_field = FieldSchema(
    name="id",                # 字段名
    dtype=DataType.INT64,     # 数据类型
    is_primary=True,          # 是否为主键
    auto_id=False,            # 是否自动生成ID
    description="主键ID"
)

# 向量字段
vector_field = FieldSchema(
    name="embedding",
    dtype=DataType.FLOAT_VECTOR,  # 浮点向量类型
    dim=768,                      # 向量维度
    description="文本embedding向量"
)

# 标量字段
text_field = FieldSchema(
    name="text",
    dtype=DataType.VARCHAR,   # 变长字符串
    max_length=65535,         # 最大长度
    description="原始文本"
)

source_field = FieldSchema(
    name="source",
    dtype=DataType.VARCHAR,
    max_length=256,
    description="数据来源"
)

# 构建Schema
schema = CollectionSchema(
    fields=[pk_field, vector_field, text_field, source_field],
    description="RAG知识库文档集合",
    enable_dynamic_field=True  # 允许动态字段
)

from pymilvus import Collection

# 方式一：使用 Schema 创建
collection = Collection(
    name="rag_documents",  # 集合名称
    schema=schema,         # Schema对象
    using="default",       # 使用的连接别名
    shards_num=2,          # 分片数量（影响写入并行度）
)

# 方式二：使用 MilvusClient 快速创建（简化API）
from pymilvus import MilvusClient

client = MilvusClient(uri="http://localhost:19530")
client.create_collection(
    collection_name="rag_documents",
    dimension=768,         # 向量维度（自动生成schema）
    metric_type="COSINE",  # 距离度量类型
)

# 获取已存在的 Collection
collection = Collection("rag_documents")

# 获取集合信息
info = collection.describe()
# 返回: {'name': 'rag_documents', 'description': '...', 'fields': [...], ...}

# 获取行数
count = collection.num_entities

# 删除 Collection
collection.drop()

# 列出所有 Collection
from pymilvus import utility
collections = utility.list_collections()

# 为向量字段创建索引
index_params = {
    "index_type": "IVF_FLAT",    # 索引类型
    "metric_type": "COSINE",     # 距离度量: L2, IP, COSINE
    "params": {                   # 索引参数
        "nlist": 128,             # 聚类中心数量
    }
}

collection.create_index(
    field_name="embedding",   # 向量字段名
    index_params=index_params,
    index_name="embedding_idx" # 索引名称
)

index_params = {
    "index_type": "IVF_FLAT",
    "metric_type": "L2",
    "params": {"nlist": 128}  # 聚类中心数量，通常设为 4*sqrt(N)
}

index_params = {
    "index_type": "IVF_SQ8",
    "metric_type": "L2",
    "params": {"nlist": 128}
}

index_params = {
    "index_type": "IVF_PQ8",
    "metric_type": "L2",
    "params": {
        "nlist": 128,    # 聚类数
        "m": 8,          # 乘积量化的子空间数（m必须能整除dim）
        "nbits": 8       # 每个子空间的量化位数
    }
}

index_params = {
    "index_type": "HNSW",
    "metric_type": "COSINE",
    "params": {
        "M": 16,           # 每个节点的最大连接数
        "efConstruction": 256  # 构建时的搜索宽度
    }
}

index_params = {
    "index_type": "ANNOY",
    "metric_type": "L2",
    "params": {"n_trees": 8}  # 树的数量
}

index_params = {
    "index_type": "SCANN",
    "metric_type": "L2",
    "params": {
        "nlist": 128,
        "with_raw_data": True   # 是否保存原始数据
    }
}

# 构建索引（显式触发）
collection.create_index(field_name="embedding", index_params=index_params)

# 查看索引信息
indexes = collection.indexes
for idx in indexes:
    print(f"索引名: {idx.index_name}")
    print(f"字段: {idx.field_name}")
    print(f"参数: {idx.params}")

# 删除索引
collection.drop_index(index_name="embedding_idx")

import numpy as np

# 准备数据
data = [
    {
        "id": i,
        "embedding": np.random.rand(768).tolist(),
        "text": f"这是第{i}条文档内容",
        "source": f"doc_{i % 10}"
    }
    for i in range(1000)
]

# 批量插入
result = collection.insert(data=data)

# 返回结果
print(result.insert_count)  # 插入的行数
print(result.primary_keys)  # 主键列表

# 也可以分字段插入
result = collection.insert(
    data=[
        [1, 2, 3],                              # id 列表
        [[0.1, 0.2, ...], [0.3, 0.4, ...], ...], # embedding 列表
        ["文本1", "文本2", "文本3"],               # text 列表
    ]
)

# 按主键删除
result = collection.delete(expr="id in [1, 2, 3]")

# 按条件删除
result = collection.delete(expr='source == "doc_0"')

# 删除结果
print(result.delete_count)

# 如果主键存在则更新，不存在则插入
data = [
    {
        "id": 1,
        "embedding": [0.1, 0.2, ...],
        "text": "更新后的文本",
        "source": "updated_doc"
    }
]

result = collection.upsert(data=data)
print(result.insert_count)  # 新插入的数量
print(result.upsert_count)  # 更新的数量

# 按条件查询标量数据
results = collection.query(
    expr='source == "doc_0" and id > 50',  # 过滤表达式
    output_fields=["id", "text", "source"],  # 返回字段
    limit=10,                                 # 返回数量限制
)

# 返回格式: [{"id": 51, "text": "...", "source": "doc_0"}, ...]

# 基本向量搜索
search_params = {
    "metric_type": "COSINE",  # 距离度量
    "params": {
        "nprobe": 16,         # IVF系列：搜索的聚类中心数
        # HNSW 使用 "ef": 64
    }
}

results = collection.search(
    data=[[0.1, 0.2, ...]],       # 查询向量列表（支持批量）
    anns_field="embedding",        # 向量字段名
    param=search_params,           # 搜索参数
    limit=10,                      # 返回Top-K结果数
    expr='source == "doc_0"',      # 过滤表达式（可选）
    output_fields=["text", "source"],  # 返回的标量字段
)

# 解析搜索结果
for hits in results:
    for hit in hits:
        print(f"ID: {hit.id}, 距离: {hit.distance}, 文本: {hit.entity.get('text')}")

from pymilvus import AnnSearchRequest

# 定义第一个搜索请求（稠密向量）
dense_search_param = {
    "data": [[0.1, 0.2, ...]],
    "anns_field": "dense_embedding",
    "param": {"metric_type": "COSINE", "params": {"nprobe": 16}},
    "limit": 10,
}
dense_request = AnnSearchRequest(**dense_search_param)

# 定义第二个搜索请求（稀疏向量/BM25）
sparse_search_param = {
    "data": [{"token_1": 0.5, "token_2": 0.3}],  # 稀疏向量
    "anns_field": "sparse_embedding",
    "param": {"metric_type": "IP"},
    "limit": 10,
}
sparse_request = AnnSearchRequest(**sparse_search_param)

# 执行混合搜索（加权融合）
from pymilvus import WeightedRanker

results = collection.hybrid_search(
    reqs=[dense_request, sparse_request],
    ranker=WeightedRanker(0.7, 0.3),  # 稠密向量权重0.7，稀疏向量权重0.3
    limit=10,
    output_fields=["text", "source"],
)

# 解析结果
for hits in results:
    for hit in hits:
        print(f"ID: {hit.id}, 分数: {hit.distance}")

from pymilvus import RRFRanker

# 使用 RRF 融合
results = collection.hybrid_search(
    reqs=[dense_request, sparse_request],
    ranker=RRFRanker(k=60),
    limit=10,
)

# 创建分区
collection.create_partition(
    name="partition_2024",  # 分区名称
    description="2024年的数据"
)

# 列出所有分区
partitions = collection.partitions
for p in partitions:
    print(f"分区名: {p.name}, 行数: {p.num_entities}")

# 插入数据到指定分区
data = [{"id": 1, "embedding": [...], "text": "..."}]
collection.insert(data=data, partition_name="partition_2024")

# 在指定分区中搜索
results = collection.search(
    data=[[0.1, 0.2, ...]],
    anns_field="embedding",
    param=search_params,
    limit=10,
    partition_names=["partition_2024"],  # 仅在指定分区搜索
)

# 删除分区
collection.drop_partition(name="partition_2024")

# 加载 Collection 到内存（搜索前必须执行）
collection.load(
    replica_number=1,  # 副本数，增加副本可提高QPS
    timeout=None,      # 超时时间
)

# 查看加载状态
load_state = utility.load_state("rag_documents")
# 可能的值: Loaded, Loading, NotLoad

# 释放 Collection 的内存（不需要搜索时释放以节省资源）
collection.release()

# 刷新数据到持久化存储（确保插入的数据可被搜索）
collection.flush()
# 注意：flush 是一个耗时操作，Milvus 也会自动flush

from pymilvus import connections, Collection, CollectionSchema, FieldSchema, DataType
from openai import OpenAI
import numpy as np

# ========== 初始化 ==========
connections.connect(host="localhost", port="19530")
openai_client = OpenAI(api_key="your-api-key")

# ========== 定义 Collection ==========
fields = [
    FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
    FieldSchema(name="text", dtype=DataType.VARCHAR, max_length=65535),
    FieldSchema(name="source", dtype=DataType.VARCHAR, max_length=256),
    FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=1536),
]
schema = CollectionSchema(fields=fields, description="RAG知识库")
collection = Collection(name="rag_kb", schema=schema)

# 创建 HNSW 索引
index_params = {
    "index_type": "HNSW",
    "metric_type": "COSINE",
    "params": {"M": 16, "efConstruction": 256}
}
collection.create_index(field_name="embedding", index_params=index_params)
collection.load()

# ========== 文档入库 ==========
def embed_texts(texts):
    """调用 OpenAI API 生成文本向量"""
    response = openai_client.embeddings.create(
        model="text-embedding-3-small",
        input=texts
    )
    return [item.embedding for item in response.data]

def ingest_documents(documents):
    """将文档切片后入库"""
    # 文档切片（简单按字符数切片）
    chunks = []
    for doc in documents:
        text = doc["text"]
        chunk_size = 500
        for i in range(0, len(text), chunk_size):
            chunks.append({
                "text": text[i:i+chunk_size],
                "source": doc["source"]
            })

    # 批量生成向量
    texts = [c["text"] for c in chunks]
    embeddings = embed_texts(texts)

    # 插入 Milvus
    data = [
        {"text": c["text"], "source": c["source"], "embedding": e}
        for c, e in zip(chunks, embeddings)
    ]
    collection.insert(data=data)
    collection.flush()
    print(f"入库完成，共 {len(data)} 个切片")

# ========== 检索与生成 ==========
def rag_query(question, top_k=5):
    """RAG查询流程"""
    # 1. 问题向量化
    question_embedding = embed_texts([question])[0]

    # 2. 向量检索
    search_params = {"metric_type": "COSINE", "params": {"ef": 64}}
    results = collection.search(
        data=[question_embedding],
        anns_field="embedding",
        param=search_params,
        limit=top_k,
        output_fields=["text", "source"],
    )

    # 3. 组装上下文
    context = "\n".join([
        f"[来源: {hit.entity.get('source')}]\n{hit.entity.get('text')}"
        for hit in results[0]
    ])

    # 4. LLM 生成回答
    response = openai_client.chat.completions.create(
        model="gpt-4",
        messages=[
            {"role": "system", "content": "基于以下参考资料回答用户问题。如果参考资料中没有相关信息，请说明。"},
            {"role": "user", "content": f"参考资料:\n{context}\n\n问题: {question}"}
        ]
    )
    return response.choices[0].message.content

# 使用示例
answer = rag_query("什么是向量数据库？")
print(answer)

# 创建支持图文混合的 Collection
fields = [
    FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
    FieldSchema(name="content_type", dtype=DataType.VARCHAR, max_length=32),
    FieldSchema(name="description", dtype=DataType.VARCHAR, max_length=65535),
    FieldSchema(name="image_url", dtype=DataType.VARCHAR, max_length=1024),
    FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=512),
]
schema = CollectionSchema(fields=fields, description="多模态知识库")
collection = Collection(name="multimodal_kb", schema=schema)

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

# 搜索时可按内容类型过滤
results = collection.search(
    data=[query_embedding],
    anns_field="embedding",
    param={"metric_type": "L2", "params": {"nprobe": 16}},
    limit=10,
    expr='content_type == "image"',  # 仅搜索图片
    output_fields=["description", "image_url"],
)

from pymilvus import MilvusClient

client = MilvusClient(uri="http://localhost:19530")

# 快速创建 Collection
client.create_collection(
    collection_name="quick_start",
    dimension=768,
    metric_type="COSINE",
    auto_id=True,
)

# 插入数据
data = [
    {"id": i, "vector": np.random.rand(768).tolist(), "text": f"doc_{i}"}
    for i in range(1000)
]
client.insert(collection_name="quick_start", data=data)

# 搜索
results = client.search(
    collection_name="quick_start",
    data=[np.random.rand(768).tolist()],
    limit=5,
    output_fields=["text"],
    search_params={"metric_type": "COSINE", "params": {"nprobe": 10}},
)

# 查询
results = client.query(
    collection_name="quick_start",
    filter="id > 500",
    output_fields=["id", "text"],
    limit=10,
)

# PyMilvus 的两层 API 设计

# 底层 ORM API（细粒度控制）
from pymilvus import connections, Collection, FieldSchema, DataType
connections.connect(...)
collection = Collection(name="...", schema=schema)
collection.insert(data)
collection.search(...)

# 高层 MilvusClient API（简化操作）
from pymilvus import MilvusClient
client = MilvusClient(uri="...")
client.insert(collection_name="...", data=data)
client.search(collection_name="...", data=...)

# HNSW 参数推荐
hnsw_params = {
    "M": 16,              # 通用场景推荐16，高召回场景推荐32-64
    "efConstruction": 256  # 推荐至少200，越高构建质量越好
}
# 搜索时
search_params = {"ef": 64}  # ef >= limit，推荐 2-4 倍 limit

# IVF 参数推荐
ivf_params = {
    "nlist": 128  # 推荐 4*sqrt(N)，N为向量数
}
# 搜索时
search_params = {"nprobe": 16}  # nprobe/nlist 在 10%-20% 较合适

# 原因：数据还在 Growing Segment 中或未 flush
# 解决方案：
collection.insert(data)
collection.flush()  # 显式刷盘（生产环境建议让 Milvus 自动flush）
# 重新 load（如果之前已 load，需要等待或重新 load）

# IVF 系列：增大 nprobe
search_params = {"metric_type": "COSINE", "params": {"nprobe": 32}}  # 原来可能是 8 或 16

# HNSW：增大 ef
search_params = {"metric_type": "COSINE", "params": {"ef": 128}}  # 原来可能是 32 或 64

# 1. 使用压缩索引（IVF_SQ8 替代 IVF_FLAT，或使用 IVF_PQ8）
# 2. 释放不常用的 Collection
collection.release()
# 3. 减少副本数
collection.load(replica_number=1)

# 批量插入的最佳实践
batch_size = 5000
for i in range(0, len(all_data), batch_size):
    batch = all_data[i:i+batch_size]
    collection.insert(batch)
    print(f"已插入 {min(i+batch_size, len(all_data))}/{len(all_data)}")

# 最后一次 flush
collection.flush()

from langchain_community.vectorstores import Milvus
from langchain_openai import OpenAIEmbeddings

# 创建 Milvus 向量存储
vectorstore = Milvus.from_documents(
    documents=documents,       # LangChain Document 列表
    embedding=OpenAIEmbeddings(),
    collection_name="langchain_docs",
    connection_args={"host": "localhost", "port": "19530"},
)

# 相似度搜索
results = vectorstore.similarity_search(
    query="什么是向量数据库？",
    k=5,
)

# 作为检索器使用
retriever = vectorstore.as_retriever(search_kwargs={"k": 5})

from llama_index.vector_stores.milvus import MilvusVectorStore
from llama_index.core import StorageContext, VectorStoreIndex

# 创建 Milvus 向量存储
vector_store = MilvusVectorStore(
    uri="http://localhost:19530",
    collection_name="llamaindex_docs",
    dim=1536,
)

# 构建索引
storage_context = StorageContext.from_defaults(vector_store=vector_store)
index = VectorStoreIndex.from_documents(
    documents,
    storage_context=storage_context,
)

# 查询
query_engine = index.as_query_engine()
response = query_engine.query("什么是向量数据库？")
print(response)