bisheng-langchain 0.3.0rc0__py3-none-any.whl → 0.3.1__py3-none-any.whl

bisheng_langchain/rag/bisheng_rag_pipeline_v2_cohere_raw_prompting.py

@@ -0,0 +1,376 @@
+import argparse
+import copy
+import time
+import inspect
+import os
+import httpx
+import pandas as pd
+import yaml
+import math
+from tqdm import tqdm
+from loguru import logger
+from collections import defaultdict
+from bisheng_langchain.retrievers import EnsembleRetriever
+from bisheng_langchain.vectorstores import ElasticKeywordsSearch, Milvus
+from langchain.docstore.document import Document
+from langchain.chains.question_answering import load_qa_chain
+from langchain_community.chat_models.cohere import ChatCohere
+from bisheng_langchain.rag.init_retrievers import (
+    BaselineVectorRetriever,
+    KeywordRetriever,
+    MixRetriever,
+    SmallerChunksVectorRetriever,
+)
+from bisheng_langchain.rag.scoring.ragas_score import RagScore
+from bisheng_langchain.rag.extract_info import extract_title
+from bisheng_langchain.rag.utils import import_by_type, import_class
+
+
+class BishengRagPipeline:
+
+    def __init__(self, yaml_path) -> None:
+        self.yaml_path = yaml_path
+        with open(self.yaml_path, 'r') as f:
+            self.params = yaml.safe_load(f)
+
+        # init data
+        self.origin_file_path = self.params['data']['origin_file_path']
+        self.question_path = self.params['data']['question']
+        self.save_answer_path = self.params['data']['save_answer']
+
+        # init embeddings
+        embedding_params = self.params['embedding']
+        embedding_object = import_by_type(_type='embeddings', name=embedding_params['type'])
+        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
+            )
+        else:
+            embedding_params.pop('type')
+            self.embeddings = embedding_object(**embedding_params)
+
+        # init llm
+        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)
+
+        # milvus
+        self.vector_store = Milvus(
+            embedding_function=self.embeddings,
+            connection_args={
+                "host": self.params['milvus']['host'],
+                "port": self.params['milvus']['port'],
+            },
+        )
+
+        # es
+        self.keyword_store = ElasticKeywordsSearch(
+            index_name='default_es',
+            elasticsearch_url=self.params['elasticsearch']['url'],
+            ssl_verify=self.params['elasticsearch']['ssl_verify'],
+        )
+
+        # init retriever
+        retriever_list = []
+        retrievers = self.params['retriever']['retrievers']
+        for retriever in retrievers:
+            retriever_type = retriever.pop('type')
+            retriever_params = {
+                'vector_store': self.vector_store,
+                'keyword_store': self.keyword_store,
+                'splitter_kwargs': retriever['splitter'],
+                'retrieval_kwargs': retriever['retrieval'],
+            }
+            retriever_list.append(self._post_init_retriever(retriever_type=retriever_type, **retriever_params))
+        self.retriever = EnsembleRetriever(retrievers=retriever_list)
+
+    def _post_init_retriever(self, retriever_type, **kwargs):
+        retriever_classes = {
+            'KeywordRetriever': KeywordRetriever,
+            'BaselineVectorRetriever': BaselineVectorRetriever,
+            'MixRetriever': MixRetriever,
+            'SmallerChunksVectorRetriever': SmallerChunksVectorRetriever,
+        }
+        if retriever_type not in retriever_classes:
+            raise ValueError(f'Unknown retriever type: {retriever_type}')
+
+        input_kwargs = {}
+        splitter_params = kwargs.pop('splitter_kwargs')
+        for key, value in splitter_params.items():
+            splitter_obj = import_by_type(_type='textsplitters', name=value.pop('type'))
+            input_kwargs[key] = splitter_obj(**value)
+
+        retrieval_params = kwargs.pop('retrieval_kwargs')
+        for key, value in retrieval_params.items():
+            input_kwargs[key] = value
+
+        input_kwargs['vector_store'] = kwargs.pop('vector_store')
+        input_kwargs['keyword_store'] = kwargs.pop('keyword_store')
+
+        retriever_class = retriever_classes[retriever_type]
+        return retriever_class(**input_kwargs)
+
+    def file2knowledge(self):
+        """
+        file to knowledge
+        """
+        df = pd.read_excel(self.question_path)
+        if ('文件名' not in df.columns) or ('知识库名' not in df.columns):
+            raise Exception(f'文件名 or 知识库名 not in {self.question_path}.')
+
+        loader_params = self.params['loader']
+        loader_object = import_by_type(_type='documentloaders', name=loader_params.pop('type'))
+
+        all_questions_info = df.to_dict('records')
+        collectionname2filename = defaultdict(set)
+        for info in all_questions_info:
+            # 存入set，去掉重复的文件名
+            collectionname2filename[info['知识库名']].add(info['文件名'])
+
+        for collection_name in tqdm(collectionname2filename):
+            all_file_paths = []
+            for file_name in collectionname2filename[collection_name]:
+                file_path = os.path.join(self.origin_file_path, file_name)
+                if not os.path.exists(file_path):
+                    raise Exception(f'{file_path} not exists.')
+                # file path可以是文件夹或者单个文件
+                if os.path.isdir(file_path):
+                    # 文件夹包含多个文件
+                    all_file_paths.extend(
+                        [os.path.join(file_path, name) for name in os.listdir(file_path) if not name.startswith('.')]
+                    )
+                else:
+                    # 单个文件
+                    all_file_paths.append(file_path)
+
+            # 当前知识库需要存储的所有文件
+            collection_name = f"{collection_name}_{self.params['retriever']['suffix']}"
+            for index, each_file_path in enumerate(all_file_paths):
+                logger.info(f'each_file_path: {each_file_path}')
+                loader = loader_object(
+                    file_name=os.path.basename(each_file_path), file_path=each_file_path, **loader_params
+                )
+                documents = loader.load()
+
+                # # load from text
+                # if each_file_path.endswith('.pdf'):
+                #     with open(each_file_path.replace('.pdf', '.txt'), 'r') as f:
+                #         content = f.read()
+                #     documents = [Document(page_content=content, metadata={'source': os.path.basename(each_file_path)})]
+                
+                logger.info(f'documents: {len(documents)}, page_content: {len(documents[0].page_content)}')
+                if len(documents[0].page_content) == 0:
+                    logger.error(f'{each_file_path} page_content is empty.')
+
+                # add aux infoerror
+                add_aux_info = self.params['retriever'].get('add_aux_info', False)
+                if add_aux_info:
+                    for doc in documents:
+                        try:
+                            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}")
+                            title = ''
+                        doc.metadata['title'] = title
+
+                vector_drop_old = self.params['milvus']['drop_old'] if index == 0 else False
+                keyword_drop_old = self.params['elasticsearch']['drop_old'] if index == 0 else False
+                for idx, retriever in enumerate(self.retriever.retrievers):
+                    retriever.add_documents(documents, collection_name, vector_drop_old, add_aux_info=add_aux_info)
+
+    def retrieval_and_rerank(self, question, collection_name, max_content=100000):
+        """
+        retrieval and rerank
+        """
+        collection_name = f"{collection_name}_{self.params['retriever']['suffix']}"
+        # EnsembleRetriever直接检索召回会默认去重
+        docs = self.retriever.get_relevant_documents(query=question, collection_name=collection_name)
+        logger.info(f'retrieval docs origin: {len(docs)}')
+
+        # delete duplicate
+        if self.params['post_retrieval']['delete_duplicate']:
+            logger.info(f'origin docs: {len(docs)}')
+            all_contents = []
+            docs_no_dup = []
+            for index, doc in enumerate(docs):
+                doc_content = doc.page_content
+                if doc_content in all_contents:
+                    continue
+                all_contents.append(doc_content)
+                docs_no_dup.append(doc)
+            docs = docs_no_dup
+            logger.info(f'delete duplicate docs: {len(docs)}')
+
+        # rerank
+        if self.params['post_retrieval']['with_rank'] and len(docs):
+            if not hasattr(self, 'ranker'):
+                rerank_params = self.params['post_retrieval']['rerank']
+                rerank_type = rerank_params.pop('type')
+                rerank_object = import_class(f'bisheng_langchain.rag.rerank.{rerank_type}')
+                self.ranker = rerank_object(**rerank_params)
+            docs = getattr(self, 'ranker').sort_and_filter(question, docs)
+
+        # delete redundancy according to max_content 
+        doc_num, doc_content_sum = 0, 0
+        for doc in docs:
+            doc_content_sum += len(doc.page_content)
+            if doc_content_sum > max_content:
+                break
+            doc_num += 1
+        docs = docs[:doc_num]
+        logger.info(f'retrieval docs after delete redundancy: {len(docs)}')
+
+        # 按照文档的source和chunk_index排序，保证上下文的连贯性和一致性
+        if self.params['post_retrieval'].get('sort_by_source_and_index', False):
+            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 load_documents(self, file_name, max_content=100000):
+        """
+        直接加载文档，如果文档过长，直接截断处理；
+        max_content: max content len of llm
+        """
+        file_path = os.path.join(self.origin_file_path, file_name)
+        if not os.path.exists(file_path):
+            raise Exception(f'{file_path} not exists.')
+        if os.path.isdir(file_path):
+            raise Exception(f'{file_path} is a directory.')
+
+        loader_params = copy.deepcopy(self.params['loader'])
+        loader_object = import_by_type(_type='documentloaders', name=loader_params.pop('type'))
+        loader = loader_object(file_name=file_name, file_path=file_path, **loader_params)
+
+        documents = loader.load()
+        logger.info(f'documents: {len(documents)}, page_content: {len(documents[0].page_content)}')
+        for doc in documents:
+            doc.page_content = doc.page_content[:max_content]
+        return documents
+
+    def question_answering(self):
+        """
+        question answer over knowledge
+        """
+        df = pd.read_excel(self.question_path)
+        all_questions_info = df.to_dict('records')
+        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
+        if not isinstance(self.llm, ChatCohere):
+            qa_chain = load_qa_chain(
+                llm=self.llm, chain_type=self.params['generate']['chain_type'], prompt=prompt, verbose=False
+            )
+        file2docs = dict()
+        for questions_info in tqdm(all_questions_info):
+            question = questions_info['问题']
+            file_name = questions_info['文件名']
+            collection_name = questions_info['知识库名']
+
+            # if question != '请分析江苏中设集团股份有限公司2021年重大关联交易的情况。':
+            #     continue
+
+            if self.params['generate']['with_retrieval']:
+                # retrieval and rerank
+                docs = self.retrieval_and_rerank(question, collection_name, max_content=self.params['generate']['max_content'])
+            else:
+                # load document
+                if file_name not in file2docs:
+                    docs = self.load_documents(file_name, max_content=self.params['generate']['max_content'])
+                    file2docs[file_name] = docs
+                else:
+                    docs = file2docs[file_name]
+
+            if isinstance(self.llm, ChatCohere):
+                try:
+                    # cohere rag
+                    # messages = prompt.format_prompt(question=question).to_messages()
+                    # ans = self.llm.invoke(messages, source_documents=docs)
+                    # rag_answer = ans.content
+
+                    # cohere rag by raw prompt
+                    documents = ''
+                    for i, doc in enumerate(docs):
+                        # documents += f'Document: {i}\n'
+                        # documents  += 'text:' + doc.page_content + '\n\n'
+
+                        documents  += doc.page_content + '\n\n'
+                    messages = prompt.format_prompt(question=question, documents=documents).to_messages()
+                    ans = self.llm.invoke(messages, raw_prompting=True)
+                    rag_answer = ans.content.replace('Answer: ', '')
+                except Exception as e:
+                    logger.error(f'question: {question}\nerror: {e}')
+                    ans = {'output_text': str(e)}
+                    rag_answer = ans['output_text']
+            else:
+                # question answer
+                try:
+                    ans = qa_chain({"input_documents": docs, "question": question}, return_only_outputs=True)
+                except Exception as e:
+                    logger.error(f'question: {question}\nerror: {e}')
+                    ans = {'output_text': str(e)}
+                rag_answer = ans['output_text']
+                
+            # context = '\n\n'.join([doc.page_content for doc in docs])
+            # content = prompt.format(context=context, question=question)
+
+            # for rate_limit
+            # time.sleep(3)
+            logger.info(f'question: {question}\nans: {rag_answer}\n')
+            questions_info['rag_answer'] = rag_answer
+            # questions_info['rag_context'] = '\n----------------\n'.join([doc.page_content for doc in docs])
+            # questions_info['rag_context'] = content
+
+        df = pd.DataFrame(all_questions_info)
+        df.to_excel(self.save_answer_path, index=False)
+
+    def score(self):
+        """
+        score
+        """
+        metric_params = self.params['metric']
+        if metric_params['type'] == 'bisheng-ragas':
+            score_params = {
+                'excel_path': self.save_answer_path,
+                'save_path': os.path.dirname(self.save_answer_path),
+                'question_column': metric_params['question_column'],
+                'gt_column': metric_params['gt_column'],
+                'answer_column': metric_params['answer_column'],
+                'query_type_column': metric_params.get('query_type_column', None),
+                'contexts_column': metric_params.get('contexts_column', None),
+                'metrics': metric_params['metrics'],
+                'batch_size': metric_params['batch_size'],
+                'gt_split_column': metric_params.get('gt_split_column', None),
+                'whether_gtsplit': metric_params.get('whether_gtsplit', False), # 是否需要模型对gt进行要点拆分
+            }
+            rag_score = RagScore(**score_params)
+            rag_score.score()
+        else:
+            # todo: 其他评分方法
+            pass
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='Process some integers.')
+    # 添加参数
+    parser.add_argument('--mode', type=str, default='qa', help='upload or qa or score')
+    parser.add_argument('--params', type=str, default='config/test/baseline_s2b.yaml', help='bisheng rag params')
+    # 解析参数
+    args = parser.parse_args()
+
+    rag = BishengRagPipeline(args.params)
+
+    if args.mode == 'upload':
+        rag.file2knowledge()
+    elif args.mode == 'qa':
+        rag.question_answering()
+    elif args.mode == 'score':
+        rag.score()

