agent with memory

.env

@@ -11,7 +11,7 @@ TARGET_SOURCE_CHUNKS=4
 
 #API token keys
 HUGGINGFACEHUB_API_TOKEN=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
-OPENAI_API_KEY=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
+OPENAI_API_KEY=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
 
 #api app
 APP_HOST=127.0.0.1

app.py

@@ -16,6 +16,7 @@ from ui.htmlTemplates import css, bot_template, user_template, source_template
 from config import MODELS, DATASETS
 
 from qaPipeline import QAPipeline
+import qaPipeline_functions
 from faissDb import create_faiss
 
 # loads environment variables
@@ -25,7 +26,12 @@ load_dotenv()
 isHuggingFaceHubEnabled = os.environ.get('ENABLE_HUGGINGFSCE_HUB_MODELS')
 isOpenAiApiEnabled = os.environ.get('ENABLE_OPENAI_API_MODELS')
 
+st.set_page_config(page_title="Chat with data",
+                       page_icon=":books:")
+st.write(css, unsafe_allow_html=True)
+
 qaPipeline = QAPipeline()
+# qaPipeline = qaPipeline_functions
 
 def initialize_session_state():
     # Initialise all session state variables with defaults
@@ -109,13 +115,22 @@ def side_bar():
 
 def chat_body():
     st.header("Chat with your own data:")
+    with st.form('chat_body'):
+
+        user_question = st.text_input(
+            "Ask a question about your documents:",
+            placeholder="enter question",
+            key='user_question',
+            # on_change=submit_user_question,
+        )
+
+        submitted = st.form_submit_button(
+            "Submit",
+            # on_click=update_parameters_change
+            )
     
-    user_question = st.text_input(
-        "Ask a question about your documents:",
-        placeholder="enter question",
-        key='user_question',
-        on_change=submit_user_question,
-    )
+        if submitted:
+            submit_user_question()
 
     # if user_question:
     #     submit_user_question()
@@ -128,18 +143,15 @@ def submit_user_question():
         user_question = st.session_state.user_question
         # st.success(user_question)
         handle_userinput(user_question)
-        st.session_state.user_question=''
+        # st.session_state.user_question=''
 
 
 def main():
    
-    st.set_page_config(page_title="Chat with data",
-                       page_icon=":books:")
-    st.write(css, unsafe_allow_html=True)
-
     initialize_session_state()
     
     side_bar()
+
     chat_body()
 
                 
@@ -157,7 +169,7 @@ def parameters_change_button(chat_model, show_source):
     time.sleep(1) # Wait for 3 seconds
     alert.empty() # Clear the alert
 
-   
+# @st.cache_data
 def get_answer_from_backend(query, model, dataset):
     # response = qaPipeline.run(query=query, model=model, dataset=dataset)
     response = qaPipeline.run_agent(query=query, model=model, dataset=dataset)
@@ -167,7 +179,7 @@ def get_answer_from_backend(query, model, dataset):
 def show_query_response(query, response, show_source_files):
     docs = []
     if isinstance(response, dict):
-        answer, docs = response['result'], response['source_documents']
+        answer, docs = response['answer'], response['source_documents']
     else:
         answer = response
 
@@ -217,9 +229,15 @@ def handle_userinput(query):
 
     except Exception as e:
         # logger.error(f"Answer retrieval failed with {e}")
-        st.error(f"Error : {e}")#, icon=":books:")
+        st.error(f"Streamlit handle_userinput Error : {e}")#, icon=":books:")
         return
 
 
 if __name__ == "__main__":
     main()
+
+# initialize_session_state()
+
+# side_bar()
+
+# chat_body()

app2.py

