用 ChromaDB 做向量检索与 RAG：从嵌入到实战查询

cosine_similarity(A, B) = (A · B) / (||A|| × ||B||)

import chromadb

client = chromadb.PersistentClient(path="./chroma_db")

# 默认 cosine，也可以选 l2（欧氏距离）或 ip（内积）
collection = client.create_collection(
    name="tech_docs",
    metadata={"hnsw:space": "cosine"}  # 可选 "l2" 或 "ip"
)

# pip install chromadb sentence-transformers

import chromadb

# 使用内存客户端，跑完即消失；换成 PersistentClient 可持久化
client = chromadb.Client()

collection = client.get_or_create_collection(
    name="python_tips",
    metadata={"hnsw:space": "cosine"}
)

# 添加文档——ChromaDB 会自动调用默认嵌入模型生成向量
docs = [
    "Python 的 list comprehension 比手动 for 循环更快更简洁",
    "用 functools.lru_cache 可以给纯函数加缓存，避免重复计算",
    "asyncio.gather 能并发运行多个协程，适合 IO 密集场景",
    "dataclass 替代手写 __init__，减少样板代码",
    "pathlib.Path 比 os.path 更面向对象，路径操作更直观",
]

ids = [f"tip_{i}" for i in range(len(docs))]

collection.add(
    documents=docs,
    ids=ids
)

# 查询：找和"缓存性能优化"最相关的两条文档
results = collection.query(
    query_texts=["如何优化函数的缓存和性能"],
    n_results=2
)

print("匹配文档：", results["documents"][0])
print("距离值：", results["distances"][0])
print("ID：", results["ids"][0])

# 先添加带元数据的文档
collection.add(
    documents=["FastAPI 的依赖注入系统非常灵活"],
    ids=["tip_5"],
    metadatas=[{"category": "web", "level": "advanced"}]
)

# 查询时同时过滤元数据
results = collection.query(
    query_texts=["Web 框架的高级用法"],
    n_results=3,
    where={"level": "advanced"}  # 只返回高级内容
)

# pip install openai chromadb

import chromadb
from openai import OpenAI

# 初始化
openai_client = OpenAI()  # 确保已设置 OPENAI_API_KEY 环境变量
chroma_client = chromadb.Client()
collection = chroma_client.get_or_create_collection("rag_demo")

# 知识库内容
knowledge = [
    "ChromaDB 默认使用 all-MiniLM-L6-v2 嵌入模型，输出 384 维向量",
    "ChromaDB 的 query 方法支持 where 和 where_document 两种过滤",
    "余弦距离 = 1 - 余弦相似度，值越小表示越相似",
    "PersistentClient 把数据存到本地磁盘，Client 只存内存",
]

collection.add(
    documents=knowledge,
    ids=[f"k_{i}" for i in range(len(knowledge))]
)

def rag_answer(question: str, n_contexts: int = 2) -> str:
    # 第一步：检索
    retrieved = collection.query(
        query_texts=[question],
        n_results=n_contexts
    )
    contexts = "\n".join(retrieved["documents"][0])

    # 第二步：构建 prompt
    prompt = f"""根据以下参考资料回答问题。如果资料中没有相关信息，明确说明。

参考资料：
{contexts}

问题：{question}
回答："""

    # 第三步：调用 LLM
    response = openai_client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{"role": "user", "content": prompt}],
        temperature=0.3
    )
    return response.choices[0].message.content

# 测试
answer = rag_answer("ChromaDB 用什么距离度量？")
print(answer)