Update ui_utils.py

ui_utils.py

@@ -1,6 +1,11 @@
 import streamlit as st
 
 def display_initial_buttons():
+    if "upload_flow" not in st.session_state:
+        st.session_state.upload_flow = False
+    if "query_flow" not in st.session_state:
+        st.session_state.query_flow = False
+
     if st.button("Upload new CSV"):
         st.session_state.upload_flow = True
         st.session_state.query_flow = False
@@ -9,19 +14,17 @@ def display_initial_buttons():
         st.session_state.query_flow = True
         st.session_state.upload_flow = False
 
-
 def display_class_dropdown(client):
     if st.session_state.upload_flow:
         existing_classes = [cls["class"] for cls in client.schema.get()["classes"]]
         class_options = existing_classes + ["New Class"]
-        selected_class = st.selectbox("Select a class or create a new one:", class_options)
-    if st.session_state.query_flow:
+        return st.selectbox("Select a class or create a new one:", class_options)
+    elif st.session_state.query_flow:
         existing_classes = [cls["class"] for cls in client.schema.get()["classes"]]
         class_options = existing_classes + ["Query all data"]
-        selected_class = st.selectbox("Select a class or query all data:", class_options)
+        return st.selectbox("Select a class or query all data:", class_options)
 
-
-def handle_new_class_selection():
+def handle_new_class_selection(selected_class):
     if selected_class == "New Class":
         class_name = st.text_input("Enter the new class name:")
         class_description = st.text_input("Enter a description for the class:")
@@ -30,39 +33,16 @@ def handle_new_class_selection():
                 # Call function to create new class schema in Weaviate
                 create_new_class_schema(class_name, class_description)
 
-    if selected_class != "Query all data":
-        # Display data preview for selected class
-        display_data_preview(selected_class)
-        if st.button("Confirm Data Selection"):
-            # Set the selected class as the data source for TAPAS queries
-            st.session_state.data_source = selected_class
-
-    else:
-        # Display all class schemas
-        display_all_class_schemas()
-        if st.button("Confirm Data Selection"):
-            # Set all data as the source for TAPAS queries
-            st.session_state.data_source = "all_data"
-
-
-def csv_upload_and_ingestion():
+def csv_upload_and_ingestion(selected_class):
     csv_file = st.file_uploader("Upload a CSV file", type=["csv"])
     if csv_file:
         if st.button("Confirm CSV upload"):
             # Call function to ingest CSV data into Weaviate
             ingest_csv_to_weaviate(csv_file, selected_class)
 
-
 def display_query_input():
     question = st.text_input("Enter your question:")
     if question:
         if st.button("Submit Query"):
             # Call function to query TAPAS with selected data and entered question
-            query_tapas_with_weaviate_data(selected_class, question)
-
-    question = st.text_input("Enter your question:")
-    if question:
-        if st.button("Submit Query"):
-            # Call function to query TAPAS with selected data source and entered question
             query_tapas_with_weaviate_data(st.session_state.data_source, question)
-