Store download file name in session and display the button again on app reload

chat_app.py

@@ -1,3 +1,4 @@
+import datetime
 import logging
 import pathlib
 import random
@@ -15,7 +16,9 @@ from langchain_core.runnables.history import RunnableWithMessageHistory
 from global_config import GlobalConfig
 from helpers import llm_helper, pptx_helper
 
+
 APP_TEXT = json5.loads(open(GlobalConfig.APP_STRINGS_FILE, 'r', encoding='utf-8').read())
+DOWNLOAD_FILE_KEY = 'download_file_name'
 # langchain.debug = True
 # langchain.verbose = True
 
@@ -82,7 +85,6 @@ def set_up_chat_ui():
         template = in_file.read()
 
     prompt = ChatPromptTemplate.from_template(template)
-
     chain = prompt | llm
 
     chain_with_history = RunnableWithMessageHistory(
@@ -103,6 +105,11 @@ def set_up_chat_ui():
         # st.chat_message(msg.type).markdown(msg.content)
         st.chat_message(msg.type).code(msg.content, language='json')
 
+    # The download button disappears on clicking (anywhere) because of app reload
+    # So, display it again
+    if DOWNLOAD_FILE_KEY in st.session_state:
+        _display_download_button(st.session_state[DOWNLOAD_FILE_KEY])
+
     progress_bar.progress(100, text='Done!')
     progress_bar.empty()
 
@@ -123,28 +130,9 @@ def set_up_chat_ui():
         # The content has been generated as JSON
         # There maybe trailing ``` at the end of the response -- remove them
         # To be careful: ``` may be part of the content as well when code is generated
-        str_len = len(response)
-        response_cleaned = response
-
+        response_cleaned = _clean_json(response)
         progress_bar_pptx.progress(50, 'Analyzing response...')
 
-        try:
-            idx = response.rindex('```')
-            logger.debug('str_len: %d, idx of ```: %d', str_len, idx)
-
-            if idx + 3 == str_len:
-                # The response ends with ``` -- most likely the end of JSON response string
-                response_cleaned = response[:idx]
-            elif idx + 3 < str_len:
-                # Looks like there are some more content beyond the last ```
-                # In the best case, it would be some additional plain-text response from the LLM
-                # and is unlikely to contain } or ] that are present in JSON
-                if '}' not in response[idx + 3:]:  # the remainder of the text
-                    response_cleaned = response[:idx]
-        except ValueError:
-            # No ``` found
-            pass
-
         # Now create the PPT file
         progress_bar_pptx.progress(75, 'Creating the slide deck...give it a moment')
         generate_slide_deck(response_cleaned)
@@ -159,28 +147,78 @@ def generate_slide_deck(json_str: str) -> List:
     :return: A list of all slide headers and the title.
     """
 
-    # # Get a unique name for the file to save -- use the session ID
-    # ctx = st_sr.get_script_run_ctx()
-    # session_id = ctx.session_id
-    # timestamp = time.time()
-    # output_file_name = f'{session_id}_{timestamp}.pptx'
-
-    temp = tempfile.NamedTemporaryFile(delete=False, suffix='.pptx')
-    path = pathlib.Path(temp.name)
+    if DOWNLOAD_FILE_KEY in st.session_state:
+        path = pathlib.Path(st.session_state[DOWNLOAD_FILE_KEY])
+        logger.debug('DOWNLOAD_FILE_KEY found in session')
+    else:
+        temp = tempfile.NamedTemporaryFile(delete=False, suffix='.pptx')
+        path = pathlib.Path(temp.name)
+        st.session_state[DOWNLOAD_FILE_KEY] = str(path)
+        logger.debug('DOWNLOAD_FILE_KEY not found in session')
 
-    logger.info('Creating PPTX file...')
+    logger.debug('Creating PPTX file: %s...', st.session_state[DOWNLOAD_FILE_KEY])
     all_headers = pptx_helper.generate_powerpoint_presentation(
         json_str,
         slides_template=pptx_template,
         output_file_path=path
     )
 
-    with open(path, 'rb') as f:
-        st.download_button('Download PPTX file ⇩', f, file_name='Presentation.pptx')
+    _display_download_button(path)
 
     return all_headers
 
 
+def _clean_json(json_str: str) -> str:
+    """
+    Attempt to clean a JSON response string from the LLM by removing the trailing ```
+    and any text beyond that. May not be always accurate.
+
+    :param json_str: The input string in JSON format.
+    :return: The "cleaned" JSON string.
+    """
+
+    str_len = len(json_str)
+    response_cleaned = json_str
+
+    try:
+        idx = json_str.rindex('```')
+        logger.debug(
+            'Fixing JSON response: str_len: %d, idx of ```: %d',
+            str_len, idx
+        )
+
+        if idx + 3 == str_len:
+            # The response ends with ``` -- most likely the end of JSON response string
+            response_cleaned = json_str[:idx]
+        elif idx + 3 < str_len:
+            # Looks like there are some more content beyond the last ```
+            # In the best case, it would be some additional plain-text response from the LLM
+            # and is unlikely to contain } or ] that are present in JSON
+            if '}' not in json_str[idx + 3:]:  # the remainder of the text
+                response_cleaned = json_str[:idx]
+    except ValueError:
+        # No ``` found
+        pass
+
+    return response_cleaned
+
+
+def _display_download_button(file_path: pathlib.Path):
+    """
+    Display a download button to download a slide deck.
+
+    :param file_path: The path of the .pptx file.
+    """
+
+    with open(file_path, 'rb') as download_file:
+        st.download_button(
+            'Download PPTX file ⬇️',
+            data=download_file,
+            file_name='Presentation.pptx',
+            key=datetime.datetime.now()
+        )
+
+
 def main():
     """
     Trigger application run.