Update app.py

app.py

@@ -9,6 +9,7 @@ import markdown
 import zipfile
 import tempfile
 from datetime import datetime
+import re
 
 def export_to_zip(images, conversations):
     """
@@ -81,12 +82,40 @@ def base64_to_image(
         print(f"Base64解码失败: {str(e)}")
         return None
 
+def extract_images_from_messages(messages):
+    """
+    从消息中提取所有图像
+    
+    Args:
+        messages: 消息JSON数据
+    
+    Returns:
+        提取的图像列表和更新后的消息
+    """
+    images = []
+    
+    for message in messages:
+        if 'content' in message and isinstance(message['content'], list):
+            for content_item in message['content']:
+                if content_item.get('type') == 'image_url':
+                    image_url = content_item.get('image_url', {}).get('url', '')
+                    if image_url.startswith('data:'):
+                        # 提取Base64图像
+                        image = base64_to_image(image_url)
+                        if image:
+                            images.append(image)
+    
+    return images, messages
+
 def process_message(file_path):
     try:
         # 读取JSON文件
         with open(file_path, "r", encoding="utf-8") as f:
             messages = json.load(f)
         
+        # 提取图像
+        images, messages = extract_images_from_messages(messages)
+        
         # 创建HTML输出
         html_output = '<div style="color: black;">'  # 添加一个包裹所有内容的div，设置文本颜色为黑色
         
@@ -110,8 +139,6 @@ def process_message(file_path):
                     # 将Markdown文本转换为HTML
                     md_text = content_item['text']
                     html_text = markdown.markdown(md_text, extensions=['fenced_code', 'codehilite'])
-                    # html_text = markdown.markdown(md_text)
-                    # html_text = md_text
                     html_output += f'<div style="color: black;">{html_text}</div>'
                 
                 elif content_type == "image_url":
@@ -125,22 +152,31 @@ def process_message(file_path):
             html_output += '</div>'
         
         html_output += '</div>'  # 关闭最外层div
-        return html_output
+        return html_output, images, messages
     
     except Exception as e:
-        return f"<div style='color: red;'>Error processing file: {str(e)}</div>"
+        return f"<div style='color: red;'>Error processing file: {str(e)}</div>", [], None
 
 def upload_and_process(file):
     if file is None:
-        return "请上传一个JSON文件"
+        return "请上传一个JSON文件", [], None
     
-    return process_message(file.name)
+    html_output, images, messages = process_message(file.name)
+    return html_output, images, messages
 
 def use_example():
     # 使用示例文件
     example_path = "test_message_gpt.json"
     return process_message(example_path)
 
+def handle_export(images, conversations):
+    """处理导出请求"""
+    if not images or conversations is None:
+        return None
+    
+    zip_path = export_to_zip(images, conversations)
+    return zip_path
+
 # 确保示例文件存在
 def setup_example_file():
     # 这里我们需要创建示例文件，因为我们没有实际的内容
@@ -196,8 +232,8 @@ with gr.Blocks(title="ChatGPT Conversation Visualizer", css="div.prose * {color:
 
     # 添加导出按钮
     with gr.Row():
-        export_btn = gr.Button("Export to ZIP")
-        download_file = gr.File(label="Download ZIP")
+        export_btn = gr.Button("导出为ZIP")
+        download_file = gr.File(label="下载ZIP文件")
 
     # 存储当前结果的状态变量
     current_images = gr.State([])
@@ -206,13 +242,19 @@ with gr.Blocks(title="ChatGPT Conversation Visualizer", css="div.prose * {color:
     visualize_button.click(
         fn=upload_and_process,
         inputs=[file_input],
-        outputs=[output]
+        outputs=[output, current_images, current_json]
     )
     
     example_button.click(
         fn=use_example,
         inputs=[],
-        outputs=[output]
+        outputs=[output, current_images, current_json]
+    )
+    
+    export_btn.click(
+        fn=handle_export,
+        inputs=[current_images, current_json],
+        outputs=[download_file]
     )
 
 # 启动Gradio应用