bisheng-langchain 0.3.1.1__py3-none-any.whl → 0.3.2__py3-none-any.whl

bisheng_langchain/chains/qa_generation/prompt_v2.py

@@ -0,0 +1,155 @@
+from langchain_core.prompts.chat import (
+    ChatPromptTemplate,
+    HumanMessagePromptTemplate,
+    SystemMessagePromptTemplate,
+)
+
+
+SEED_QUESTION_SYSTEM = SystemMessagePromptTemplate.from_template(
+    """\
+您的任务是遵循以下规则从给定的上下文中提出一个问题，规则如下：
+
+    1. 即使在没有给定上下文的情况下，问题也应该对人类有意义。
+    2. 应该可以从给定上下文中完全回答问题。
+    3. 问题应该来自包含重要信息的上下文部分。它也可以来自表格、段落、或者代码等。
+    4. 回答问题时不应包含任何链接。
+    5. 问题的难度应该是中等的。
+    6. 问题必须是合理的，并且必须能被人理解和回答。
+    7. 不要在问题中使用“提供的上下文”等短语。
+    8. 避免使用可以分解成多个问题的“和”字样来构建问题。
+    9. 如果上下文是中文，那么问题也应该是中文的。
+
+Examples:
+context:武汉达梦数据库股份有限公司 招股说明书 （申报稿） 1-1-226 表中作出恰当列报。 2、研发费用 2021年度、 2020年度、 2019 年度，达梦数据 研发费用金额分别 为11,786.99 万元、 9,660.26 万元、 6,255.86万元， 各年度研发费用占营 业收入的比例分别为 15.86 % 、 21.46 %、20.74 %。 由于研发投入金额及其占当期 营业收入的比例是 达梦数据 的关键 指标之一，可能存在因为核算不准 确而导致的错报风险。因此， 中天 运会计师 将研发费用的归集和核算 确定为关键审计事项。 针对研发费用的真实性与准确性，会计师执行的 重要审计程序主要包括： （1）了解与研发费用相关的关键内部控制，评价 这些控制的设计，确定其是否得到执行，并对相关内 部控制的运行有效性进行测试； （2）获取研发项目立项、审批资料，抽查重要研 发项目的过程文档，判断研发项目的真实性； （3）获取研发费用按项目、性质分类明细表，分
+question:达梦2021年的研发费用占营业收入的比例是多少？
+
+context:武汉达梦数据库股份有限公司 招股说明书 （申报稿） 1-1-329 （2）存货周转率 公司与同行业可比公司存货周转率对比情况如下： 公司简称 2021年度 2020年度 2019年度 中望软件 6.93 5.62 10.66 星环科技 3.38 3.21 2.24 金山办公 212.60 175.46 162.91 平均值 74.30 61.43 58.60 本公司 1.13 0.57 0.87 数据来源：可比公司招股说明书、定期报告。 报告期各期， 公司存货周转率显著低于同行业可比公司存货周转率平均水平， 主要是因为公司将未验收的数据及行业解决方案项目所发生的累 计成本均作为 存货核算。报告期各期末，公司存在 “湖北省司法行政数据中心项目 ”、“政法云 大数据中心基础设施服务及大数据中心软件采购 项目”等金额较大且实施周期较 长的数据及行业解决方案项目，导致年末存货金额较大。
+question:达梦2021年的存货周转率相较于前一年有何变化？
+"""  # noqa: E501
+)
+
+
+SEED_QUESTION_HUMAN = HumanMessagePromptTemplate.from_template(
+"""
+context:{context}
+question:
+"""
+)
+
+
+SEED_QUESTION_CHAT_PROMPT = ChatPromptTemplate.from_messages(
+    [
+        SEED_QUESTION_SYSTEM,
+        SEED_QUESTION_HUMAN
+    ]
+)
+
+
+SCORE_CONTEXT_SYSTEM = SystemMessagePromptTemplate.from_template(
+"""Evaluate the provided context and assign a numerical score between 0 and 10 based on the following criteria:
+1. Award a high score to context that thoroughly delves into and explains concepts.
+2. Assign a lower score to context that contains excessive references, acknowledgments, external links, personal information, or other non-essential elements.
+
+And you should only output the score.
+
+Examples:
+Context:
+01-2022.04.30 贷方发生额共 计 2535.43 万元，户名；X 贸易有限公司；\n③根据用款企业提供的增值税纳税申报表来看，2021 年度用款企业年累计开票额为\n7826.48 万元，年累计应纳税合计 95.32 万元，年累计已纳税额 86.23 万元；截止至 2022 年 3 月，用款企业累计开票额为 1986.54 万元，累计应纳税合计19.54 万元，累计已纳税额\n20.23 万元。\n根据核算用款企业的银行流水及企业会计记账系统，剔除借款人往来转账款，估算用款 企业年营业额约在 6000 万元左右(纳税申报营业额)，全部营业收入约 20000 万元左右，借 款人所在 X 贸易有限公司综合毛
+利润率约为 35%，净利润约 20%左右。\n\n| 资产种类 | 坐落 | 产权人 | 建筑面积 | 现价值 | 贷款余额 | 资产净值 |\n| --- | --- | --- | --- | --- | --- | --- |\n| 房产 | HN 省 YY 市 PP 小区 5#2-101 | A | 240.20 | 365.23 万 | 165.
+Score: 4
+
+Context:
+认缴出资额 200 万元 实缴出资额 200 万元 持股比例 20% |\n| 企业所属商圈 | 无 | 是否为已准入商圈 | 是□ 否 ☑ |\n(1) 企业经营历史及现状说明\nX 贸易有限公司 (下称“用款企业”) 注册成立于 2015 年 11 月，统一社会信用代码1234567890ACBDEFGH，法定代表人 A，公司注册地址位于 M 市 N 区 JF 路 20 号 NJ 大厦 18 楼1807 室，实际办公地址位于 M 市 N 区 K 广场 C 座 19 楼 1901、1906、1908、1910、1912、1914，办公面积为 880.51 ㎡，经营场所为用款企业租赁房产，租赁期限，现阶段年租金 73 万余元。\n用款企业是著名品牌“XYZ”的运营公司，是
+以经营短袜、连裤袜、 内衣、家居服、配饰为主要品类的亲体织物公司，致力于为年轻消费群体提供“一站式”多品类亲体织物购物 体验。 作为织物文化的传播者和输出者，用款企业秉承一贯的高品质与原创精神，依托中国 研发团队，创领多项核心技术，不断建立并升级健康织物行业标准，目前拥有实用新型专利 6 项，发明专利 1 项，注册商标 30 余个，为品牌的商标保护构建了全面的商标防御体系。\n“XYZ”品牌创立于 2006 年，于 2009 年正式进入中国市场，在成立 10 年的时间里，在全国共有 400 余家店面，运营主要有以下三种模式：\n①直营模式：目前用款企业
+管控的直营店有 100 家左右，其中在 M 地区共有 9 家直营店，分别为 Y1 店、Y2 店、Y3 店、Y4 店、Y5 店、Y6 店、Y7 店、Y8 店、Y9 店。经查看用款企业相关财务系统并截屏 ，用款企业 2021 年度 、2022 年 1-4 月直营店营业收入合计分别为7623.45 万元、1987.23 万元，M 地区 9 家直营店收入合计分别为 1238.67 万元、302.54 万元。根据数据测算直营部分毛利润率65%。
+Score: 7
+"""  # noqa: E501
+)
+
+
+SCORE_CONTEXT_HUMAN = HumanMessagePromptTemplate.from_template(
+"""
+Context:
+{context}
+Score:
+"""  # noqa: E501
+)
+
+
+SCORE_CONTEXT_CHAT_PROMPT = ChatPromptTemplate.from_messages(
+    [
+        SCORE_CONTEXT_SYSTEM,
+        SCORE_CONTEXT_HUMAN
+    ]
+)
+
+
+FILTER_QUESTION_SYSTEM = SystemMessagePromptTemplate.from_template(
+    """\
+Determine if the given question can be clearly understood even when presented without any additional context. Specify reason and verdict is a valid json format.
+
+Examples:
+question: What is the discovery about space?
+{{
+    "reason":"The question is too vague and does not specify which discovery about space it is referring to."
+    "verdit":"No"
+}}
+
+question: What caused the Great Depression?
+{{
+    "reason":"The question is specific and refers to a well-known historical economic event, making it clear and answerable.",
+    "verdict":"Yes"
+}}
+
+question: What is the keyword that best describes the paper's focus in natural language understanding tasks?
+{{
+    "reason": "The question mentions a 'paper' in it without referring it's name which makes it unclear without it",
+    "verdict": "No"
+}}
+
+question: Who wrote 'Romeo and Juliet'?
+{{
+    "reason": "The question is clear and refers to a specific work by name therefore it is clear",
+    "verdict": "Yes"
+}}
+
+question: What did the study mention?
+{{
+    "reason": "The question is vague and does not specify which study it is referring to",
+    "verdict": "No"
+}}
+
+question: What is the focus of the REPLUG paper?
+{{
+    "reason": "The question refers to a specific work by it's name hence can be understood", 
+    "verdict": "Yes"
+}}
+"""  # noqa: E501
+)
+
+
+FILTER_QUESTION_HUMAN = HumanMessagePromptTemplate.from_template(
+    """\
+question:{question}
+"""  # noqa: E501
+)
+
+
+FILTER_QUESTION_CHAT_PROMPT = ChatPromptTemplate.from_messages(
+    [
+        FILTER_QUESTION_SYSTEM,
+        FILTER_QUESTION_HUMAN
+    ]
+)
+
+
+ANSWER_FORMULATE = HumanMessagePromptTemplate.from_template(
+    """\
+Answer the question using the information from the given context. 
+
+context:{context}
+
+question:{question}
+answer:
+"""  # noqa: E501
+)

