well-defined export format

app.py

@@ -8,6 +8,7 @@ from PIL import Image
 from settings_mgr import generate_download_settings_js, generate_upload_settings_js
 from llm import LLM, log_to_console
 from code_exec import eval_restricted_script
+from chat_export import import_history, get_export_js
 from botocore.config import Config
 
 dump_controls = False
@@ -146,45 +147,21 @@ def bot(message, history, aws_access, aws_secret, aws_token, system_prompt, temp
     except Exception as e:
         raise gr.Error(f"Error: {str(e)}")
 
-def import_history(history, file):
-    with open(file.name, mode="rb") as f:
-        content = f.read()
-
-        if isinstance(content, bytes):
-            content = content.decode('utf-8', 'replace')
-        else:
-            content = str(content)
-
-    # Deserialize the JSON content
-    import_data = json.loads(content)
-
-    # Check if 'history' key exists for backward compatibility
-    if 'history' in import_data:
-        history = import_data['history']
-        system_prompt_value = import_data.get('system_prompt', '')  # Set default if not present
-    else:
-        # Assume it's an old format with only history data
-        history = import_data
-        system_prompt_value = ''
-
-    # Process the history to handle image data
-    processed_history = []
-    for pair in history:
-        processed_pair = []
-        for message in pair:
-            if isinstance(message, dict) and 'file' in message and 'data' in message['file']:
-                # Create a gradio.Image from the base64 data
-                image_data = base64.b64decode(message['file']['data'].split(',')[1])
-                img = Image.open(io.BytesIO(image_data))
-                gr_image = gr.Image(img)
-                processed_pair.append(gr_image)
-
-                gr.Warning("Reusing images across sessions is limited to one conversation turn")
-            else:
-                processed_pair.append(message)
-        processed_history.append(processed_pair)
-
-    return processed_history, system_prompt_value
+def import_history_guarded(aws_access, aws_secret, aws_token, region, history, file):
+    # check credentials first
+    try:
+        sess = boto3.Session(
+            aws_access_key_id = aws_access,
+            aws_secret_access_key = aws_secret,
+            aws_session_token = aws_token,
+            region_name = region)
+        br = sess.client(service_name="bedrock")
+        br.list_foundation_models(byProvider="invalid")
+    except Exception as e:
+        raise gr.Error(f"Bedrock login error: {str(e)}")
+
+    # actual import
+    return import_history(history, file)
 
 def export_history(h, s):
     pass
@@ -253,7 +230,7 @@ with gr.Blocks(delete_cache=(86400, 86400)) as demo:
         dl_settings_button.click(None, controls, js=generate_download_settings_js("amz_chat_settings.bin", control_ids))
         ul_settings_button.click(None, None, None, js=generate_upload_settings_js(control_ids))
 
-    chat = gr.ChatInterface(fn=bot, multimodal=True, additional_inputs=controls, autofocus = False)
+    chat = gr.ChatInterface(fn=bot, multimodal=True, additional_inputs=controls, autofocus = False, type = "messages")
     chat.textbox.file_count = "multiple"
     chatbot = chat.chatbot
     chatbot.show_copy_button = True
@@ -268,53 +245,7 @@ with gr.Blocks(delete_cache=(86400, 86400)) as demo:
     with gr.Accordion("Import/Export", open = False):
         import_button = gr.UploadButton("History Import")
         export_button = gr.Button("History Export")
-        export_button.click(export_history, [chatbot, system_prompt], js="""
-            async (chat_history, system_prompt) => {
-                console.log('Chat History:', JSON.stringify(chat_history, null, 2));
-
-                async function fetchAndEncodeImage(url) {
-                    const response = await fetch(url);
-                    const blob = await response.blob();
-                    return new Promise((resolve, reject) => {
-                        const reader = new FileReader();
-                        reader.onloadend = () => resolve(reader.result);
-                        reader.onerror = reject;
-                        reader.readAsDataURL(blob);
-                    });
-                }
-
-                const processedHistory = await Promise.all(chat_history.map(async (pair) => {
-                    return await Promise.all(pair.map(async (message) => {
-                        if (message && message.file && message.file.url) {
-                            const base64Image = await fetchAndEncodeImage(message.file.url);
-                            return {
-                                ...message,
-                                file: {
-                                    ...message.file,
-                                    data: base64Image
-                                }
-                            };
-                        }
-                        return message;
-                    }));
-                }));
-
-                const export_data = {
-                    history: processedHistory,
-                    system_prompt: system_prompt
-                };
-                const history_json = JSON.stringify(export_data);
-                const blob = new Blob([history_json], {type: 'application/json'});
-                const url = URL.createObjectURL(blob);
-                const a = document.createElement('a');
-                a.href = url;
-                a.download = 'chat_history.json';
-                document.body.appendChild(a);
-                a.click();
-                document.body.removeChild(a);
-                URL.revokeObjectURL(url);
-            }
-        """)
+        export_button.click(lambda: None, [chatbot, system_prompt], js=get_export_js())
         dl_button = gr.Button("File download")
         dl_button.click(lambda: None, [chatbot], js="""
             (chat_history) => {
@@ -370,6 +301,7 @@ with gr.Blocks(delete_cache=(86400, 86400)) as demo:
                 }
             }
         """)
-        import_button.upload(import_history, inputs=[chatbot, import_button], outputs=[chatbot, system_prompt])
-
+        import_button.upload(import_history_guarded, 
+                            inputs=[aws_access, aws_secret, aws_token, region, chatbot, import_button], 
+                            outputs=[chatbot, system_prompt])
 demo.queue(default_concurrency_limit = None).launch()

chat_export.py

@@ -0,0 +1,209 @@
+import json
+import base64
+import os, io
+import mimetypes
+from PIL import Image
+import gradio as gr
+
+def import_history(history, file):
+    if os.path.getsize(file.name) > 100e6:
+        raise ValueError("History larger than 100 MB")
+
+    with open(file.name, mode="rb") as f:
+        content = f.read().decode('utf-8', 'replace')
+
+    import_data = json.loads(content)
+    
+    # Handle different import formats
+    if 'messages' in import_data:
+        # New OpenAI-style format
+        messages = import_data['messages']
+        system_prompt_value = ''
+        chat_history = []
+        
+        msg_num = 1
+        for msg in messages:
+            if msg['role'] == 'system':
+                system_prompt_value = msg['content']
+                continue
+                
+            if msg['role'] == 'user':
+                content = msg['content']
+                if isinstance(content, list):
+                    for item in content:
+                        if item.get('type', '') == 'image_url':
+                            # Create gr.Image from data URI
+                            image_data = base64.b64decode(item['image_url']['url'].split(',')[1])
+                            img = Image.open(io.BytesIO(image_data))
+                            chat_history.append({
+                                "role": msg['role'],
+                                "content": gr.Image(value=img)
+                            })
+                        elif item.get('type', '') == 'file':
+                            # Handle file content with gr.File
+                            fname = os.path.basename(item['file'].get('name', f'download{msg_num}'))
+                            dir_path = os.path.dirname(file.name)
+                            temp_path = os.path.join(dir_path, fname)
+                            file_data = base64.b64decode(item['file']['url'].split(',')[1])
+                            if (len(file_data) > 15e6):
+                                raise ValueError(f"file content `{fname}` larger than 15 MB")
+                            
+                            with open(temp_path, "wb") as tempf:
+                                tempf.write(file_data)
+                            chat_history.append({
+                                "role": msg['role'],
+                                "content": gr.File(value=temp_path, 
+                                                 label=fname)
+                            })
+                        else:
+                            chat_history.append(msg)
+                else:
+                    chat_history.append(msg)
+                    
+            elif msg['role'] == 'assistant':
+                chat_history.append(msg)
+
+            msg_num = msg_num + 1
+            
+    else:
+        # Legacy format handling
+        if 'history' in import_data:
+            legacy_history = import_data['history']
+            system_prompt_value = import_data.get('system_prompt', '')
+        else:
+            legacy_history = import_data
+            system_prompt_value = ''
+        
+        chat_history = []
+        # Convert tuple/pair format to messages format
+        for pair in legacy_history:
+            if pair[0]:  # User message
+                if isinstance(pair[0], dict) and 'file' in pair[0]:
+                    if 'data' in pair[0]['file']:
+                        # Legacy format with embedded data
+                        file_data = pair[0]['file']['data']
+                        mime_type = file_data.split(';')[0].split(':')[1]
+                        
+                        if mime_type.startswith('image/'):
+                            image_data = base64.b64decode(file_data.split(',')[1])
+                            img = Image.open(io.BytesIO(image_data))
+                            chat_history.append({
+                                "role": "user",
+                                "content": gr.Image(value=img)
+                            })
+                        else:
+                            fname = pair[0]['file'].get('name', 'download')
+                            dir_path = os.path.dirname(file.name)
+                            temp_path = os.path.join(dir_path, fname)
+                            file_data = base64.b64decode(file_data.split(',')[1])
+                            
+                            with open(temp_path, "wb") as tempf:
+                                tempf.write(file_data)
+                            chat_history.append({
+                                "role": "user",
+                                "content": gr.File(value=temp_path, 
+                                                 label=fname)
+                            })
+                    else:
+                        # Keep as-is but convert to message format
+                        chat_history.append({
+                            "role": "user",
+                            "content": pair[0]
+                        })
+                else:
+                    chat_history.append({
+                        "role": "user",
+                        "content": pair[0]
+                    })
+            
+            if pair[1]:  # Assistant message
+                chat_history.append({
+                    "role": "assistant",
+                    "content": pair[1]
+                })
+
+    return chat_history, system_prompt_value
+
+def get_export_js():
+    return """
+    async (chat_history, system_prompt) => {
+        let messages = [];
+        
+        if (system_prompt) {
+            messages.push({
+                "role": "system",
+                "content": system_prompt
+            });
+        }
+
+        async function processFile(file_url) {
+            const response = await fetch(file_url);
+            const blob = await response.blob();
+            return new Promise((resolve) => {
+                const reader = new FileReader();
+                reader.onloadend = () => resolve({
+                    data: reader.result,
+                    type: blob.type
+                });
+                reader.onerror = (error) => resolve(null);
+                reader.readAsDataURL(blob);
+            });
+        }
+
+        for (let message of chat_history) {
+            if (!message.role || !message.content) continue;
+
+            if (message.content && typeof message.content === 'object') {
+                if (message.content.file) {
+                    try {
+                        const file_data = await processFile(message.content.file.url);
+                        if (!file_data) continue;
+
+                        if (file_data.type.startsWith('image/')) {
+                            messages.push({
+                                "role": message.role,
+                                "content": [{
+                                    "type": "image_url",
+                                    "image_url": {
+                                        "url": file_data.data
+                                    }
+                                }]
+                            });
+                        } else {
+                            const fileLink = document.querySelector(`a[data-testid="chatbot-file"][download][href*="${message.content.file.url.split('/').pop()}"]`);
+                            const fileName = fileLink ? fileLink.getAttribute('download') : (message.content.file.name || "download");
+                            
+                            messages.push({
+                                "role": message.role,
+                                "content": [{
+                                    "type": "file",
+                                    "file": {
+                                        "url": file_data.data,
+                                        "name": fileName,
+                                        "mime_type": file_data.type
+                                    }
+                                }]
+                            });
+                        }
+                    } catch (error) {}
+                }
+            } else {
+                messages.push({
+                    "role": message.role,
+                    "content": message.content
+                });
+            }
+        }
+
+        const export_data = { messages };
+        const blob = new Blob([JSON.stringify(export_data)], {type: 'application/json'});
+        const url = URL.createObjectURL(blob);
+        const a = document.createElement('a');
+        a.href = url;
+        a.download = 'chat_history.json';
+        document.body.appendChild(a);
+        a.click();
+        document.body.removeChild(a);
+        URL.revokeObjectURL(url);
+    }
+    """

llm.py

@@ -32,28 +32,36 @@ class LLM:
         # AWS API requires strict user, assi, user, ... sequence
         lastTypeHuman = False
 
-        for human, assi in history:
-            if human:
+        for msg in history:
+            if msg['role'] == "user":
                 if lastTypeHuman:
                     last_msg = messages.pop()
                     user_msg_parts = last_msg["content"]
                 else:
                     user_msg_parts = []
 
-                if isinstance(human, tuple):
-                    user_msg_parts.extend(self._process_file(human[0]))
-                elif isinstance(human, gradio.Image):
-                    user_msg_parts.extend(self._process_file(human.value["path"]))
+                content = msg['content']
+                if isinstance(content, gradio.File):
+                    user_msg_parts.extend(self._process_file(content.value['path']))
+                elif isinstance(content, gradio.Image):
+                    user_msg_parts.extend(self._process_file(content.value["path"]))
                 else:
-                    user_msg_parts.extend([{"text": human}])
+                    user_msg_parts.extend([{"text": content}])
 
                 messages.append({"role": "user", "content": user_msg_parts})
                 lastTypeHuman = True
-            if assi:
-                messages.append({"role": "assistant", "content": [{"text": assi}]})
+            else:
+                messages.append({
+                    "role": "assistant",
+                    "content":[{"text": msg['content']}]
+                })
                 lastTypeHuman = False
         
-        user_msg_parts = []
+        if lastTypeHuman:
+            last_msg = messages.pop()
+            user_msg_parts = last_msg["content"]
+        else:
+            user_msg_parts = []
         if message["text"]:
             user_msg_parts.append({"text": message["text"]})
         if message["files"]:

requirements.txt

@@ -1,6 +1,6 @@
 gradio == 5.1
 langchain
-boto3>1.34.54
+boto3>1.35.68
 lxml
 PyMuPDF
 RestrictedPython