""" Python Backend API to chat with private data 08/14/2023 D.M. Theekshana Samaradiwakara """ import os import time from dotenv import load_dotenv from langchain.chains import RetrievalQA from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler from langchain.llms import GPT4All from langchain.llms import HuggingFaceHub from langchain.chat_models import ChatOpenAI # from langchain.retrievers._query.base import SelfQueryRetriever # from langchain.chains.query_constructor.base import AttributeInfo # from chromaDb import load_store from faissDb import load_FAISS_store from langchain.agents import ZeroShotAgent, Tool, AgentExecutor from langchain.prompts import PromptTemplate from langchain.chains import LLMChain, ConversationalRetrievalChain from conversationBufferWindowMemory import ConversationBufferWindowMemory from langchain.memory import ReadOnlySharedMemory load_dotenv() #gpt4 all model gpt4all_model_path = os.environ.get('GPT4ALL_MODEL_PATH') model_n_ctx = os.environ.get('MODEL_N_CTX') model_n_batch = int(os.environ.get('MODEL_N_BATCH',8)) target_source_chunks = int(os.environ.get('TARGET_SOURCE_CHUNKS',4)) openai_api_key = os.environ.get('OPENAI_API_KEY') verbose = os.environ.get('VERBOSE') # activate/deactivate the streaming StdOut callback for LLMs callbacks = [StreamingStdOutCallbackHandler()] memory = ConversationBufferWindowMemory( memory_key="chat_history", input_key="question", return_messages=True, k=3 ) readonlymemory = ReadOnlySharedMemory(memory=memory) print("\n\n> Initializing QAPipeline:") global llm_name llm_name = 'None' global llm llm = 'None' global dataset_name dataset_name = 'None' global vectorstore vectorstore = 'None' qa_chain = None agent = None def run(query, model, dataset): if (llm_name != model) or (dataset_name != dataset) or (qa_chain == None): set_model(model) set_vectorstore(dataset) set_qa_chain() # Get the answer from the chain start = time.time() res = qa_chain(query) # answer, docs = res['result'],res['source_documents'] end = time.time() # Print the result print("\n\n> Question:") print(query) print(f"\n> Answer (took {round(end - start, 2)} s.):") print( res) return res def run_agent(query, model, dataset): try: if (llm_name != model) or (dataset_name != dataset) or (agent == None): set_model(model) set_vectorstore(dataset) set_qa_chain_with_agent() # Get the answer from the chain start = time.time() res = agent(query) # answer, docs = res['result'],res['source_documents'] end = time.time() # Print the result print("\n\n> Question:") print(query) print(f"\n> Answer (took {round(end - start, 2)} s.):") print( res) return res["output"] except Exception as e: # logger.error(f"Answer retrieval failed with {e}") print(f"> QAPipeline run_agent Error : {e}")#, icon=":books:") return def set_model(model_type): if model_type != llm_name: global llm match model_type: case "gpt4all": # llm = GPT4All(model=gpt4all_model_path, n_ctx=model_n_ctx, backend='gptj', n_batch=model_n_batch, callbacks=callbacks, verbose=verbose) llm = GPT4All(model=gpt4all_model_path, max_tokens=model_n_ctx, backend='gptj', n_batch=model_n_batch, callbacks=callbacks, verbose=verbose) # llm = HuggingFaceHub(repo_id="nomic-ai/gpt4all-j", model_kwargs={"temperature":0.001, "max_length":1024}) case "google/flan-t5-xxl": llm = HuggingFaceHub(repo_id="google/flan-t5-xxl", model_kwargs={"temperature":0.001, "max_length":1024}) case "tiiuae/falcon-7b-instruct": llm = HuggingFaceHub(repo_id=model_type, model_kwargs={"temperature":0.001, "max_length":1024}) case "openai": llm = ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0) case _default: # raise exception if model_type is not supported raise Exception(f"Model type {model_type} is not supported. Please choose a valid one") # global llm_name llm_name = model_type def set_vectorstore( dataset): if dataset != dataset_name: # vectorstore = load_store(dataset) global vectorstore vectorstore = load_FAISS_store() print("\n\n> vectorstore loaded:") dataset_name = dataset def set_qa_chain(): global qa_chain qa_chain = RetrievalQA.from_chain_type( llm=llm, chain_type="stuff", retriever = vectorstore.as_retriever(), # retriever = vectorstore.as_retriever(search_kwargs={"k": target_source_chunks} return_source_documents= True ) def set_qa_chain_with_agent(): try: # Define a custom prompt general_qa_template = ( """You can have a general conversation with the users like greetings. Continue the conversation and only answer questions related to banking sector like financial and legal. If you dont know the answer say you dont know, dont try to makeup answers. Conversation: {chat_history} Question: {question} """ ) general_qa_chain_prompt = PromptTemplate(input_variables=["question", "chat_history"], template=general_qa_template) general_qa_chain = LLMChain( llm=llm, prompt=general_qa_chain_prompt, verbose=True, memory=readonlymemory, # use the read-only memory to prevent the tool from modifying the memory ) general_qa_chain_tool = Tool( name="general qa", func= general_qa_chain.run, description='''useful for when you need to have a general conversation with the users like greetings or to answer general purpose questions related to banking sector like financial and legal. Input should be a fully formed question.''', return_direct=True, ) # Define a custom prompt retrieval_qa_template = ( """ please answer the question based on the chat history and context with the latest information. You have provided context information below related to central bank acts published in various years. The content of a bank act can updated by a bank act from a latest year. If you dont know the answer say you dont know, dont try to makeup answers. Conversation: {chat_history} Context: {context} Question : {question} """ ) retrieval_qa_chain_prompt = PromptTemplate( input_variables=["question", "context", "chat_history"], template=retrieval_qa_template ) bank_regulations_qa = ConversationalRetrievalChain.from_llm( llm=llm, chain_type="stuff", retriever = vectorstore.as_retriever(), # retriever = vectorstore.as_retriever(search_kwargs={"k": target_source_chunks} return_source_documents= True, get_chat_history=lambda h : h, combine_docs_chain_kwargs={"prompt": retrieval_qa_chain_prompt}, verbose=True, memory=readonlymemory, # use the read-only memory to prevent the tool from modifying the memory ) bank_regulations_qa_tool = Tool( name="bank regulations", func= lambda question: bank_regulations_qa({"question": question}), description='''useful for when you need to answer questions about financial and legal information issued from central bank regarding banks and bank regulations. Input should be a fully formed question.''', return_direct=True, ) tools = [ bank_regulations_qa_tool, general_qa_chain_tool ] prefix = """Have a conversation with a human, answering the following questions as best you can. You have access to the following tools:""" suffix = """Begin!" {chat_history} Question: {question} {agent_scratchpad}""" agent_prompt = ZeroShotAgent.create_prompt( tools, prefix=prefix, suffix=suffix, input_variables=["question", "chat_history", "agent_scratchpad"], ) llm_chain = LLMChain(llm=llm, prompt=agent_prompt) zeroShotAgent = ZeroShotAgent( llm_chain=llm_chain, tools=tools, verbose=True, ) agent_chain = AgentExecutor.from_agent_and_tools( agent=zeroShotAgent, tools=tools, verbose=True, memory=memory, handle_parsing_errors=True, ) global agent agent = agent_chain print(f"\n> agent_chain created") except Exception as e: # logger.error(f"Answer retrieval failed with {e}") print(f"> QAPipeline set_qa_chain_with_agent Error : {e}")#, icon=":books:") return