概述

Chroma 是一个开源的向量数据库，专为 AI 应用设计，支持存储 embeddings 和元数据，实现高效的相似度检索。它可以用作 AI Agent 的长期记忆系统，使 Agent 能够跨会话记住和检索相关信息。

安装与配置

安装

pip install chromadb chromadb-server

启动模式

嵌入式模式（默认，开发用）：

import chromadb
client = chromadb.Client()

持久化模式（生产用）：

import chromadb
client = chromadb.PersistentClient(path="./chroma_data")

客户端-服务器模式：

# 启动服务器
chroma run --path /path/to/data --port 8000

import chromadb
client = chromadb.HttpClient(host="localhost", port=8000)

集合管理

创建/获取集合

# 创建集合（如果不存在）
collection = client.get_or_create_collection(
    name="knowledge_base",
    metadata={"description": "知识库集合"}
)

# 获取集合
collection = client.get_collection(name="knowledge_base")

# 列出所有集合
collections = client.list_collections()
for c in collections:
    print(f"- {c.name}: {c.count()} 条文档")

删除集合

client.delete_collection(name="old_collection")

文档操作

添加文档

# 简单添加
collection.add(
    documents=[
        "Python 是一种高级编程语言",
        "JavaScript 主要用于 Web 开发",
        "Go 语言以并发性能著称"
    ],
    ids=["doc1", "doc2", "doc3"],
    metadatas=[
        {"language": "programming", "level": "beginner"},
        {"language": "web", "level": "beginner"},
        {"language": "system", "level": "intermediate"}
    ]
)

自定义 Embedding 函数

from sentence_transformers import SentenceTransformer

# 使用自定义 embedding 模型
model = SentenceTransformer("all-MiniLM-L6-v2")

def embed_texts(texts):
    return model.encode(texts).tolist()

collection.add(
    documents=["文档内容..."],
    ids=["doc1"],
    embedding_function=embed_texts  # 自定义 embedding
)

更新和删除

# 更新文档
collection.update(
    ids=["doc1"],
    documents=["更新后的内容"],
    metadatas=[{"updated": True}]
)

# 删除文档
collection.delete(ids=["doc1"])

相似度检索

基础检索

# 按文本检索
results = collection.query(
    query_texts=["什么编程语言适合初学者"],
    n_results=3  # 返回前 3 个最相似结果
)

print(results["documents"])  # 文档内容
print(results["distances"])    # 距离分数
print(results["metadatas"])    # 元数据

元数据过滤

# 过滤条件
results = collection.query(
    query_texts=["高效编程语言"],
    n_results=5,
    where={"language": "programming"},  # 元数据过滤
    where_document={"$contains": "性能"}  # 文档内容过滤
)

距离计算方式

# 创建时指定距离函数
collection = client.get_or_create_collection(
    name="my_collection",
    metadata={"hnsw:space": "cosine"}  # cosine/l2/ip
)

Agent 记忆集成示例

import chromadb
from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain_core.tools import tool

# 初始化
client = chromadb.PersistentClient(path="./agent_memory")
memory_collection = client.get_or_create_collection("agent_memory")
embeddings = OpenAIEmbeddings()

@tool
def remember(topic: str, content: str) -> str:
    """存储重要信息到记忆
    
    Args:
        topic: 信息主题
        content: 要记住的内容
    """
    vector = embeddings.embed_query(content)
    memory_collection.add(
        documents=[content],
        embeddings=[vector],
        ids=[f"mem_{topic}"],
        metadatas=[{"topic": topic, "content": content}]
    )
    return f"已记住关于 '{topic}' 的信息"

@tool
def recall(query: str) -> str:
    """从记忆中检索相关信息
    
    Args:
        query: 搜索查询
    """
    query_vector = embeddings.embed_query(query)
    results = memory_collection.query(
        query_embeddings=[query_vector],
        n_results=3
    )
    if not results["documents"][0]:
        return "记忆中未找到相关信息"
    return "\n".join([
        f"- {doc}" 
        for doc in results["documents"][0]
    ])

# Agent 工具列表
tools = [remember, recall]

# 使用示例
# remember.invoke({"topic": "用户偏好", "content": "用户喜欢简洁的代码风格"})
# recall.invoke({"query": "用户的代码风格偏好是什么"})

常见问题

Q1: Chroma 默认使用什么 Embedding 模型？

• 默认使用 SentenceTransformer 的 all-MiniLM-L6-v2 模型
• 首次使用时会自动下载
• 可以通过 embedding_function 参数指定自定义模型

Q2: 如何选择距离函数？

• cosine：余弦相似度，适合方向相似性
• l2：欧几里得距离，适合数值大小
• ip：内积，适合未归一化的向量

Q3: 如何处理大规模数据？

• 使用客户端-服务器模式便于扩展
• 配置 HNSW 索引参数优化检索性能
• 定期清理过期数据

参考资料

• Chroma 官方文档
• Chroma GitHub
• Sentence Transformers