cohere-rerank.py

# CohereRerank wrapper to use as a document compressor

from typing import Any, Dict, Sequence

from pydantic import root_validator

from langchain.retrievers.document_compressors.base import BaseDocumentCompressor
from langchain.schema import Document
from langchain.utils import get_from_dict_or_env


class CohereRerank(BaseDocumentCompressor):
    client: Any
    top_n: int = 3
    model: str = "rerank-english-v2.0"

    @root_validator()
    def validate_environment(cls, values: Dict) -> Dict:
        """Validate that api key and python package exists in environment."""
        cohere_api_key = get_from_dict_or_env(
            values, "cohere_api_key", "COHERE_API_KEY"
        )
        try:
            import cohere

            values["client"] = cohere.Client(cohere_api_key)
        except ImportError:
            raise ValueError(
                "Could not import cohere python package. "
                "Please install it with `pip install cohere`."
            )
        return values

    def compress_documents(
        self, documents: Sequence[Document], query: str
    ) -> Sequence[Document]:
        doc_list = list(documents)
        _docs = [d.page_content for d in doc_list]
        results = self.client.rerank(
            model=self.model, query=query, documents=_docs, top_n=self.top_n
        )
        final_results = []
        for r in results:
            doc = doc_list[r.index]
            doc.metadata["relevance_score"] = r.relevance_score
            final_results.append(doc)
        return final_results

    async def acompress_documents(
        self, documents: Sequence[Document], query: str
    ) -> Sequence[Document]:
        raise NotImplementedError

Thanks, @hwchase17: I keep getting this error when following https://python.langchain.com/en/latest/modules/indexes/retrievers/examples/cohere-reranker.html exactly:

pydantic.error_wrappers.ValidationError: 1 validation error for CohereRerank client field required (type=value_error.missing)

@cnevelle-blueprint This gist is out of date. See: langchain-ai/langchain@84cfa76
Hopefully this error will be solved for you now because of the Pydantic logic update.