bisheng-langchain 0.3.7.1__py3-none-any.whl → 0.4.0.dev1__py3-none-any.whl

bisheng_langchain/rag/bisheng_rag_chain.py

@@ -1,25 +1,20 @@
 """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 bisheng_langchain.vectorstores import ElasticKeywordsSearch, Milvus
+from langchain.chains.base import Chain
+from langchain_core.callbacks import (AsyncCallbackManagerForChainRun, CallbackManagerForChainRun,
+                                      Callbacks)
 from langchain_core.language_models import BaseLanguageModel
+from langchain_core.prompts import (ChatPromptTemplate, HumanMessagePromptTemplate,
+                                    SystemMessagePromptTemplate)
 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. 
+# 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}"""
@@ -29,7 +24,6 @@ from .bisheng_rag_tool import BishengRAGTool
 # ]
 # DEFAULT_QA_PROMPT = ChatPromptTemplate.from_messages(messages)
 
-
 system_template_general = """你是一个准确且可靠的知识库问答助手，能够借助上下文知识回答问题。你需要根据以下的规则来回答问题：
 1. 如果上下文中包含了正确答案，你需要根据上下文进行准确的回答。但是在回答前，你需要注意，上下文中的信息可能存在事实性错误，如果文档中存在和事实不一致的错误，请根据事实回答。
 2. 如果上下文中不包含答案，就说你不知道，不要试图编造答案。
@@ -51,15 +45,13 @@ 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:
+    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"
-    )
+    bisheng_rag_tool: BishengRAGTool = Field(default_factory=BishengRAGTool,
+                                             description='RAG tool')
 
     class Config:
         """Configuration for this pydantic object."""
@@ -84,7 +76,7 @@ class BishengRetrievalQA(Chain):
         """
         _output_keys = [self.output_key]
         if self.return_source_documents:
-            _output_keys = _output_keys + ["source_documents"]
+            _output_keys = _output_keys + ['source_documents']
         return _output_keys
 
     @classmethod
@@ -100,15 +92,13 @@ class BishengRetrievalQA(Chain):
         return_source_documents: bool = False,
         **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