bisheng_langchain/rag/bisheng_rag_tool.py

@@ -0,0 +1,288 @@
+import time
+import os
+import yaml
+import httpx
+from typing import Any, Dict, Tuple, Type, Union, Optional
+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.chains.question_answering import load_qa_chain
+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
+
+
+class MultArgsSchemaTool(Tool):
+
+    def _to_args_and_kwargs(self, tool_input: Union[str, Dict]) -> Tuple[Tuple, Dict]:
+        # For backwards compatibility, if run_input is a string,
+        # pass as a positional argument.
+        if isinstance(tool_input, str):
+            return (tool_input, ), {}
+        else:
+            return (), tool_input
+
+
+class BishengRAGTool:
+
+    def __init__(
+        self,
+        vector_store: Optional[Milvus] = None,
+        keyword_store: Optional[ElasticKeywordsSearch] = None,
+        llm: Optional[LanguageModelLike] = None,
+        collection_name: Optional[str] = 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')
+        self.collection_name = collection_name
+        
+        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)
+        self.params['generate']['max_content'] = max_content
+        self.params['post_retrieval']['sort_by_source_and_index'] = sort_by_source_and_index
+        
+        # init milvus
+        if vector_store:
+            self.vector_store = vector_store
+        else:
+            # init embeddings
+            embedding_params = self.params['embedding']
+            embedding_object = import_by_type(_type='embeddings', name=embedding_params['type'])
+            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
+                )
+            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'],
+                },
+            )
+        
+        # init keyword store
+        if keyword_store:
+            self.keyword_store = keyword_store
+        else:
+            self.keyword_store = ElasticKeywordsSearch(
+                index_name='default_es',
+                elasticsearch_url=self.params['elasticsearch']['url'],
+                ssl_verify=self.params['elasticsearch']['ssl_verify'],
+            )
+        
+        # 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 = []
+        retrievers = self.params['retriever']['retrievers']
+        for retriever in retrievers:
+            retriever_type = retriever.pop('type')
+            retriever_params = {
+                'vector_store': self.vector_store,
+                'keyword_store': self.keyword_store,
+                'splitter_kwargs': retriever['splitter'],
+                'retrieval_kwargs': retriever['retrieval'],
+            }
+            retriever_list.append(self._post_init_retriever(retriever_type=retriever_type, **retriever_params))
+        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}')
+        else:
+            prompt = None
+        self.qa_chain = load_qa_chain(
+            llm=self.llm, 
+            chain_type=self.params['generate']['chain_type'], 
+            prompt=prompt, 
+            verbose=False
+        )
+    
+    def _post_init_retriever(self, retriever_type, **kwargs):
+        retriever_classes = {
+            'KeywordRetriever': KeywordRetriever,
+            'BaselineVectorRetriever': BaselineVectorRetriever,
+            'MixRetriever': MixRetriever,
+            'SmallerChunksVectorRetriever': SmallerChunksVectorRetriever,
+        }
+        if retriever_type not in retriever_classes:
+            raise ValueError(f'Unknown retriever type: {retriever_type}')
+
+        input_kwargs = {}
+        splitter_params = kwargs.pop('splitter_kwargs')
+        for key, value in splitter_params.items():
+            splitter_obj = import_by_type(_type='textsplitters', name=value.pop('type'))
+            input_kwargs[key] = splitter_obj(**value)
+
+        retrieval_params = kwargs.pop('retrieval_kwargs')
+        for key, value in retrieval_params.items():
+            input_kwargs[key] = value
+
+        input_kwargs['vector_store'] = kwargs.pop('vector_store')
+        input_kwargs['keyword_store'] = kwargs.pop('keyword_store')
+
+        retriever_class = retriever_classes[retriever_type]
+        return retriever_class(**input_kwargs)
+
+    def file2knowledge(self, file_path, drop_old=True):
+        """
+        file to knowledge
+        """
+        loader_params = self.params['loader']
+        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
+        )
+        documents = loader.load()
+        logger.info(f'documents: {len(documents)}, page_content: {len(documents[0].page_content)}')
+        if len(documents[0].page_content) == 0:
+            logger.error(f'{file_path} page_content is empty.')
+
+        # add aux info
+        add_aux_info = self.params['retriever'].get('add_aux_info', False)
+        if add_aux_info:
+            for doc in documents:
+                try:
+                    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}")
+                    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
+            )
+    
+    def retrieval_and_rerank(self, query):
+        """
+        retrieval and rerank
+        """
+        # EnsembleRetriever直接检索召回会默认去重
+        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 
+        doc_num, doc_content_sum = 0, 0
+        for doc in docs:
+            doc_content_sum += len(doc.page_content)
+            if doc_content_sum > self.params['generate']['max_content']:
+                break
+            doc_num += 1
+        docs = docs[:doc_num]
+        logger.info(f'retrieval docs after delete redundancy: {len(docs)}')
+
+        # 按照文档的source和chunk_index排序，保证上下文的连贯性和一致性
+        if self.params['post_retrieval'].get('sort_by_source_and_index', False):
+            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) -> str:
+        docs = self.retrieval_and_rerank(query)
+        try:
+            ans = self.qa_chain({"input_documents": docs, "question": query}, return_only_outputs=True)
+        except Exception as e:
+            logger.error(f'question: {query}\nerror: {e}')
+            ans = {'output_text': str(e)}
+        rag_answer = ans['output_text']
+        return rag_answer
+    
+    async def arun(self, query: str) -> str:
+        rag_answer = self.run(query)
+        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.')
+
+        return MultArgsSchemaTool(name=name,
+                                  description=description,
+                                  func=cls(**kwargs).run,
+                                  coroutine=cls(**kwargs).arun,
+                                  args_schema=InputArgs)
+    
+
+if __name__ == '__main__':
+    # rag_tool = BishengRAGTool(collection_name='rag_finance_report_0_test')
+    # rag_tool.file2knowledge(file_path='/home/public/rag_benchmark_finance_report/金融年报财报的来源文件/2021-04-23__金宇生物技术股份有限公司__600201__生物股份__2020年__年度报告.pdf')
+
+    from langchain.chat_models import ChatOpenAI
+    from langchain.embeddings import OpenAIEmbeddings
+    # embedding
+    embeddings = OpenAIEmbeddings(model='text-embedding-ada-002')
+    # llm
+    llm = ChatOpenAI(model='gpt-4-1106-preview', temperature=0.01)
+    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',
+            },
+    )
+    # es
+    keyword_store = ElasticKeywordsSearch(
+        index_name=collection_name,
+        elasticsearch_url='http://110.16.193.170:50062/es',
+        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
+    )
+    print(tool.run('能否根据2020年金宇生物技术股份有限公司的年报，给我简要介绍一下报告期内公司的社会责任工作情况？'))
+
+    # tool = BishengRAGTool.get_rag_tool(
+    #     name='rag_knowledge_retrieve', 
+    #     description='金融年报财报知识库问答',
+    #     collection_name='rag_finance_report_0_benchmark_caibao_1000_source_title'
+    # )
+    # print(tool.run('能否根据2020年金宇生物技术股份有限公司的年报，给我简要介绍一下报告期内公司的社会责任工作情况？'))