04 - Chroma 与 FAISS 向量数据库实战

学习目标

学完本节你将能够:

  1. 理解向量数据库的核心概念和工作原理
  2. 掌握 Chroma 和 FAISS 的使用方法和差异
  3. 实现元数据过滤(metadata filtering)
  4. 为设备维修系统选择合适的向量存储方案

一、为什么需要向量数据库?

在 RAG 系统中,知识库文档被向量化后需要一个高效的存储和检索方案:

传统数据库                     向量数据库
┌──────────┐                ┌──────────────────┐
│ 精确匹配  │                │ 语义相似度搜索     │
│ WHERE x=5 │                │ 找最像的 K 个向量   │
│ B-Tree索引 │                │ ANN 近似最近邻索引  │
└──────────┘                └──────────────────┘

二、主流向量数据库对比

特性 Chroma FAISS Milvus Qdrant
类型 嵌入式 DB 索引库 分布式 DB 分布式 DB
安装难度 ⭐ 简单 ⭐ 简单 ⭐⭐⭐ 复杂 ⭐⭐ 中等
元数据过滤 ✅ 内置 ❌ 需自行实现 ✅ 强大 ✅ 强大
持久化 ✅ 自动 需手动 ✅ 自动 ✅ 自动
Python API 优秀 优秀 良好 优秀
适合规模 <100万 <1000万 不限 不限
推荐场景 原型/中小项目 高性能搜索 生产环境 生产环境

对于设备维修系统的选择建议

  • 原型开发 / 中小规模:Chroma(简单、自带元数据过滤)
  • 大规模 / 高性能需求:FAISS(速度最快)+ 自建元数据过滤
  • 生产环境 / 多用户:Milvus 或 Qdrant

三、Chroma 快速上手

Chroma 是最易用的向量数据库,几行代码即可创建知识库:

import chromadb

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

# 创建集合(类似数据库中的表)
collection = client.create_collection(
    name="maintenance_kb",
    metadata={"description": "设备维修知识库"}
)

# 添加文档
collection.add(
    documents=["空压机排气温度过高,需要检查冷却器"],
    metadatas=[{"device_type": "空压机", "fault_type": "温度异常"}],
    ids=["doc_001"]
)

# 查询
results = collection.query(
    query_texts=["压缩机温度太高"],
    n_results=3
)

四、FAISS 快速上手

FAISS 是 Facebook 开源的向量搜索库,速度极快:

import faiss
import numpy as np

# 创建索引
dimension = 384
index = faiss.IndexFlatL2(dimension)  # 精确 L2 距离

# 添加向量
vectors = np.random.randn(100, dimension).astype('float32')
index.add(vectors)

# 搜索
query = np.random.randn(1, dimension).astype('float32')
distances, indices = index.search(query, k=5)

五、元数据过滤(Metadata Filtering)

元数据过滤是在向量搜索之前先用结构化条件缩小范围:

用户查询: "空压机温度高怎么办"
元数据过滤: device_type="空压机" AND priority="高"
                │
                ▼
        先过滤出空压机相关文档(假设 100 条→20 条)
                │
                ▼
        再在这 20 条中做向量相似度搜索
                │
                ▼
        返回最相关的 top-K 结果

好处:

  • 减少搜索空间,提高速度
  • 提高结果相关性(排除无关设备的文档)
  • 实现权限控制(不同角色看到不同文档)

六、设备维修系统架构

                    ┌─────────────────────┐
                    │   设备维修知识库       │
                    │  (维修手册/工单/日志)  │
                    └─────────┬───────────┘
                              │
                    ┌─────────▼───────────┐
                    │   文档解析 + 分块     │
                    └─────────┬───────────┘
                              │
                    ┌─────────▼───────────┐
                    │   Embedding 向量化   │
                    └─────────┬───────────┘
                              │
               ┌──────────────┴──────────────┐
               │                             │
    ┌──────────▼──────────┐    ┌─────────────▼────────┐
    │  Chroma (开发/原型)  │    │  FAISS (高性能搜索)   │
    │  + 元数据过滤        │    │  + 自定义过滤逻辑     │
    └──────────┬──────────┘    └─────────────┬────────┘
               │                             │
               └──────────────┬──────────────┘
                              │
                    ┌─────────▼───────────┐
                    │    RAG 检索 + 生成   │
                    └─────────────────────┘

