codekingpro/portable-devtools

1"""Document compressor."""2 3from __future__ import annotations4 5from abc import ABC, abstractmethod6from typing import TYPE_CHECKING7 8from pydantic import BaseModel9 10from langchain_core.runnables import run_in_executor11 12if TYPE_CHECKING:13    from collections.abc import Sequence14 15    from langchain_core.callbacks import Callbacks16    from langchain_core.documents import Document17 18 19class BaseDocumentCompressor(BaseModel, ABC):20    """Base class for document compressors.21 22    This abstraction is primarily used for post-processing of retrieved documents.23 24    `Document` objects matching a given query are first retrieved.25 26    Then the list of documents can be further processed.27 28    For example, one could re-rank the retrieved documents using an LLM.29 30    !!! note31        Users should favor using a `RunnableLambda` instead of sub-classing from this32        interface.33 34    """35 36    @abstractmethod37    def compress_documents(38        self,39        documents: Sequence[Document],40        query: str,41        callbacks: Callbacks | None = None,42    ) -> Sequence[Document]:43        """Compress retrieved documents given the query context.44 45        Args:46            documents: The retrieved `Document` objects.47            query: The query context.48            callbacks: Optional `Callbacks` to run during compression.49 50        Returns:51            The compressed documents.52 53        """54 55    async def acompress_documents(56        self,57        documents: Sequence[Document],58        query: str,59        callbacks: Callbacks | None = None,60    ) -> Sequence[Document]:61        """Async compress retrieved documents given the query context.62 63        Args:64            documents: The retrieved `Document` objects.65            query: The query context.66            callbacks: Optional `Callbacks` to run during compression.67 68        Returns:69            The compressed documents.70 71        """72        return await run_in_executor(73            None, self.compress_documents, documents, query, callbacks74        )75