Execute PPT, poster, and PR separately

app.py

@@ -54,71 +54,143 @@ def save_api_key(api_key):
     except Exception as e:
         return f"❌ 出错: {str(e)}", get_debug_info()
 
-def run_mineru_parsing():
-    """执行 PDF 解析并捕获完整日志"""
+
+def run_mineru_parsing_and_dag_gen():
+    """执行 PDF 解析并捕获完整日志，随后执行DAG生成流程"""
     if not os.path.exists(PAPERS_DIR) or not any(f.endswith('.pdf') for f in os.listdir(PAPERS_DIR)):
         return "❌ 未发现 PDF 文件", get_debug_info(), "No execution logs."
 
     try:
+        # ================= 第一步：执行 Mineru 解析 =================
         env = os.environ.copy()
         env["MINERU_FORMULA_ENABLE"] = "false"
         env["MINERU_TABLE_ENABLE"] = "false"
         env["MINERU_DEVICE_MODE"] = "cpu"
         env["MINERU_VIRTUAL_VRAM_SIZE"] = "8"
 
-        command = ["mineru", "-p", PAPERS_DIR, "-o", OUTPUT_DIR]
+        command_mineru = ["mineru", "-p", PAPERS_DIR, "-o", OUTPUT_DIR]
         
-        result = subprocess.run(
-            command, 
+        result_mineru = subprocess.run(
+            command_mineru, 
             env=env, 
             capture_output=True, 
             text=True,
             timeout=300 
         )
 
-        full_log = f"--- STDOUT ---\n{result.stdout}\n\n--- STDERR ---\n{result.stderr}"
+        # 记录 Mineru 的日志
+        full_log = f"--- Mineru STDOUT ---\n{result_mineru.stdout}\n\n--- Mineru STDERR ---\n{result_mineru.stderr}\n"
+        
+        # 如果解析失败，直接返回，不执行后续 DAG 生成
+        if result_mineru.returncode != 0:
+            status = f"❌ Mineru 解析失败 (Exit Code: {result_mineru.returncode})"
+            return status, get_debug_info(), full_log
+
+        # ================= 第二步：执行 DAG 生成 =================
+        # 使用 sys.executable 确保使用当前虚拟环境的 Python 解释器
+        command_dag = [sys.executable, "gen_dag.py"]
+        
+        result_dag = subprocess.run(
+            command_dag, 
+            capture_output=True, 
+            text=True,
+            timeout=300  # 可以根据你 DAG 生成的时间需求调整超时
+        )
+
+        # 追加 DAG 生成的日志
+        full_log += f"\n--- DAG Gen STDOUT ---\n{result_dag.stdout}\n\n--- DAG Gen STDERR ---\n{result_dag.stderr}"
         
-        if result.returncode == 0:
-            status = "✅ PDF解析完成"
+        if result_dag.returncode == 0:
+            status = "✅ PDF解析与DAG生成全部完成"
         else:
-            status = f"❌ 解析失败 (Exit Code: {result.returncode})"
+            status = f"❌ DAG生成失败 (Exit Code: {result_dag.returncode})"
             
         return status, get_debug_info(), full_log
 
+    except subprocess.TimeoutExpired as e:
+        error_log = f"Execution timed out:\n{str(e)}"
+        return "❌ 运行超时", get_debug_info(), error_log
     except Exception as e:
         error_log = f"Exception occurred during execution:\n{str(e)}"
         return "❌ 运行异常", get_debug_info(), error_log
 
-def run_final_generation():
-    """执行 python main.py 并压缩结果"""
+def run_final_generation(task_type="all"):
+    """
+    执行对应的生成脚本并压缩结果（支持并行执行）
+    task_type 支持: 'ppt', 'poster', 'pr', 'all'
+    """
     if not os.path.exists(OUTPUT_DIR):
         return "❌ 请先执行第二步解析", get_debug_info(), "No output folder found.", None
 
-    try:
-        # 1. 运行 python main.py
-        command = [sys.executable, "main.py"]
+    # 根据传入的 task_type 决定要运行哪些脚本
+    scripts_to_run = []
+    if task_type == "ppt":
+        scripts_to_run = ["gen_ppt.py"]
+    elif task_type == "poster":
+        scripts_to_run = ["gen_poster.py"]
+    elif task_type == "pr":
+        scripts_to_run = ["gen_pr.py"]
+    elif task_type == "all":
+        scripts_to_run = ["gen_ppt.py", "gen_poster.py", "gen_pr.py"]
+    else:
+        return "❌ 未知任务类型", get_debug_info(), "Invalid task_type.", None
+
+    full_log = f"🚀 准备启动 {len(scripts_to_run)} 个任务...\n"
+    success = True
+
+    # 定义单个脚本的执行包装器
+    def execute_script(script):
+        command = [sys.executable, script]
         result = subprocess.run(
             command,
             capture_output=True,
             text=True,
-            timeout=600 # 适当增加超时时间
+            timeout=600 # 每个脚本独立的超时时间
         )
+        return script, result
 
-        full_log = f"--- STDOUT ---\n{result.stdout}\n\n--- STDERR ---\n{result.stderr}"
+    try:
+        # 使用 ThreadPoolExecutor 并行执行脚本
+        with concurrent.futures.ThreadPoolExecutor(max_workers=len(scripts_to_run)) as executor:
+            # 提交所有任务
+            future_to_script = {executor.submit(execute_script, s): s for s in scripts_to_run}
+            
+            # as_completed 会在某个任务完成时立刻生成结果
+            for future in concurrent.futures.as_completed(future_to_script):
+                script_name = future_to_script[future]
+                try:
+                    # 获取该任务的执行结果
+                    _, result = future.result()
+                    
+                    full_log += f"\n================ ✅ 执行完成: {script_name} ================\n"
+                    full_log += f"--- STDOUT ---\n{result.stdout}\n\n--- STDERR ---\n{result.stderr}\n"
+                    
+                    # 检查此任务是否失败
+                    if result.returncode != 0:
+                        success = False
+                        full_log += f"❌ [错误] {script_name} 返回非零退出码 (Exit Code: {result.returncode})\n"
+                
+                except subprocess.TimeoutExpired as e:
+                    success = False
+                    full_log += f"\n================ ❌ 任务超时: {script_name} ================\n{str(e)}\n"
+                except Exception as e:
+                    success = False
+                    full_log += f"\n================ ❌ 任务异常: {script_name} ================\n{str(e)}\n"
 
-        if result.returncode != 0:
-            return f"❌ 生成失败 (Exit Code: {result.returncode})", get_debug_info(), full_log, None
+        # 如果有任何一个脚本执行失败，直接返回，不打包压缩
+        if not success:
+            return f"❌ {task_type.upper()} 包含失败任务，请检查日志", get_debug_info(), full_log, None
 
-        # 2. 压缩 mineru_outputs 文件夹
-        # shutil.make_archive 会自动加上 .zip 后缀，所以 base_name 不带后缀
+        # 所有脚本都运行成功后，压缩 mineru_outputs 文件夹
         zip_base_name = ZIP_OUTPUT_PATH.replace(".zip", "")
         shutil.make_archive(zip_base_name, 'zip', OUTPUT_DIR)
 
-        return "✅ 最终生成并压缩完成", get_debug_info(), full_log, ZIP_OUTPUT_PATH
+        success_msg = f"✅ {task_type.upper()} 生成并压缩完成"
+        return success_msg, get_debug_info(), full_log, ZIP_OUTPUT_PATH
 
     except Exception as e:
-        error_log = f"Exception occurred during final generation:\n{str(e)}"
-        return "❌ 最终生成异常", get_debug_info(), error_log, None
+        error_log = full_log + f"\n[全局异常] Exception occurred:\n{str(e)}"
+        return "❌ 最终生成发生全局异常", get_debug_info(), error_log, None
 
 # --- UI ---
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
@@ -136,12 +208,18 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
                 
             with gr.Group():
                 gr.Markdown("### 2. 执行解析")
-                parse_btn = gr.Button("🚀 Run Mineru (CPU Mode)", variant="primary")
+                parse_btn = gr.Button("🚀 Run Mineru & DAG Gen", variant="primary")
                 parse_status = gr.Textbox(label="运行状态")
 
             with gr.Group():
                 gr.Markdown("### 3. 最终生成")
-                gen_btn = gr.Button("🔨 执行 main.py 并打包", variant="primary")
+                gr.Markdown("请选择要生成的内容类型：")
+                with gr.Row():
+                    gen_ppt_btn = gr.Button("📊 生成 PPT")
+                    gen_poster_btn = gr.Button("🖼️ 生成 Poster")
+                    gen_pr_btn = gr.Button("📰 生成 PR")
+                gen_all_btn = gr.Button("🔨 生成全部 (ALL)", variant="primary")
+                
                 gen_status = gr.Textbox(label="生成状态")
                 download_file = gr.File(label="下载压缩后的结果", interactive=False)
 
@@ -158,20 +236,20 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
             )
             refresh_btn = gr.Button("🔄 刷新状态")
 
