/frontend/demo_light/pages_util/CreateNewArticle.py

import os
import time

import demo_util
import streamlit as st
from demo_util import DemoFileIOHelper, DemoTextProcessingHelper, DemoUIHelper


def create_new_article_page():
    demo_util.clear_other_page_session_state(page_index=3)

    if "page3_write_article_state" not in st.session_state:
        st.session_state["page3_write_article_state"] = "not started"

    if st.session_state["page3_write_article_state"] == "not started":

        _, search_form_column, _ = st.columns([2, 5, 2])
        with search_form_column:
            with st.form(key='search_form'):
                # Text input for the search topic
                DemoUIHelper.st_markdown_adjust_size(content="Enter the topic you want to learn in depth:",
                                                     font_size=18)
                st.session_state["page3_topic"] = st.text_input(label='page3_topic', label_visibility="collapsed")
                pass_appropriateness_check = True

                # Submit button for the form
                submit_button = st.form_submit_button(label='Research')
                # only start new search when button is clicked, not started, or already finished previous one
                if submit_button and st.session_state["page3_write_article_state"] in ["not started", "show results"]:
                    if not st.session_state["page3_topic"].strip():
                        pass_appropriateness_check = False
                        st.session_state["page3_warning_message"] = "topic could not be empty"

                    st.session_state["page3_topic_name_cleaned"] = st.session_state["page3_topic"].replace(
                        ' ', '_').replace('/', '_')
                    if not pass_appropriateness_check:
                        st.session_state["page3_write_article_state"] = "not started"
                        alert = st.warning(st.session_state["page3_warning_message"], icon="⚠️")
                        time.sleep(5)
                        alert.empty()
                    else:
                        st.session_state["page3_write_article_state"] = "initiated"

    if st.session_state["page3_write_article_state"] == "initiated":
        current_working_dir = os.path.join(demo_util.get_demo_dir(), "DEMO_WORKING_DIR")
        if not os.path.exists(current_working_dir):
            os.makedirs(current_working_dir)

        if "runner" not in st.session_state:
            demo_util.set_storm_runner()
        st.session_state["page3_current_working_dir"] = current_working_dir
        st.session_state["page3_write_article_state"] = "pre_writing"

    if st.session_state["page3_write_article_state"] == "pre_writing":
        status = st.status("I am brain**STORM**ing now to research the topic. (This may take 2-3 minutes.)")
        st_callback_handler = demo_util.StreamlitCallbackHandler(status)
        with status:
            # STORM main gen outline
            st.session_state["runner"].run(
                topic=st.session_state["page3_topic"],
                do_research=True,
                do_generate_outline=True,
                do_generate_article=False,
                do_polish_article=False,
                callback_handler=st_callback_handler
            )
            conversation_log_path = os.path.join(st.session_state["page3_current_working_dir"],
                                                 st.session_state["page3_topic_name_cleaned"], "conversation_log.json")
            demo_util._display_persona_conversations(DemoFileIOHelper.read_json_file(conversation_log_path))
            st.session_state["page3_write_article_state"] = "final_writing"
            status.update(label="brain**STORM**ing complete!", state="complete")

    if st.session_state["page3_write_article_state"] == "final_writing":
        # polish final article
        with st.status(
                "Now I will connect the information I found for your reference. (This may take 4-5 minutes.)") as status:
            st.info('Now I will connect the information I found for your reference. (This may take 4-5 minutes.)')
            st.session_state["runner"].run(topic=st.session_state["page3_topic"], do_research=False,
                                           do_generate_outline=False,
                                           do_generate_article=True, do_polish_article=True, remove_duplicate=False)
            # finish the session
            st.session_state["runner"].post_run()

            # update status bar
            st.session_state["page3_write_article_state"] = "prepare_to_show_result"
            status.update(label="information snythesis complete!", state="complete")

    if st.session_state["page3_write_article_state"] == "prepare_to_show_result":
        _, show_result_col, _ = st.columns([4, 3, 4])
        with show_result_col:
            if st.button("show final article"):
                st.session_state["page3_write_article_state"] = "completed"
                st.rerun()

    if st.session_state["page3_write_article_state"] == "completed":
        # display polished article
        current_working_dir_paths = DemoFileIOHelper.read_structure_to_dict(
            st.session_state["page3_current_working_dir"])
        current_article_file_path_dict = current_working_dir_paths[st.session_state["page3_topic_name_cleaned"]]
        demo_util.display_article_page(selected_article_name=st.session_state["page3_topic_name_cleaned"],
                                       selected_article_file_path_dict=current_article_file_path_dict,
                                       show_title=True, show_main_article=True)