157 lines
7.2 KiB
Python
157 lines
7.2 KiB
Python
import streamlit as st
|
|
from flow import create_initial_processing_flow, create_finalization_flow
|
|
|
|
st.title("PocketFlow HITL with Streamlit")
|
|
|
|
# Initialize session state variables if they don't exist
|
|
if 'stage' not in st.session_state:
|
|
st.session_state.stage = "initial"
|
|
st.session_state.error_message = None
|
|
st.session_state.task_input = ""
|
|
# Flow-related data will be added directly as needed
|
|
print("Initialized session state.")
|
|
|
|
# --- Helper Function to Reset State ---
|
|
def reset_state():
|
|
# Keep essential Streamlit state keys if necessary, or clear selectively
|
|
keys_to_clear = [k for k in st.session_state.keys() if k not in ['stage', 'error_message', 'task_input']]
|
|
for key in keys_to_clear:
|
|
del st.session_state[key]
|
|
|
|
st.session_state.stage = "initial"
|
|
st.session_state.error_message = None
|
|
st.session_state.task_input = ""
|
|
print("Reset session state (keeping core stage/error keys).")
|
|
|
|
# --- Display Area for Shared Data (now the entire session state) ---
|
|
with st.expander("Show Session State (Shared Data)"):
|
|
# Convert to dict for clean JSON display, excluding internal Streamlit keys if desired
|
|
display_state = {k: v for k, v in st.session_state.items() if not k.startswith("_")}
|
|
st.json(display_state)
|
|
|
|
# --- Stage: Initial Input ---
|
|
if st.session_state.stage == "initial":
|
|
st.header("1. Submit Data for Processing")
|
|
# Use st.session_state.task_input directly for the text area's value
|
|
task_input_value = st.text_area("Enter data to process:", value=st.session_state.task_input, height=150)
|
|
|
|
if st.button("Submit"):
|
|
if not task_input_value.strip():
|
|
st.error("Please enter some data to process.")
|
|
else:
|
|
print(f"Submit button clicked. Input: '{task_input_value[:50]}...'")
|
|
# Store input directly in session state
|
|
st.session_state.task_input = task_input_value
|
|
st.session_state.error_message = None
|
|
# Clear previous results if any
|
|
if "processed_output" in st.session_state: del st.session_state.processed_output
|
|
if "final_result" in st.session_state: del st.session_state.final_result
|
|
if "input_used_by_process" in st.session_state: del st.session_state.input_used_by_process
|
|
|
|
try:
|
|
with st.spinner("Processing initial task..."):
|
|
initial_flow = create_initial_processing_flow()
|
|
# Pass the entire session state as shared data
|
|
initial_flow.run(st.session_state)
|
|
|
|
# Check if processing was successful (output exists directly in session state)
|
|
if "processed_output" in st.session_state:
|
|
st.session_state.stage = "awaiting_review"
|
|
print("Initial processing complete. Moving to 'awaiting_review' stage.")
|
|
st.rerun()
|
|
else:
|
|
st.session_state.error_message = "Processing failed to produce an output."
|
|
print("Error: Processing failed, no output found.")
|
|
# Keep stage as initial to allow retry/correction
|
|
|
|
except Exception as e:
|
|
st.session_state.error_message = f"An error occurred during initial processing: {e}"
|
|
print(f"Exception during initial processing: {e}")
|
|
# Keep stage as initial
|
|
|
|
# --- Stage: Awaiting Review ---
|
|
elif st.session_state.stage == "awaiting_review":
|
|
st.header("2. Review Processed Output")
|
|
# Get processed output directly from session state
|
|
processed_output = st.session_state.get("processed_output", "Error: Processed output not found!")
|
|
|
|
st.subheader("Output to Review:")
|
|
st.markdown(f"```\n{str(processed_output)}\n```") # Display as markdown code block
|
|
|
|
col1, col2, _ = st.columns([1, 1, 5]) # Layout buttons
|
|
with col1:
|
|
if st.button("Approve"):
|
|
print("Approve button clicked.")
|
|
st.session_state.error_message = None
|
|
try:
|
|
with st.spinner("Finalizing result..."):
|
|
finalization_flow = create_finalization_flow()
|
|
# Pass the entire session state
|
|
finalization_flow.run(st.session_state)
|
|
|
|
# Check for final result directly in session state
|
|
if "final_result" in st.session_state:
|
|
st.session_state.stage = "completed"
|
|
print("Approval processed. Moving to 'completed' stage.")
|
|
st.rerun()
|
|
else:
|
|
st.session_state.error_message = "Finalization failed to produce a result."
|
|
print("Error: Finalization failed, no final_result found.")
|
|
# Stay in review stage and show error.
|
|
|
|
except Exception as e:
|
|
st.session_state.error_message = f"An error occurred during finalization: {e}"
|
|
print(f"Exception during finalization: {e}")
|
|
# Stay in review stage
|
|
st.rerun() # Rerun to show error message
|
|
|
|
with col2:
|
|
if st.button("Reject"):
|
|
print("Reject button clicked.")
|
|
st.session_state.error_message = None # Clear previous errors
|
|
# Go back to initial stage to allow modification/resubmission
|
|
st.session_state.stage = "initial"
|
|
# Keep the rejected output visible in the input field for modification
|
|
st.session_state.task_input = st.session_state.get("processed_output", st.session_state.task_input)
|
|
# Clear the processed output so it doesn't linger
|
|
if "processed_output" in st.session_state: del st.session_state.processed_output
|
|
if "final_result" in st.session_state: del st.session_state.final_result
|
|
st.info("Task rejected. Modify the input below and resubmit.")
|
|
print("Task rejected. Moving back to 'initial' stage.")
|
|
st.rerun()
|
|
|
|
# --- Stage: Completed ---
|
|
elif st.session_state.stage == "completed":
|
|
st.header("3. Task Completed")
|
|
# Get final result directly from session state
|
|
final_result = st.session_state.get("final_result", "Error: Final result not found!")
|
|
st.subheader("Final Result:")
|
|
st.success("Task approved and completed successfully!")
|
|
st.text_area("", value=str(final_result), height=200, disabled=True)
|
|
|
|
if st.button("Start Over"):
|
|
print("Start Over button clicked.")
|
|
reset_state()
|
|
st.rerun()
|
|
|
|
# --- Stage: Rejected ---
|
|
elif st.session_state.stage == "rejected_final":
|
|
st.header("3. Task Rejected")
|
|
st.error("The processed output was rejected.")
|
|
# Get rejected output directly from session state
|
|
rejected_output = st.session_state.get("processed_output", "")
|
|
if rejected_output:
|
|
st.text_area("Rejected Output:", value=str(rejected_output), height=150, disabled=True)
|
|
|
|
if st.button("Start Over"):
|
|
print("Start Over button clicked.")
|
|
reset_state()
|
|
st.rerun()
|
|
|
|
# --- Display Error Messages ---
|
|
if st.session_state.error_message:
|
|
st.error(st.session_state.error_message)
|
|
|
|
# --- Add a button to reset state anytime (for debugging) ---
|
|
# st.sidebar.button("Reset State", on_click=reset_state) # Removed sidebar
|