Source code for langchain.retrievers.document_compressors.embeddings_filter
from typing import Callable, Dict, Optional, Sequence
import numpy as np
from langchain.callbacks.manager import Callbacks
from langchain.document_transformers.embeddings_redundant_filter import (
    _get_embeddings_from_stateful_docs,
    get_stateful_documents,
)
from langchain.pydantic_v1 import root_validator
from langchain.retrievers.document_compressors.base import (
    BaseDocumentCompressor,
)
from langchain.schema import Document
from langchain.schema.embeddings import Embeddings
from langchain.utils.math import cosine_similarity
[docs]class EmbeddingsFilter(BaseDocumentCompressor):
    """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 = cosine_similarity
    """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: Optional[int] = 20
    """The number of relevant documents to return. Can be set to None, in which case
    `similarity_threshold` must be specified. Defaults to 20."""
    similarity_threshold: Optional[float]
    """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."""
    class Config:
        """Configuration for this pydantic object."""
        arbitrary_types_allowed = True
    @root_validator()
    def validate_params(cls, values: Dict) -> Dict:
        """Validate similarity parameters."""
        if values["k"] is None and values["similarity_threshold"] is None: