listwise_rerank.py — langchain Source File

Source Code

"""Filter that uses an LLM to rerank documents listwise and select top-k."""

from collections.abc import Sequence
from typing import Any

from langchain_core.callbacks import Callbacks
from langchain_core.documents import BaseDocumentCompressor, Document
from langchain_core.language_models import BaseLanguageModel
from langchain_core.prompts import BasePromptTemplate, ChatPromptTemplate
from langchain_core.runnables import Runnable, RunnableLambda, RunnablePassthrough
from pydantic import BaseModel, ConfigDict, Field

_default_system_tmpl = """{context}

Sort the Documents by their relevance to the Query."""
_DEFAULT_PROMPT = ChatPromptTemplate.from_messages(
    [("system", _default_system_tmpl), ("human", "{query}")],
)


def _get_prompt_input(input_: dict) -> dict[str, Any]:
    """Return the compression chain input."""
    documents = input_["documents"]
    context = ""
    for index, doc in enumerate(documents):
        context += f"Document ID: {index}\n```{doc.page_content}```\n\n"
    document_range = "empty list"
    if len(documents) > 0:
        document_range = f"Document ID: 0, ..., Document ID: {len(documents) - 1}"
    context += f"Documents = [{document_range}]"
    return {"query": input_["query"], "context": context}


def _parse_ranking(results: dict) -> list[Document]:
    ranking = results["ranking"]
    docs = results["documents"]
    return [docs[i] for i in ranking.ranked_document_ids]


class LLMListwiseRerank(BaseDocumentCompressor):
    """Document compressor that uses `Zero-Shot Listwise Document Reranking`.

    Adapted from: https://arxiv.org/pdf/2305.02156.pdf

    `LLMListwiseRerank` uses a language model to rerank a list of documents based on
    their relevance to a query.

    !!! note
        Requires that underlying model implement `with_structured_output`.

    Example usage:
        ```python
        from langchain_classic.retrievers.document_compressors.listwise_rerank import (
            LLMListwiseRerank,
        )
        from langchain_core.documents import Document
        from langchain_openai import ChatOpenAI

        documents = [
            Document("Sally is my friend from school"),
// ... (87 more lines)