bisheng_langchain/document_loaders/elem_unstrcutured_loader.py

@@ -128,16 +128,43 @@ class ElemUnstructuredLoaderV0(BasePDFLoader):
         super().__init__(file_path)
 
     def load(self) -> List[Document]:
+        page_content, metadata = self.get_text_metadata()
+        doc = Document(page_content=page_content, metadata=metadata)
+        return [doc]
+
+    def get_text_metadata(self):
         b64_data = base64.b64encode(open(self.file_path, 'rb').read()).decode()
         payload = dict(filename=os.path.basename(self.file_name), b64_data=[b64_data], mode='text')
         payload.update({'start': self.start, 'n': self.n})
         payload.update(self.extra_kwargs)
-        resp = requests.post(self.unstructured_api_url, headers=self.headers, json=payload).json()
-
-        if 200 != resp.get('status_code'):
-            logger.info(f'not return resp={resp}')
-
-        page_content = resp['text']
-        meta = {'source': self.file_name}
-        doc = Document(page_content=page_content, metadata=meta)
-        return [doc]
+        resp = requests.post(self.unstructured_api_url, headers=self.headers, json=payload)
+        # 说明文件解析成功
+        if resp.status_code == 200 and resp.json().get('status_code') == 200:
+            res = resp.json()
+            return res['text'], {'source': self.file_name}
+        # 说明文件解析失败，pdf文件直接返回报错
+        if self.file_name.endswith('.pdf'):
+            raise Exception(f'file text {os.path.basename(self.file_name)} failed resp={resp.text}')
+        # 非pdf文件，先将文件转为pdf格式，让后再执行partition模式解析文档
+        # 把文件转为pdf
+        resp = requests.post(self.unstructured_api_url, headers=self.headers, json={
+            'filename': os.path.basename(self.file_name),
+            'b64_data': [b64_data],
+            'mode': 'topdf',
+        })
+        if resp.status_code != 200 or resp.json().get('status_code') != 200:
+            raise Exception(f'file topdf {os.path.basename(self.file_name)} failed resp={resp.text}')
+        # 解析pdf文件
+        payload['mode'] = 'partition'
+        payload['b64_data'] = [resp.json()['b64_pdf']]
+        payload['filename'] = os.path.basename(self.file_name) + '.pdf'
+        resp = requests.post(self.unstructured_api_url, headers=self.headers, json=payload)
+        if resp.status_code != 200 or resp.json().get('status_code') != 200:
+            raise Exception(f'file partition {os.path.basename(self.file_name)} failed resp={resp.text}')
+        res = resp.json()
+        partitions = res['partitions']
+        if not partitions:
+            raise Exception(f'file partition empty {os.path.basename(self.file_name)} resp={resp.text}')
+        # 拼接结果为文本
+        content, _ = merge_partitions(partitions)
+        return content, {'source': self.file_name}