-        )
+        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,
@@ -134,8 +124,12 @@ class BishengRetrievalQA(Chain):
         """
         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}
+            answer, docs = self.bisheng_rag_tool.run(
+                question,
+                return_only_outputs=False,
+                run_manager=run_manager,
+            )
+            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}
@@ -160,7 +154,7 @@ class BishengRetrievalQA(Chain):
 
         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}
+            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_tool.py

@@ -1,27 +1,25 @@
-import time
 import os
-import yaml
-import httpx
-from typing import Any, Dict, Tuple, Type, Union, Optional
+from typing import Any, Dict, Optional, Tuple, Union
 
-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
+import httpx
+import yaml
+from bisheng_langchain.rag.extract_info import extract_title
+from bisheng_langchain.rag.init_retrievers import (BaselineVectorRetriever, KeywordRetriever,
+                                                   MixRetriever, SmallerChunksVectorRetriever)
+from bisheng_langchain.rag.utils import import_by_type, import_class
 from bisheng_langchain.retrievers import EnsembleRetriever
 from bisheng_langchain.vectorstores import ElasticKeywordsSearch, Milvus
-from bisheng_langchain.rag.init_retrievers import (
-    BaselineVectorRetriever,
-    KeywordRetriever,
-    MixRetriever,
-    SmallerChunksVectorRetriever,
-)
-from bisheng_langchain.rag.utils import import_by_type, import_class
-from bisheng_langchain.rag.extract_info import extract_title
+from langchain.chains.llm import LLMChain
+from langchain.chains.question_answering import load_qa_chain
+from langchain.chains.combine_documents import create_stuff_documents_chain
+from langchain_core.callbacks import CallbackManagerForChainRun
+from langchain_core.language_models.base import LanguageModelLike
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.pydantic_v1 import BaseModel, Field
+from langchain_core.runnables import RunnableConfig
+from langchain_core.tools import BaseTool, Tool
+from langchain_core.vectorstores import VectorStoreRetriever
+from loguru import logger
 
 
 class MultArgsSchemaTool(Tool):
@@ -37,26 +35,27 @@ class MultArgsSchemaTool(Tool):
 
 class BishengRAGTool:
 
-    def __init__(
-        self,
-        vector_store: Optional[Milvus] = None,
-        keyword_store: Optional[ElasticKeywordsSearch] = None,
-        llm: Optional[LanguageModelLike] = None,
-        collection_name: Optional[str] = None,
-        QA_PROMPT: Optional[ChatPromptTemplate] = None,
-        **kwargs
-    ) -> None:
+    def __init__(self,
+                 vector_store: Optional[Milvus] = None,
+                 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):
-            raise ValueError('collection_name must be provided if keyword_store or vector_store is not provided')
+            raise ValueError(
+                'collection_name must be provided if keyword_store or vector_store is not provided'
+            )
         self.collection_name = collection_name
-        
-        yaml_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'config/baseline_v2.yaml')
+
+        yaml_path = os.path.join(os.path.dirname(os.path.abspath(__file__)),
+                                 'config/baseline_v2.yaml')
         with open(yaml_path, 'r', encoding='utf-8') as f:
             self.params = yaml.safe_load(f)
-        
+
         # update params
-        max_content = kwargs.get("max_content", 15000)
-        sort_by_source_and_index = kwargs.get("sort_by_source_and_index", True)
+        max_content = kwargs.get('max_content', 15000)
+        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
 
@@ -68,11 +67,12 @@ class BishengRAGTool:
             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']),
+                                      **llm_params)
             else:
                 llm_params.pop('type')
                 self.llm = llm_object(**llm_params)
-        
+
         # init milvus
         if vector_store:
             # if vector_store is retriever, get vector_store instance
@@ -87,26 +87,27 @@ class BishengRAGTool:
             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']),
+                    **embedding_params)
             else:
                 embedding_params.pop('type')
                 self.embeddings = embedding_object(**embedding_params)
-            
+
             self.vector_store = Milvus(
                 embedding_function=self.embeddings,
                 connection_args={
-                    "host": self.params['milvus']['host'],
-                    "port": self.params['milvus']['port'],
+                    'host': self.params['milvus']['host'],
+                    'port': self.params['milvus']['port'],
                 },
             )
-        
+
         # init keyword store
         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')
+                extract_key_prompt = import_class(
+                    'bisheng_langchain.rag.prompts.EXTRACT_KEY_PROMPT')
                 llm_chain = LLMChain(llm=self.llm, prompt=extract_key_prompt)
             else:
                 llm_chain = None
@@ -128,10 +129,11 @@ class BishengRAGTool:
                 'splitter_kwargs': retriever['splitter'],
                 'retrieval_kwargs': retriever['retrieval'],
             }
-            retriever_list.append(self._post_init_retriever(retriever_type=retriever_type, **retriever_params))
+            retriever_list.append(
+                self._post_init_retriever(retriever_type=retriever_type, **retriever_params))
         self.retriever = EnsembleRetriever(retrievers=retriever_list)
 
-        # init qa chain    
+        # init qa chain
         if QA_PROMPT:
             prompt = QA_PROMPT
         else:
@@ -140,13 +142,8 @@ class BishengRAGTool:
                 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'], 
-            prompt=prompt, 
-            verbose=False
-        )
-    
+        self.qa_chain = create_stuff_documents_chain(llm=self.llm, prompt=prompt)
+
     def _post_init_retriever(self, retriever_type, **kwargs):
         retriever_classes = {
             'KeywordRetriever': KeywordRetriever,
@@ -181,9 +178,9 @@ class BishengRAGTool:
         loader_object = import_by_type(_type='documentloaders', name=loader_params.pop('type'))
 
         logger.info(f'file_path: {file_path}')
-        loader = loader_object(
-            file_name=os.path.basename(file_path), file_path=file_path, **loader_params
-        )
+        loader = loader_object(file_name=os.path.basename(file_path),
+                               file_path=file_path,
+                               **loader_params)
         documents = loader.load()
         logger.info(f'documents: {len(documents)}, page_content: {len(documents[0].page_content)}')
         if len(documents[0].page_content) == 0:
@@ -197,30 +194,26 @@ class BishengRAGTool:
                     title = extract_title(llm=self.llm, text=doc.page_content)
                     logger.info(f'extract title: {title}')
                 except Exception as e:
-                    logger.error(f"Failed to extract title: {e}")
+                    logger.error(f'Failed to extract title: {e}')
                     title = ''
                 doc.metadata['title'] = title
 
         for idx, retriever in enumerate(self.retriever.retrievers):
-            retriever.add_documents(
-                documents, 
-                self.collection_name, 
-                drop_old=drop_old, 
-                add_aux_info=add_aux_info
-            )
-    
+            retriever.add_documents(documents,
+                                    self.collection_name,
+                                    drop_old=drop_old,
+                                    add_aux_info=add_aux_info)
+
     def retrieval_and_rerank(self, query):
         """
         retrieval and rerank
         """
         # EnsembleRetriever直接检索召回会默认去重
-        docs = self.retriever.get_relevant_documents(
-            query=query, 
-            collection_name=self.collection_name
-        )
+        docs = self.retriever.get_relevant_documents(query=query,
+                                                     collection_name=self.collection_name)
         logger.info(f'retrieval docs origin: {len(docs)}')
 
-        # delete redundancy according to max_content 
+        # delete redundancy according to max_content
         doc_num, doc_content_sum = 0, 0
         for doc in docs:
             doc_content_sum += len(doc.page_content)
@@ -235,28 +228,37 @@ class BishengRAGTool:
             logger.info('sort chunks by source and chunk_index')
             docs = sorted(docs, key=lambda x: (x.metadata['source'], x.metadata['chunk_index']))
         return docs
-    
-    def run(self, query, return_only_outputs=True) -> Any:
+
+    def run(self,
+            query,
+            return_only_outputs=True,
+            run_manager: Optional[CallbackManagerForChainRun] = None) -> Any:
         docs = self.retrieval_and_rerank(query)
         try:
-            ans = self.qa_chain({"input_documents": docs, "question": query}, return_only_outputs=return_only_outputs)
+            kwargs = {}
+            if run_manager:
+                kwargs['config'] = RunnableConfig(callbacks=[run_manager])
+            ans = self.qa_chain.invoke(
+                {
+                    'context': docs,
+                    'question': query
+                }, **kwargs
+            )
         except Exception as e:
-            logger.error(f'question: {query}\nerror: {e}')
-            ans = {'output_text': str(e)}
+            logger.exception(f'question: {query}\nerror: {e}')
+            ans = str(e)
         if return_only_outputs:
-            rag_answer = ans['output_text']
-            return rag_answer
+            return ans
         else:
-            rag_answer = ans['output_text']
-            input_documents = ans['input_documents']
-            return rag_answer, input_documents
-    
+            return ans, docs
+
     async def arun(self, query: str, return_only_outputs=True) -> str:
         rag_answer = self.run(query, return_only_outputs)
         return rag_answer
-    
+
     @classmethod
     def get_rag_tool(cls, name, description, **kwargs: Any) -> BaseTool:
+
         class InputArgs(BaseModel):
             query: str = Field(description='question asked by the user.')
 
@@ -265,7 +267,7 @@ class BishengRAGTool:
                                   func=cls(**kwargs).run,
                                   coroutine=cls(**kwargs).arun,
                                   args_schema=InputArgs)
-    
+
 
 if __name__ == '__main__':
     # rag_tool = BishengRAGTool(collection_name='rag_finance_report_0_test')
@@ -280,32 +282,30 @@ if __name__ == '__main__':
     collection_name = 'rag_finance_report_0_benchmark_caibao_1000_source_title'
     # milvus
     vector_store = Milvus(
-            collection_name=collection_name,
-            embedding_function=embeddings,
-            connection_args={
-                "host": '110.16.193.170',
-                "port": '50062',
-            },
+        collection_name=collection_name,
+        embedding_function=embeddings,
+        connection_args={
+            'host': '110.16.193.170',
+            'port': '50062',
+        },
     )
     # es
     keyword_store = ElasticKeywordsSearch(
         index_name=collection_name,
         elasticsearch_url='http://110.16.193.170:50062/es',
-        ssl_verify={'basic_auth': ["elastic", "oSGL-zVvZ5P3Tm7qkDLC"]},
+        ssl_verify={'basic_auth': ['elastic', 'oSGL-zVvZ5P3Tm7qkDLC']},
     )
 
-    tool = BishengRAGTool.get_rag_tool(
-        name='rag_knowledge_retrieve', 
-        description='金融年报财报知识库问答',
-        vector_store=vector_store, 
-        keyword_store=keyword_store, 
-        llm=llm
-    )
+    tool = BishengRAGTool.get_rag_tool(name='rag_knowledge_retrieve',
+                                       description='金融年报财报知识库问答',
+                                       vector_store=vector_store,
+                                       keyword_store=keyword_store,
+                                       llm=llm)
     print(tool.run('能否根据2020年金宇生物技术股份有限公司的年报，给我简要介绍一下报告期内公司的社会责任工作情况？'))
 
     # tool = BishengRAGTool.get_rag_tool(
-    #     name='rag_knowledge_retrieve', 
+    #     name='rag_knowledge_retrieve',
     #     description='金融年报财报知识库问答',
     #     collection_name='rag_finance_report_0_benchmark_caibao_1000_source_title'
     # )
-    # print(tool.run('能否根据2020年金宇生物技术股份有限公司的年报，给我简要介绍一下报告期内公司的社会责任工作情况？'))
+    # print(tool.run('能否根据2020年金宇生物技术股份有限公司的年报，给我简要介绍一下报告期内公司的社会责任工作情况？'))

bisheng_langchain/rag/extract_info.py

@@ -1,5 +1,3 @@
-import httpx
-from langchain.chat_models import ChatOpenAI
 from bisheng_langchain.chat_models import ChatQWen
 from langchain.chains import LLMChain
 from langchain.prompts.chat import (

bisheng_langchain/rag/init_retrievers/baseline_vector_retriever.py

@@ -1,20 +1,15 @@
-import os
-import uuid
-from loguru import logger
-from typing import Any, Dict, Iterable, List, Optional
+from typing import Any, List, Optional
 
-from bisheng_langchain.vectorstores.milvus import Milvus
+from langchain.text_splitter import TextSplitter
 from langchain_core.documents import Document
 from langchain_core.pydantic_v1 import Field
 from langchain_core.retrievers import BaseRetriever
-from langchain_core.vectorstores import VectorStore
-
-from langchain.callbacks.manager import CallbackManagerForRetrieverRun
-from langchain.text_splitter import TextSplitter
+from loguru import logger
 
 
 class BaselineVectorRetriever(BaseRetriever):
-    vector_store: Milvus
+
+    vector_store: Any
     text_splitter: TextSplitter
     search_type: str = 'similarity'
     search_kwargs: dict = Field(default_factory=dict)
@@ -27,13 +22,14 @@ class BaselineVectorRetriever(BaseRetriever):
         **kwargs,
     ) -> None:
         split_docs = self.text_splitter.split_documents(documents)
-        logger.info(f"BaselineVectorRetriever: split document into {len(split_docs)} chunks")
+        logger.info(f'BaselineVectorRetriever: split document into {len(split_docs)} chunks')
         for chunk_index, split_doc in enumerate(split_docs):
             if 'chunk_bboxes' in split_doc.metadata:
                 split_doc.metadata.pop('chunk_bboxes')
             split_doc.metadata['chunk_index'] = chunk_index
             if kwargs.get('add_aux_info', False):
-                split_doc.page_content = split_doc.metadata["source"] + '\n' + split_doc.metadata["title"] + '\n' + split_doc.page_content
+                split_doc.page_content = split_doc.metadata['source'] + '\n' + split_doc.metadata[
+                    'title'] + '\n' + split_doc.page_content
 
         connection_args = self.vector_store.connection_args
         embedding_function = self.vector_store.embedding_func

bisheng_langchain/rag/init_retrievers/keyword_retriever.py

@@ -1,22 +1,14 @@
-import os
-import uuid
-from loguru import logger
-from dataclasses import dataclass
-from typing import Any, Dict, Iterable, List, Optional
+from typing import Any, List, Optional
 
-from bisheng_langchain.vectorstores import ElasticKeywordsSearch
-from bisheng_langchain.vectorstores.milvus import Milvus
+from langchain.text_splitter import TextSplitter
 from langchain_core.documents import Document
 from langchain_core.pydantic_v1 import Field
 from langchain_core.retrievers import BaseRetriever
-from langchain_core.vectorstores import VectorStore
-
-from langchain.callbacks.manager import CallbackManagerForRetrieverRun
-from langchain.text_splitter import TextSplitter
+from loguru import logger
 
 
 class KeywordRetriever(BaseRetriever):
-    keyword_store: ElasticKeywordsSearch
+    keyword_store: Any
     text_splitter: TextSplitter
     search_type: str = 'similarity'
     search_kwargs: dict = Field(default_factory=dict)
@@ -29,13 +21,14 @@ class KeywordRetriever(BaseRetriever):
         **kwargs,
     ) -> None:
         split_docs = self.text_splitter.split_documents(documents)
-        logger.info(f"KeywordRetriever: split document into {len(split_docs)} chunks")
+        logger.info(f'KeywordRetriever: split document into {len(split_docs)} chunks')
         for chunk_index, split_doc in enumerate(split_docs):
             if 'chunk_bboxes' in split_doc.metadata:
                 split_doc.metadata.pop('chunk_bboxes')
             split_doc.metadata['chunk_index'] = chunk_index
             if kwargs.get('add_aux_info', False):
-                split_doc.page_content = split_doc.metadata["source"] + '\n' + split_doc.metadata["title"] + '\n' + split_doc.page_content
+                split_doc.page_content = split_doc.metadata['source'] + '\n' + split_doc.metadata[
+                    'title'] + '\n' + split_doc.page_content
 
         elasticsearch_url = self.keyword_store.elasticsearch_url
         ssl_verify = self.keyword_store.ssl_verify
@@ -58,8 +51,7 @@ class KeywordRetriever(BaseRetriever):
                 index_name=collection_name,
                 elasticsearch_url=self.keyword_store.elasticsearch_url,
                 ssl_verify=self.keyword_store.ssl_verify,
-                llm_chain=self.keyword_store.llm_chain
-            )
+                llm_chain=self.keyword_store.llm_chain)
         if self.search_type == 'similarity':
             result = self.keyword_store.similarity_search(query, **self.search_kwargs)
         return result

bisheng_langchain/rag/init_retrievers/mix_retriever.py

@@ -1,17 +1,14 @@
-from typing import Any, Dict, Iterable, List, Optional
+from typing import Any, List, Optional
 
 from bisheng_langchain.vectorstores import ElasticKeywordsSearch
-from bisheng_langchain.vectorstores.milvus import Milvus
+from langchain.text_splitter import TextSplitter
 from langchain_core.documents import Document
 from langchain_core.pydantic_v1 import Field
 from langchain_core.retrievers import BaseRetriever
 
-from langchain.schema import BaseRetriever, Document
-from langchain.text_splitter import TextSplitter
-
 
 class MixRetriever(BaseRetriever):
-    vector_store: Milvus
+    vector_store: Any
     keyword_store: ElasticKeywordsSearch
     vector_text_splitter: TextSplitter
     keyword_text_splitter: TextSplitter
@@ -34,14 +31,16 @@ class MixRetriever(BaseRetriever):
                 split_doc.metadata.pop('chunk_bboxes')
             split_doc.metadata['chunk_index'] = chunk_index
             if kwargs.get('add_aux_info', False):
-                split_doc.page_content = split_doc.metadata["source"] + '\n' + split_doc.metadata["title"] + '\n' + split_doc.page_content
+                split_doc.page_content = split_doc.metadata['source'] + '\n' + split_doc.metadata[
+                    'title'] + '\n' + split_doc.page_content
         keyword_split_docs = self.keyword_text_splitter.split_documents(documents)
         for chunk_index, split_doc in enumerate(keyword_split_docs):
             if 'chunk_bboxes' in split_doc.metadata:
                 split_doc.metadata.pop('chunk_bboxes')
             split_doc.metadata['chunk_index'] = chunk_index
             if kwargs.get('add_aux_info', False):
-                split_doc.page_content = split_doc.metadata["source"] + '\n' + split_doc.metadata["title"] + '\n' + split_doc.page_content
+                split_doc.page_content = split_doc.metadata['source'] + '\n' + split_doc.metadata[
+                    'title'] + '\n' + split_doc.page_content
 
         self.keyword_store.from_documents(
             keyword_split_docs,
@@ -70,15 +69,15 @@ class MixRetriever(BaseRetriever):
                 index_name=collection_name,
                 elasticsearch_url=self.keyword_store.elasticsearch_url,
                 ssl_verify=self.keyword_store.ssl_verify,
-                llm_chain=self.keyword_store.llm_chain
-            )
+                llm_chain=self.keyword_store.llm_chain)
             self.vector_store = self.vector_store.__class__(
                 collection_name=collection_name,
                 embedding_function=self.vector_store.embedding_func,
                 connection_args=self.vector_store.connection_args,
             )
         if self.search_type == 'similarity':
-            keyword_docs = self.keyword_store.similarity_search(query, **self.keyword_search_kwargs)
+            keyword_docs = self.keyword_store.similarity_search(query,
+                                                                **self.keyword_search_kwargs)
             vector_docs = self.vector_store.similarity_search(query, **self.vector_search_kwargs)
             if self.combine_strategy == 'keyword_front':
                 return keyword_docs + vector_docs
@@ -94,10 +93,10 @@ class MixRetriever(BaseRetriever):
                 combine_docs.extend(vector_docs[min_len:])
                 return combine_docs
             else:
-                raise ValueError(
-                    f'Expected combine_strategy to be one of '
-                    f'(keyword_front, vector_front, mix),'
-                    f'instead found {self.combine_strategy}'
-                )
+                raise ValueError(f'Expected combine_strategy to be one of '
+                                 f'(keyword_front, vector_front, mix),'
+                                 f'instead found {self.combine_strategy}')
         else:
-            raise ValueError(f'Expected search_type to be one of (similarity), instead found {self.search_type}')
+            raise ValueError(
+                f'Expected search_type to be one of (similarity), instead found {self.search_type}'
+            )