Download output as a zip file

src/plan/app_text2plan.py

@@ -16,6 +16,8 @@ from src.plan.speedvsdetail import SpeedVsDetailEnum
 from src.prompt.prompt_catalog import PromptCatalog
 from src.utils.get_env_as_string import get_env_as_string
 from src.utils.time_since_last_modification import time_since_last_modification
+import shutil
+import zipfile
     
 # Global variables
 active_proc = None
@@ -87,6 +89,26 @@ class MarkdownBuilder:
 
     def to_markdown(self):
         return "\n".join(self.rows)
+
+def create_zip_archive(directory_to_zip: str) -> str:
+    """
+    Creates a zip archive of the given directory, excluding 'log.txt'.
+    Returns the path to the zip file, or None if an error occurred.
+    """
+    zip_file_path = os.path.join(os.path.dirname(directory_to_zip), os.path.basename(directory_to_zip) + ".zip")
+    try:
+        with zipfile.ZipFile(zip_file_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
+            for root, _, files in os.walk(directory_to_zip):
+                for file in files:
+                    if file == "log.txt":
+                        continue  # Skip the log file
+
+                    file_path = os.path.join(root, file)
+                    zipf.write(file_path, os.path.relpath(file_path, directory_to_zip))
+        return zip_file_path
+    except Exception as e:
+        print(f"Error creating zip archive: {e}")
+        return None
     
 def run_planner(submit_or_retry_button: str, plan_prompt: str, llm_model: str, speedvsdetail: str):
     """
@@ -95,7 +117,8 @@ def run_planner(submit_or_retry_button: str, plan_prompt: str, llm_model: str, s
     1. Prepare the run directory.
     2. Updates environment with RUN_ID and LLM_MODEL.
     3. Launches the Python pipeline script.
-    5. Polls the output dir (every second) to yield a Markdown-formatted file listing.
+    4. Polls the output dir (every second) to yield a Markdown-formatted file listing.
+    5. Creates a zip archive (every 10 seconds) with the latest output.
        If a stop is requested or if the pipeline-complete file is detected, it exits.
     """
     global active_proc, stop_event, latest_run_dir, latest_run_id
@@ -145,6 +168,8 @@ def run_planner(submit_or_retry_button: str, plan_prompt: str, llm_model: str, s
     env["SPEED_VS_DETAIL"] = speedvsdetail
 
     start_time = time.perf_counter()
+    last_zip_time = time.time()  # Initialize the last zip creation time
+    most_recent_zip_file = None # Store the path to the latest created zip
 
     # Launch the pipeline as a separate Python process.
     command = [sys.executable, "-m", MODULE_PATH_PIPELINE]
@@ -192,7 +217,10 @@ def run_planner(submit_or_retry_button: str, plan_prompt: str, llm_model: str, s
             markdown_builder.status("Process terminated by user.")
             markdown_builder.path_to_run_dir(absolute_path_to_run_dir)
             markdown_builder.list_files(run_path)
-            yield markdown_builder.to_markdown()
+            zip_file_path = create_zip_archive(run_path)
+            if zip_file_path:
+                most_recent_zip_file = zip_file_path
+            yield markdown_builder.to_markdown(), most_recent_zip_file
             break
 
         last_update = ceil(time_since_last_modification(run_path))
@@ -200,7 +228,16 @@ def run_planner(submit_or_retry_button: str, plan_prompt: str, llm_model: str, s
         markdown_builder.status(f"Working. {duration} seconds elapsed. Last output update was {last_update} seconds ago.")
         markdown_builder.path_to_run_dir(absolute_path_to_run_dir)
         markdown_builder.list_files(run_path)
-        yield markdown_builder.to_markdown()
+
+        # Check if it's time to create a new zip archive (every 10 seconds)
+        current_time = time.time()
+        if current_time - last_zip_time >= 10:
+            zip_file_path = create_zip_archive(run_path)
+            if zip_file_path: # Only update the path if the creation succeeded
+                most_recent_zip_file = zip_file_path
+            last_zip_time = current_time
+
+        yield markdown_builder.to_markdown(), most_recent_zip_file
 
         # If the pipeline complete file is found, finish streaming.
         if has_pipeline_complete_file(run_path):
@@ -230,7 +267,13 @@ def run_planner(submit_or_retry_button: str, plan_prompt: str, llm_model: str, s
     markdown_builder.status(f"{status_message} {duration} seconds elapsed.")
     markdown_builder.path_to_run_dir(absolute_path_to_run_dir)
     markdown_builder.list_files(run_path)
-    yield markdown_builder.to_markdown()
+
+    # Create zip archive
+    zip_file_path = create_zip_archive(run_path)
+    if zip_file_path:
+        most_recent_zip_file = zip_file_path
+
+    yield markdown_builder.to_markdown(), most_recent_zip_file
 
 def stop_planner():
     """
@@ -287,6 +330,7 @@ with gr.Blocks(title="PlanExe") as demo_text2plan:
 
                 output_markdown = gr.Markdown("Output will appear here...")
                 status_markdown = gr.Markdown("Status messages will appear here...")
+                download_output = gr.File(label="Download latest output (excluding log.txt) as zip")
 
             with gr.Column(scale=1, min_width=300):
                 examples = gr.Examples(
@@ -320,10 +364,10 @@ with gr.Blocks(title="PlanExe") as demo_text2plan:
 """)
 
     # Start the planning assistant with streaming updates.
-    submit_btn.click(fn=run_planner, inputs=[submit_btn, prompt_input, model_radio, speedvsdetail_radio], outputs=output_markdown)
+    submit_btn.click(fn=run_planner, inputs=[submit_btn, prompt_input, model_radio, speedvsdetail_radio], outputs=[output_markdown, download_output])
     # The stop button simply calls stop_planner and displays its status message.
     stop_btn.click(fn=stop_planner, outputs=status_markdown)
-    retry_btn.click(fn=run_planner, inputs=[retry_btn, prompt_input, model_radio, speedvsdetail_radio], outputs=output_markdown)
+    retry_btn.click(fn=run_planner, inputs=[retry_btn, prompt_input, model_radio, speedvsdetail_radio], outputs=[output_markdown, download_output])
     open_dir_btn.click(fn=open_output_dir, outputs=status_markdown)
 
 if __name__ == "__main__":