AIAgent零基础入门: Day13 向量库Chroma、FAISS、Milvus 选哪个?

pip install langchain==1.4.3 langchain-openai==1.6.6 langchain-text-splitters==1.1.2 langchain-chroma==1.1.0 langchain-milvus==0.4.0 milvus-lite langchain-community==0.4.2 python-dotenv

OPENAI_API_KEY=sk-xxxxOPENAI_BASE_URL=https://api.openai.com/v1OPENAI_MODEL_NAME=gpt-4o-miniEMBEDDING_MODEL=text-embedding-3-small

# uv add langchain==1.4.3 langchain-openai==1.6.6 langchain-text-splitters==1.1.2 langchain-chroma==1.1.0 langchain-milvus==0.4.0 milvus-lite langchain-community==0.4.2 python-dotenv# ---------------- 0. 依赖 ----------------import os, re, time, warnings, zlibfrom dotenv import load_dotenvfrom langchain_core.documents import Documentfrom langchain_core.embeddings import Embeddingsfrom langchain_openai import OpenAIEmbeddingsfrom langchain_text_splitters import RecursiveCharacterTextSplitterload_dotenv()                                                  # 读 .env：填了 Key 就自动切真模型# ---------------- 1. 语料切分：沿用 Day12 的报销制度，多给每块挂一个 category ----------------DOC = """【员工报销制度 v3.2】1. 差旅住宿标准：一线城市每晚不超过 600 元，二线城市 450 元，其他城市 350 元。2. 交通标准：高铁二等座，飞机经济舱；连续乘车超过 6 小时可升级高铁一等座。3. 报销时效：费用发生后 30 天内提交，跨年度费用必须在次年 1 月 10 日前提交。4. 审批流程：单笔 1000 元以下由直属主管审批；1000 至 5000 元由部门负责人审批；   超过 5000 元需要财务总监审批，且必须附上级书面说明。5. 发票要求：必须提供增值税发票原件，电子发票需打印并附支付凭证。6. 违规处理：虚报费用一经查实，全额追回并取消当年评优资格。"""CATEGORIES = ["差旅", "交通", "时效", "审批", "发票", "违规"]def build_chunks() -> list[Document]:"""切分时给每块挂上 category 元数据——能不能按元数据过滤，是库和索引的分水岭"""    splitter = RecursiveCharacterTextSplitter(        chunk_size=60, chunk_overlap=15,                       # 中文必须自定义分隔符        separators=["\n\n", "\n", "。", "；", "，", " "],    )    chunks = splitter.split_documents([Document(page_content=DOC)])for c in chunks:        nums = re.findall(r"(\d+)\.\s", c.page_content)        # 按「1. 2. 3.」认领类目if nums and int(nums[0]) <= len(CATEGORIES):            c.metadata["category"] = CATEGORIES[int(nums[0]) - 1]return [c for c in chunks if c.metadata.get("category")]class LexicalEmbeddings(Embeddings):"""没 Key 时的本地替身：字符 bigram + 哈希桶，256 维，输出 L2 归一化向量"""    DIM = 256def _vec(self, t: str) -> list[float]:        v = [0.0] * self.DIM        s = re.sub(r"\s+", "", t.lower())for i in range(len(s) - 1):            v[zlib.crc32(s[i:i + 2].encode("utf-8")) % self.DIM] += 1.0        norm = sum(x * x for x in v) ** 0.5or1.0# 归一化：余弦与 L2 排序一致return [x / norm for x in v]def embed_documents(self, texts: list[str]) -> list[list[float]]:return [self._vec(t) for t in texts]def embed_query(self, text: str) -> list[float]:return self._vec(text)def build_embeddings() -> Embeddings:if os.getenv("OPENAI_API_KEY"):return OpenAIEmbeddings(model=os.getenv("EMBEDDING_MODEL", "text-embedding-3-small"),                                base_url=os.getenv("OPENAI_BASE_URL") orNone)return LexicalEmbeddings()# ---------------- 2. 四套存储，一套接口 ----------------from langchain.agents import create_agentfrom langchain.tools import toolfrom langchain_chroma import Chromafrom langchain_core.language_models.fake_chat_models import FakeMessagesListChatModelfrom langchain_core.messages import AIMessagefrom langchain_core.vectorstores import InMemoryVectorStorefrom langgraph.checkpoint.memory import InMemorySaverclass Store:"""统一门面：上层（Agent/Retriever）只见 search()，换库只改下面的 build()"""def __init__(self, name, store, note, can_filter):        self.name, self.store, self.note, self.can_filter = name, store, note, can_filterdef search(self, q: str, k: int = 3, where: dict | None = None):"""where 统一写 dict；各库方言不同——Chroma 收 dict，Milvus 收布尔表达式串"""        kw = {}if where and self.can_filter:if self.name == "Milvus":                kw["expr"] = " and ".join(f"{key} == {val!r}"for key, val in where.items())else:                kw["filter"] = wherereturn self.store.similarity_search(q, k=k, **kw)def build(backend: str, chunks: list[Document], emb: Embeddings) -> Store:"""后端工厂：新增一种向量库 = 加一个 if 分支，其它代码一动不动"""if backend == "memory":        s = InMemoryVectorStore(emb)        s.add_documents(chunks)                                # 进程内字典，进程结束即消失return Store("InMemory", s, "无落盘，重启即丢", False)if backend == "chroma":# 嵌入式库：persist_directory 一给就自动落盘；换成 host/port 就是服务模式        s = Chroma.from_documents(chunks, emb, collection_name="policy",                                  persist_directory="./chroma_policy")return Store("Chroma", s, "自动落盘 ./chroma_policy", True)if backend == "faiss":        warnings.filterwarnings("ignore", category=DeprecationWarning)   # 官方集成包正在迁移from langchain_community.vectorstores import FAISS        s = FAISS.from_documents(chunks, emb)        s.save_local("./faiss_policy")                          # 只有 index.faiss + index.pkl        s = FAISS.load_local("./faiss_policy", emb, allow_dangerous_deserialization=True)return Store("FAISS", s, "纯索引，只有 .faiss/.pkl", False)if backend == "milvus":from langchain_milvus import Milvus# Milvus Lite：本地单文件版，uri 换成 http://host:19530 就直连分布式集群        s = Milvus.from_documents(chunks, emb, collection_name="policy",                                  connection_args={"uri": "./milvus_policy.db"}, drop_old=True)return Store("Milvus", s, "Lite 本地库 / 改 uri 上集群", True)    raise ValueError(f"未知后端：{backend}")# ---------------- 3. 离线模型替身（无 Key 时按脚本回放） ----------------class OfflineChat(FakeMessagesListChatModel):def bind_tools(self, tools, **kwargs):return self                                            # 补上基类缺失的 bind_toolsdef build_model():if os.getenv("OPENAI_API_KEY"):from langchain_openai import ChatOpenAIreturn ChatOpenAI(model=os.getenv("OPENAI_MODEL_NAME", "gpt-4o-mini"),                          base_url=os.getenv("OPENAI_BASE_URL") orNone, temperature=0)return OfflineChat(responses=[        AIMessage(content="", tool_calls=[{"name": "search_policy","args": {"query": "虚报费用 追回 评优 资格"}, "id": "c1"}]),        AIMessage(content="根据第 6 条：虚报费用一经查实，全额追回并取消当年评优资格。"),        AIMessage(content="", tool_calls=[{"name": "search_policy","args": {"query": "电子发票 打印 支付凭证"}, "id": "c2"}]),        AIMessage(content="根据第 5 条：必须提供增值税发票原件，电子发票需打印并附支付凭证。"),    ])def build_agent(store: Store):@tooldef search_policy(query: str) -> str:"""在员工报销制度中检索相关条款，返回最相关的原文片段。"""return"\n\n".join(f"[{i}] {d.page_content.strip()}"for i, d in enumerate(store.search(query), 1))return create_agent(model=build_model(), tools=[search_policy],                        system_prompt="你是报销制度助手，先检索再作答，无依据就说「资料未提及」。",                        checkpointer=InMemorySaver())CASES = [("住一晚最多能报多少？", "差旅"), ("报销时效是多久？", "时效"),         ("电子发票要怎么处理？", "发票"), ("超过五千谁审批？", "审批"),         ("虚报费用会怎样？", "违规"), ("高铁一等座什么条件？", "交通")]def main() -> None:    chunks = build_chunks()    emb = build_embeddings()    reused = os.path.exists("./chroma_policy")                 # 第二次运行：库已经在盘上了print(f"[1] 语料：{len(chunks)} 条 chunk，向量维度 {len(emb.embed_query('测试'))}"f" | 磁盘上{'已有' if reused else '还没有'}落盘数据")    stores = {}for b in ["memory", "chroma", "faiss", "milvus"]:try:            stores[b] = build(b, chunks, emb)except Exception as exc:                                # 缺依赖就跳过，不打断整体演示print(f"[2] 建库 {b:<12} 跳过：{type(exc).__name__}（pip 装对应包后可复现）")            continue        tag = "复用盘上数据"if reused and b != "memory"else"重新灌库"print(f"[2] 建库 {stores[b].name:<12} {tag} | {stores[b].note}")print(f"\n[3] 同一个问题在四个库上的 Top1（{len(CASES)} 条）")    hit = {b: 0for b in stores}for q, want in CASES:        row = []for b, s in stores.items():            ok = want in s.search(q, 1)[0].page_content            hit[b] += ok            row.append(f"{s.name}:{'√' if ok else '×'}")print(f"    {q:<12}" + "  ".join(row))print("    → " + "  ".join(f"{stores[b].name} {v}/{len(CASES)}"for b, v in hit.items()))print("\n[4] 延迟对比（每库 20 次查询；6 条数据，绝对数字无意义，看量级）")for b, s in stores.items():        t0 = time.perf_counter()for _ in range(20):            s.search(CASES[0][0], 3)print(f"    {s.name:<12} 平均 {(time.perf_counter() - t0) / 20 * 1000:.2f} ms")print("\n[5] 元数据过滤：一句「报销标准」会命中哪些类目？（各取 Top2）")    q = "报销标准是什么"for b, s in stores.items():        where = {"category": "差旅"} if s.can_filter elseNone        cats = [d.metadata.get("category", "-") for d in s.search(q, 2, where)]print(f"    {s.name:<12} {'按[差旅]过滤' if where else '不支持过滤':<8} 命中类目 {cats}")print("\n[6] 生产形态怎么切：Chroma 服务模式 / Milvus 集群，只改连接参数")print("    Chroma(collection_name='policy', host='localhost', port=8000, embedding_function=emb)")print("    Milvus(embedding_function=emb, connection_args={'uri': 'http://localhost:19530'})")    agent = build_agent(stores["chroma"])    cfg = {"configurable": {"thread_id": "day13-demo"}}print()for question in ["虚报费用被查出来会怎么样？", "那电子发票呢？"]:print(f"[问] {question}")        res = agent.invoke({"messages": [{"role": "user", "content": question}]}, cfg)print(f"[答] {res['messages'][-1].content}\n")if __name__ == "__main__":    main()

