Update app.py

app.py

@@ -1,5 +1,6 @@
 import streamlit as st
-from streamlit_chat import message
+#from streamlit_chat import message
+from streamlit_option_menu import option_menu
 
 import os
 from langchain.llms import HuggingFaceHub # for calling HuggingFace Inference API (free for our use case)
@@ -26,15 +27,17 @@ warnings.filterwarnings("ignore", category=DeprecationWarning)
 
 # os.environ['HUGGINGFACEHUB_API_TOKEN'] = 'your_api_key' # for using HuggingFace Inference API
 
-
 from langchain.callbacks.base import BaseCallbackHandler
+
+# callback is needed to print intermediate steps of agent reasoning in the chatbot
+# i.e. when action is taken, when tool is called, when tool call is complete etc.
 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_llm_new_token(self, token, **kwargs) -> None: # HuggingFaceHub() cannot stream unfortunately!
+    #     self.tokens.append(token)
+    #     print(token)      
 
     def on_agent_action(self, action, **kwargs):
         """Run on agent action."""
@@ -74,7 +77,7 @@ class MyCallbackHandler(BaseCallbackHandler):
     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..."
+        tool_output = f"Tool Output for Me: {output}  \n  \nI am processing the output from the tool..."
         st.session_state.messages.append(
             {"role": "assistant", "content": tool_output}
         )
@@ -114,7 +117,6 @@ if 'countries_to_scrape' not in st.session_state:
 # 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)
@@ -155,7 +157,6 @@ countries = [
     "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
@@ -264,6 +265,7 @@ def retrieve_answer_for_country(query_and_country: str) -> str: # TODO, change d
             return_source_documents=True # returned in result['source_documents']
         )
         result = qa(query)
+        st.session_state['source_documents'].append(f"Documents retrieved for agent query '{query}' for country '{country}'.") 
         st.session_state['source_documents'].append(result['source_documents']) # let user know what source docs are used
         return result['result']
     
@@ -305,14 +307,17 @@ def generic_chat_llm(query: str) -> str:
 @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.
+    Use this tool as a final step, only after you have used other tools to obtain all the information you need.
     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."""
+    return f"""Once again, check through all your previous observations to answer the user query. 
+    Make sure every part of the query is addressed by the context, or that you have at least tried to do so.
+    Make sure you have not forgotten to address anything in the query.
+    If you still need more details, you can use another tool to find out more if you have not tried using the same tool with the necessary input earlier. 
+    If you have enough information, use your reasoning to answer them to the best of your ability. 
+    Give as much elaboration in your answer as possible but they MUST be from the earlier context.
+    Do not give details that cannot be found in the earlier context."""
 
 retrieve_answer_for_country.callbacks = [my_callback_handler]
 compare.callbacks = [my_callback_handler]
@@ -333,77 +338,94 @@ agent = initialize_agent(
     # max_iterations=10
 )
 
+# original menu options
+if "menu" not in st.session_state:
+    st.session_state["menu"] = ["Chatbot", 'Source Documents  \n (for Last Query, Click Only After Full Execution)', 'Settings']
 
-
-# Create a header element
-st.header("Chat")
+with st.sidebar:
+    selected = option_menu("Main Menu", st.session_state["menu"], 
+        icons=['house', 'gear', 'gear'], menu_icon="cast", default_index=0)
 
 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)
+if selected == "Chatbot":
+    st.header("Chat")
+    
+    # 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? example qn, "}
+        ]
+
+    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"):
+        # remove source documents option from menu while query is running
+
+        st.session_state['source_documents'] = [f"User query: '{user_query}'"] # reset source documents list
+
+        formatted_user_query = f":blue[{user_query}]"
+        # Add our input to the session state
+        st.session_state.messages.append(
+            {"role": "user", "content": formatted_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 our input to the chat window
+        with st.chat_message("user"):
+            st.markdown(formatted_user_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)
+        # 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 necessary information to answer your query."
 
-    # 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": action_plan_message}
+        )
+        # Add the response to the chat window
+        with st.chat_message("assistant"):
+            st.markdown(action_plan_message)
 
-    # Add the response to the session state
-    st.session_state.messages.append(
-        {"role": "assistant", "content": response}
-    )
+        # 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":blue[The answer to your query is:] {results['output']}"
 
-    # Add the response to the chat window
-    with st.chat_message("assistant"):
-        st.markdown(response)
+        # 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")
 
+    # with col2:
+    #     st.write("hi")
 
 
+if selected == "Source Documents  \n (for Last Query, Click Only After Full Execution)":
+    st.header("Source Documents for Last Query")
+    try:
+        st.subheader(st.session_state['source_documents'][0])
+        for doc in st.session_state['source_documents'][1:]:
+            st.write(doc)
+    except:
+        st.write("No source documents retrieved yet. Please run a user query before coming back to this page.")