Update functions.py

functions.py

@@ -26,25 +26,44 @@ from langgraph.graph import START, END, StateGraph
 
 
 
-async def predict_custom_agent_answer(example: dict):
+async def handle_userinput(user_question, custom_graph):
+    # Add the user's question to the chat history and display it in the UI
+    st.session_state.messages.append({"role": "user", "content": user_question})
+    st.chat_message("user").write(user_question)
+
+    # Generate a unique thread ID for the graph's state
     config = {"configurable": {"thread_id": str(uuid.uuid4())}}
-    
+
     try:
         # Invoke the custom graph with the input question
-        state_dict = await  custom_graph.ainvoke(
-            {"question": example["input"], "steps": []}, config
+        state_dict = await custom_graph.ainvoke(
+            {"question": user_question, "steps": []}, config
         )
+
+        # Retrieve the documents from the graph's state (if available)
+        docs = state_dict.get("documents", [])
         
-        # Check if 'generation' exists and is not empty
+        # Display the retrieved documents in the sidebar
+        with st.sidebar:
+            st.subheader("Your documents")
+            with st.spinner("Processing"):
+                for doc in docs:
+                    st.write(f"Document: {doc.page_content}")  # Assuming doc.page_content contains the text
+
+        # Check if a response (generation) was produced by the graph
         if 'generation' in state_dict and state_dict['generation']:
-            return {"response": state_dict["generation"], "steps": state_dict["steps"]}
+            response = state_dict["generation"]
+
+            # Add the assistant's response to the chat history and display it
+            st.session_state.messages.append({"role": "assistant", "content": response})
+            st.chat_message("assistant").write(response)
         else:
-            # Raise an exception if the generation is missing or empty
-            print("Your question violates toxicity rules or contains sensitive information.")
-    
+            # Handle cases where no valid generation is present
+            st.chat_message("assistant").write("Your question violates toxicity rules or contains sensitive information.")
+
     except Exception as e:
-        # Raise the exception with a custom message
-        print("An error occurred: Try to change the question.") 
+        # Display an error message in case of failure
+        st.chat_message("assistant").write("An error occurred: Try to change the question.")