1
# ===== 使用 Chroma（最轻量、最适合学习） =====
2
import chromadb
3
from chromadb.utils import embedding_functions
4

5
# ===== 第1步：创建客户端和集合 =====
6
# 持久化模式：数据存磁盘，下次重启还在
7
client = chromadb.PersistentClient(path="./chroma_data")
8

9
# 使用默认的all-MiniLM-L6-v2 Embedding模型（免费、轻量）
10
sentence_transformer_ef = embedding_functions.SentenceTransformerEmbeddingFunction(
11
    model_name="all-MiniLM-L6-v2"
12
)
13

14
# 创建或获取集合（类似数据库的表）
15
collection = client.get_or_create_collection(
16
    name="my_knowledge_base",
17
    embedding_function=sentence_transformer_ef,
18
    metadata={"description": "我的个人知识库"}
19
)
20
# 解释：collection ≈ 数据库中的一张表
21
# embedding_function决定了怎么把文本变成向量
22

23
# ===== 第2步：添加文档 =====
24
documents = [
25
    "Python是一种广泛使用的解释型高级编程语言",
26
    "FastAPI是一个现代、快速（高性能）的Web框架，用于构建API",
27
    "向量数据库专门存储和检索高维向量，常用于语义搜索",
28
    "RAG（检索增强生成）通过检索外部知识来增强LLM的回答",
29
    "LangGraph是一个用于构建有状态Agent的框架"
30
]
31
ids = [f"doc_{i}" for i in range(len(documents))]
32

33
collection.add(
34
    documents=documents,
35
    ids=ids,
36
    metadatas=[{"source": "知识库", "type": "技术"} for _ in documents]
37
)
38
# 解释：add时自动调用embedding_function将文档转为向量存储
39
# metadata可附带额外信息，用于过滤（如只搜某类文档）
40

41
# ===== 第3步：语义搜索 =====
42
query = "什么框架可以用来构建API？"
43
results = collection.query(
44
    query_texts=[query],   # 自动把查询文本转为向量
45
    n_results=2,           # 返回最相似的2条
46
    include=["documents", "distances", "metadatas"]
47
)
48

49
print("查询结果：")
50
for i, doc in enumerate(results["documents"][0]):
51
    distance = results["distances"][0][i]
52
    print(f"  [{distance:.4f}] {doc}")
53
# 解释：distance是余弦距离，值越小表示越相似
54
# 输出：[0.2345] FastAPI是一个现代、快速（高性能）的Web框架，用于构建API
55

56
# ===== 第4步：带元数据过滤的搜索 =====
57
results = collection.query(
58
    query_texts=["编程语言"],
59
    n_results=5,
60
    where={"type": "技术"},          # 只搜type="技术"的文档
61
    where_document={"$contains": "Python"}  # 只搜包含"Python"的文档
62
)
63
# 解释：where是针对metadata的过滤，where_document是对文档内容的过滤
64

65
# ===== 第5步：更新和删除 =====
66
collection.update(ids=["doc_0"], documents=["Python是最流行的AI开发语言"])
67
collection.delete(ids=["doc_3"])