如何将LlamaIndex SubDocSummaryPack生成的嵌入持久化到本地ChromaDB

核心思路 #

SubDocSummaryPack处理文档后会生成带层级的父/子节点集合（父块含自动生成的摘要，子块是细粒度分块），我们需要提取这些节点，结合Chroma的存储上下文构建索引，即可完成嵌入的持久化。同时通过检查Chroma集合的内容，避免文档无变更时重复执行分块和嵌入操作。

修改后的完整代码 #

from llama_index.core import SimpleDirectoryReader
import os
from llama_index.packs.subdoc_summary import SubDocSummaryPack
from llama_index.llms.openai import OpenAI
from llama_index.embeddings.openai import OpenAIEmbedding
import chromadb
from llama_index.core import VectorStoreIndex
from llama_index.vector_stores.chroma import ChromaVectorStore
from llama_index.core import StorageContext

os.environ["OPENAI_API_KEY"] = "%API_KEY%"

# 初始化ChromaDB
db = chromadb.PersistentClient(path="./chroma")
chroma_collection = db.get_or_create_collection("quickstart")
vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
storage_context = StorageContext.from_defaults(vector_store=vector_store)

# 检查Chroma集合是否已有数据，避免重复处理
if chroma_collection.count() == 0:
    # 仅当集合为空时，重新加载并处理文档
    documents = SimpleDirectoryReader("data").load_data()
    
    subdoc_summary_pack = SubDocSummaryPack(
        documents,
        parent_chunk_size=8192,
        child_chunk_size=512,
        llm=OpenAI(model="gpt-3.5-turbo"),
        embed_model=OpenAIEmbedding(),
    )
    
    # 提取SubDocSummaryPack处理后的所有节点（父节点+子节点）
    processed_nodes = subdoc_summary_pack.get_nodes()
    
    # 构建索引并自动将嵌入写入ChromaDB
    index = VectorStoreIndex(
        processed_nodes,
        storage_context=storage_context,
        embed_model=OpenAIEmbedding()
    )
else:
    # 直接从ChromaDB加载已有的索引
    index = VectorStoreIndex.from_vector_store(
        vector_store,
        embed_model=OpenAIEmbedding()
    )

# 后续即可用index进行查询等操作
# query_engine = index.as_query_engine()
# response = query_engine.query("你的问题")

关键说明 #

• 提取处理后的节点：通过subdoc_summary_pack.get_nodes()获取SubDocSummaryPack生成的所有层级节点，这是后续构建索引的核心数据。
• 持久化判断逻辑：利用chroma_collection.count()检查集合是否为空，为空才重新执行文档加载、分块和嵌入生成，否则直接加载已有索引。
• 索引复用：下次启动时，只要ChromaDB的./chroma目录未被删除，且文档无变更，就会直接复用已存储的嵌入，无需重复消耗API资源。

文档变更时的更新策略 #

如果后续文档有修改，只需删除ChromaDB的./chroma目录，或者删除目标集合（db.delete_collection("quickstart")），重新运行代码即可生成新的嵌入并持久化。

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