@@ -0,0 +1,184 @@
+"""
+Python Backend API to chat with private data  
+
+08/16/2023
+D.M. Theekshana Samaradiwakara
+"""
+
+import os
+import time
+import streamlit as st
+from streamlit.logger import get_logger
+
+logger = get_logger(__name__)
+
+from ui.htmlTemplates import css, bot_template, user_template, source_template
+from config import MODELS, DATASETS
+
+from qaPipeline import QAPipeline
+from faissDb import create_faiss
+
+# loads environment variables
+from dotenv import load_dotenv
+load_dotenv()
+
+isHuggingFaceHubEnabled = os.environ.get('ENABLE_HUGGINGFSCE_HUB_MODELS')
+isOpenAiApiEnabled = os.environ.get('ENABLE_OPENAI_API_MODELS')
+
+st.set_page_config(page_title="Chat with data",
+                       page_icon=":books:")
+st.write(css, unsafe_allow_html=True)
+
+
+SESSION_DEFAULTS = {
+    "model": MODELS["DEFAULT"],
+    "dataset": DATASETS["DEFAULT"],
+    "chat_history": None,
+    "is_parameters_changed":False,
+    "show_source_files": False,
+    "user_question":'',
+}
+
+for k, v in SESSION_DEFAULTS.items():
+    if k not in st.session_state:
+        st.session_state[k] = v
+
+
+with st.sidebar:
+    st.subheader("Chat parameters")
+
+    with st.form('param_form'):
+
+        chat_model = st.selectbox(
+            "Chat model",
+            MODELS,
+            key="chat_model",
+            help="Select the LLM model for the chat",
+            # on_change=update_parameters_change,
+        )
+
+        st.session_state.dataset =  "DEFAULT"
+
+        show_source = st.checkbox(
+            label="show source files",
+            key="show_source",
+            help="Select this to show relavant source files for the query",
+        )
+
+        submitted = st.form_submit_button(
+            "Submit",
+            # on_click=parameters_change_button,
+            # args=[chat_model, show_source]
+            )
+        
+        # submitted = st.button(
+        # "Submit",
+        # # on_click=parameters_change_button,
+        # # args=[chat_model, show_source]
+        # )
+
+        if submitted:
+            st.session_state.model = chat_model
+            st.session_state.dataset = "DEFAULT"
+            st.session_state.show_source_files = show_source
+            st.session_state.is_parameters_changed = False
+
+            alert =  st.success("chat parameters updated")
+            time.sleep(1) # Wait for 3 seconds
+            alert.empty() # Clear the alert
+                    
+        st.markdown("\n")
+
+    # if st.button("Create FAISS db"):
+    #     try:
+    #         with st.spinner('creating faiss vector store'):
+    #             create_faiss()
+    #             st.success('faiss saved')
+    #     except Exception as e:
+    #         st.error(f"Error : {e}")#, icon=":books:")
+    #         return
+            
+    st.markdown(
+        "### How to use\n"
+        "1. Select the chat model\n"  # noqa: E501
+        "2. Select \"show source files\" to show the source files related to the answer.📄\n"
+        "3. Ask a question about the documents💬\n"
+    )
+
+
+
+st.header("Chat with your own data:")
+@st.experimental_singleton  # 👈 Add the caching decorator
+def load_QaPipeline():
+    print('> QAPipeline loaded for front end')
+    return QAPipeline()
+
+qaPipeline = load_QaPipeline()
+# qaPipeline = QAPipeline()
+with st.form('chat_body'):
+
+
+    user_question = st.text_input(
+        "Ask a question about your documents:",
+        placeholder="enter question",
+        key='user_question',
+        # on_change=submit_user_question,
+    )
+
+    submitted = st.form_submit_button(
+        "Submit",
+        # on_click=submit_user_question
+        )
+
+    if submitted:
+        with st.spinner("Processing"):
+            user_question = st.session_state.user_question
+            # st.success(user_question)
+            query = user_question
+            # st.session_state.user_question=''
+   
+            # Get the answer from the chain
+            try:
+                if (not query) or (query.strip() == ''):
+                    st.error("Please enter a question!")
+                    st.stop()
+
+                model = MODELS[st.session_state.model]
+                dataset = DATASETS[st.session_state.dataset]
+                show_source_files = st.session_state.show_source_files
+
+                # Try to access openai and deeplake
+                print(f">\n model: {model} \n dataset : {dataset} \n show_source_files : {show_source_files}")
+
+                # response = qaPipeline.run(query=query, model=model, dataset=dataset)
+                response = qaPipeline.run_agent(query=query, model=model, dataset=dataset)
+
+            
+                docs = []
+                if isinstance(response, dict):
+                    answer, docs = response['answer'], response['source_documents']
+                else:
+                    answer = response
+
+                st.write(user_template.replace(
+                            "{{MSG}}", query), unsafe_allow_html=True)
+                st.write(bot_template.replace(
+                            "{{MSG}}",  answer ), unsafe_allow_html=True)
+                
+                if show_source_files:
+                    # st.write(source_template.replace(
+                    #         "{{MSG}}",  "source files" ), unsafe_allow_html=True)
+
+                    if len(docs)>0 :
+                        st.markdown("#### source files : ")
+                        for source in docs:
+                            # st.info(source.metadata)
+                            with st.expander(source.metadata["source"]):
+                                st.markdown(source.page_content)
+                        
+                    # st.write(response)
+                
+            except Exception as e:
+                # logger.error(f"Answer retrieval failed with {e}")
+                st.error(f"Error : {e}")#, icon=":books:")
+

conversationBufferWindowMemory.py

