Update app.py

app.py

@@ -1,11 +1,10 @@
 import streamlit as st
 import os
 from pathlib import Path
-from llama_index.core.query_engine.router_query_engine import RouterQueryEngine
 from llama_index.core.selectors import LLMSingleSelector
 from llama_index.core.tools import QueryEngineTool
-from llama_index.core import SummaryIndex, VectorStoreIndex
-from llama_index.core import VectorStoreIndex, Settings
+from llama_index.core import VectorStoreIndex
+from llama_index.core import Settings
 from llama_index.core import SimpleDirectoryReader
 from llama_index.llms.groq import Groq
 from llama_index.embeddings.huggingface import HuggingFaceEmbedding
@@ -13,6 +12,9 @@ from typing import Tuple
 from llama_index.core import StorageContext, load_index_from_storage
 from llama_index.core.objects import ObjectIndex
 from llama_index.core.agent import ReActAgent
+import time
+import sys
+import io
 
 # Function to process files and create document tools
 def create_doc_tools(document_fp: str, doc_name: str, verbose: bool = True) -> Tuple[QueryEngineTool,]:
@@ -45,8 +47,11 @@ def find_tex_files(directory: str):
     tex_files.sort()
     return tex_files
 
+
+
 # Main app function
 def main():
+    
     st.title("AMGPT, Powered by LlamaIndex")
 
     # API Key input
@@ -56,59 +61,98 @@ def main():
 
     llm = Groq(model="mixtral-8x7b-32768")
 
-    if "history" not in st.session_state:
-        st.session_state.history = []
-
-    if apikey:
-        directory = '/home/user/app/rag_docs_final_review_tex_merged'
-        tex_files = find_tex_files(directory)
-
-        paper_to_tools_dict = {}
-        for paper in tex_files:
-            path = Path(paper)
-            vector_tool = create_doc_tools(doc_name=path.stem, document_fp=path)
-            paper_to_tools_dict[path.stem] = [vector_tool]
-
-        initial_tools = [t for paper in tex_files for t in paper_to_tools_dict[Path(paper).stem]]
-
-        obj_index = ObjectIndex.from_objects(
-            initial_tools,
-            index_cls=VectorStoreIndex,
-        )
-
-        obj_retriever = obj_index.as_retriever(similarity_top_k=6)
+    with st.sidebar:
+        verbose_toggle = st.toggle("Verbose") # get verbose or only LLM response
+        reset = st.button('Reset Chat!') # reset the chat
 
-        context = """You are an agent designed to answer scientific queries over a set of given documents.
-        Please always use the tools provided to answer a question. Do not rely on prior knowledge.
-        """
 
-        agent = ReActAgent.from_tools(
-            tool_retriever=obj_retriever,
-            llm=llm,
-            verbose=True,
-            context=context
-        )
-
-        user_prompt = st.text_input("Enter your question")
-
-        if user_prompt:
-            st.session_state.history.append({"user": user_prompt})
-
-            with st.spinner("Processing..."):
-                response = agent.query(user_prompt)
-                st.session_state.history.append({"agent": response})
-
-            # Display the latest response
-            st.markdown(f"### Query Response:\n{response}")
-
-        # Display chat history
-        if st.session_state.history:
-            st.markdown("## Chat History")
-            for chat in st.session_state.history:
-                if "user" in chat:
-                    st.markdown(f"**User:** {chat['user']}")
-                if "agent" in chat:
-                    st.markdown(f"**Agent:** {chat['agent']}")
+    if apikey:
+        if "tools_loaded" not in st.session_state:
+            try:
+                directory = '/home/user/app/rag_docs_final_review_tex_merged'
+                tex_files = find_tex_files(directory)
+
+                paper_to_tools_dict = {}
+                for paper in tex_files:
+                    path = Path(paper)
+                    vector_tool = create_doc_tools(doc_name=path.stem, document_fp=path)
+                    paper_to_tools_dict[path.stem] = [vector_tool]
+        
+                initial_tools = [t for paper in tex_files for t in paper_to_tools_dict[Path(paper).stem]]
+        
+                obj_index = ObjectIndex.from_objects(
+                    initial_tools,
+                    index_cls=VectorStoreIndex,
+                )
+        
+                obj_retriever = obj_index.as_retriever(similarity_top_k=6)
+        
+        
+                context = """You are an agent designed to answer scientific queries over a set of given documents.
+                Please always use the tools provided to answer a question. Do not rely on prior knowledge.
+                """
+        
+                agent = ReActAgent.from_tools(
+                    tool_retriever=obj_retriever,
+                    llm=llm,
+                    verbose=True,
+                    context=context
+                )
+
+                # store session state variables
+                st.session_state["tools_loaded"] = True
+                st.session_state["agent"] = agent
+            except Exception as e:
+                st.error(e)
+    
+    
+    
+    
+    
+    if "messages" not in st.session_state or reset==True:
+        st.session_state["messages"] = [{"role": "assistant", "content": "How can I help you?"}]
+    
+    for msg in st.session_state.messages:
+        st.chat_message(msg["role"]).write(msg["content"])
+    
+    if prompt := st.chat_input():
+    
+        # if the user started chatting without setting the OPENAI API KEY
+        if not apikey:
+            st.info("Please add your Groq API key to continue.")
+            st.stop()
+    
+        st.session_state.messages.append({"role": "user", "content": prompt})
+        st.chat_message("user").write(prompt)
+        try:
+            with st.spinner('Wait for output...'):
+                # Redirect stdout
+                original_stdout = sys.stdout
+                sys.stdout = io.StringIO()
+    
+                # query the agent
+                response = st.session_state.agent.query(prompt)
+    
+                # Get the captured output and restore stdout
+                output = sys.stdout.getvalue()
+                sys.stdout = original_stdout
+    
+                # format the received verbose output
+                verbose = ''
+                for output_string in output.split('==='):
+                    verbose+=output_string
+                    verbose+='\n'
+    
+                # assistant response
+                msg = f'{verbose}' if verbose_toggle else f'{response.response[10:]}'
+            
+            # write the response
+            st.session_state.messages.append({"role": "assistant", "content": msg})
+            st.chat_message("assistant").write(msg)
+        except Exception as e:
+            st.error(e)
 
 if __name__ == "__main__":
     main()
+
+