七、动手实践

配套脚本:scripts/04_vector_db_demo.py

该脚本演示:

  1. 用 Chroma 创建维修知识库并添加文档
  2. 用 FAISS 创建相同数据的索引
  3. 元数据过滤查询
  4. Chroma vs FAISS 结果对比
  5. 搜索性能测试

运行方式

pip install chromadb faiss-cpu sentence-transformers numpy
python scripts/04_vector_db_demo.py

八、FAISS 索引类型选择

索引类型 说明 适用规模
IndexFlatL2 精确 L2 距离(暴力搜索) <10 万
IndexFlatIP 精确内积 <10 万
IndexIVFFlat 倒排文件索引(近似搜索) 10-100 万
IndexHNSW 分层导航小世界图 10-1000 万

对于设备维修系统(通常不超过 10 万条记录),IndexFlatL2 已足够。


九、Chroma 持久化存储

# 内存模式(开发用)
client = chromadb.Client()
# 持久化模式(生产用)
client = chromadb.PersistentClient(path="./chroma_db")
# 重启后数据自动恢复
collection = client.get_collection("maintenance_kb")

练习

  1. 在脚本中增加 5 条新的设备维修文档,观察搜索结果变化
  2. 实现一个组合查询:同时按设备类型和故障优先级过滤
  3. 思考:如果维修知识库有 100 万条记录,应该选择哪个向量数据库?

验证清单

  • 能说出 Chroma 和 FAISS 的核心区别
  • 理解元数据过滤的作用
  • 成功运行 04_vector_db_demo.py 并理解输出
  • 能为设备维修系统选择合适的向量存储方案

code

# -*- coding: utf-8 -*-
"""
04 - Chroma 与 FAISS 向量数据库实战
设备维修养护系统 - 第二阶段学习材料

pip install chromadb faiss-cpu sentence-transformers numpy
"""

import time
import numpy as np

# ============================================================
# 设备维修知识库样本数据
# ============================================================

