Add optional image upload support with better UI layout

app.py

@@ -27,45 +27,49 @@ def image_to_base64(image):
     
     return None
 
-def process_message(message, history, system_prompt, temperature, max_tokens, top_p):
-    """处理消息并生成响应"""
+def process_message(message, history, image, system_prompt, temperature, max_tokens, top_p):
+    """处理消息并生成响应，支持可选的图片输入"""
     print(f"[DEBUG] Processing message: {message[:100] if message else 'None'}...")
+    print(f"[DEBUG] Has image: {image is not None}")
     
-    if not message:
-        print("[DEBUG] No message provided, skipping")
+    if not message and not image:
+        print("[DEBUG] No message or image provided, skipping")
         yield history
         return
     
     if not STEP_API_KEY:
         print("[DEBUG] No API key configured")
-        history.append([message, "❌ API key not configured. Please add STEP_API_KEY in Settings."])
+        error_msg = "❌ API key not configured. Please add STEP_API_KEY in Settings."
+        display_msg = f"[Image] {message}" if image and message else "[Image]" if image else message
+        history.append([display_msg, error_msg])
         yield history
         return
     
-    # 检查是否有图片（通过检查消息中是否有图片标签）
-    image_pattern = r'<img[^>]+src="([^">]+)"'
-    image_match = re.search(image_pattern, message)
+    # 处理图片
     image_content = None
-    text_content = message
+    text_content = message or ""
     
-    if image_match:
-        # 提取图片路径
-        image_path = image_match.group(1)
-        text_content = re.sub(image_pattern, '', message).strip()
-        
+    if image:
         # 转换图片为base64
         try:
-            if image_path.startswith('data:'):
-                # 已经是base64格式
-                image_content = image_path.split(',')[1]
-            else:
-                image_content = image_to_base64(image_path)
+            image_content = image_to_base64(image)
             print(f"[DEBUG] Image processed successfully")
         except Exception as e:
             print(f"[DEBUG] Failed to process image: {e}")
+            history.append([message or "[Image]", f"❌ Failed to process image: {str(e)}"])
+            yield history
+            return
+    
+    # 构造显示消息
+    if image and message:
+        display_message = f"🖼️ [Image] {message}"
+    elif image:
+        display_message = "🖼️ [Image]"
+    else:
+        display_message = message
     
     # 添加到历史
-    history.append([message, ""])
+    history.append([display_message, ""])
     yield history
     
     # 构建API消息
@@ -240,15 +244,23 @@ with gr.Blocks(title="Step-3 Chat", theme=gr.themes.Soft()) as demo:
             
             # 输入区域
             with gr.Row():
-                with gr.Column(scale=10):
+                with gr.Column(scale=8):
                     msg = gr.Textbox(
                         label="Message",
-                        placeholder="Type your message here... (You can drag and drop images)",
+                        placeholder="Type your message here...",
                         lines=2,
                         max_lines=10,
                         show_label=False,
                         elem_id="message-textbox"
                     )
+                with gr.Column(scale=2):
+                    image_input = gr.Image(
+                        label="Upload Image (Optional)",
+                        type="filepath",
+                        height=100,
+                        interactive=True,
+                        show_label=True
+                    )
                 with gr.Column(scale=1, min_width=100):
                     submit_btn = gr.Button("Send", variant="primary")
             
@@ -307,20 +319,22 @@ with gr.Blocks(title="Step-3 Chat", theme=gr.themes.Soft()) as demo:
                 )
     
     # 事件处理函数
-    def user_submit(message, history):
+    def user_submit(message, history, image):
         """用户提交消息时的处理"""
         print(f"[DEBUG] user_submit called with message: {message[:50] if message else 'None'}...")
-        if message:
-            # 保存消息内容用于后续处理
-            return gr.update(value=""), history, message
-        return gr.update(value=message), history, message
+        print(f"[DEBUG] user_submit called with image: {image is not None}")
+        if message or image:
+            # 清空输入，保存消息和图片用于后续处理
+            return gr.update(value=""), history, gr.update(value=None), message, image
+        return gr.update(value=message), history, gr.update(value=image), message, image
     
-    def bot_response(history, saved_message, system_prompt, temperature, max_tokens, top_p):
+    def bot_response(history, saved_message, saved_image, system_prompt, temperature, max_tokens, top_p):
         """生成机器人响应"""
         print(f"[DEBUG] bot_response called with saved_message: {saved_message[:50] if saved_message else 'None'}...")
-        if saved_message:
+        print(f"[DEBUG] bot_response called with saved_image: {saved_image is not None}")
+        if saved_message or saved_image:
             # 使用生成器处理消息
-            for updated_history in process_message(saved_message, history, system_prompt, temperature, max_tokens, top_p):
+            for updated_history in process_message(saved_message, history, saved_image, system_prompt, temperature, max_tokens, top_p):
                 yield updated_history
         else:
             yield history
@@ -337,38 +351,39 @@ with gr.Blocks(title="Step-3 Chat", theme=gr.themes.Soft()) as demo:
             return new_history, last_message
         return history, ""
     
-    # 创建一个隐藏的组件来存储消息
+    # 创建隐藏的组件来存储消息和图片
     saved_msg = gr.State("")
+    saved_img = gr.State(None)
     
     # 提交消息 - Enter键
     msg.submit(
         user_submit,
-        [msg, chatbot],
-        [msg, chatbot, saved_msg],
+        [msg, chatbot, image_input],
+        [msg, chatbot, image_input, saved_msg, saved_img],
         queue=False
     ).then(
         bot_response,
-        [chatbot, saved_msg, system_prompt, temperature, max_tokens, top_p],
+        [chatbot, saved_msg, saved_img, system_prompt, temperature, max_tokens, top_p],
         chatbot
     )
     
     # 提交消息 - Send按钮
     submit_btn.click(
         user_submit,
-        [msg, chatbot],
-        [msg, chatbot, saved_msg],
+        [msg, chatbot, image_input],
+        [msg, chatbot, image_input, saved_msg, saved_img],
         queue=False
     ).then(
         bot_response,
-        [chatbot, saved_msg, system_prompt, temperature, max_tokens, top_p],
+        [chatbot, saved_msg, saved_img, system_prompt, temperature, max_tokens, top_p],
         chatbot
     )
     
     # 清空对话
     clear_btn.click(
-        lambda: ([], ""),
+        lambda: ([], "", None),
         None,
-        [chatbot, msg]
+        [chatbot, msg, image_input]
     )
     
     # 撤销最后一条
@@ -385,7 +400,7 @@ with gr.Blocks(title="Step-3 Chat", theme=gr.themes.Soft()) as demo:
         [chatbot, saved_msg]
     ).then(
         bot_response,
-        [chatbot, saved_msg, system_prompt, temperature, max_tokens, top_p],
+        [chatbot, saved_msg, saved_img, system_prompt, temperature, max_tokens, top_p],
         chatbot
     )