如何仅针对特定PDF文件提取相关的ID和元数据，而不是针对整个chromadb集合？

Question

我正在用Langchain和Python开发一个聊天应用。这个应用的想法是，用户可以提交一些PDF文件，然后聊天模型会根据这些文件进行训练，用户可以向模型提问。我们把这些文件的嵌入信息存储在Chromadb向量数据库里。所以，这实际上是一个基于RAG（检索增强生成）的解决方案。

现在，创建和存储嵌入信息的功能都正常，聊天功能也很好。不过，我在嵌入信息中存储了一些自定义的元数据和一些ID。相关的代码如下：

def read_docs(pdf_file):
    pdf_loader = PyPDFLoader(pdf_file)
    pdf_documents = pdf_loader.load()

    text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
    documents = text_splitter.split_documents(pdf_documents)
    
    return documents

def generate_and_store_embeddings(documents, pdf_file, user_id):
    client = chromadb.PersistentClient(path="./trained_db")
    collection = client.get_or_create_collection("PDF_Embeddings", embedding_function=embedding_functions.OpenAIEmbeddingFunction(api_key=config["OPENAI_API_KEY"], model_name=configs.EMBEDDINGS_MODEL))
    now = datetime.now()

    #custom metadata and ids I want to store along with the embeddings for each pdf
    metadata = {"source": pdf_file.filename, "user": str(user_id), 'created_at': 
                          now.strftime("%d/%m/%Y %H:%M:%S")}
    ids = [str(uuid.uuid4()) for _ in range(len(documents))]

    try:
        vectordb = Chroma.from_documents(
                    documents,         
                    embedding=OpenAIEmbeddings(openai_api_key=config["OPENAI_API_KEY"], 
                    model=configs.EMBEDDINGS_MODEL),
                    persist_directory='./trained_db',
                    collection_name = collection.name, 
                    client = client,
                    ids = ids,
                    collection_metadata = {item: value for (item, value) in metadata.items()}
                )
        vectordb.persist()
        
    except Exception as err:
        print(f"An error occured: {err=}, {type(err)=}")
        return {"answer": "An error occured while generating embeddings. Please check terminal 
                           for more details."}
    return vectordb

我现在想要的是，能够获取与特定PDF文件相关的那些ID和元数据，而不是获取整个集合中的所有ID和元数据。这样，当用户输入要删除嵌入信息的PDF文件时，我就可以只获取那个PDF文件的元数据和ID，这样我就能用这些ID从集合中删除该PDF文件的嵌入信息。

我知道有一个函数叫vectordb._collection.get()，但它会返回所有的ID。我也试过这个代码：print(vectordb.get(where={"source": pdf_file.filename}))，但返回的结果是：

{'ids': [], 'embeddings': None, 'metadatas': [], 'documents': [], 'uris': None, 'data': None}

Answer 1

vectordb.get(where={"source": pdf_file.filename})这个功能是用来处理每个Document的元数据的。你需要给每个Document添加source这个元数据，这样才能查询到同一个来源的所有文档。这样做的目的是为了能够获取到某个PDF的所有文档ID，以便后面可以删除它们。

def read_docs(pdf_file):
    pdf_loader = PyPDFLoader(pdf_file)
    pdf_documents = pdf_loader.load()

    text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
    documents = text_splitter.split_documents(pdf_documents)

    # add Document metadata
    for doc in documents:
        doc.metadata = {
            "id": "1",  # assign ID here
            "source": pdf_file.filename,
        }
    
    return documents

当你调用generate_and_store_embeddings(documents, pdf_file, user_id)时，Document的元数据会被保存下来。你不需要把元数据传递给Chroma.from_documents()函数中的collection_metadata参数。