基于 LangChain + LLM + 向量数据库搭建企业智能问答系统

{"response":"你好！我是人工智能助手 ChatGLM-6B，很高兴见到你，欢迎问我任何问题。",  "history":[["你好","你好！我是人工智能助手 ChatGLM-6B，很高兴见到你，欢迎问我任何问题。"]],  "status":200,  "time":"2023-03-23 21:38:40"}

{
    "llm_config": {
        "URL": "http://1**.**.*.1:8000",
        "TOKEN": "xxxxxx"
    },
    "vdb_config": {
        "VDB_URL": "http://10.**.**.**:80",
        "VDB_USERNAME": "r****",
        "VDB_KEY":"tO**********xxxon**********T",
        "DATABASE_NAME":"LangChainDatabase",
        "COLLECTION_NAME":"LangChainCollection"
    },
    "embedding": {
        "model_id": "bge-base-zh",
        "model_dimension": 768
    },
    "query_topk": 4,
    "prompt_template": "Use the following pieces of information to answer the user's question.\\n If you don't know the answer, just say that you don't know, don't try to make up an answer.\\n Context: {context}\\nQuestion: {question}\\n Only return the helpful answer below and nothing else.\\nHelpful answer:"
}

# 将本地的文档“/path/to/local/username.PDF”scp到云服务的“/vdbproject/data"路径下，实际操作中这两个路径需要替换为您的实际路径
# 上传单个PDF文档
scp -P xxx /path/to/local/filename.pdf username@server_ip:/vdbproject/data
# 批量上传多个PDF文档
scp -P xxx /path/to/local/*.pdf username@server_ip:/vdbproject/data

from langchain.document_loaders import PyPDFLoader  # 这里以PDF格式文档为例，实际过程中如果使用其他文档格式，需要进行适配
import os, pdfplumber, tempfile
import argparse
import json
from langchain.embeddings.fake import FakeEmbeddings
from langchain.text_splitter import CharacterTextSplitter
from langchain.vectorstores import TencentVectorDB
from langchain.vectorstores.tencentvectordb import ConnectionParams
from langchain.vectorstores.tencentvectordb import IndexParams
from langchain.llms import ChatGLM
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain
import time
import requests


# 创建一个 VectorDB 和 LLM 模型的对话
class ChatLLMbot:
    def __init__(self, config,clear_db,no_vector_store) -> None:
        self.config = config
        if no_vector_store == False:
            self.vector_db = self.connect_vectorstore(clear_db)
        self.llm = self.connect_llm()


    # 连接到 LLM 模型
    def connect_llm(self):
        print("Start connecting to LLM.")
        endpoint_url = self.config['llm_config']['URL']      
        
        # 配置 LLM 参数
        llm = ChatGLM(
            endpoint_url=endpoint_url,
            max_token=80000,
            top_p=0.9,
            model_kwargs={"sample_model_args": False}
        )

        return llm


    # 连接到 VectorDB
    def connect_vectorstore(self,clear_db):
        print("Start connecting to VectorDB.")
        VDB_URL = self.config['vdb_config']['VDB_URL']
        VDB_USERNAME = self.config['vdb_config']['VDB_USERNAME']
        VDB_KEY = self.config['vdb_config']['VDB_KEY']
        DATABSE_NAME = self.config['vdb_config']['DATABASE_NAME']
        COLLECTION_NAME = self.config['vdb_config']['COLLECTION_NAME']

        # 为 VectorDB 建立连接参数
        conn_params = ConnectionParams(
            url=VDB_URL,
            key=VDB_KEY,
            username=VDB_USERNAME,
            timeout=20
        )

        # 创建 Embedding 对象，如下示例中使用了虚拟的 Embedding，您在使用过程中需要替换为真实的 Embedding 服务参数。
        embeddings = FakeEmbeddings(size=128)
        vector_db = TencentVectorDB(
            embedding = embeddings,
            connection_params=conn_params,
            index_params = IndexParams (128),
            database_name = DATABSE_NAME,
            collection_name = COLLECTION_NAME,
            drop_old = clear_db
        )

        return vector_db


    # 读取文档   
    def load_data(self, files: list[str]):
        documents = []
        for fname in files:
            loader = PyPDFLoader(fname)
            documents += loader.load()

        # 分割文档
        text_splitter = CharacterTextSplitter(
            chunk_size=1000, chunk_overlap=100)
        documents = text_splitter.split_documents(documents)
        self.vector_db.add_documents(documents)
 
       
    # 使用 VectorDB+LLM 检索
    def query(self, query: str, use_vdb: bool = True) -> str:
        context = ''
        if use_vdb:
            answer_from_vdb = self.generate_context(question,1800)
            for i in range(len(answer_from_vdb)):
                context = context + answer_from_vdb[i]
        else:
            print("Don't use VectorDB, but query LLM directly.")
        # 返回查询结果给 LLM 模型
        answer = self.query_to_llm(context, query)  
        return answer
      
          
    # 根据提问匹配上下文
    def generate_context(self, query: str,max_context_length: int) -> str:
        print("Start querying VectorDB with query: " + query)
        # 使用向量数据库做相似性检索
        docs = self.vector_db.similarity_search(
            question, k=self.config['query_topk'])
        # 限制发给大模型的上下文文本总长度
        current_context_length = 0
        ret = []
        for doc in docs:
            if len(doc.page_content) + \\
                    current_context_length > max_context_length:
                continue
            current_context_length += len(doc.page_content)
            ret.append(doc.page_content)
        return ret


    # 将 VectorDB 匹配的结果通过 prompt 发送给 LLM
    def query_to_llm(self, context: str, query: str) -> str:
        template = self.config['prompt_template']
        prompt = PromptTemplate(template=template, input_variables=["context", "question"])
        print("Start querying LLM with prompt.")
        start_time = time.time()
        llm_chain = LLMChain(prompt=prompt, llm=self.llm)
        # 使用 LLM 模型进行预测
        answer = llm_chain.predict(context=context,question=question)
        end_time = time.time()
        print(
            "Get response from LLM success. Cost Time: {:.2f}s".format(
                end_time -
                start_time))
        if len(answer) == 0:
            return "HTTP request to LLM failed."

        return answer


# 命令行参数解析
if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        prog='chatbot',
        description='llm+vdb chatbot command line interface')
    parser.add_argument('-l', '--load', action='store_true',
                        help='generate embeddings and update the vector database.')
    parser.add_argument('-c', '--clear', action='store_true',
                        help='clear all data in vector store')
    parser.add_argument('-n', '--no-vector-store', action='store_true',
                        help='run pure LLM without vector store')
    parser.add_argument('--config', help='input configuration json file',default='./config/config.json')
    
    args = parser.parse_args()

    # 检查是否配置 JSON 文件并加载对应配置
    if args.config:
        if os.path.exists(args.config):
            with open(args.config) as f:
                config = json.load(f)
                #print(config)
                bot = ChatLLMbot(config,args.clear,args.no_vector_store)
                if args.load :
                    DIR_PATH = os.path.dirname(os.path.realpath(__file__))
                    files = [os.path.join(DIR_PATH, 'data', x)
                            for x in os.listdir(os.path.join(DIR_PATH, 'data'))]
                    print(f'Start loading files: {files}')
                    bot.load_data(files)
                    exit(0)


                # 开始提问
                while True:
                    print("Please enter a Question: ")
                    question = input()
                    if(args.no_vector_store):
                        answer = bot.query(question,False)
                        print('LLM answer:\\n ' + answer)
                    else:
                        answer = bot.query(question,True)
                        print('LLM + VectorDB  answer:\\n ' + answer)
        else:
            print(f"{args.config} is not existed.")
    else :
        print("The config json file must be set.")