qaPipeline_functions.py

"""
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