Class AdvancedVectorStoreRetriever

document_qa/langchain.py:19–71

Retriever that can enrich documents with similarity scores and embeddings. Extends LangChain's ``VectorStoreRetriever`` with a ``"similarity_with_embeddings"`` search type. When used, each returned document's ``metadata`` dict gains ``__similarity`` (float) and ``__embeddings`` (li

class AdvancedVectorStoreRetriever(VectorStoreRetriever):
    """Retriever that can enrich documents with similarity scores and embeddings.

    Extends LangChain's ``VectorStoreRetriever`` with a
    ``"similarity_with_embeddings"`` search type.  When used, each
    returned document's ``metadata`` dict gains ``__similarity`` (float)
    and ``__embeddings`` (list[float]) keys.
    """

    allowed_search_types: ClassVar[Collection[str]] = (
        "similarity",
        "similarity_score_threshold",
        "mmr",
        "similarity_with_embeddings",
    )

    def _get_relevant_documents(self, query: str, *, run_manager: CallbackManagerForRetrieverRun) -> List[Document]:
        """Fetch relevant documents for the configured search type.

        Supports all standard search types plus
        ``"similarity_with_embeddings"`` which attaches score and
        embedding vector metadata to each document.

        Args:
            query: The search query string.
            run_manager: LangChain callback manager.

        Returns:
            list[Document]: Retrieved documents, optionally enriched
            with similarity scores and embeddings.
        """

        if self.search_type == "similarity_with_embeddings":
            docs_scores_and_embeddings = self.vectorstore.advanced_similarity_search(query, **self.search_kwargs)

            for doc, score, embeddings in docs_scores_and_embeddings:
                if "__embeddings" not in doc.metadata.keys():
                    doc.metadata["__embeddings"] = embeddings
                if "__similarity" not in doc.metadata.keys():
                    doc.metadata["__similarity"] = score

            docs = [doc for doc, _, _ in docs_scores_and_embeddings]
        elif self.search_type == "similarity_score_threshold":
            docs_and_similarities = self.vectorstore.similarity_search_with_relevance_scores(query, **self.search_kwargs)
            for doc, similarity in docs_and_similarities:
                if "__similarity" not in doc.metadata.keys():
                    doc.metadata["__similarity"] = similarity

            docs = [doc for doc, _ in docs_and_similarities]
        else:
            docs = super()._get_relevant_documents(query, run_manager=run_manager)

        return docs


class AdvancedVectorStore(VectorStore):