1from enum import Enum2from typing import Any34from langchain_core.callbacks import (5    AsyncCallbackManagerForRetrieverRun,6    CallbackManagerForRetrieverRun,7)8from langchain_core.documents import Document9from langchain_core.retrievers import BaseRetriever10from langchain_core.stores import BaseStore, ByteStore11from langchain_core.vectorstores import VectorStore12from pydantic import Field, model_validator13from typing_extensions import override1415from langchain_classic.storage._lc_store import create_kv_docstore161718class SearchType(str, Enum):19    """Enumerator of the types of search to perform."""2021    similarity = "similarity"22    """Similarity search."""23    similarity_score_threshold = "similarity_score_threshold"24    """Similarity search with a score threshold."""25    mmr = "mmr"26    """Maximal Marginal Relevance reranking of similarity search."""272829class MultiVectorRetriever(BaseRetriever):30    """Retriever that supports multiple embeddings per parent document.3132    This retriever is designed for scenarios where documents are split into33    smaller chunks for embedding and vector search, but retrieval returns34    the original parent documents rather than individual chunks.3536    It works by:37    - Performing similarity (or MMR) search over embedded child chunks38    - Collecting unique parent document IDs from chunk metadata39    - Fetching and returning the corresponding parent documents from the docstore4041    This pattern is commonly used in RAG pipelines to improve answer grounding42    while preserving full document context.43    """4445    vectorstore: VectorStore46    """The underlying `VectorStore` to use to store small chunks47    and their embedding vectors"""4849    byte_store: ByteStore | None = None50    """The lower-level backing storage layer for the parent documents"""5152    docstore: BaseStore[str, Document]53    """The storage interface for the parent documents"""5455    id_key: str = "doc_id"5657    search_kwargs: dict = Field(default_factory=dict)58    """Keyword arguments to pass to the search function."""5960    search_type: SearchType = SearchType.similarity61    """Type of search to perform (similarity / mmr)"""6263    @model_validator(mode="before")64    @classmethod65    def _shim_docstore(cls, values: dict) -> Any:66        byte_store = values.get("byte_store")67        docstore = values.get("docstore")68        if byte_store is not None:69            docstore = create_kv_docstore(byte_store)70        elif docstore is None:71            msg = "You must pass a `byte_store` parameter."72            raise ValueError(msg)73        values["docstore"] = docstore74        return values7576    @override77    def _get_relevant_documents(78        self,79        query: str,80        *,81        run_manager: CallbackManagerForRetrieverRun,82    ) -> list[Document]:83        """Get documents relevant to a query.8485        Args:86            query: String to find relevant documents for87            run_manager: The callbacks handler to use88        Returns:89            List of relevant documents.90        """91        if self.search_type == SearchType.mmr:92            sub_docs = self.vectorstore.max_marginal_relevance_search(93                query,94                **self.search_kwargs,95            )96        elif self.search_type == SearchType.similarity_score_threshold:97            sub_docs_and_similarities = (98                self.vectorstore.similarity_search_with_relevance_scores(99                    query,100                    **self.search_kwargs,101                )102            )103            sub_docs = [sub_doc for sub_doc, _ in sub_docs_and_similarities]104        else:105            sub_docs = self.vectorstore.similarity_search(query, **self.search_kwargs)106107        # We do this to maintain the order of the IDs that are returned108        ids = []109        for d in sub_docs:110            if self.id_key in d.metadata and d.metadata[self.id_key] not in ids:111                ids.append(d.metadata[self.id_key])112        docs = self.docstore.mget(ids)113        return [d for d in docs if d is not None]114115    @override116    async def _aget_relevant_documents(117        self,118        query: str,119        *,120        run_manager: AsyncCallbackManagerForRetrieverRun,121    ) -> list[Document]:122        """Asynchronously get documents relevant to a query.123124        Args:125            query: String to find relevant documents for126            run_manager: The callbacks handler to use127        Returns:128            List of relevant documents.129        """130        if self.search_type == SearchType.mmr:131            sub_docs = await self.vectorstore.amax_marginal_relevance_search(132                query,133                **self.search_kwargs,134            )135        elif self.search_type == SearchType.similarity_score_threshold:136            sub_docs_and_similarities = (137                await self.vectorstore.asimilarity_search_with_relevance_scores(138                    query,139                    **self.search_kwargs,140                )141            )142            sub_docs = [sub_doc for sub_doc, _ in sub_docs_and_similarities]143        else:144            sub_docs = await self.vectorstore.asimilarity_search(145                query,146                **self.search_kwargs,147            )148149        # We do this to maintain the order of the IDs that are returned150        ids = []151        for d in sub_docs:152            if self.id_key in d.metadata and d.metadata[self.id_key] not in ids:153                ids.append(d.metadata[self.id_key])154        docs = await self.docstore.amget(ids)155        return [d for d in docs if d is not None]