core[minor]: move document compressor base (#17910)

This commit is contained in:
Harrison Chase
2024-02-26 17:20:50 -08:00
committed by GitHub
parent b3f4de38ae
commit d7c607ca00
4 changed files with 41 additions and 30 deletions

View File

@@ -1,38 +1,15 @@
from abc import ABC, abstractmethod
from inspect import signature
from typing import List, Optional, Sequence, Union
from langchain_core.documents import BaseDocumentTransformer, Document
from langchain_core.pydantic_v1 import BaseModel
from langchain_core.runnables.config import run_in_executor
from langchain_core.documents import (
BaseDocumentCompressor,
BaseDocumentTransformer,
Document,
)
from langchain.callbacks.manager import Callbacks
class BaseDocumentCompressor(BaseModel, ABC):
"""Base class for document compressors."""
@abstractmethod
def compress_documents(
self,
documents: Sequence[Document],
query: str,
callbacks: Optional[Callbacks] = None,
) -> Sequence[Document]:
"""Compress retrieved documents given the query context."""
async def acompress_documents(
self,
documents: Sequence[Document],
query: str,
callbacks: Optional[Callbacks] = None,
) -> Sequence[Document]:
"""Compress retrieved documents given the query context."""
return await run_in_executor(
None, self.compress_documents, documents, query, callbacks
)
class DocumentCompressorPipeline(BaseDocumentCompressor):
"""Document compressor that uses a pipeline of Transformers."""