Update app.py

app.py

@@ -6,6 +6,43 @@ import base64
 from io import BytesIO
 import gradio as gr
 import markdown
+import zipfile
+import tempfile
+from datetime import datetime
+
+def export_to_zip(images, conversations):
+    """
+    将图像和对话数据导出为ZIP文件
+    
+    Args:
+        images: 提取的图像列表
+        conversations: 对话JSON数据
+    
+    Returns:
+        生成的ZIP文件路径
+    """
+    # 创建临时目录
+    temp_dir = tempfile.mkdtemp()
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    zip_filename = os.path.join(temp_dir, f"export_{timestamp}.zip")
+    
+    # 创建ZIP文件
+    with zipfile.ZipFile(zip_filename, 'w') as zipf:
+        # 保存图像
+        for i, img in enumerate(images):
+            img_path = os.path.join(temp_dir, f"image_{i}.png")
+            img.save(img_path)
+            zipf.write(img_path, f"images/image_{i}.png")
+            os.remove(img_path)  # 删除临时图像文件
+        
+        # 保存对话数据
+        json_path = os.path.join(temp_dir, "conversations.json")
+        with open(json_path, 'w', encoding='utf-8') as f:
+            json.dump(conversations, f, ensure_ascii=False, indent=4)
+        zipf.write(json_path, "conversations.json")
+        os.remove(json_path)  # 删除临时JSON文件
+    
+    return zip_filename
 
 def base64_to_image(
     base64_str: str, 
@@ -156,6 +193,15 @@ with gr.Blocks(title="ChatGPT Conversation Visualizer", css="div.prose * {color:
     
     with gr.Row():
         output = gr.HTML(label="对话内容")
+
+    # 添加导出按钮
+    with gr.Row():
+        export_btn = gr.Button("Export to ZIP")
+        download_file = gr.File(label="Download ZIP")
+
+    # 存储当前结果的状态变量
+    current_images = gr.State([])
+    current_json = gr.State(None)
     
     visualize_button.click(
         fn=upload_and_process,