Source code for langchain_community.retrievers.remote_retriever

from typing import List, Optional

import aiohttp
import requests
from langchain_core.callbacks import (
    AsyncCallbackManagerForRetrieverRun,
    CallbackManagerForRetrieverRun,
)
from langchain_core.documents import Document
from langchain_core.retrievers import BaseRetriever


class RemoteLangChainRetriever(BaseRetriever):
    """`LangChain API` 检索器。"""

    url: str
    """远程LangChain API的URL。"""
    headers: Optional[dict] = None
    """用于请求的标头。"""
    input_key: str = "message"
    """请求中用于输入的键。"""
    response_key: str = "response"
    """用于请求中响应的键。"""
    page_content_key: str = "page_content"
    """用于响应中页面内容的键。"""
    metadata_key: str = "metadata"
    """用于响应中元数据的键。"""

    def _get_relevant_documents(
        self, query: str, *, run_manager: CallbackManagerForRetrieverRun
    ) -> List[Document]:
        response = requests.post(
            self.url, json={self.input_key: query}, headers=self.headers
        )
        result = response.json()
        return [
            Document(
                page_content=r[self.page_content_key], metadata=r[self.metadata_key]
            )
            for r in result[self.response_key]
        ]

    async def _aget_relevant_documents(
        self, query: str, *, run_manager: AsyncCallbackManagerForRetrieverRun
    ) -> List[Document]:
        async with aiohttp.ClientSession() as session:
            async with session.request(
                "POST", self.url, headers=self.headers, json={self.input_key: query}
            ) as response:
                result = await response.json()
        return [
            Document(
                page_content=r[self.page_content_key], metadata=r[self.metadata_key]
            )
            for r in result[self.response_key]
        ]