Source code for langchain.retrievers.parent_document_retriever

import uuid
from typing import List, Optional, Sequence

from langchain_core.documents import Document
from langchain_text_splitters import TextSplitter

from langchain.retrievers import MultiVectorRetriever


class ParentDocumentRetriever(MultiVectorRetriever):
    """获取小块，然后检索它们的父文档。

    在检索文档时，通常存在冲突的需求：

    1. 您可能希望有小型文档，以便它们的嵌入可以最准确地反映它们的含义。如果太长，那么嵌入可能会失去意义。
    2. 您希望有足够长的文档，以保留每个块的上下文。

    ParentDocumentRetriever 通过拆分和存储小块数据来平衡这一点。在检索过程中，它首先获取小块，然后查找这些块的父ID，并返回那些更大的文档。

    请注意，“父文档”指的是小块来源的文档。这既可以是整个原始文档，也可以是更大的块。

    示例：

        .. code-block:: python

            from langchain_community.embeddings import OpenAIEmbeddings
            from langchain_community.vectorstores import Chroma
            from langchain_text_splitters import RecursiveCharacterTextSplitter
            from langchain.storage import InMemoryStore

            # 此文本拆分器用于创建父文档
            parent_splitter = RecursiveCharacterTextSplitter(chunk_size=2000, add_start_index=True)
            # 此文本拆分器用于创建子文档
            # 它应该创建比父文档更小的文档
            child_splitter = RecursiveCharacterTextSplitter(chunk_size=400, add_start_index=True)
            # 用于索引子块的向量存储
            vectorstore = Chroma(embedding_function=OpenAIEmbeddings())
            # 用于父文档的存储层
            store = InMemoryStore()

            # 初始化检索器
            retriever = ParentDocumentRetriever(
                vectorstore=vectorstore,
                docstore=store,
                child_splitter=child_splitter,
                parent_splitter=parent_splitter,
            )"""  # noqa: E501

    child_splitter: TextSplitter
    """用于创建子文档的文本分割器。"""

    """用于跟踪父ID的关键。这将存储在子文档的元数据中。"""
    parent_splitter: Optional[TextSplitter] = None
    """用于创建父文档的文本分割器。
    如果没有指定，则父文档将是传入的原始文档。"""

    child_metadata_fields: Optional[Sequence[str]] = None
    """要保留在子文档中的元数据字段。如果为None，则保留所有父文档的元数据。"""

    def add_documents(
        self,
        documents: List[Document],
        ids: Optional[List[str]] = None,
        add_to_docstore: bool = True,
    ) -> None:
        """将文档添加到文档存储库和向量存储库中。

参数:
    documents: 要添加的文档列表
    ids: 文档的可选id列表。如果提供应与文档列表长度相同。
        如果父文档已经在文档存储库中，且不想重新添加到文档存储库，则可以提供此参数。
        如果未提供，则将使用随机UUID作为id。
    add_to_docstore: 是否将文档添加到文档存储库的布尔值。
        如果提供了`ids`，则此参数可以为false。如果文档已经在文档存储库中，且不想重新添加，则可能希望将其设置为False。
"""
        if self.parent_splitter is not None:
            documents = self.parent_splitter.split_documents(documents)
        if ids is None:
            doc_ids = [str(uuid.uuid4()) for _ in documents]
            if not add_to_docstore:
                raise ValueError(
                    "If ids are not passed in, `add_to_docstore` MUST be True"
                )
        else:
            if len(documents) != len(ids):
                raise ValueError(
                    "Got uneven list of documents and ids. "
                    "If `ids` is provided, should be same length as `documents`."
                )
            doc_ids = ids

        docs = []
        full_docs = []
        for i, doc in enumerate(documents):
            _id = doc_ids[i]
            sub_docs = self.child_splitter.split_documents([doc])
            if self.child_metadata_fields is not None:
                for _doc in sub_docs:
                    _doc.metadata = {
                        k: _doc.metadata[k] for k in self.child_metadata_fields
                    }
            for _doc in sub_docs:
                _doc.metadata[self.id_key] = _id
            docs.extend(sub_docs)
            full_docs.append((_id, doc))
        self.vectorstore.add_documents(docs)
        if add_to_docstore:
            self.docstore.mset(full_docs)