Update app.py

app.py

@@ -4,6 +4,8 @@ from typing import Optional, Union
 from PIL import Image
 import base64
 from io import BytesIO
+import gradio as gr
+import markdown
 
 def base64_to_image(
     base64_str: str, 
@@ -20,10 +22,6 @@ def base64_to_image(
     
     Returns:
         PIL.Image.Image 对象，解码失败时返回None
-        
-    Examples:
-        >>> img = base64_to_image("data:image/png;base64,iVBORw0KGg...")
-        >>> img = base64_to_image("iVBORw0KGg...", remove_prefix=False)
     """
     try:
         # 1. 处理Base64前缀
@@ -46,20 +44,127 @@ def base64_to_image(
         print(f"Base64解码失败: {str(e)}")
         return None
 
+def process_message(file_path):
+    try:
+        # 读取JSON文件
+        with open(file_path, "r", encoding="utf-8") as f:
+            messages = json.load(f)
+        
+        # 创建HTML输出
+        html_output = ""
+        
+        for message_item in messages:
+            role = message_item['role']
+            content = message_item['content']
+            
+            # 根据角色设置样式
+            if role == "user" or role == "human":
+                html_output += f'<div style="background-color: #f0f0f0; padding: 10px; margin: 10px 0; border-radius: 10px;"><strong>User:</strong><br>'
+            elif role == "assistant":
+                html_output += f'<div style="background-color: #e6f7ff; padding: 10px; margin: 10px 0; border-radius: 10px;"><strong>Assistant:</strong><br>'
+            else:
+                html_output += f'<div style="background-color: #f9f9f9; padding: 10px; margin: 10px 0; border-radius: 10px;"><strong>{role.capitalize()}:</strong><br>'
+            
+            # 处理内容
+            for content_item in content:
+                content_type = content_item['type']
+                
+                if content_type == "text":
+                    # 将Markdown文本转换为HTML
+                    md_text = content_item['text']
+                    html_text = markdown.markdown(md_text, extensions=['fenced_code', 'codehilite'])
+                    html_output += html_text
+                
+                elif content_type == "image_url":
+                    content_value = content_item['image_url']['url']
+                    # 如果是base64图片
+                    if content_value.startswith("data:"):
+                        html_output += f'<img src="{content_value}" style="max-width: 100%; margin: 10px 0;">'
+                    else:
+                        html_output += f'<img src="{content_value}" style="max-width: 100%; margin: 10px 0;">'
+            
+            html_output += '</div>'
+        
+        return html_output
+    
+    except Exception as e:
+        return f"<div style='color: red;'>Error processing file: {str(e)}</div>"
 
-file_path = "./test_message_gpt.json"
-messages = json.load(open(file_path, "r", encoding="utf-8"))
+def upload_and_process(file):
+    if file is None:
+        return "请上传一个JSON文件"
+    
+    return process_message(file.name)
 
-for message_item in messages:
-    role = message_item['role']
-    content = message_item['content']
+def use_example():
+    # 使用示例文件
+    example_path = "test_message_gpt.json"
+    return process_message(example_path)
 
-    for content_item in content:
-        content_type = content_item['type']
-        if content_type == "text":
-            content_value = content_item['text']
-        elif conent_type == "image_url":
-            content_value = content_item['image_url']['url']
-            image = base64_to_image(content_value)
+# 确保示例文件存在
+def setup_example_file():
+    # 这里我们需要创建示例文件，因为我们没有实际的内容
+    # 在实际应用中，你应该将原始的test_message_gpt.json文件放在Space的根目录下
+    example_path = "test_message_gpt.json"
+    
+    # 如果文件不存在，创建一个简单的示例
+    if not os.path.exists(example_path):
+        example_messages = [
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "text",
+                        "text": "你好，请介绍一下自己"
+                    }
+                ]
+            },
+            {
+                "role": "assistant",
+                "content": [
+                    {
+                        "type": "text",
+                        "text": "你好！我是一个AI助手。我可以帮助回答问题、提供信息、进行对话等。我被设计用来协助用户完成各种任务，从简单的问答到更复杂的讨论。\n\n我可以处理文本信息，也能理解和描述图像内容。虽然我有一些限制，但我会尽力提供有用、准确和有帮助的回应。\n\n有什么我可以帮助你的吗？"
+                    }
+                ]
+            }
+        ]
+        
+        with open(example_path, "w", encoding="utf-8") as f:
+            json.dump(example_messages, f, ensure_ascii=False, indent=2)
+
+# 设置示例文件
+setup_example_file()
+
+# 创建Gradio界面
+with gr.Blocks(title="ChatGPT Conversation Visualizer") as demo:
+    gr.Markdown("# ChatGPT 对话可视化工具")
+    gr.Markdown("上传一个包含ChatGPT对话记录的JSON文件，或使用示例文件查看可视化结果")
+    
+    with gr.Row():
+        file_input = gr.File(label="上传JSON文件", file_types=[".json"])
+    
+    with gr.Row():
+        col1, col2 = gr.Column(), gr.Column()
+        with col1:
+            visualize_button = gr.Button("可视化上传的对话")
+        with col2:
+            example_button = gr.Button("使用示例文件")
+    
+    with gr.Row():
+        output = gr.HTML(label="对话内容")
+    
+    visualize_button.click(
+        fn=upload_and_process,
+        inputs=[file_input],
+        outputs=[output]
+    )
+    
+    example_button.click(
+        fn=use_example,
+        inputs=[],
+        outputs=[output]
+    )
 
-为我修改这段代码，用gradio包装一下，这里的file是一个chatgpt对话过程的message，我希望你用gradio可视化出来。注意，这段代码在huggingface的space上运行
+# 启动Gradio应用
+demo.launch()