Source code for langchain_community.document_compressors.jina_rerank

from __future__ import annotations

from copy import deepcopy
from typing import Any, Dict, List, Optional, Sequence, Union

import requests
from langchain_core.callbacks import Callbacks
from langchain_core.documents import BaseDocumentCompressor, Document
from langchain_core.pydantic_v1 import Extra, root_validator
from langchain_core.utils import get_from_dict_or_env

JINA_API_URL: str = "https://api.jina.ai/v1/rerank"


class JinaRerank(BaseDocumentCompressor):
    """使用`Jina Rerank API`的文档压缩器。"""

    session: Any = None
    """使用Requests会话与API进行通信。"""
    top_n: Optional[int] = 3
    """返回的文档数量。"""
    model: str = "jina-reranker-v1-base-en"
    """用于重新排序的模型。"""
    jina_api_key: Optional[str] = None
    """Jina API密钥。必须直接指定或通过环境变量JINA_API_KEY指定。"""
    user_agent: str = "langchain"
    """用于发出请求的应用程序的标识符。"""

    class Config:
        """此pydantic对象的配置。"""

        extra = Extra.forbid
        arbitrary_types_allowed = True

    @root_validator(pre=True)
    def validate_environment(cls, values: Dict) -> Dict:
        """验证环境中是否存在API密钥。"""
        jina_api_key = get_from_dict_or_env(values, "jina_api_key", "JINA_API_KEY")
        user_agent = values.get("user_agent", "langchain")
        session = requests.Session()
        session.headers.update(
            {
                "Authorization": f"Bearer {jina_api_key}",
                "Accept-Encoding": "identity",
                "Content-type": "application/json",
                "user-agent": user_agent,
            }
        )
        values["session"] = session
        return values

    def rerank(
        self,
        documents: Sequence[Union[str, Document, dict]],
        query: str,
        *,
        model: Optional[str] = None,
        top_n: Optional[int] = -1,
        max_chunks_per_doc: Optional[int] = None,
    ) -> List[Dict[str, Any]]:
        """返回按与提供的查询相关性排序的文档的有序列表。

参数：
    query：用于重新排序的查询。
    documents：要重新排序的文档序列。
    model：用于重新排序的模型。默认为self.model。
    top_n：要返回的结果数量。如果为None，则返回所有结果。默认为self.top_n。
    max_chunks_per_doc：从文档中派生的最大块数。
"""  # noqa: E501
        if len(documents) == 0:  # to avoid empty api call
            return []
        docs = [
            doc.page_content if isinstance(doc, Document) else doc for doc in documents
        ]
        model = model or self.model
        top_n = top_n if (top_n is None or top_n > 0) else self.top_n
        data = {
            "query": query,
            "documents": docs,
            "model": model,
            "top_n": top_n,
        }

        resp = self.session.post(
            JINA_API_URL,
            json=data,
        ).json()

        if "results" not in resp:
            raise RuntimeError(resp["detail"])

        results = resp["results"]
        result_dicts = []
        for res in results:
            result_dicts.append(
                {"index": res["index"], "relevance_score": res["relevance_score"]}
            )
        return result_dicts

    def compress_documents(
        self,
        documents: Sequence[Document],
        query: str,
        callbacks: Optional[Callbacks] = None,
    ) -> Sequence[Document]:
        """使用Jina的Rerank API 压缩文档。

参数:
    documents: 需要压缩的文档序列。
    query: 用于压缩文档的查询。
    callbacks: 在压缩过程中运行的回调函数。

返回:
    压缩后的文档序列。
"""
        compressed = []
        for res in self.rerank(documents, query):
            doc = documents[res["index"]]
            doc_copy = Document(doc.page_content, metadata=deepcopy(doc.metadata))
            doc_copy.metadata["relevance_score"] = res["relevance_score"]
            compressed.append(doc_copy)
        return compressed