Update app.py

app.py

@@ -204,7 +204,6 @@ def run_agent(
     temperature: float=0.7,
     agent_type: Literal["Tool Calling", "ReAct"]="Tool Calling",
 ) -> Union[str, None]:
-
     """
     Generate text based on user queries.
     Args:
@@ -216,40 +215,42 @@ def run_agent(
         agent_type: 'Tool Calling' or 'ReAct'
     Return:
         generated text
-    The chat prompt and message history are stored in
-    st.session_state variables.
     """
 
     try:
-        # Ensure retriever tool is included in tool
+        # Ensure retriever tool is included when "Retrieval" is selected
         if "Retrieval" in st.session_state.tool_names[0]:
             if st.session_state.retriever_tool:
-                if st.session_state.retriever_tool not in tools:
+                retriever_tool_name = "retriever"  # Ensure naming consistency
+                if retriever_tool_name not in [tool.name for tool in tools]:
                     tools.append(st.session_state.retriever_tool)
-                st.write("**Retriever tool has been successfully added to the tools list.**")
+                    st.write(f"✅ **{retriever_tool_name} tool has been added successfully.**")
             else:
-                st.error("Retriever tool is not initialized. Please create a vector store first.")
+                st.error("❌ Retriever tool is not initialized. Please create a vector store first.")
+                return None  # Exit early to avoid broken tool usage
+
+        # Debugging: Print final tools list
+        st.write("**Final Tools Being Used:**", [tool.name for tool in tools])
 
-        # Debugging: Print tools being used
-        st.write("**Tools being used by the agent:**", [tool.name for tool in tools])
         if "retriever" in [tool.name for tool in tools]:
-            st.write("✅ Retriever tool is included in the tools list.")
-        else:
-            st.error("❌ Retriever tool is missing. Check the tool setup.")
+            st.success("✅ Retriever tool is confirmed and ready for use.")
+        elif "Retrieval" in st.session_state.tool_names[0]:
+            st.warning("⚠️ 'Retrieval' was selected but the retriever tool is missing!")
 
-        st.write("**Tools passed to run_agent:**", [tool.name for tool in tools])
+        # Initialize the LLM model
         llm = get_chat_model(model, temperature, [StreamHandler(st.empty())])
         if llm is None:
-            st.error(f"Unsupported model: {model}", icon="🚨")
+            st.error(f"❌ Unsupported model: {model}", icon="🚨")
             return None
         
+        # Prepare chat history
         if agent_type == "Tool Calling":
             chat_history = st.session_state.history
         else:
             chat_history = message_history_to_string()
 
         history_query = {"chat_history": chat_history, "input": query}
-        
+
         # Generate message content
         message_with_no_image = st.session_state.chat_prompt.invoke(history_query)
         message_content = message_with_no_image.messages[0].content
@@ -286,7 +287,6 @@ def run_agent(
         return None
 
 
-
 def openai_create_image(
     description: str, model: str="dall-e-3", size: str="1024x1024"
 ) -> Optional[str]:
@@ -379,10 +379,14 @@ def get_vector_store(uploaded_files: List[UploadedFile]) -> Optional[FAISS]:
 def get_retriever() -> None:
     """
     Upload document(s), create a vector store, prepare a retriever tool,
-    save the tool to the variable st.session_state.retriever_tool
+    save the tool to the variable st.session_state.retriever_tool.
     """
+
+    # Section Title
     st.write("")
     st.write("**Query Document(s)**")
+
+    # File Upload Input
     uploaded_files = st.file_uploader(
         label="Upload an article",
         type=["txt", "pdf", "docx"],
@@ -391,29 +395,42 @@ def get_retriever() -> None:
         key="document_upload_" + str(st.session_state.uploader_key),
     )
 
-    if uploaded_files and st.button(label="Create the vector store"):
-        st.info("Creating the vector store and initializing the retriever tool...")
-        vector_store = get_vector_store(uploaded_files)
+    # Check if files are uploaded
+    if uploaded_files:
+        # Use a unique button key to avoid duplicate presses
+        if st.button(label="Create the vector store", key=f"create_vector_{st.session_state.uploader_key}"):
+            st.info("Creating the vector store and initializing the retriever tool...")
+
+            # Attempt to create the vector store
+            vector_store = get_vector_store(uploaded_files)
 
-        if vector_store:
-            uploaded_file_names = [file.name for file in uploaded_files]
-            st.session_state.vector_store_message = (
-                f"Vector store for :blue[[{', '.join(uploaded_file_names)}]] is ready!"
-            )
-            retriever = vector_store.as_retriever()
-            st.session_state.retriever_tool = create_retriever_tool(
-                retriever,
-                name="retriever",
-                description="Search uploaded documents for information when queried.",
-            )
+            if vector_store:
+                uploaded_file_names = [file.name for file in uploaded_files]
+                st.session_state.vector_store_message = (
+                    f"Vector store for :blue[[{', '.join(uploaded_file_names)}]] is ready!"
+                )
 
-            # Dynamically add Retrieval to the tool_names if not already present
-            if "Retrieval" not in st.session_state.tool_names[0]:
-                st.session_state.tool_names[0].append("Retrieval")
+                # Initialize retriever and create tool
+                retriever = vector_store.as_retriever()
+                st.session_state.retriever_tool = create_retriever_tool(
+                    retriever,
+                    name="retriever",
+                    description="Search uploaded documents for information when queried.",
+                )
+
+                # Add "Retrieval" to the tools list if not already present
+                if "Retrieval" not in st.session_state.tool_names[0]:
+                    st.session_state.tool_names[0].append("Retrieval")
+
+                st.success("✅ Retriever tool has been successfully initialized and is ready to use.")
+
+                # Debugging output
+                st.write("**Current Tools:**", st.session_state.tool_names[0])
+            else:
+                st.error("❌ Failed to create vector store. Please check the uploaded files (supported formats: txt, pdf, docx).")
+    else:
+        st.info("Please upload document(s) to create the vector store.")
 
-            st.success("Retriever tool has been successfully initialized and is ready to use.")
-        else:
-            st.error("Failed to create vector store. Please check the uploaded files.")