如何通过Python与LangChain从ChromaDB提取带ID的全量文档及嵌入向量

1. 加载已持久化的ChromaDB实例 #

先连接到你之前保存的ChromaDB存储，确保嵌入函数和持久化路径与创建时完全一致：

from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores import Chroma

# 初始化嵌入模型（和创建ChromaDB时的配置保持一致）
embeddings = OpenAIEmbeddings()
# 加载已持久化的ChromaDB实例
db = Chroma(persist_directory='db', embedding_function=embeddings)

2. 提取所有文档、嵌入向量、元数据及ID #

用ChromaDB的get()方法一次性提取所需数据，通过include参数指定要获取的字段集合：

# 提取包含ID、文档内容、嵌入向量、元数据的全部记录
all_records = db.get(include=["documents", "embeddings", "metadatas"])

# 拆分各字段列表
doc_ids = all_records["ids"]
documents = all_records["documents"]
embedding_vectors = all_records["embeddings"]
metadata_list = all_records["metadatas"]

# 整理为便于后续处理的字典列表
formatted_data = [
    {
        "_id": doc_id,
        "content": content,
        "embedding": vector,
        "metadata": meta
    }
    for doc_id, content, vector, meta in zip(doc_ids, documents, embedding_vectors, metadata_list)
]

3. 回存数据到MongoDB #

使用pymongo库连接MongoDB，将整理好的数据批量插入目标集合：

from pymongo import MongoClient

# 连接MongoDB（替换为你的实际连接配置，如带认证、远程地址等）
client = MongoClient("mongodb://localhost:27017/")
mongo_db = client["your_target_database"]
target_collection = mongo_db["your_target_collection"]

# 批量插入数据
target_collection.insert_many(formatted_data)

4. 用BERTopic进行主题分类 #

直接复用已有的OpenAI嵌入向量可以避免重复计算，提升效率：

from bertopic import BERTopic

# 初始化BERTopic，禁用内置嵌入模型（使用预计算的OpenAI向量）
topic_model = BERTopic(embedding_model="none")
# 执行主题分类，传入文档内容和预计算的嵌入向量
topics, probabilities = topic_model.fit_transform(documents, embedding_vectors)

# 将主题信息关联到原数据
for idx, item in enumerate(formatted_data):
    item["topic_id"] = topics[idx]
    item["topic_probability"] = probabilities[idx]
    # 获取可读性强的主题标签（格式如"topic_0: xxx"）
    item["topic_label"] = topic_model.get_topic_info().loc[topics[idx], "Name"]

# 更新MongoDB中的文档，添加主题分类结果
for item in formatted_data:
    target_collection.update_one(
        {"_id": item["_id"]},
        {"$set": {
            "topic_id": item["topic_id"],
            "topic_probability": item["topic_probability"],
            "topic_label": item["topic_label"]
        }}
    )

关键注意事项 #

• 确保OpenAI API密钥已通过环境变量（OPENAI_API_KEY）或直接传入OpenAIEmbeddings()完成配置。
• 如果不需要复用OpenAI嵌入，可删除embedding_model="none"参数，让BERTopic自动使用默认的all-MiniLM-L6-v2模型计算嵌入。
• MongoDB连接字符串需根据你的部署环境调整，比如添加用户名密码、指定端口等。

内容的提问来源于stack exchange，提问作者user791793