@@ -0,0 +1,118 @@
+from abc import ABC
+from typing import Any, Dict, Optional, Tuple
+# import json
+
+from langchain.memory.chat_message_histories.in_memory import ChatMessageHistory
+from langchain.memory.utils import get_prompt_input_key
+from langchain.pydantic_v1 import Field
+from langchain.schema import BaseChatMessageHistory, BaseMemory
+
+from typing import List, Union
+
+# from langchain.memory.chat_memory import BaseChatMemory
+from langchain.schema.messages import BaseMessage, get_buffer_string
+
+
+class BaseChatMemory(BaseMemory, ABC):
+    """Abstract base class for chat memory."""
+
+    chat_memory: BaseChatMessageHistory = Field(default_factory=ChatMessageHistory)
+    output_key: Optional[str] = None
+    input_key: Optional[str] = None
+    return_messages: bool = False
+
+    def _get_input_output(
+        self, inputs: Dict[str, Any], outputs: Dict[str, str]
+    ) -> Tuple[str, str]:
+        
+
+        if self.input_key is None:
+            prompt_input_key = get_prompt_input_key(inputs, self.memory_variables)
+        else:
+            prompt_input_key = self.input_key
+
+        if self.output_key is None:
+            """
+            output for agent with LLM chain tool                     = {answer} 
+            output for agent with ConversationalRetrievalChain tool  = {'question', 'chat_history', 'answer','source_documents'}
+            """
+        
+            LLM_key = 'output'
+            Retrieval_key = 'answer'
+            if isinstance(outputs[LLM_key], dict):
+                Retrieval_dict = outputs[LLM_key]
+                if Retrieval_key in Retrieval_dict.keys():
+                    #output keys are 'answer' , 'source_documents'
+                    output = Retrieval_dict[Retrieval_key]
+                else:
+                    raise ValueError(f"output key: {LLM_key} not a valid dictionary")
+
+            else:
+                #otherwise output key will be 'output'
+                output_key = list(outputs.keys())[0]
+                output = outputs[output_key]
+
+            # if len(outputs) != 1:
+            #     raise ValueError(f"One output key expected, got {outputs.keys()}")
+                
+            
+        else:
+            output_key = self.output_key
+            output = outputs[output_key]
+
+        return inputs[prompt_input_key], output
+
+    def save_context(self, inputs: Dict[str, Any], outputs: Dict[str, str]) -> None:
+        """Save context from this conversation to buffer."""
+        input_str, output_str = self._get_input_output(inputs, outputs)
+        self.chat_memory.add_user_message(input_str)
+        self.chat_memory.add_ai_message(output_str)
+
+    def clear(self) -> None:
+        """Clear memory contents."""
+        self.chat_memory.clear()
+
+
+
+
+
+class ConversationBufferWindowMemory(BaseChatMemory):
+    """Buffer for storing conversation memory inside a limited size window."""
+
+    human_prefix: str = "Human"
+    ai_prefix: str = "AI"
+    memory_key: str = "history"  #: :meta private:
+    k: int = 5
+    """Number of messages to store in buffer."""
+
+    @property
+    def buffer(self) -> Union[str, List[BaseMessage]]:
+        """String buffer of memory."""
+        return self.buffer_as_messages if self.return_messages else self.buffer_as_str
+
+    @property
+    def buffer_as_str(self) -> str:
+        """Exposes the buffer as a string in case return_messages is True."""
+        messages = self.chat_memory.messages[-self.k * 2 :] if self.k > 0 else []
+        return get_buffer_string(
+            messages,
+            human_prefix=self.human_prefix,
+            ai_prefix=self.ai_prefix,
+        )
+
+    @property
+    def buffer_as_messages(self) -> List[BaseMessage]:
+        """Exposes the buffer as a list of messages in case return_messages is False."""
+        return self.chat_memory.messages[-self.k * 2 :] if self.k > 0 else []
+
+    @property
+    def memory_variables(self) -> List[str]:
+        """Will always return list of memory variables.
+
+        :meta private:
+        """
+        return [self.memory_key]
+
+    def load_memory_variables(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
+        """Return history buffer."""
+        return {self.memory_key: self.buffer}

qaPipeline.py

@@ -23,10 +23,12 @@ from langchain.chat_models import ChatOpenAI
 # from chromaDb import load_store
 from faissDb import load_FAISS_store
 
-from langchain.agents import initialize_agent, Tool
-from langchain.agents import AgentType
+from langchain.agents import ZeroShotAgent, Tool, AgentExecutor
+
 from langchain.prompts import PromptTemplate
-from langchain.chains import LLMChain
+from langchain.chains import LLMChain, ConversationalRetrievalChain
+from conversationBufferWindowMemory import ConversationBufferWindowMemory
+from langchain.memory import ReadOnlySharedMemory
 
 load_dotenv()
 
@@ -43,10 +45,35 @@ 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)
+
+class Singleton:
+   __instance = None
+   @staticmethod 
+   def getInstance():
+      """ Static access method. """
+      if Singleton.__instance == None:
+         Singleton()
+      return Singleton.__instance
+   def __init__(self):
+      """ Virtually private constructor. """
+      if Singleton.__instance != None:
+         raise Exception("This class is a singleton!")
+      else:
+         Singleton.__instance = QAPipeline()
+
 class QAPipeline:
 
     def __init__(self):