MAINTENANCE_DOCS = [
    {
        "id": "doc_001",
        "text": "空压机排气温度过高:检查冷却器翅片是否积灰,用压缩空气清洁;检查温控阀是否正常工作;确认润滑油量充足。",
        "metadata": {"device_type": "空压机", "fault_type": "温度异常", "priority": "高", "date": "2024-03-15"}
    },
    {
        "id": "doc_002",
        "text": "空压机排气量不足:更换空气过滤器滤芯;检查皮带张紧度;检测管路是否有泄漏;检查进气阀工作状态。",
        "metadata": {"device_type": "空压机", "fault_type": "性能下降", "priority": "中", "date": "2024-03-10"}
    },
    {
        "id": "doc_003",
        "text": "空压机润滑油乳化:油中水分含量超标,需要更换润滑油并检查油气分离器。建议每2000小时更换一次润滑油。",
        "metadata": {"device_type": "空压机", "fault_type": "润滑问题", "priority": "中", "date": "2024-02-28"}
    },
    {
        "id": "doc_004",
        "text": "注塑机加热圈损坏:关闭电源等待冷却,拆卸旧加热圈,清洁安装面,安装新加热圈并检查绝缘电阻。一般8000-12000小时更换。",
        "metadata": {"device_type": "注塑机", "fault_type": "加热异常", "priority": "高", "date": "2024-03-18"}
    },
    {
        "id": "doc_005",
        "text": "注塑机合模力不足导致产品飞边:检查合模油缸密封件是否磨损,检查合模压力设定值,必要时更换全套密封件。",
        "metadata": {"device_type": "注塑机", "fault_type": "合模异常", "priority": "高", "date": "2024-03-20"}
    },
    {
        "id": "doc_006",
        "text": "注塑机螺杆磨损:塑化能力下降,产品出现未熔颗粒。需要拆卸螺杆检查磨损程度,磨损超过0.5mm应更换。",
        "metadata": {"device_type": "注塑机", "fault_type": "塑化异常", "priority": "中", "date": "2024-02-15"}
    },
    {
        "id": "doc_007",
        "text": "冷却塔风机皮带断裂:停机锁定电源,松开电机底座,更换新皮带,调整张紧度至按压下沉10-15mm,试运行30分钟。",
        "metadata": {"device_type": "冷却塔", "fault_type": "皮带故障", "priority": "高", "date": "2024-03-22"}
    },
    {
        "id": "doc_008",
        "text": "冷却塔出水温度偏高:检查填料层是否堵塞,清洗或更换填料;检查风机转速是否正常;检查进水量是否过大。",
        "metadata": {"device_type": "冷却塔", "fault_type": "温度异常", "priority": "中", "date": "2024-03-05"}
    },
    {
        "id": "doc_009",
        "text": "输送带跑偏:调整尾部滚筒张紧螺栓,跑偏向哪侧就紧哪侧;检查托辊是否转动灵活;检查输送带接头是否平整。",
        "metadata": {"device_type": "输送带", "fault_type": "跑偏", "priority": "中", "date": "2024-03-12"}
    },
    {
        "id": "doc_010",
        "text": "输送带打滑:张紧力不足或滚筒表面磨损。增加张紧力,清洁滚筒表面,必要时更换包胶滚筒。",
        "metadata": {"device_type": "输送带", "fault_type": "打滑", "priority": "低", "date": "2024-03-08"}
    },
    {
        "id": "doc_011",
        "text": "液压系统压力不足:检查油位是否充足;检查溢流阀设定压力;检查液压泵是否磨损;检查管路有无泄漏。",
        "metadata": {"device_type": "液压系统", "fault_type": "压力异常", "priority": "高", "date": "2024-03-25"}
    },
    {
        "id": "doc_012",
        "text": "液压缸漏油:密封件老化导致。需要更换活塞密封和杆密封,同时检查缸筒内壁是否有划痕。使用原厂密封件套装。",
        "metadata": {"device_type": "液压系统", "fault_type": "泄漏", "priority": "高", "date": "2024-03-28"}
    },
    {
        "id": "doc_013",
        "text": "液压油温度过高:检查冷却器是否工作正常,清洗冷却器;检查油液粘度是否合适;减少连续高负荷运行时间。",
        "metadata": {"device_type": "液压系统", "fault_type": "温度异常", "priority": "中", "date": "2024-02-20"}
    },
]


def load_embedding_model():
    """加载 Embedding 模型"""
    try:
        from sentence_transformers import SentenceTransformer
        print("📦 加载 Embedding 模型 (all-MiniLM-L6-v2)...")
        model = SentenceTransformer("all-MiniLM-L6-v2")
        return model
    except ImportError:
        print("⚠️  请安装: pip install sentence-transformers")
        return None


def get_embeddings(model, texts):
    """获取文本向量"""
    if model:
        return model.encode(texts)
    else:
        # 模拟向量
        np.random.seed(42)
        return np.random.randn(len(texts), 384).astype("float32")


# ============================================================
# 第一部分:Chroma 演示
# ============================================================