[1] 语料：6 条 chunk，向量维度 256 | 磁盘上还没有落盘数据[2] 建库 InMemory     重新灌库 | 无落盘，重启即丢[2] 建库 Chroma       重新灌库 | 自动落盘 ./chroma_policy[2] 建库 FAISS        重新灌库 | 纯索引，只有 .faiss/.pkl[2] 建库 Milvus       重新灌库 | Lite 本地库 / 改 uri 上集群[3] 同一个问题在四个库上的 Top1（6 条）    住一晚最多能报多少？  InMemory:×  Chroma:×  FAISS:×  Milvus:×    报销时效是多久？    InMemory:√  Chroma:√  FAISS:√  Milvus:√    电子发票要怎么处理？  InMemory:√  Chroma:√  FAISS:√  Milvus:√    超过五千谁审批？    InMemory:×  Chroma:×  FAISS:×  Milvus:×    虚报费用会怎样？    InMemory:√  Chroma:√  FAISS:√  Milvus:√    高铁一等座什么条件？  InMemory:√  Chroma:√  FAISS:√  Milvus:√    → InMemory 4/6  Chroma 4/6  FAISS 4/6  Milvus 4/6[4] 延迟对比（每库 20 次查询；6 条数据，绝对数字无意义，看量级）    InMemory     平均 0.18 ms    Chroma       平均 0.89 ms    FAISS        平均 0.06 ms    Milvus       平均 1.32 ms[5] 元数据过滤：一句「报销标准」会命中哪些类目？（各取 Top2）    InMemory     不支持过滤    命中类目 ['交通', '审批']    Chroma       按[差旅]过滤  命中类目 ['差旅']    FAISS        不支持过滤    命中类目 ['交通', '审批']    Milvus       按[差旅]过滤  命中类目 ['差旅'][6] 生产形态怎么切：Chroma 服务模式 / Milvus 集群，只改连接参数    Chroma(collection_name='policy', host='localhost', port=8000, embedding_function=emb)    Milvus(embedding_function=emb, connection_args={'uri': 'http://localhost:19530'})[问] 虚报费用被查出来会怎么样？[答] 根据第 6 条：虚报费用一经查实，全额追回并取消当年评优资格。[问] 那电子发票呢？[答] 根据第 5 条：必须提供增值税发票原件，电子发票需打印并附支付凭证。

[1] 语料：6 条 chunk，向量维度 256 | 磁盘上已有落盘数据[2] 建库 Chroma       复用盘上数据 | 自动落盘 ./chroma_policy[2] 建库 FAISS        复用盘上数据 | 纯索引，只有 .faiss/.pkl[2] 建库 Milvus       复用盘上数据 | Lite 本地库 / 改 uri 上集群