-        
+
+        print("\n\n> Initializing QAPipeline:")
         self.llm_name = None
         self.llm = None
 
@@ -56,6 +83,7 @@ class QAPipeline:
         self.qa_chain = None
         self.agent = None
 
+
     def run(self,query, model, dataset):
 
         if (self.llm_name != model) or (self.dataset_name != dataset) or (self.qa_chain == None):
@@ -79,24 +107,31 @@ class QAPipeline:
     
     def run_agent(self,query, model, dataset):
 
-        if (self.llm_name != model) or (self.dataset_name != dataset) or (self.agent == None):
-            self.set_model(model)
-            self.set_vectorstore(dataset)
-            self.set_qa_chain_with_agent()
+        try:
 
-        # Get the answer from the chain
-        start = time.time()
-        res = self.agent(query)
-        # answer, docs = res['result'],res['source_documents']
-        end = time.time()
+            if (self.llm_name != model) or (self.dataset_name != dataset) or (self.agent == None):
+                self.set_model(model)
+                self.set_vectorstore(dataset)
+                self.set_qa_chain_with_agent()
 
-        # Print the result
-        print("\n\n> Question:")
-        print(query)
-        print(f"\n> Answer (took {round(end - start, 2)} s.):")
-        print( res)
+            # Get the answer from the chain
+            start = time.time()
+            res = self.agent(query)
+            # answer, docs = res['result'],res['source_documents']
+            end = time.time()
 
-        return res["output"]
+            # 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 
 
 
     
@@ -139,67 +174,115 @@ class QAPipeline:
 
     def set_qa_chain_with_agent(self):
 
-        # Define a custom prompt
-        general_qa_template = (
-            """You are the AI assistant of the Boardpac company which provide services for companies board members.
-            You can have a general conversation with the users like greetings.
-            But 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.
-            each answer should start with code word BoardPac AI (Conversation):
-            Question: {question}
-            """
-        )
-
-        general_qa_chain_prompt = PromptTemplate.from_template(general_qa_template)
-        general_qa_chain = LLMChain(llm=self.llm, prompt=general_qa_chain_prompt)
-
-        # Define a custom prompt
-        retrieval_qa_template = (
-            """You are the AI assistant of the Boardpac company which provide services for companies board members.
-            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.
-            {context}
-            Given this information, please answer the question with the latest information. 
-            If you dont know the answer say you dont know, dont try to makeup answers. 
-            each answer should start with code word BoardPac AI (Retrieval):
-            Question: {question}
-            """
-        )
-        retrieval_qa_chain_prompt = PromptTemplate.from_template(retrieval_qa_template)
-
-        bank_regulations_qa = RetrievalQA.from_chain_type(
-            llm=self.llm,
-            chain_type="stuff",
-            retriever = self.vectorstore.as_retriever(), 
-            # retriever = self.vectorstore.as_retriever(search_kwargs={"k": target_source_chunks}
-            return_source_documents= True,
-            input_key="question",
-            chain_type_kwargs={"prompt": retrieval_qa_chain_prompt},
-        )
-
-        tools = [
-            Tool(
+        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=self.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=self.llm,
+                chain_type="stuff",
+                retriever = self.vectorstore.as_retriever(), 
+                # retriever = self.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 query: bank_regulations_qa({"question": query}),
+                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,
-            ),
+            )
 
-            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,
-            ),  
-        ]
-
-        self.agent = initialize_agent(
-            tools, 
-            self.llm, 
-            agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, 
-            verbose=True,
-            max_iterations=3,
-        )
+            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=self.llm, prompt=agent_prompt)
+
+            agent = ZeroShotAgent(
+                llm_chain=llm_chain, 
+                tools=tools, 
+                verbose=True,
+            )
+
+            agent_chain = AgentExecutor.from_agent_and_tools(
+                agent=agent, 
+                tools=tools, 
+                verbose=True,
+                memory=memory,
+                handle_parsing_errors=True,
+            )
+
+            self.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 

qaPipeline_functions.py

@@ -0,0 +1,278 @@
+"""
+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