app.py

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
from local_storage.entities import PersonalDetails, LocationDetails, ContractorDetails
from local_storage.ls_manager import save_details, get_detail
from prompts.prompts_manager import PromptsManager
from enums import Questions as Q
from repository.repository import build_repo_from_environment, get_repository
from repository import ModelRoles, Model
from utils.parsing_utils import check_for_missing_answers

user_msg = "Please describe what you need to do. To get the best results try to answer all the following questions:"

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)
        self.repository = (build_repo_from_environment(self.pm.system_prompt) or
                           get_repository("testing",
                                          Model("fakeModel", ModelRoles("a", "b", "c"))))

    @staticmethod
    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 + "\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
                    self.set_current_step(Steps.PARSING_ANSWERS)
                    st.rerun()

    def build_ui_for_parsing_answers(self):
        self._build_base_ui()
        with st.status("parsing user input for tags"):
            tags = find_tags_regex.findall(ss["user_input"])
            details = [get_detail(t) for t in tags]
        with st.status("initialising LLM"):
            self.repository.init()
        with st.status("waiting for LLM"):
            answer = self.repository.send_prompt(self.pm.verify_user_input_prompt(ss["user_input"], details))
            st.write(f"answers from LLM: {answer['content']}")
        with st.status("Checking for missing answers"):
            answers = LlmParser.parse_verification_prompt_answers(answer['content'], details)
            ss["answers"] = answers
            if len(answers) != len(Q):
                self.set_current_step(Steps.PARSING_ERROR)
                st.rerun()
        ss["missing_answers"] = check_for_missing_answers(ss["answers"])
        if not ss.get("missing_answers"):
            self.set_current_step(Steps.FIND_CATEGORIES)
        else:
            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.value].lower(), key=ma)
            submitted = st.form_submit_button("Submit answers")
            if submitted:
                for ma in ss["missing_answers"]:
                    ss["answers"][ma] = ss[ma]
                self.set_current_step(Steps.FIND_CATEGORIES)
                st.rerun()

    def build_ui_for_check_category(self):
        self._build_base_ui()
        with st.status("finding the work categories applicable to your work"):
            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
            self.set_current_step(Steps.VALIDATE_DATA)
            st.rerun()

    def build_ui_for_form_created(self):
        self._build_base_ui()
        st.download_button("download form", ss["pdf_form"],
                           file_name=ss["pdf_form_filename"], mime="application/pdf")
        start_over_button = st.button("Start over")
        if start_over_button:
            del ss["step"]
            del ss["pdf_form"]
            del ss["pdf_form_filename"]
            if "signature" in ss:
                del ss["signature"]
            st.rerun()

    def _integrate_llm_answers_with_user_corrections(self):
        for i in range(len(Q)):
            ss["answers"][Q(i)] = ss[f"fq_{Q(i).name}"]

        for details_key, func in [("your_details", self._get_personal_details),
                                  ("location_details", self._get_location_details),
                                  ("contractor_details", self._get_contractor_details)]:
            details = func(details_key)
            if details:
                key = ss[details_key]  # get the name under which this data should be saved
                save_details(details, key)
        self.set_current_step(Steps.FIND_CATEGORIES)
        st.rerun()

    def _create_pdf_form(self):
        with st.status("categories found, creating PDF form"):
            form_data, filename = build_form_data_from_answers(ss["answers"], ss["categories"],
                                                               ss.get("signature"))
            pdf_form = write_pdf_form(form_data)
            pdf_form_filename = filename
            ss["pdf_form"] = pdf_form
            ss["pdf_form_filename"] = pdf_form_filename
            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,
                                    "Find work categories")

    def build_ui_to_confirm_form_data(self):
        self._build_validation_form(True, self._create_pdf_form,
                                    "Create work permit request")


    @staticmethod
    def _get_personal_details(personal_details_key) -> PersonalDetails | None:
        key_ = ss.get(personal_details_key)
        if key_:
            details = PersonalDetails(ss[f"fq_{Q.FULL_NAME.name}"], ss[f"fq_{Q.YOUR_EMAIL.name}"], ss[f"fq_{Q.CONTACT_NUMBER.name}"])
            return details
        return None

    @staticmethod
    def _get_location_details(location_details_key) -> LocationDetails | None:
        if ss.get(location_details_key):
            return LocationDetails(ss[f"fq_{Q.OWNER_OR_TENANT.name}"], ss[f"fq_{Q.COMMUNITY.name}"],
                                   ss[f"fq_{Q.BUILDING.name}"], ss[f"fq_{Q.UNIT_APT_NUMBER.name}"])
        return None

    @staticmethod
    def _get_contractor_details(contractor_details_key) -> ContractorDetails | None:
        if ss.get(contractor_details_key):
            return ContractorDetails(ss[f"fq_{Q.COMPANY_NAME.name}"], ss[f"fq_{Q.COMPANY_NUMBER.name}"],
                                     ss[f"fq_{Q.COMPANY_EMAIL.name}"])
        return None



    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:
                with col:
                    form_.text_input(self.pm.questions_to_field_labels()[user_data], value=ss.get("answers", {})
                                     .get(user_data), key=f"fq_{user_data.name}")
            if add_save:
                with col:
                    form_.text_input("Save as", key=title.replace(" ", "_"))

        self._build_base_ui()
        f = st.form("Please check the following information and correct fix any inaccuracies")
        col1, col2 = f.columns(2)
        build_form_fragment(f, col1, "your details", True, Q.FULL_NAME, Q.CONTACT_NUMBER, Q.YOUR_EMAIL)
        build_form_fragment(f, col2, "work details", False, Q.WORK_TO_DO, Q.START_DATE, Q.END_DATE)
        build_form_fragment(f, col1, "location details",  True, Q.COMMUNITY, Q.BUILDING, Q.UNIT_APT_NUMBER,
                            Q.OWNER_OR_TENANT)
        build_form_fragment(f, col2, "contractor details", True, Q.COMPANY_NAME, Q.COMPANY_NUMBER, Q.COMPANY_EMAIL)
        if show_categories:
            for k, wc in work_categories.items():
                f.checkbox(label=wc, key=k, value=k in ss["categories"])

        submit_data = f.form_submit_button(label=submit_button_label)
        if submit_data:
            onsubmit()


um = UIManager()

def use_streamlit():

    if um.get_current_step() == Steps.INITIAL_STATE.value :
        um.build_ui_for_initial_state(user_msg)
    elif um.get_current_step() == Steps.PARSING_ANSWERS.value:
        um.build_ui_for_parsing_answers()
    elif um.get_current_step() == Steps.PARSING_ERROR.value:
        um.build_ui_for_validate_data_after_correction()
    elif um.get_current_step() == Steps.ASK_AGAIN.value:
        um.build_ui_for_ask_again()
    elif um.get_current_step() == Steps.FIND_CATEGORIES.value:
        um.build_ui_for_check_category()
    elif um.get_current_step() == Steps.VALIDATE_DATA.value:
        um.build_ui_to_confirm_form_data()
    elif um.get_current_step() == Steps.FORM_CREATED.value:
        um.build_ui_for_form_created()


use_streamlit()