bisheng_langchain/document_loaders/parsers/ellm_client.py

@@ -1,13 +1,13 @@
 # import base64
 import copy
-import base64
+from typing import Optional
+
 import requests
-from typing import Any, Iterator, List, Mapping, Optional, Union
 
 
 class ELLMClient(object):
-    def __init__(self,
-                 api_base_url: Optional[str] = None):
+
+    def __init__(self, api_base_url: Optional[str] = None):
         self.ep = api_base_url
         self.client = requests.Session()
         self.timeout = 10000
@@ -26,8 +26,8 @@ class ELLMClient(object):
                 'ellm': 'ELLM'
             },
             'form': {
-                'det': 'mrcnn-v5.1',
-                'recog': 'transformer-v2.8-gamma-faster',
+                'det': 'general_text_det_v2.0',
+                'recog': 'general_text_reg_nb_v1.0_faster',
                 'ellm': 'ELLM'
             },
             'hand': {
@@ -48,9 +48,7 @@ class ELLMClient(object):
         req_data = {'data': [b64_image], 'param': params}
 
         try:
-            r = self.client.post(url=self.ep,
-                                 json=req_data,
-                                 timeout=self.timeout)
+            r = self.client.post(url=self.ep, json=req_data, timeout=self.timeout)
             return r.json()
         except Exception as e:
             return {'status_code': 400, 'status_message': str(e)}

bisheng_langchain/document_loaders/universal_kv.py

@@ -47,6 +47,7 @@ def transpdf2png(pdf_file):
 class UniversalKVLoader(BaseLoader):
     """Extract key-value from pdf or image.
     """
+
     def __init__(self,
                  file_path: str,
                  ellm_model_url: str = None,
@@ -83,7 +84,7 @@ class UniversalKVLoader(BaseLoader):
 
             kv_results = defaultdict(list)
             for key, value in key_values.items():
-                kv_results[key] = value['text']
+                kv_results[key].extend([v['text'] for v in value])
 
             content = json.dumps(kv_results, indent=2, ensure_ascii=False)
             file_name = os.path.basename(self.file_path)
@@ -95,7 +96,7 @@ class UniversalKVLoader(BaseLoader):
             pdf_images = transpdf2png(self.file_path)
 
             kv_results = defaultdict(list)
-            for pdf_name in pdf_images:
+            for index, pdf_name in enumerate(pdf_images):
                 page = int(pdf_name.split('page_')[-1])
                 if page > self.max_pages:
                     continue
@@ -110,7 +111,7 @@ class UniversalKVLoader(BaseLoader):
                     raise ValueError(f'universal kv load failed: {resp}')
 
                 for key, value in key_values.items():
-                    kv_results[key].extend(value['text'])
+                    kv_results[key].extend([v['text'] for v in value])
 
             content = json.dumps(kv_results, indent=2, ensure_ascii=False)
             file_name = os.path.basename(self.file_path)

bisheng_langchain/gpts/tools/api_tools/openapi.py

@@ -13,7 +13,7 @@ class OpenApiTools(APIToolBase):
         return self.url + self.params["path"]
 
     def get_request_method(self):
-        return self.params["method"]
+        return self.params["method"].lower()
 
     def get_params_json(self, **kwargs):
         params_define = {}
@@ -59,11 +59,11 @@ class OpenApiTools(APIToolBase):
         if method == "get":
             resp = self.client.get(path, params=params)
         elif method == 'post':
-            resp = self.client.post(path, params=params, json=self.params)
+            resp = self.client.post(path, params=params, json=json_data)
         elif method == 'put':
-            resp = self.client.put(path, params=params, json=self.params)
+            resp = self.client.put(path, params=params, json=json_data)
         elif method == 'delete':
-            resp = self.client.delete(path, params=params, json=self.params)
+            resp = self.client.delete(path, params=params, json=json_data)
         else:
             raise Exception(f"http method is not support: {method}")
         if resp.status_code != 200:
@@ -81,11 +81,11 @@ class OpenApiTools(APIToolBase):
         if method == "get":
             resp = await self.async_client.aget(path, params=params)
         elif method == 'post':
-            resp = await self.async_client.apost(path, params=params, json=self.params)
+            resp = await self.async_client.apost(path, params=params, json=json_data)
         elif method == 'put':
-            resp = await self.async_client.aput(path, params=params, json=self.params)
+            resp = await self.async_client.aput(path, params=params, json=json_data)
         elif method == 'delete':
-            resp = await self.async_client.adelete(path, params=params, json=self.params)
+            resp = await self.async_client.adelete(path, params=params, json=json_data)
         else:
             raise Exception(f"http method is not support: {method}")
         return resp

bisheng_langchain/rag/__init__.py

@@ -1,5 +1,7 @@
 from bisheng_langchain.rag.bisheng_rag_tool import BishengRAGTool
+from bisheng_langchain.rag.bisheng_rag_chain import BishengRetrievalQA
 
 __all__ = [
     "BishengRAGTool",
+    "BishengRetrievalQA"
 ]

bisheng_langchain/rag/bisheng_rag_chain.py

@@ -0,0 +1,164 @@
+"""Chain for question-answering against a vector database."""
+from __future__ import annotations
+
+import inspect
+from abc import abstractmethod
+from typing import Any, Dict, List, Optional
+
+from langchain_core.callbacks import (
+    AsyncCallbackManagerForChainRun,
+    CallbackManagerForChainRun,
+    Callbacks
+)
+from langchain_core.prompts import PromptTemplate, BasePromptTemplate, ChatPromptTemplate, HumanMessagePromptTemplate, SystemMessagePromptTemplate
+from langchain_core.language_models import BaseLanguageModel
+from langchain_core.pydantic_v1 import Extra, Field
+from bisheng_langchain.vectorstores import ElasticKeywordsSearch, Milvus
+
+from langchain.chains.base import Chain
+from .bisheng_rag_tool import BishengRAGTool
+
+
+# system_template = """Use the following pieces of context to answer the user's question. 
+# If you don't know the answer, just say that you don't know, don't try to make up an answer.
+# ----------------
+# {context}"""
+# messages = [
+#     SystemMessagePromptTemplate.from_template(system_template),
+#     HumanMessagePromptTemplate.from_template("{question}"),
+# ]
+# DEFAULT_QA_PROMPT = ChatPromptTemplate.from_messages(messages)
+
+
+system_template_general = """你是一个准确且可靠的知识库问答助手，能够借助上下文知识回答问题。你需要根据以下的规则来回答问题：
+1. 如果上下文中包含了正确答案，你需要根据上下文进行准确的回答。但是在回答前，你需要注意，上下文中的信息可能存在事实性错误，如果文档中存在和事实不一致的错误，请根据事实回答。
+2. 如果上下文中不包含答案，就说你不知道，不要试图编造答案。
+3. 你需要根据上下文给出详细的回答，不要试图偷懒，不要遗漏括号中的信息，你必须回答的尽可能详细。
+"""
+human_template_general = """
+上下文：
+{context}
+
+问题：
+{question}
+"""
+messages_general = [
+    SystemMessagePromptTemplate.from_template(system_template_general),
+    HumanMessagePromptTemplate.from_template(human_template_general),
+]
+DEFAULT_QA_PROMPT = ChatPromptTemplate.from_messages(messages_general)
+
+
+class BishengRetrievalQA(Chain):
+    """Base class for question-answering chains."""
+
+    """Chain to use to combine the documents."""
+    input_key: str = "query"  #: :meta private:
+    output_key: str = "result"  #: :meta private:
+    return_source_documents: bool = False
+    """Return the source documents or not."""
+    bisheng_rag_tool: BishengRAGTool = Field(
+        default_factory=BishengRAGTool, description="RAG tool"
+    )
+
+    class Config:
+        """Configuration for this pydantic object."""
+
+        extra = Extra.forbid
+        arbitrary_types_allowed = True
+        allow_population_by_field_name = True
+
+    @property
+    def input_keys(self) -> List[str]:
+        """Input keys.
+
+        :meta private:
+        """
+        return [self.input_key]
+
+    @property
+    def output_keys(self) -> List[str]:
+        """Output keys.
+
+        :meta private:
+        """
+        _output_keys = [self.output_key]
+        if self.return_source_documents:
+            _output_keys = _output_keys + ["source_documents"]
+        return _output_keys
+
+    @classmethod
+    def from_llm(
+        cls,
+        llm: BaseLanguageModel,
+        vector_store: Milvus,
+        keyword_store: ElasticKeywordsSearch,
+        QA_PROMPT: ChatPromptTemplate = DEFAULT_QA_PROMPT,
+        max_content: int = 15000,
+        sort_by_source_and_index: bool = False,
+        callbacks: Callbacks = None,
+        **kwargs: Any,
+    ) -> BishengRetrievalQA:
+        bisheng_rag_tool = BishengRAGTool(
+            vector_store=vector_store, 
+            keyword_store=keyword_store,
+            llm=llm,
+            QA_PROMPT=QA_PROMPT,
+            max_content=max_content,
+            sort_by_source_and_index=sort_by_source_and_index,
+            **kwargs
+        )
+        return cls(
+            bisheng_rag_tool=bisheng_rag_tool,
+            callbacks=callbacks,
+            **kwargs,
+        )
+
+    def _call(
+        self,
+        inputs: Dict[str, Any],
+        run_manager: Optional[CallbackManagerForChainRun] = None,
+    ) -> Dict[str, Any]:
+        """Run get_relevant_text and llm on input query.
+
+        If chain has 'return_source_documents' as 'True', returns
+        the retrieved documents as well under the key 'source_documents'.
+
+        Example:
+        .. code-block:: python
+
+        res = indexqa({'query': 'This is my query'})
+        answer, docs = res['result'], res['source_documents']
+        """
+        question = inputs[self.input_key]
+        if self.return_source_documents:
+            answer, docs = self.bisheng_rag_tool.run(question, return_only_outputs=False)
+            return {self.output_key: answer, "source_documents": docs}
+        else:
+            answer = self.bisheng_rag_tool.run(question, return_only_outputs=True)
+            return {self.output_key: answer}
+
+    async def _acall(
+        self,
+        inputs: Dict[str, Any],
+        run_manager: Optional[AsyncCallbackManagerForChainRun] = None,
+    ) -> Dict[str, Any]:
+        """Run get_relevant_text and llm on input query.
+
+        If chain has 'return_source_documents' as 'True', returns
+        the retrieved documents as well under the key 'source_documents'.
+
+        Example:
+        .. code-block:: python
+
+        res = indexqa({'query': 'This is my query'})
+        answer, docs = res['result'], res['source_documents']
+        """
+        question = inputs[self.input_key]
+
+        if self.return_source_documents:
+            answer, docs = await self.bisheng_rag_tool.arun(question, return_only_outputs=False)
+            return {self.output_key: answer, "source_documents": docs}
+        else:
+            answer = await self.bisheng_rag_tool.arun(question, return_only_outputs=True)
+            return {self.output_key: answer}

bisheng_langchain/rag/bisheng_rag_pipeline_v2.py

@@ -44,7 +44,9 @@ class BishengRagPipeline:
         if embedding_params['type'] == 'OpenAIEmbeddings' and embedding_params['openai_proxy']:
             embedding_params.pop('type')
             self.embeddings = embedding_object(
-                http_client=httpx.Client(proxies=embedding_params['openai_proxy']), **embedding_params
+                http_client=httpx.Client(proxies=embedding_params['openai_proxy']),
+                http_async_client=httpx.AsyncClient(proxies=embedding_params['openai_proxy']),
+                **embedding_params
             )
         else:
             embedding_params.pop('type')
@@ -55,7 +57,11 @@ class BishengRagPipeline:
         llm_object = import_by_type(_type='llms', name=llm_params['type'])
         if llm_params['type'] == 'ChatOpenAI' and llm_params['openai_proxy']:
             llm_params.pop('type')
-            self.llm = llm_object(http_client=httpx.Client(proxies=llm_params['openai_proxy']), **llm_params)
+            self.llm = llm_object(
+                http_client=httpx.Client(proxies=llm_params['openai_proxy']),
+                http_async_client=httpx.AsyncClient(proxies=llm_params['openai_proxy']),
+                **llm_params
+            )
         else:
             llm_params.pop('type')
             self.llm = llm_object(**llm_params)

bisheng_langchain/rag/bisheng_rag_tool.py

@@ -3,10 +3,14 @@ import os
 import yaml
 import httpx
 from typing import Any, Dict, Tuple, Type, Union, Optional
+
+from langchain_core.vectorstores import VectorStoreRetriever
 from loguru import logger
 from langchain_core.tools import BaseTool, Tool
 from langchain_core.pydantic_v1 import BaseModel, Extra, Field, root_validator
 from langchain_core.language_models.base import LanguageModelLike
+from langchain_core.prompts import ChatPromptTemplate
+from langchain.chains.llm import LLMChain
 from langchain.chains.question_answering import load_qa_chain
 from bisheng_langchain.retrievers import EnsembleRetriever
 from bisheng_langchain.vectorstores import ElasticKeywordsSearch, Milvus
@@ -39,6 +43,7 @@ class BishengRAGTool:
         keyword_store: Optional[ElasticKeywordsSearch] = None,
         llm: Optional[LanguageModelLike] = None,
         collection_name: Optional[str] = None,
+        QA_PROMPT: Optional[ChatPromptTemplate] = None,
         **kwargs
     ) -> None:
         if collection_name is None and (keyword_store is None or vector_store is None):
@@ -54,10 +59,27 @@ class BishengRAGTool:
         sort_by_source_and_index = kwargs.get("sort_by_source_and_index", True)
         self.params['generate']['max_content'] = max_content
         self.params['post_retrieval']['sort_by_source_and_index'] = sort_by_source_and_index
+
+        # init llm
+        if llm:
+            self.llm = llm
+        else:
+            llm_params = self.params['chat_llm']
+            llm_object = import_by_type(_type='llms', name=llm_params['type'])
+            if llm_params['type'] == 'ChatOpenAI' and llm_params['openai_proxy']:
+                llm_params.pop('type')
+                self.llm = llm_object(http_client=httpx.Client(proxies=llm_params['openai_proxy']), **llm_params)
+            else:
+                llm_params.pop('type')
+                self.llm = llm_object(**llm_params)
         
         # init milvus
         if vector_store:
-            self.vector_store = vector_store
+            # if vector_store is retriever, get vector_store instance
+            if isinstance(vector_store, VectorStoreRetriever):
+                self.vector_store = vector_store.vectorstore
+            else:
+                self.vector_store = vector_store
         else:
             # init embeddings
             embedding_params = self.params['embedding']
@@ -83,24 +105,17 @@ class BishengRAGTool:
         if keyword_store:
             self.keyword_store = keyword_store
         else:
+            if self.params['elasticsearch'].get('extract_key_by_llm', False):
+                extract_key_prompt = import_class(f'bisheng_langchain.rag.prompts.EXTRACT_KEY_PROMPT')
+                llm_chain = LLMChain(llm=self.llm, prompt=extract_key_prompt)
+            else:
+                llm_chain = None
             self.keyword_store = ElasticKeywordsSearch(
                 index_name='default_es',
                 elasticsearch_url=self.params['elasticsearch']['url'],
                 ssl_verify=self.params['elasticsearch']['ssl_verify'],
+                llm_chain=llm_chain,
             )
-        
-        # init llm
-        if llm:
-            self.llm = llm
-        else:
-            llm_params = self.params['chat_llm']
-            llm_object = import_by_type(_type='llms', name=llm_params['type'])
-            if llm_params['type'] == 'ChatOpenAI' and llm_params['openai_proxy']:
-                llm_params.pop('type')
-                self.llm = llm_object(http_client=httpx.Client(proxies=llm_params['openai_proxy']), **llm_params)
-            else:
-                llm_params.pop('type')
-                self.llm = llm_object(**llm_params)
 
         # init retriever
         retriever_list = []
@@ -117,11 +132,14 @@ class BishengRAGTool:
         self.retriever = EnsembleRetriever(retrievers=retriever_list)
 
         # init qa chain    
-        if 'prompt_type' in self.params['generate']:
-            prompt_type = self.params['generate']['prompt_type']
-            prompt = import_class(f'bisheng_langchain.rag.prompts.{prompt_type}')
+        if QA_PROMPT:
+            prompt = QA_PROMPT
         else:
-            prompt = None
+            if 'prompt_type' in self.params['generate']:
+                prompt_type = self.params['generate']['prompt_type']
+                prompt = import_class(f'bisheng_langchain.rag.prompts.{prompt_type}')
+            else:
+                prompt = None
         self.qa_chain = load_qa_chain(
             llm=self.llm, 
             chain_type=self.params['generate']['chain_type'], 
@@ -218,18 +236,23 @@ class BishengRAGTool:
             docs = sorted(docs, key=lambda x: (x.metadata['source'], x.metadata['chunk_index']))
         return docs
     
-    def run(self, query) -> str:
+    def run(self, query, return_only_outputs=True) -> Any:
         docs = self.retrieval_and_rerank(query)
         try:
-            ans = self.qa_chain({"input_documents": docs, "question": query}, return_only_outputs=True)
+            ans = self.qa_chain({"input_documents": docs, "question": query}, return_only_outputs=return_only_outputs)
         except Exception as e:
             logger.error(f'question: {query}\nerror: {e}')
             ans = {'output_text': str(e)}
-        rag_answer = ans['output_text']
-        return rag_answer
+        if return_only_outputs:
+            rag_answer = ans['output_text']
+            return rag_answer
+        else:
+            rag_answer = ans['output_text']
+            input_documents = ans['input_documents']
+            return rag_answer, input_documents
     
-    async def arun(self, query: str) -> str:
-        rag_answer = self.run(query)
+    async def arun(self, query: str, return_only_outputs=True) -> str:
+        rag_answer = self.run(query, return_only_outputs)
         return rag_answer
     
     @classmethod

bisheng_langchain/rag/config/baseline_caibao_v2.yaml

@@ -13,7 +13,7 @@ elasticsearch:
   ssl_verify:
     basic_auth: ["elastic", "oSGL-zVvZ5P3Tm7qkDLC"]
   drop_old: True
-  extract_key_by_llm: False 
+  extract_key_by_llm: False
 
 embedding:
   type: 'OpenAIEmbeddings'