Update app.py

app.py

@@ -13,6 +13,42 @@ from typing import Optional, Union
 import gradio as gr
 import markdown
 from openai import AzureOpenAI
+import zipfile
+import tempfile
+
+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, 
@@ -428,12 +464,45 @@ def create_demo():
                 extracted_images = gr.Gallery(label="Extracted Images")
             with gr.Column(scale=1):
                 extracted_json = gr.JSON(label="Extracted JSON")
+                
+        # 添加导出按钮
+        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)
+
         # 处理提交
+        def handle_submit(model, key, url, q, img):
+            html, images, json_data = o3_chat(model, key, url, q, img)
+            return html, images, json_data, images, json_data
+            
+        # # 处理提交
+        # submit_btn.click(
+        #     fn=o3_chat,
+        #     inputs=[model_name, api_key, base_url, question, image_input],
+        #     outputs=[output, extracted_images, extracted_json]
+        # )
+
         submit_btn.click(
-            fn=o3_chat,
+            fn=handle_submit,
             inputs=[model_name, api_key, base_url, question, image_input],
-            outputs=[output, extracted_images, extracted_json]
+            outputs=[output, extracted_images, extracted_json, current_images, current_json]
+        )
+
+        # 处理导出
+        def handle_export(images, conversations):
+            if not images or conversations is None:
+                return None
+            zip_path = export_to_zip(images, conversations)
+            return zip_path
+
+        export_btn.click(
+            fn=handle_export,
+            inputs=[current_images, current_json],
+            outputs=[download_file]
         )
         
         # 示例部分