Source code for langchain.chains.qa_with_sources.retrieval
"""在索引上使用来源进行问答。"""
from typing import Any, Dict, List
from langchain_core.callbacks import (
AsyncCallbackManagerForChainRun,
CallbackManagerForChainRun,
)
from langchain_core.documents import Document
from langchain_core.pydantic_v1 import Field
from langchain_core.retrievers import BaseRetriever
from langchain.chains.combine_documents.stuff import StuffDocumentsChain
from langchain.chains.qa_with_sources.base import BaseQAWithSourcesChain
[docs]class RetrievalQAWithSourcesChain(BaseQAWithSourcesChain):
"""在索引上使用来源进行问答。"""
retriever: BaseRetriever = Field(exclude=True)
"""连接的索引。"""
reduce_k_below_max_tokens: bool = False
"""根据令牌限制减少从存储中返回的结果数量"""
max_tokens_limit: int = 3375
"""根据令牌限制从存储返回的文档,
仅对StuffDocumentChain强制执行,如果reduce_k_below_max_tokens设置为true。"""
def _reduce_tokens_below_limit(self, docs: List[Document]) -> List[Document]:
num_docs = len(docs)
if self.reduce_k_below_max_tokens and isinstance(
self.combine_documents_chain, StuffDocumentsChain
):
tokens = [
self.combine_documents_chain.llm_chain._get_num_tokens(doc.page_content)
for doc in docs
]
token_count = sum(tokens[:num_docs])
while token_count > self.max_tokens_limit:
num_docs -= 1
token_count -= tokens[num_docs]
return docs[:num_docs]
def _get_docs(
self, inputs: Dict[str, Any], *, run_manager: CallbackManagerForChainRun
) -> List[Document]:
question = inputs[self.question_key]
docs = self.retriever.invoke(
question, config={"callbacks": run_manager.get_child()}
)
return self._reduce_tokens_below_limit(docs)
async def _aget_docs(
self, inputs: Dict[str, Any], *, run_manager: AsyncCallbackManagerForChainRun
) -> List[Document]:
question = inputs[self.question_key]
docs = await self.retriever.ainvoke(
question, config={"callbacks": run_manager.get_child()}
)
return self._reduce_tokens_below_limit(docs)
@property
def _chain_type(self) -> str:
"""返回链的类型。"""
return "retrieval_qa_with_sources_chain"