Source code for langchain.chains.qa_with_sources.vector_db

"""使用向量数据库进行带来源的问答。"""

import warnings
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, root_validator
from langchain_core.vectorstores import VectorStore

from langchain.chains.combine_documents.stuff import StuffDocumentsChain
from langchain.chains.qa_with_sources.base import BaseQAWithSourcesChain


class VectorDBQAWithSourcesChain(BaseQAWithSourcesChain):
    """使用向量数据库进行带来源的问答。"""

    vectorstore: VectorStore = Field(exclude=True)
    """连接到的向量数据库。"""
    k: int = 4
    """从商店返回的结果数量"""
    reduce_k_below_max_tokens: bool = False
    """根据令牌限制减少从存储返回的结果数量"""
    max_tokens_limit: int = 3375
    """根据令牌限制从存储返回的文档，仅对StuffDocumentChain生效，如果reduce_k_below_max_tokens设置为true。"""
    search_kwargs: Dict[str, Any] = Field(default_factory=dict)
    """额外的搜索参数。"""

    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.vectorstore.similarity_search(
            question, k=self.k, **self.search_kwargs
        )
        return self._reduce_tokens_below_limit(docs)

    async def _aget_docs(
        self, inputs: Dict[str, Any], *, run_manager: AsyncCallbackManagerForChainRun
    ) -> List[Document]:
        raise NotImplementedError("VectorDBQAWithSourcesChain does not support async")

    @root_validator()
    def raise_deprecation(cls, values: Dict) -> Dict:
        warnings.warn(
            "`VectorDBQAWithSourcesChain` is deprecated - "
            "please use `from langchain.chains import RetrievalQAWithSourcesChain`"
        )
        return values

    @property
    def _chain_type(self) -> str:
        return "vector_db_qa_with_sources_chain"