import streamlit as st from streamlit_chat import message import os from langchain.llms import HuggingFaceHub # for calling HuggingFace Inference API (free for our use case) from langchain.embeddings import HuggingFaceEmbeddings # to let program know what embeddings the vector store was embedded in earlier # to set up the agent and tools which will be used to answer questions later from langchain.agents import initialize_agent from langchain.agents import tool # decorator so each function will be recognized as a tool from langchain.chains.retrieval_qa.base import RetrievalQA # to answer questions from vector store retriever # from langchain.chains.question_answering import load_qa_chain # to further customize qa chain if needed from langchain.vectorstores import Chroma # vector store for retriever import ast # to parse user string input to list for one of the tools (agent tools do not support 2 inputs) #from langchain.memory import ConversationBufferMemory # not used as of now import pickle # for loading the bm25 retriever from langchain.retrievers import EnsembleRetriever # to use chroma and # for defining a generic LLMChain as a generic chat tool (if needed) from langchain.prompts import PromptTemplate from langchain.chains import LLMChain import warnings warnings.filterwarnings("ignore", category=FutureWarning) warnings.filterwarnings("ignore", category=DeprecationWarning) # os.environ['HUGGINGFACEHUB_API_TOKEN'] = 'your_api_key' # for using HuggingFace Inference API from langchain.callbacks.base import BaseCallbackHandler class MyCallbackHandler(BaseCallbackHandler): def __init__(self): self.tokens = [] def on_llm_new_token(self, token, **kwargs) -> None: # HuggingFaceHub() cannot stream self.tokens.append(token) print(token) def on_agent_action(self, action, **kwargs): """Run on agent action.""" print("\n\nnew action", action) thought = action.log.replace('\n', ' \n') # so streamlit will recognize as newline tool_called = action.tool # tool_input = action.tool_input calling_tool = f"I am calling the '{tool_called}' tool and waiting for it to give me a result..." st.session_state.messages.extend( [{"role": "assistant", "content": thought}, {"role": "assistant", "content": calling_tool}] ) # Add the response to the chat window with st.chat_message("assistant"): st.markdown(thought) st.markdown(calling_tool) # def on_agent_finish(self, finish, **kwargs): # """Run on agent end.""" # #print("\n\nEnd", finish) # finish_string = finish.log.replace('\n', ' \n') # so streamlit will recognize as newline # st.session_state.messages.append( # {"role": "assistant", "content": finish_string} # ) # with st.chat_message("assistant"): # st.markdown(finish_string) # def on_llm_start(self, serialized, prompts, **kwargs): # """Run when LLM starts running.""" # print("LLM Start: ", prompts) # def on_llm_end(self, response, **kwargs): # """Run when LLM ends running.""" # print(response) def on_tool_end(self, output, **kwargs): """Run when tool ends running.""" #print("\n\nTool End: ", output) tool_output = f"Tool Output: {output} \n \nI am processing the output from the tool..." st.session_state.messages.append( {"role": "assistant", "content": tool_output} ) with st.chat_message("assistant"): st.markdown(tool_output) my_callback_handler = MyCallbackHandler() # # Set the webpage title # st.set_page_config( # page_title="Your own AI-Chat!", # layout="wide" # ) # llm for HuggingFace Inference API # model = "mistralai/Mistral-7B-Instruct-v0.2" model = "mistralai/Mixtral-8x7B-Instruct-v0.1" # with st.spinner('Downloading pre-built Chroma and BM25 vector stores'): # chroma_db = Chroma(persist_directory=persist_directory,embedding_function=hf_embeddings) # Document config if 'chunk_size' not in st.session_state: st.session_state['chunk_size'] = 1000 # choose one of [500, 600, 700, 800, 900, 1000, 1250, 1500, 1750, 2000, 2250, 2500, 2750, 3000] if 'chunk_overlap' not in st.session_state: st.session_state['chunk_overlap'] = 100 # choose one of [50, 100, 150, 200] # scraping results using DuckDuckGo if 'top_n_results' not in st.session_state: st.session_state['top_n_results'] = 10 # this is for returning top n search results using DuckDuckGo if 'countries_to_scrape' not in st.session_state: st.session_state['countries_to_scrape'] = [] # this is for returning top n search results using DuckDuckGo # in main app, add configuration for user to scrape new data from DuckDuckGo # in main app, add configuration for user to upload PDF to override country's existing policies in vectorstore # Retriever config if 'chroma_n_similar_documents' not in st.session_state: st.session_state['chroma_n_similar_documents'] = 5 # number of chunks returned by chroma vector store retriever (semantic) if 'bm25_n_similar_documents' not in st.session_state: st.session_state['bm25_n_similar_documents'] = 5 # number of chunks returned by bm25 retriever (keyword) if 'retriever_config' not in st.session_state: st.session_state['retriever_config'] = 'ensemble' # choose one of ['semantic', 'keyword', 'ensemble'] if 'keyword_retriever_weight' not in st.session_state: st.session_state['keyword_retriever_weight'] = 0.3 # choose between 0 and 1, only when using ensemble if 'source_documents' not in st.session_state: st.session_state['source_documents'] = [] # this is to store all source documents for a particular search # LLM config if 'temperature' not in st.session_state: st.session_state['temperature'] = 0.25 if 'max_new_tokens' not in st.session_state: st.session_state['max_new_tokens'] = 500 # max tokens generated by LLM # This is the list of countries present in the vector store, since the vector store is previously prepared as they take very long to prepare # This is for the RetrievalQA tool later to check, because even if the country given to it is not in the vector store, # it would still filter the vector store with this country and give an empty result, instead of giving an error. # We have to manually return the error to let the agent using the tool know. # The countries were reduced to just 6 as the time taken to get the embeddings to build up the chunks is too long. # However, having more countries **will not affect** the quality of the answers in comparing between 2 countries in the RAG application # as the RAG only picks out document chunks for the 2 countries of interest. countries = [ "Australia", "China", "Japan", "Malaysia", "Singapore", "Germany", ] @st.cache_data # only going to get once def get_llm(temp = st.session_state['temperature'], tokens = st.session_state['max_new_tokens']): # This is an inference endpoint API from huggingface, the model is not run locally, it is run on huggingface # It is a free API that is very good for deploying online for quick testing without users having to deploy a local LLM llm = HuggingFaceHub(repo_id=model, model_kwargs={ 'temperature':temp, "max_new_tokens":tokens }, ) return llm llm = get_llm(st.session_state['temperature'], tokens = st.session_state['max_new_tokens']) @st.cache_data # only going to get once def get_embeddings(): with st.spinner(f'Getting HuggingFaceEmbeddings'): # We use HuggingFaceEmbeddings() as it is open source and free to use. # Initialize the default hf model for embedding the tokenized texts into vectors with semantic meanings hf_embeddings = HuggingFaceEmbeddings() return hf_embeddings hf_embeddings = get_embeddings() # Chromadb vector stores have already been pre-created for the countries above for each of the different chunk sizes and overlaps, and zipped up, # to save time when experimenting as the embeddings take a long time to generate. # The existing stores will be pulled using from google drive above when app starts. When using the existing vector stores, # just need to change the name of the persist directory when selecting the different chunk sizes and overlaps. # Later in the main app if the user choose to scrape new data, or override with their own PDF, a new chromadb would be created. # This step will take some time if not os.path.exists("bm25.zip"): with st.spinner(f'Downloading bm25 retriever for all chunk sizes and overlaps, will take some time'): os.system("gdown https://drive.google.com/uc?id=1q-hNnyyBA8tKyF3vR69nkwCk9kJj7WHi") if not os.path.exists("chromadb.zip"): with st.spinner(f'Downloading chromadb retrievers for all chunk sizes and overlaps, will take some time'): os.system("gdown https://drive.google.com/uc?id=1zad6tgYm2o5M9E2dTLQqmm6GoI8kxNC3") if not os.path.exists("bm25/"): with st.spinner(f'Unzipping bm25 retriever for all chunk sizes and overlaps, will take some time'): os.system("unzip bm25.zip") if not os.path.exists("chromadb/"): with st.spinner(f'Unzipping chromadb retrievers for all chunk sizes and overlaps, will take some time'): os.system("unzip chromadb.zip") persist_directory = f"chromadb/chromadb_esg_countries_chunk_{st.session_state['chunk_size']}_overlap_{st.session_state['chunk_overlap']}" with st.spinner(f'Setting up pre-built chroma vector store'): chroma_db = Chroma(persist_directory=persist_directory,embedding_function=hf_embeddings) # Initialize BM25 Retriever # Unlike Chroma (semantic) BM25 is a keyword-based algorithm that performs well on queries containing keywords without capturing the semantic meaning of the query terms, # hence there is no need to embed the text with HuggingFaceEmbeddings and it is relatively faster to set up. # The retrievers with different chunking sizes and overlaps and countries were created in advanced and saved as pickle files and pulled using !wget. # Need to initialize one BM25Retriever for each country so the search results later in the main app can be limited to just a particular country. # (Chroma DB gives an option to filter metadata for just a particular country during the retrieval processbut BM25 does not because it makes use of external ranking library.) # A separate retriever was saved for each country. bm25_retrievers = {} # to store retrievers for different countries with st.spinner(f'Setting up pre-built bm25 retrievers'): for country in countries: bm25_filename = f"bm25/bm25_esg_countries_{country}_chunk_{st.session_state['chunk_size']}_overlap_{st.session_state['chunk_overlap']}.pickle" with open(bm25_filename, 'rb') as handle: bm25_retriever = pickle.load(handle) bm25_retrievers[country] = bm25_retriever # Tools for LLM to Use # The most important tool is the first one, which uses a RetrievalQA chain to answer a question about a specific country's ESG policies, # e.g. carbon emissions policy of Singapore. # By calling this tool multiple times, the agent is able to look at the responses from this tool for both countries and compare them. # This is far better than just retrieving relevant chunks for the user's query and throw everything to a single RetrievalQA chain to process # Multi input tools are not available, hence we have to prompt the agent to give an input list as a string # then use ast.literal_eval to convert it back into a list @tool def retrieve_answer_for_country(query_and_country: str) -> str: # TODO, change diff chain type diff version answers, change """Gives answer to a query about a single country's public ESG policy. The input list should be of the following format: [query, country] The first element of the list is the user query, surrounded by double quotes. The second element is the full name of the country involved, surrounded by double quotes, for example "Singapore". The 2 inputs are separated by a comma. Do not write a list comprehension. The 2 inputs, together, are surrounded by square brackets as it is a list. Do not put multiple countries into the input at once. Instead use this tool multiple times, one time for each country. If you have multiple queries to ask about a country, break the query into separate parts and use this tool multiple times, one for each query. """ try: query_and_country_list = ast.literal_eval(query_and_country) query = query_and_country_list[0] country = query_and_country_list[1].capitalize() # in case LLM did not capitalize first letter as filtering for metadata is case sensitive if not country in countries: return """The country that you input into the tool cannot be found. If you did not make a mistake and the country that you input is indeed what the user asked, then there is no record for the country and no answer can be obtained.""" # different retrievers bm = bm25_retrievers[country] # keyword based bm.k = st.session_state['bm25_n_similar_documents'] chroma = chroma_db.as_retriever(search_kwargs={'filter': {'country':country}, 'k': st.session_state['chroma_n_similar_documents']}) # semantic # ensemble (below) reranks results from both retrievers ensemble = EnsembleRetriever(retrievers=[bm, chroma], weights=[st.session_state['keyword_retriever_weight'], 1 - st.session_state['keyword_retriever_weight']]) retrievers = {'ensemble': ensemble, 'semantic': chroma, 'keyword': bm} qa = RetrievalQA.from_chain_type( llm=llm, chain_type='stuff', retriever=retrievers[st.session_state['retriever_config']], # selected retriever based on user config return_source_documents=True # returned in result['source_documents'] ) result = qa(query) st.session_state['source_documents'].append(result['source_documents']) # let user know what source docs are used return result['result'] except Exception as e: return f"""There is an error using this tool: {e}. Check if you have input anything wrongly and try again. Remember the 2 inputs, query and country, must both be surrounded by double quotes. The 2 inputs, together, are surrounded by square brackets as it is a list.""" # if a user tries to casually chat with the agent chatbot, the LLM will be able to use this tool to reply instead # this is optional, better to let user's know the chatbot is not for casual chatting @tool def generic_chat_llm(query: str) -> str: """Use this tool for general queries and casual chat. Forward the user input directly into this tool, do not come up with your own input. This tool IS NOT FOR MAKING COMPARISONS of anything. This tool IS NOT FOR FINDING ESG POLICY of any country! It is only for casual chat! Do not use this tool unnecessarily! """ try: # Second Generic Tool prompt = PromptTemplate( input_variables=["query"], template="{query}" ) llm_chain = LLMChain(llm=llm, prompt=prompt) return llm_chain.run(query) except Exception as e: return f"""There is an error using this tool: {e}. Check if you have input anything wrongly and try again. If you have already tried 2 times, do not try anymore, there is no response for your input. Move on to the next step of your plan.""" # sometimes the agent will suddenly ask for a 'compare' tool even though it was not given this tool # hence I have decided to give it this tool that gives a prompt to remind it to look at past information # and decide whether it is time to darw a conclusion # tools cannot have no input, hence I let the agent input a 'query' parameter even though it is not used # having the query as input let the LLM 'recall' what is being asked # instead of it being lost all the way at the start of the ReAct process @tool def compare(query:str) -> str: """Use this tool to give you hints and instructions on how you can compare between policies of countries. Use this tool only at one of your final steps, do not use it at the start. When putting the query into this tool, look at the entire query that the user has asked at the start, do not leave any details in the query out. """ return f"""Look at all your previous observations to answer the user query. Use as much relevant information as possible but only from your previous thoughts and observations. If you need more details, you can use a tool to find out more. If you have enough information, use your reasoning to answer them to the best of your ability. Give as much detail as you want in your answer.""" retrieve_answer_for_country.callbacks = [my_callback_handler] compare.callbacks = [my_callback_handler] generic_chat_llm.callbacks = [my_callback_handler] agent = initialize_agent( [retrieve_answer_for_country, compare], # tools #[retrieve_answer_for_country, generic_chat_llm, compare], llm=llm, agent="zero-shot-react-description", # this is good verbose=False, handle_parsing_errors=True, return_intermediate_steps=True, callbacks=[my_callback_handler] # memory=ConversationBufferMemory( # memory_key="chat_history", return_messages=True # ), # max_iterations=10 ) # Create a header element st.header("Chat") col1, col2 = st.columns(2) # with col1: # Store the conversation in the session state. # Used to render the chat conversation. # Initialize it with the first message for users to be greeted with if "messages" not in st.session_state: st.session_state.messages = [ {"role": "assistant", "content": "How may I help you today?"} ] if "current_response" not in st.session_state: st.session_state.current_response = "" # Loop through each message in the session state and render it as a chat message. for message in st.session_state.messages: with st.chat_message(message["role"]): st.markdown(message["content"]) # We initialize the quantized LLM from a local path. # Currently most parameters are fixed but we can make them # configurable. #llm_chain = create_chain(retriever) # We take questions/instructions from the chat input to pass to the LLM if user_query := st.chat_input("Your message here", key="user_input"): # Add our input to the session state st.session_state.messages.append( {"role": "user", "content": user_query} ) # Add our input to the chat window with st.chat_message("user"): st.markdown(user_query) # Let user know agent is planning the actions action_plan_message = "Please wait while I plan out a best set of actions to obtain the information and answer your query." # Add the response to the session state st.session_state.messages.append( {"role": "assistant", "content": action_plan_message} ) # Add the response to the chat window with st.chat_message("assistant"): st.markdown(action_plan_message) # Pass our input to the llm chain and capture the final responses. # It is worth noting that the Stream Handler is already receiving the # streaming response as the llm is generating. We get our response # here once the llm has finished generating the complete response. results = agent(user_query) response = f"The answer to your query is: {results['output']}" # Add the response to the session state st.session_state.messages.append( {"role": "assistant", "content": response} ) # Add the response to the chat window with st.chat_message("assistant"): st.markdown(response) # with col2: # st.write("hi")