def demo_chroma(embeddings, docs):
    """Chroma 向量数据库演示"""
    import chromadb

    print("\n" + "=" * 70)
    print("  Chroma 向量数据库演示")
    print("=" * 70)

    # 创建客户端(内存模式)
    client = chromadb.Client()

    # 创建集合
    collection = client.create_collection(
        name="maintenance_kb",
        metadata={"description": "设备维修养护知识库"}
    )
    print(f"  ✅ 创建集合: {collection.name}")

    # 添加文档
    texts = [d["text"] for d in docs]
    metadatas = [d["metadata"] for d in docs]
    ids = [d["id"] for d in docs]

    collection.add(
        embeddings=embeddings.tolist(),
        documents=texts,
        metadatas=metadatas,
        ids=ids
    )
    print(f"  ✅ 已添加 {collection.count()} 条文档")

    # --- 基础查询 ---
    print("\n  🔍 查询 1: 「空压机温度过高怎么处理」(无过滤)")
    results = collection.query(
        query_embeddings=[embeddings[0].tolist()],  # 用第一条文档模拟查询
        n_results=3
    )
    for i, (doc, meta, dist) in enumerate(zip(
            results["documents"][0],
            results["metadatas"][0],
            results["distances"][0]), 1):
        print(f"    {i}. [{meta['device_type']}] {doc[:60]}... (距离: {dist:.4f})")

    # --- 元数据过滤查询 ---
    print("\n  🔍 查询 2: 只搜索「注塑机」相关文档")
    results = collection.query(
        query_embeddings=[embeddings[3].tolist()],
        n_results=5,
        where={"device_type": "注塑机"}
    )
    for i, (doc, meta, dist) in enumerate(zip(
            results["documents"][0],
            results["metadatas"][0],
            results["distances"][0]), 1):
        print(f"    {i}. [{meta['device_type']}|{meta['priority']}] {doc[:50]}...")

    # --- 多条件过滤 ---
    print("\n  🔍 查询 3: 搜索「高优先级」且「设备类型为液压系统」")
    results = collection.query(
        query_embeddings=[embeddings[10].tolist()],
        n_results=5,
        where={
            "$and": [
                {"device_type": {"$eq": "液压系统"}},
                {"priority": {"$eq": "高"}}
            ]
        }
    )
    for i, (doc, meta, dist) in enumerate(zip(
            results["documents"][0],
            results["metadatas"][0],
            results["distances"][0]), 1):
        print(f"    {i}. [{meta['device_type']}|{meta['fault_type']}|{meta['priority']}] {doc[:50]}...")

    return collection


# ============================================================
# 第二部分:FAISS 演示
# ============================================================

def demo_faiss(embeddings, docs):
    """FAISS 向量数据库演示"""
    import faiss

    print("\n" + "=" * 70)
    print("  FAISS 向量数据库演示")
    print("=" * 70)

    dim = embeddings.shape[1]

    # 创建索引(使用精确 L2 距离)
    index = faiss.IndexFlatL2(dim)
    vectors = embeddings.astype("float32")
    index.add(vectors)
    print(f"  ✅ 创建 FAISS 索引,维度: {dim},文档数: {index.ntotal}")

    # --- 基础查询 ---
    print("\n  🔍 查询 1: 「空压机温度过高怎么处理」")
    query = embeddings[0:1].astype("float32")
    distances, indices = index.search(query, k=3)

    for rank, (dist, idx) in enumerate(zip(distances[0], indices[0]), 1):
        doc = docs[idx]
        print(f"    {rank}. [{doc['metadata']['device_type']}] "
              f"{doc['text'][:60]}... (距离: {dist:.4f})")

    # --- 手动元数据过滤 ---
    print("\n  🔍 查询 2: 手动过滤「注塑机」后搜索")

    # 先过滤出注塑机的索引
    target_device = "注塑机"
    filtered_indices = [i for i, d in enumerate(docs)
                        if d["metadata"]["device_type"] == target_device]
    filtered_vectors = embeddings[filtered_indices].astype("float32")

    if len(filtered_vectors) > 0:
        sub_index = faiss.IndexFlatL2(dim)
        sub_index.add(filtered_vectors)

        query = embeddings[3:4].astype("float32")  # 注塑机加热圈
        k = min(3, sub_index.ntotal)
        distances, sub_indices = sub_index.search(query, k=k)

        for rank, (dist, si) in enumerate(zip(distances[0], sub_indices[0]), 1):
            original_idx = filtered_indices[si]
            doc = docs[original_idx]
            print(f"    {rank}. [{doc['metadata']['device_type']}] "
                  f"{doc['text'][:60]}... (距离: {dist:.4f})")

    return index


