"""通过第一遍处理文档,然后在更多文档上进行细化。"""
from __future__ import annotations
from typing import Any, Dict, List, Tuple
from langchain_core.callbacks import Callbacks
from langchain_core.documents import Document
from langchain_core.prompts import BasePromptTemplate, format_document
from langchain_core.prompts.prompt import PromptTemplate
from langchain_core.pydantic_v1 import Extra, Field, root_validator
from langchain.chains.combine_documents.base import (
BaseCombineDocumentsChain,
)
from langchain.chains.llm import LLMChain
def _get_default_document_prompt() -> PromptTemplate:
return PromptTemplate(input_variables=["page_content"], template="{page_content}")
[docs]class RefineDocumentsChain(BaseCombineDocumentsChain):
"""将文档合并,首先进行第一遍处理,然后在更多文档上进行细化。
该算法首先在第一个文档上调用`initial_llm_chain`,传入该第一个文档,并生成一个名为`initial_response_name`的新变量。
然后,它循环遍历每个剩余文档。这称为“细化”步骤。
它调用`refine_llm_chain`,传入该文档以及前一个响应,前一个响应的变量名为`initial_response_name`。
示例:
.. code-block:: python
from langchain.chains import RefineDocumentsChain, LLMChain
from langchain_core.prompts import PromptTemplate
from langchain_community.llms import OpenAI
# 这控制每个文档的格式。具体来说,它将传递给`format_document` - 有关更多详细信息,请参阅该函数。
document_prompt = PromptTemplate(
input_variables=["page_content"],
template="{page_content}"
)
document_variable_name = "context"
llm = OpenAI()
# 这里的提示应该将`document_variable_name`作为输入变量
prompt = PromptTemplate.from_template(
"总结这段内容:{context}"
)
initial_llm_chain = LLMChain(llm=llm, prompt=prompt)
initial_response_name = "prev_response"
# 这里的提示应该将`document_variable_name`以及`initial_response_name`作为输入变量
prompt_refine = PromptTemplate.from_template(
"这是您的第一个摘要:{prev_response}。"
"现在根据以下内容添加内容:{context}"
)
refine_llm_chain = LLMChain(llm=llm, prompt=prompt_refine)
chain = RefineDocumentsChain(
initial_llm_chain=initial_llm_chain,
refine_llm_chain=refine_llm_chain,
document_prompt=document_prompt,
document_variable_name=document_variable_name,
initial_response_name=initial_response_name,
)"""
initial_llm_chain: LLMChain
"""用于初始文档的LLM链。"""
refine_llm_chain: LLMChain
"""用于细化时使用的LLM链。"""
document_variable_name: str
"""初始_llm_chain 中放置文档的变量名。
如果在 initial_llm_chain 中只有一个变量,则无需提供此变量。"""
initial_response_name: str
"""用于在细化时格式化初始响应的变量名称。"""
document_prompt: BasePromptTemplate = Field(
default_factory=_get_default_document_prompt
)
"""传递给`format_document`的用于格式化每个文档的提示。"""
return_intermediate_steps: bool = False
"""将细化步骤的结果返回到输出中。"""
@property
def output_keys(self) -> List[str]:
"""期望输入键。
:元数据 私有:
"""
_output_keys = super().output_keys
if self.return_intermediate_steps:
_output_keys = _output_keys + ["intermediate_steps"]
return _output_keys
class Config:
"""这个pydantic对象的配置。"""
extra = Extra.forbid
arbitrary_types_allowed = True
@root_validator(pre=True)
def get_return_intermediate_steps(cls, values: Dict) -> Dict:
"""为了向后兼容。"""
if "return_refine_steps" in values:
values["return_intermediate_steps"] = values["return_refine_steps"]
del values["return_refine_steps"]
return values
@root_validator(pre=True)
def get_default_document_variable_name(cls, values: Dict) -> Dict:
"""获取默认文档变量名称,如果未提供。"""
if "document_variable_name" not in values:
llm_chain_variables = values["initial_llm_chain"].prompt.input_variables
if len(llm_chain_variables) == 1:
values["document_variable_name"] = llm_chain_variables[0]
else:
raise ValueError(
"document_variable_name must be provided if there are "
"multiple llm_chain input_variables"
)
else:
llm_chain_variables = values["initial_llm_chain"].prompt.input_variables
if values["document_variable_name"] not in llm_chain_variables:
raise ValueError(
f"document_variable_name {values['document_variable_name']} was "
f"not found in llm_chain input_variables: {llm_chain_variables}"
)
return values
[docs] def combine_docs(
self, docs: List[Document], callbacks: Callbacks = None, **kwargs: Any
) -> Tuple[str, dict]:
"""将首先通过映射将所有内容连接起来,然后将其填充到最终链中。
参数:
docs:需要合并的文档列表
callbacks:要传递的回调函数
**kwargs:要传递给LLM调用的其他参数(例如文档之外的其他输入变量)
返回值:
返回的第一个元素是单个字符串输出。返回的第二个元素是要返回的其他键的字典。
"""
inputs = self._construct_initial_inputs(docs, **kwargs)
res = self.initial_llm_chain.predict(callbacks=callbacks, **inputs)
refine_steps = [res]
for doc in docs[1:]:
base_inputs = self._construct_refine_inputs(doc, res)
inputs = {**base_inputs, **kwargs}
res = self.refine_llm_chain.predict(callbacks=callbacks, **inputs)
refine_steps.append(res)
return self._construct_result(refine_steps, res)
[docs] async def acombine_docs(
self, docs: List[Document], callbacks: Callbacks = None, **kwargs: Any
) -> Tuple[str, dict]:
"""异步通过映射第一个链条覆盖所有内容,然后填充到最终链条中。
参数:
docs: 需要合并的文档列表
callbacks: 需要传递的回调函数
**kwargs: 需要传递给LLM调用的其他参数(例如文档之外的其他输入变量)
返回值:
返回的第一个元素是单个字符串输出。返回的第二个元素是要返回的其他键的字典。
"""
inputs = self._construct_initial_inputs(docs, **kwargs)
res = await self.initial_llm_chain.apredict(callbacks=callbacks, **inputs)
refine_steps = [res]
for doc in docs[1:]:
base_inputs = self._construct_refine_inputs(doc, res)
inputs = {**base_inputs, **kwargs}
res = await self.refine_llm_chain.apredict(callbacks=callbacks, **inputs)
refine_steps.append(res)
return self._construct_result(refine_steps, res)
def _construct_result(self, refine_steps: List[str], res: str) -> Tuple[str, dict]:
if self.return_intermediate_steps:
extra_return_dict = {"intermediate_steps": refine_steps}
else:
extra_return_dict = {}
return res, extra_return_dict
def _construct_refine_inputs(self, doc: Document, res: str) -> Dict[str, Any]:
return {
self.document_variable_name: format_document(doc, self.document_prompt),
self.initial_response_name: res,
}
def _construct_initial_inputs(
self, docs: List[Document], **kwargs: Any
) -> Dict[str, Any]:
base_info = {"page_content": docs[0].page_content}
base_info.update(docs[0].metadata)
document_info = {k: base_info[k] for k in self.document_prompt.input_variables}
base_inputs: dict = {
self.document_variable_name: self.document_prompt.format(**document_info)
}
inputs = {**base_inputs, **kwargs}
return inputs
@property
def _chain_type(self) -> str:
return "refine_documents_chain"