Source code for langchain.retrievers.re_phraser

import logging
from typing import List

from langchain_core.callbacks import (
    AsyncCallbackManagerForRetrieverRun,
    CallbackManagerForRetrieverRun,
)
from langchain_core.documents import Document
from langchain_core.language_models import BaseLLM
from langchain_core.prompts.prompt import PromptTemplate
from langchain_core.retrievers import BaseRetriever

from langchain.chains.llm import LLMChain

logger = logging.getLogger(__name__)

# Default template
DEFAULT_TEMPLATE = """You are an assistant tasked with taking a natural language \
query from a user and converting it into a query for a vectorstore. \
In this process, you strip out information that is not relevant for \
the retrieval task. Here is the user query: {question}"""

# Default prompt
DEFAULT_QUERY_PROMPT = PromptTemplate.from_template(DEFAULT_TEMPLATE)


class RePhraseQueryRetriever(BaseRetriever):
    """给定一个查询，使用LLM来重新表达它。
    然后，检索重新表达的查询的文档。"""

    retriever: BaseRetriever
    llm_chain: LLMChain

    @classmethod
    def from_llm(
        cls,
        retriever: BaseRetriever,
        llm: BaseLLM,
        prompt: PromptTemplate = DEFAULT_QUERY_PROMPT,
    ) -> "RePhraseQueryRetriever":
        """使用默认模板从llm初始化。

这里使用的提示期望一个单一的输入：`question`

参数:
    retriever: 用于查询文档的检索器
    llm: 用于使用DEFAULT_QUERY_PROMPT生成查询的llm
    prompt: 用于查询生成的提示模板

返回:
    RePhraseQueryRetriever
"""

        llm_chain = LLMChain(llm=llm, prompt=prompt)
        return cls(
            retriever=retriever,
            llm_chain=llm_chain,
        )

    def _get_relevant_documents(
        self,
        query: str,
        *,
        run_manager: CallbackManagerForRetrieverRun,
    ) -> List[Document]:
        """给定用户问题，获取相关文档。

参数:
    query: 用户问题

返回:
    重新表述问题的相关文档
"""
        response = self.llm_chain(query, callbacks=run_manager.get_child())
        re_phrased_question = response["text"]
        logger.info(f"Re-phrased question: {re_phrased_question}")
        docs = self.retriever.invoke(
            re_phrased_question, config={"callbacks": run_manager.get_child()}
        )
        return docs

    async def _aget_relevant_documents(
        self,
        query: str,
        *,
        run_manager: AsyncCallbackManagerForRetrieverRun,
    ) -> List[Document]:
        raise NotImplementedError