codekingpro/portable-devtools

1"""Document transformers."""2 3from __future__ import annotations4 5from abc import ABC, abstractmethod6from typing import TYPE_CHECKING, Any7 8from langchain_core.runnables.config import run_in_executor9 10if TYPE_CHECKING:11    from collections.abc import Sequence12 13    from langchain_core.documents import Document14 15 16class BaseDocumentTransformer(ABC):17    """Abstract base class for document transformation.18 19    A document transformation takes a sequence of `Document` objects and returns a20    sequence of transformed `Document` objects.21 22    Example:23        ```python24        class EmbeddingsRedundantFilter(BaseDocumentTransformer, BaseModel):25            embeddings: Embeddings26            similarity_fn: Callable = cosine_similarity27            similarity_threshold: float = 0.9528 29            class Config:30                arbitrary_types_allowed = True31 32            def transform_documents(33                self, documents: Sequence[Document], **kwargs: Any34            ) -> Sequence[Document]:35                stateful_documents = get_stateful_documents(documents)36                embedded_documents = _get_embeddings_from_stateful_docs(37                    self.embeddings, stateful_documents38                )39                included_idxs = _filter_similar_embeddings(40                    embedded_documents,41                    self.similarity_fn,42                    self.similarity_threshold,43                )44                return [stateful_documents[i] for i in sorted(included_idxs)]45 46            async def atransform_documents(47                self, documents: Sequence[Document], **kwargs: Any48            ) -> Sequence[Document]:49                raise NotImplementedError50        ```51    """52 53    @abstractmethod54    def transform_documents(55        self, documents: Sequence[Document], **kwargs: Any56    ) -> Sequence[Document]:57        """Transform a list of documents.58 59        Args:60            documents: A sequence of `Document` objects to be transformed.61 62        Returns:63            A sequence of transformed `Document` objects.64        """65 66    async def atransform_documents(67        self, documents: Sequence[Document], **kwargs: Any68    ) -> Sequence[Document]:69        """Asynchronously transform a list of documents.70 71        Args:72            documents: A sequence of `Document` objects to be transformed.73 74        Returns:75            A sequence of transformed `Document` objects.76        """77        return await run_in_executor(78            None, self.transform_documents, documents, **kwargs79        )80