wasted a lot of time on streamlit today, but it now seems like it is kind of working

app.py

@@ -1,7 +1,9 @@
 import re
+from enum import Enum
 
 import streamlit as st
 from streamlit import session_state as ss
+from streamlit.runtime.scriptrunner_utils.exceptions import StopException
 
 from form.form import build_form_data_from_answers, write_pdf_form, work_categories
 from llm_manager.llm_parser import LlmParser
@@ -16,6 +18,23 @@ user_msg = "Please describe what you need to do. To get the best results try to
 
 find_tags_regex = re.compile(r"(@\S)")
 
+class Steps(Enum):
+    INITIAL_STATE = 1
+    PARSING_ANSWERS = 2
+    ASK_AGAIN = 3
+    FIND_CATEGORIES = 4
+    VALIDATE_DATA = 5
+    PARSING_ERROR = 6
+    FORM_CREATED = 7
+
+    def __eq__(self, other):
+        if not isinstance(other, self.__class__):
+            return False
+        return self.value == other.value
+
+    def __hash__(self):
+        return hash(self.value)
+
 class UIManager:
     def __init__(self):
         self.pm: PromptsManager = PromptsManager(work_categories=work_categories)
@@ -24,26 +43,33 @@ class UIManager:
                                           Model("fakeModel", ModelRoles("a", "b", "c"))))
 
     @staticmethod
-    def get_current_step():
-        return ss.get("step")
+    def get_current_step() -> int:
+        try:
+            return ss.get("step") or Steps.INITIAL_STATE.value
+        except StopException:
+            return Steps.INITIAL_STATE.value
+
+    @staticmethod
+    def set_current_step(step: Steps):
+        ss["step"] = step.value
 
     @staticmethod
     def _build_base_ui():
         st.markdown("## Dubai Asset Management red tape cutter")
 
     def build_ui_for_initial_state(self, user_message):
-        help_ = user_message
+        help_ = user_message + "\n".join(self.pm.questions)
         self._build_base_ui()
         with st.form("Please describe your request"):
-            user_input = st.text_area("Your input", height=700, label_visibility="hidden", placeholder=help_,
-                                      help=help_)
-            signature = st.file_uploader("Your signature", key="file_upload")
-            ss["signature"] = signature
-            submit_button = st.form_submit_button()
-            if submit_button:
-                ss["user_input"] = user_input
-                ss["step"] = "parsing_answers"
-                st.rerun()
+                user_input = st.text_area("Your input", height=700, label_visibility="hidden", placeholder=help_,
+                                          help=help_)
+                signature = st.file_uploader("Your signature", key="file_upload")
+                ss["signature"] = signature
+                submit_button = st.form_submit_button()
+                if submit_button:
+                    ss["user_input"] = user_input
+                    self.set_current_step(Steps.PARSING_ANSWERS)
+                    st.rerun()
 
     def build_ui_for_parsing_answers(self):
         self._build_base_ui()
@@ -60,25 +86,25 @@ class UIManager:
             answers = LlmParser.parse_verification_prompt_answers(answer['content'])
             ss["answers"] = answers
             if len(answers) != len(Q):
-                ss["step"] = "parsing_error"
+                self.set_current_step(Steps.PARSING_ERROR)
                 st.rerun()
         ss["missing_answers"] = check_for_missing_answers(ss["answers"])
         if not ss.get("missing_answers"):
-            ss["step"] = "check_category"
+            self.set_current_step(Steps.FIND_CATEGORIES)
         else:
-            ss["step"] = "ask_again"
+            self.set_current_step(Steps.ASK_AGAIN)
         st.rerun()
 
     def build_ui_for_ask_again(self):
         self._build_base_ui()
         with st.form("form1"):
             for ma in ss["missing_answers"]:
-                st.text_input(self.pm.questions[ma].lower(), key=ma)
+                st.text_input(self.pm.questions[ma.value].lower(), key=ma)
             submitted = st.form_submit_button("Submit answers")
             if submitted:
                 for ma in ss["missing_answers"]:
                     ss["answers"][ma] = ss[ma]
-                ss["step"] = "check_category"
+                self.set_current_step(Steps.FIND_CATEGORIES)
                 st.rerun()
 
     def build_ui_for_check_category(self):
@@ -87,7 +113,7 @@ class UIManager:
             answer = self.repository.send_prompt(self.pm.get_work_category(ss["answers"][Q.WORK_TO_DO]))
             categories = LlmParser.parse_get_categories_answer(answer['content'])
             ss["categories"] = categories