# ============================================================
# 第三部分:性能对比
# ============================================================

def benchmark(embeddings, docs):
    """Chroma vs FAISS 性能对比"""
    import chromadb
    import faiss

    print("\n" + "=" * 70)
    print("  搜索性能对比")
    print("=" * 70)

    n = len(docs)
    dim = embeddings.shape[1]
    vectors = embeddings.astype("float32")
    query = embeddings[0:1].astype("float32")

    # --- Chroma 性能 ---
    client = chromadb.Client()
    coll = client.create_collection("bench")
    coll.add(
        embeddings=vectors.tolist(),
        documents=[d["text"] for d in docs],
        ids=[d["id"] for d in docs]
    )

    iterations = 100
    start = time.time()
    for _ in range(iterations):
        coll.query(query_embeddings=query.tolist(), n_results=5)
    chroma_time = (time.time() - start) / iterations * 1000

    # --- FAISS 性能 ---
    idx = faiss.IndexFlatL2(dim)
    idx.add(vectors)

    start = time.time()
    for _ in range(iterations):
        idx.search(query, k=5)
    faiss_time = (time.time() - start) / iterations * 1000

    print(f"\n  文档数量: {n} 条")
    print(f"  向量维度: {dim}")
    print(f"  查询次数: {iterations}")
    print(f"\n  {'方案':<15} {'平均耗时':<15} {'相对速度'}")
    print(f"  {'-'*45}")
    print(f"  {'Chroma':<15} {chroma_time:<15.2f}ms  1.0x")
    speedup = chroma_time / faiss_time if faiss_time > 0 else 1
    print(f"  {'FAISS':<15} {faiss_time:<15.2f}ms  {speedup:.1f}x")

    print(f"\n  💡 FAISS 在纯向量搜索上通常更快")
    print(f"  💡 Chroma 内置元数据过滤,综合使用更方便")


# ============================================================
# 主流程
# ============================================================

def main():
    print("=" * 70)
    print("  Chroma 与 FAISS 向量数据库实战")
    print("  设备维修养护系统 - 第二阶段学习材料")
    print("=" * 70)

    # 加载模型
    model = load_embedding_model()

    # 获取所有文档的向量
    texts = [d["text"] for d in MAINTENANCE_DOCS]
    embeddings = get_embeddings(model, texts)
    print(f"📊 已生成 {len(embeddings)} 个文档向量,维度: {embeddings.shape[1]}")

    # Chroma 演示
    try:
        demo_chroma(embeddings, MAINTENANCE_DOCS)
    except ImportError:
        print("\n⚠️  Chroma 未安装,请运行: pip install chromadb")

    # FAISS 演示
    try:
        demo_faiss(embeddings, MAINTENANCE_DOCS)
    except ImportError:
        print("\n⚠️  FAISS 未安装,请运行: pip install faiss-cpu")

    # 性能对比
    try:
        benchmark(embeddings, MAINTENANCE_DOCS)
    except ImportError:
        print("\n⚠️  需要同时安装 chromadb 和 faiss-cpu 才能运行性能对比")

    # --- 选择建议 ---
    print("\n" + "=" * 70)
    print("  向量数据库选择建议")
    print("=" * 70)
    print("""
  📌 设备维修系统推荐方案:

  🔹 开发/原型阶段 → Chroma
     优势:API 简单、自带元数据过滤、自动持久化
     适合:验证 RAG 方案可行性

  🔹 生产环境(<100万文档) → Chroma + 持久化
     优势:够用且维护成本低
     部署:chromadb.PersistentClient(path="./db")

  🔹 生产环境(>100万文档) → FAISS 或 Milvus
     优势:高性能、支持大规模数据
     部署:FAISS 作为索引库 + 自建元数据层
""")
    print("=" * 70)
    print("  演示完成!")
    print("=" * 70)


if __name__ == "__main__":
    main()

Logo

汇聚全球AI编程工具,助力开发者即刻编程。

更多推荐