-    # 逻辑绑定
+    # ================= 逻辑绑定 =================
     key_btn.click(save_api_key, inputs=key_input, outputs=[parse_status, debug_view])
     pdf_btn.click(save_pdf, inputs=pdf_input, outputs=[parse_status, debug_view])
     
     parse_btn.click(
-        fn=run_mineru_parsing, 
+        fn=run_mineru_parsing_and_dag_gen, 
         outputs=[parse_status, debug_view, cmd_logs]
     )
 
-    # 最终生成逻辑绑定
-    gen_btn.click(
-        fn=run_final_generation,
-        outputs=[gen_status, debug_view, cmd_logs, download_file]
-    )
+    # 最终生成逻辑绑定 (使用 lambda 传递固定的 task_type 参数)
+    gen_ppt_btn.click(fn=lambda: run_final_generation("ppt"), outputs=[gen_status, debug_view, cmd_logs, download_file])
+    gen_poster_btn.click(fn=lambda: run_final_generation("poster"), outputs=[gen_status, debug_view, cmd_logs, download_file])
+    gen_pr_btn.click(fn=lambda: run_final_generation("pr"), outputs=[gen_status, debug_view, cmd_logs, download_file])
+    gen_all_btn.click(fn=lambda: run_final_generation("all"), outputs=[gen_status, debug_view, cmd_logs, download_file])
     
     refresh_btn.click(get_debug_info, outputs=debug_view)