-            ss["step"] = "validate_data"
+            self.set_current_step(Steps.VALIDATE_DATA)
             st.rerun()
 
     def build_ui_for_form_created(self):
@@ -114,7 +140,7 @@ class UIManager:
             if details:
                 key = ss[details_key]  # get the name under which this data should be saved
                 save_details(details, key)
-        ss["step"] = "check_categories"
+        self.set_current_step(Steps.FIND_CATEGORIES)
         st.rerun()
 
     def _create_pdf_form(self):
@@ -125,14 +151,16 @@ class UIManager:
             pdf_form_filename = filename
             ss["pdf_form"] = pdf_form
             ss["pdf_form_filename"] = pdf_form_filename
-            ss["step"] = "form_created"
+            self.set_current_step(Steps.FORM_CREATED)
             st.rerun()
 
     def build_ui_for_validate_data_after_correction(self):
-        self._build_validation_form(False, self._integrate_llm_answers_with_user_corrections)
+        self._build_validation_form(False, self._integrate_llm_answers_with_user_corrections,
+                                    "Find work categories")
 
     def build_ui_to_confirm_form_data(self):
-        self._build_validation_form(True, self._create_pdf_form)
+        self._build_validation_form(True, self._create_pdf_form,
+                                    "Create work permit request")
 
 
     @staticmethod
@@ -159,7 +187,7 @@ class UIManager:
 
 
 
-    def _build_validation_form(self, show_categories:bool, onsubmit):
+    def _build_validation_form(self, show_categories:bool, onsubmit, submit_button_label):
         def build_form_fragment(form_, col, title, add_save, *questions):
             form_.text(title)
             for user_data in questions:
@@ -182,27 +210,28 @@ class UIManager:
             for k, wc in work_categories.items():
                 f.checkbox(label=wc, key=k, value=k in ss["categories"])
 
-        submit_data = f.form_submit_button()
+        submit_data = f.form_submit_button(label=submit_button_label)
         if submit_data:
             onsubmit()
 
-def use_streamlit():
 
-    um = UIManager()
+um = UIManager()
+
+def use_streamlit():
 
-    if not um.get_current_step():
+    if um.get_current_step() == Steps.INITIAL_STATE.value :
         um.build_ui_for_initial_state(user_msg)
-    if um.get_current_step() == "parsing_answers":
+    elif um.get_current_step() == Steps.PARSING_ANSWERS.value:
         um.build_ui_for_parsing_answers()
-    if um.get_current_step() == "parsing_error":
+    elif um.get_current_step() == Steps.PARSING_ERROR.value:
         um.build_ui_for_validate_data_after_correction()
-    if um.get_current_step() == "ask_again":
+    elif um.get_current_step() == Steps.ASK_AGAIN.value:
         um.build_ui_for_ask_again()
-    if um.get_current_step() == "check_category":
+    elif um.get_current_step() == Steps.FIND_CATEGORIES.value:
         um.build_ui_for_check_category()
-    if um.get_current_step() == "validate_form_data":
+    elif um.get_current_step() == Steps.VALIDATE_DATA.value:
         um.build_ui_to_confirm_form_data()
-    if um.get_current_step() == "form_created":
+    elif um.get_current_step() == Steps.FORM_CREATED.value:
         um.build_ui_for_form_created()
 
 

local_storage/ls_manager.py

@@ -9,7 +9,7 @@ ls = LocalStorage()
 
 def get_detail(key: str):
     for detail_type in DetailsType.values():
-        detail = json.loads((ls.getItem(detail_type.name) or {}).get(key))
+        detail = (ls.getItem(detail_type.name) or {}).get(key)
         if detail:
             type_ = detail["type"]
             if type_ == DetailsType.PERSONAL_DETAILS.name:
@@ -31,11 +31,11 @@ def save_details(details: SavedDetails, key: str):
         type_ = DetailsType.CONTRACTOR_DETAILS
     else:
         raise ValueError("Unexpected type: {}", type(details))
-    existing_data = json.loads(ls.getItem(type_.name) or "{}")
+    existing_data = (ls.getItem(type_.name) or {})
     existing_data[key] = details.to_json()
     ls.setItem(type_.name, existing_data, type_.name)
 
 
 def get_details(type_: DetailsType):
-    return json.loads((ls.getItem(type_.name) or "{}"))
+    return ls.getItem(type_.name) or {}