Upload med_streamlit.py

med_streamlit.py

@@ -49,6 +49,19 @@ def update_dataframe(records: List[Dict] | pd.DataFrame):
     st.session_state.df = new_data  # Assign the updated DataFrame
     #st.rerun()  # Trigger a rerun
 
+def undo():
+    """Undo the last operation by restoring the previous DataFrame."""
+    if st.session_state.prev_df is not None:
+        st.session_state.df = st.session_state.prev_df
+        st.session_state.prev_df = None
+
+    if st.session_state.history:
+        st.session_state.history.pop()
+    if st.session_state.explanation:
+        st.session_state.explanation = "Changes undone."
+    if st.session_state.references:
+        st.session_state.references = ""
+
 
 # Page config
 st.set_page_config(layout="wide", page_title="RP Medication Analyzer")
@@ -114,6 +127,9 @@ if "last_response" not in st.session_state:
     st.session_state.last_response = {}
 if "history" not in st.session_state:
     st.session_state.history = []  # Stores all past interactions
+if "prev_df" not in st.session_state:
+    st.session_state.prev_df = None  # Stores the previous DataFrame for undo
+
 
 # File uploader
 file = st.file_uploader("Upload an Excel file", type=["xlsx"])
@@ -129,7 +145,7 @@ if file and file != st.session_state.uploaded_file:
                 update_dataframe(data_df)
             else:
                 st.error("🚨 No 'Data' sheet found in the uploaded file. Make sure the file has it")
-
+            print(f"History: {st.session_state.history}")
         st.session_state.uploaded_file = file
         print("File uploaded successfully!")
         st.success("✅ File uploaded successfully!")
@@ -171,7 +187,6 @@ if input_text:
     else:
         # Convert dataframe to JSON for LLM processing
         json_data = st.session_state.df.to_json(orient="records")
-        print(json_data)
         with st.spinner(f"Processing request: *{input_text}*..."):
             response = None  # Ensure response is defined before use
 
@@ -191,14 +206,11 @@ if input_text:
                     api_key=api_key,
                 )
 
-            print(f"Response:{response}")
-
         # Ensure response exists before processing
         if response:
-            st.session_state.df = None
+            st.session_state.prev_df = st.session_state.df
             try:
                 parsed_response = validate_llm_response(response)
-                print(f"Parsed response: {parsed_response}")
 
                 st.session_state.last_prompt = input_text
                 st.session_state.last_response = response  # Keep full JSON response
@@ -207,7 +219,6 @@ if input_text:
                 if "error" in parsed_response:
                     st.error(parsed_response["error"])
                 else:
-                    print(f"Parsed data: {parsed_response['dataset']}")
                     update_dataframe(parsed_response["dataset"])
                     st.session_state.explanation = parsed_response["explanation"]
                     st.session_state.references = parsed_response["references"]
@@ -260,3 +271,7 @@ if st.session_state.df is not None:
         file_name="updated_dataset.xlsx",
         mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
     )
+
+st.sidebar.subheader("Undo Changes")
+if st.sidebar.button("Undo", disabled=st.session_state.prev_df is None):
+    undo()