embeddings_filter.py — langchain Source File

Source Code

from collections.abc import Callable, Sequence

from langchain_core.callbacks import Callbacks
from langchain_core.documents import BaseDocumentCompressor, Document
from langchain_core.embeddings import Embeddings
from langchain_core.utils import pre_init
from pydantic import ConfigDict, Field
from typing_extensions import override


def _get_similarity_function() -> Callable:
    try:
        from langchain_community.utils.math import cosine_similarity
    except ImportError as e:
        msg = (
            "To use please install langchain-community "
            "with `pip install langchain-community`."
        )
        raise ImportError(msg) from e
    return cosine_similarity


class EmbeddingsFilter(BaseDocumentCompressor):
    """Embeddings Filter.

    Document compressor that uses embeddings to drop documents unrelated to the query.
    """

    embeddings: Embeddings
    """Embeddings to use for embedding document contents and queries."""
    similarity_fn: Callable = Field(default_factory=_get_similarity_function)
    """Similarity function for comparing documents. Function expected to take as input
    two matrices (List[List[float]]) and return a matrix of scores where higher values
    indicate greater similarity."""
    k: int | None = 20
    """The number of relevant documents to return. Can be set to `None`, in which case
    `similarity_threshold` must be specified."""
    similarity_threshold: float | None = None
    """Threshold for determining when two documents are similar enough
    to be considered redundant. Defaults to `None`, must be specified if `k` is set
    to None."""

    model_config = ConfigDict(
        arbitrary_types_allowed=True,
    )

    @pre_init
    def validate_params(cls, values: dict) -> dict:
        """Validate similarity parameters."""
        if values["k"] is None and values["similarity_threshold"] is None:
            msg = "Must specify one of `k` or `similarity_threshold`."
            raise ValueError(msg)
        return values

    @override
    def compress_documents(
        self,
        documents: Sequence[Document],
        query: str,
        callbacks: Callbacks | None = None,
// ... (82 more lines)