Update app.py

app.py

@@ -3,6 +3,7 @@
 Multimodal chat front‑end for Amazon Bedrock Nova Premier v1
 (text and/or image ➜ assistant text) with per‑session JSONL logging.
 Logs FULL conversations to HF Dataset and images to HF repo via Git LFS.
+Supports multiple concurrent users with separate sessions.
 
 Prereqs:
     pip install gradio==5.38 boto3 pillow datasets huggingface_hub
@@ -16,8 +17,9 @@ Update REPO_ID to your HF dataset repository.
 """
 
 from __future__ import annotations
-import base64, datetime, io, json, pathlib, uuid, os
+import base64, datetime, io, json, pathlib, uuid, os, threading
 from typing import Dict, List, Optional, Tuple
+from concurrent.futures import ThreadPoolExecutor
 
 import boto3
 from botocore.config import Config
@@ -52,6 +54,10 @@ IMG_DIR = LOG_DIR / "imgs"
 LOG_DIR.mkdir(exist_ok=True)
 IMG_DIR.mkdir(exist_ok=True)
 
+# Thread pool for background operations
+executor = ThreadPoolExecutor(max_workers=4)
+dataset_lock = threading.Lock()  # Prevent concurrent dataset updates
+
 # ====== Bedrock client ====== #
 bedrock = boto3.client(
     "bedrock-runtime",
@@ -101,16 +107,16 @@ def call_bedrock(
     messages.append({"role": "assistant", "content": [{"text": reply}]})
     return reply, messages
 
-def upload_image_to_hf_repo(session_id: str, pil_img: Image.Image, message_index: int) -> Optional[str]:
+def upload_image_to_hf_repo_async(session_id: str, pil_img: Image.Image, message_index: int) -> Optional[str]:
     """Upload image to HF repo and return the repo URL."""
     if not HF_TOKEN:
-        print("Warning: HF_TOKEN not set, skipping image upload to HF repo")
         return None
         
     try:
-        # Create unique filename with message index
+        # Create unique filename with message index and random component to avoid conflicts
         ts = datetime.datetime.utcnow().strftime("%Y%m%dT%H%M%S")
-        filename = f"images/{session_id}_{message_index:03d}_{ts}.png"
+        random_suffix = str(uuid.uuid4())[:8]
+        filename = f"images/{session_id}_{message_index:03d}_{ts}_{random_suffix}.png"
         
         # Save locally first
         local_path = IMG_DIR / f"{session_id}_{message_index:03d}_{ts}.png"
@@ -128,85 +134,86 @@ def upload_image_to_hf_repo(session_id: str, pil_img: Image.Image, message_index
         
         # Return the HF repo URL
         hf_image_url = f"https://huggingface.co/datasets/{REPO_ID}/resolve/main/{filename}"
-        print(f"Image uploaded to HF repo: {hf_image_url}")
         return hf_image_url
         
     except Exception as e:
         print(f"Failed to upload image to HF repo: {e}")
         return None
 
-def save_full_conversation_to_hf_dataset(session_id: str, chat_history: List[Tuple], br_history: List[Dict], image_urls: Dict[int, str]):
-    """Save the ENTIRE conversation to HF Dataset."""
+def save_full_conversation_to_hf_dataset_async(session_id: str, chat_history: List[Tuple], br_history: List[Dict], image_urls: Dict[int, str]):
+    """Save the ENTIRE conversation to HF Dataset with thread safety."""
     if not HF_TOKEN:
-        print("Warning: HF_TOKEN not set, skipping upload to HF dataset")
         return
         
-    try:
-        # Convert chat history to structured format
-        messages = []
-        for i, (user_msg, assistant_msg) in enumerate(chat_history):
-            messages.append({
-                "message_index": i,
-                "role": "user",
-                "content": user_msg,
-                "image_url": image_urls.get(i, ""),
-                "has_image": i in image_urls
-            })
-            messages.append({
-                "message_index": i,
-                "role": "assistant", 
-                "content": assistant_msg,
-                "image_url": "",
-                "has_image": False
-            })
-        
-        # Create conversation record
-        conversation_record = {
-            "session_id": session_id,
-            "timestamp": datetime.datetime.utcnow().isoformat() + "Z",
-            "message_count": len(chat_history),
-            "total_messages": len(messages),
-            "conversation_messages": messages,
-            "bedrock_history": br_history,  # Full Bedrock conversation context
-            "images_count": len(image_urls)
-        }
-        
-        # Load existing dataset
+    def background_save():
         try:
-            existing_dataset = Dataset.load_dataset(REPO_ID, token=HF_TOKEN, split="train")
-            records = existing_dataset.to_list()
-            
-            # Check if session already exists and update it
-            session_exists = False
-            for idx, record in enumerate(records):
-                if record.get("session_id") == session_id:
-                    records[idx] = conversation_record
-                    session_exists = True
-                    print(f"Updated existing session {session_id[:8]} in dataset")
-                    break
-            
-            if not session_exists:
-                records.append(conversation_record)
-                print(f"Added new session {session_id[:8]} to dataset")
+            with dataset_lock:  # Prevent concurrent dataset updates
+                # Convert chat history to structured format
+                messages = []
+                for i, (user_msg, assistant_msg) in enumerate(chat_history):
+                    messages.append({
+                        "message_index": i,
+                        "role": "user",
+                        "content": user_msg,
+                        "image_url": image_urls.get(i, ""),
+                        "has_image": i in image_urls
+                    })
+                    messages.append({
+                        "message_index": i,
+                        "role": "assistant", 
+                        "content": assistant_msg,
+                        "image_url": "",
+                        "has_image": False
+                    })
                 
-            updated_dataset = Dataset.from_list(records)
-            
-        except Exception as load_error:
-            print(f"Could not load existing dataset (creating new): {load_error}")
-            updated_dataset = Dataset.from_list([conversation_record])
-        
-        # Push updated dataset
-        updated_dataset.push_to_hub(
-            REPO_ID,
-            token=HF_TOKEN,
-            private=True,
-            commit_message=f"Update full conversation for session {session_id[:8]} ({len(chat_history)} exchanges)"
-        )
-        
-        print(f"Full conversation logged to HF dataset: {REPO_ID}")
-        
-    except Exception as e:
-        print(f"Failed to upload full conversation to HF dataset: {e}")
+                # Create conversation record
+                conversation_record = {
+                    "session_id": session_id,
+                    "timestamp": datetime.datetime.utcnow().isoformat() + "Z",
+                    "message_count": len(chat_history),
+                    "total_messages": len(messages),
+                    "conversation_messages": messages,
+                    "bedrock_history": br_history,  # Full Bedrock conversation context
+                    "images_count": len(image_urls)
+                }
+                
+                # Load existing dataset
+                try:
+                    existing_dataset = Dataset.load_dataset(REPO_ID, token=HF_TOKEN, split="train")
+                    records = existing_dataset.to_list()
+                    
+                    # Check if session already exists and update it
+                    session_exists = False
+                    for idx, record in enumerate(records):
+                        if record.get("session_id") == session_id:
+                            records[idx] = conversation_record
+                            session_exists = True
+                            break
+                    
+                    if not session_exists:
+                        records.append(conversation_record)
+                        
+                    updated_dataset = Dataset.from_list(records)
+                    
+                except Exception as load_error:
+                    # Dataset doesn't exist yet, create new one
+                    updated_dataset = Dataset.from_list([conversation_record])
+                
+                # Push updated dataset
+                updated_dataset.push_to_hub(
+                    REPO_ID,
+                    token=HF_TOKEN,
+                    private=True,
+                    commit_message=f"Update conversation {session_id[:8]} ({len(chat_history)} exchanges)"
+                )
+                
+                print(f"Conversation {session_id[:8]} saved to HF dataset")
+                
+        except Exception as e:
+            print(f"Failed to save conversation {session_id[:8]} to HF dataset: {e}")
+    
+    # Run in background thread to avoid blocking UI
+    executor.submit(background_save)
 
 def save_local_conversation_log(session_id: str, chat_history: List[Tuple], image_urls: Dict[int, str]):
     """Save full conversation to local JSONL file."""
@@ -237,11 +244,16 @@ with gr.Blocks(title="Multimodal Chat with Full Conversation Logging") as demo:
         ## Multimodal Chat with Full Conversation Logging 📝
         Upload an image *(optional)*, ask a question, and continue the conversation.
         
+        **Multi-User Support:**
+        - 👥 Each browser tab/session gets a unique conversation ID
+        - 🔒 Conversations are isolated between users
+        - ⚡ Background logging doesn't block the UI
+        
         **Logging Features:**
         - 💾 **Full conversation history** saved after each message
         - 🔄 **Context preservation** - entire chat context maintained
         - 📸 **Image tracking** - all images linked to specific messages
-        - 🏷️ **Session management** - each browser tab gets unique session ID
+        - 🏷️ **Session management** - unique session ID per conversation
         
         **Storage:**
         - 🤗 HF Dataset: {"✅ Enabled" if HF_TOKEN else "❌ Disabled (set HF_TOKEN)"} - Repo: `{REPO_ID}`
@@ -249,7 +261,7 @@ with gr.Blocks(title="Multimodal Chat with Full Conversation Logging") as demo:
         """
     )
 
-    chatbot = gr.Chatbot(height=420)
+    chatbot = gr.Chatbot(height=420, type="tuples")  # Fix the warning
     chat_state = gr.State([])   # [(user, assistant), …] - Full chat history
     br_state = gr.State([])     # Bedrock message dicts - Full conversation context
     sess_state = gr.State("")   # UUID for this browser tab
@@ -278,6 +290,15 @@ with gr.Blocks(title="Multimodal Chat with Full Conversation Logging") as demo:
             max_lines=3
         )
 
+    # Session info display
+    with gr.Row():
+        session_info = gr.Textbox(
+            label="Session Info",
+            value="New session will be created on first message",
+            interactive=False,
+            max_lines=1
+        )
+
     # ---- main handler ---- #
     def chat(chat_log, br_history, sess_id, img_urls_dict,
              image, text,
@@ -286,8 +307,10 @@ with gr.Blocks(title="Multimodal Chat with Full Conversation Logging") as demo:
         if image is None and not text.strip():
             raise gr.Error("Upload an image or enter a message.")
 
+        # Create new session if needed
         if not sess_id:
             sess_id = str(uuid.uuid4())
+            print(f"Created new session: {sess_id}")
 
         # Call Bedrock with full conversation context
         try:
@@ -303,32 +326,40 @@ with gr.Blocks(title="Multimodal Chat with Full Conversation Logging") as demo:
         hf_img_url = None
         if image:
             message_index = len(chat_log)  # Current message index
-            hf_img_url = upload_image_to_hf_repo(sess_id, image, message_index)
-            if hf_img_url:
+            # Upload image in background (non-blocking)
+            if HF_TOKEN:
+                executor.submit(lambda: upload_image_to_hf_repo_async(sess_id, image, message_index))
+                # For immediate reference, we'll generate the expected URL pattern
+                ts = datetime.datetime.utcnow().strftime("%Y%m%dT%H%M%S")
+                hf_img_url = f"[Image uploaded to HF repo - session {sess_id[:8]}, message {message_index}]"
                 img_urls_dict[message_index] = hf_img_url
 
         # Update chat history
         display_user = text.strip() if text.strip() else "[image uploaded]"
         chat_log.append((display_user, reply))
         
-        # Save FULL conversation (local and HF)
+        # Save FULL conversation (local immediately, HF in background)
         save_local_conversation_log(sess_id, chat_log, img_urls_dict)
-        save_full_conversation_to_hf_dataset(sess_id, chat_log, new_br_history, img_urls_dict)
+        save_full_conversation_to_hf_dataset_async(sess_id, chat_log, new_br_history, img_urls_dict)
         
         # Update status message
         status_msg = f"✅ Full conversation logged for session {sess_id[:8]}\n"
         status_msg += f"📊 Total exchanges: {len(chat_log)} | Messages in context: {len(new_br_history)}"
         
         if image:
-            if hf_img_url:
-                status_msg += f"\n🖼️ Image #{len(img_urls_dict)} uploaded to HF repo"
+            if HF_TOKEN:
+                status_msg += f"\n🖼️ Image #{len([k for k in img_urls_dict.keys() if isinstance(k, int)])} uploading to HF repo..."
             else:
                 status_msg += f"\n⚠️ Image saved locally only"
         
         if not HF_TOKEN:
             status_msg += "\n❌ HF logging disabled (no token) - local only"
 
-        return chat_log, chat_log, new_br_history, sess_id, img_urls_dict, None, "", status_msg
+        # Update session info
+        session_display = f"Session: {sess_id[:8]}... | Messages: {len(chat_log)} | Images: {len([k for k in img_urls_dict.keys() if isinstance(k, int)])}"
+
+        return (chat_log, chat_log, new_br_history, sess_id, img_urls_dict, 
+                None, "", status_msg, session_display)
 
     send_btn.click(
         chat,
@@ -336,41 +367,53 @@ with gr.Blocks(title="Multimodal Chat with Full Conversation Logging") as demo:
                 img_in, txt_in,
                 max_tk, temp, top_p, top_k],
         outputs=[chatbot, chat_state, br_state, sess_state, img_urls_state, 
-                img_in, txt_in, log_status],
+                img_in, txt_in, log_status, session_info],
     )
 
     # ---- clear chat ---- #
     def reset():
-        return [], [], "", {}, None, "", "Ready to start logging full conversations..."
+        new_session_id = str(uuid.uuid4())
+        return ([], [], new_session_id, {}, None, "", 
+                "Ready to start logging full conversations...",
+                f"New session: {new_session_id[:8]}...")
 
     clear_btn.click(
         reset,
         inputs=None,
         outputs=[chatbot, chat_state, sess_state, img_urls_state, 
-                img_in, txt_in, log_status],
+                img_in, txt_in, log_status, session_info],
         queue=False,
     )
 
     # Add info about the logging structure
     gr.Markdown(
-        f"""
-        ### 📊 Conversation Log Structure
+        """
+        ### 📊 Multi-User Conversation Log Structure
         
         Each conversation is saved with:
-        - **Session ID**: Unique identifier for this browser tab/conversation
-        - **Full message history**: All user and assistant messages in order
+        - **Session ID**: Unique identifier for each conversation (8-char prefix shown in UI)
+        - **Full message history**: All user and assistant messages in chronological order
         - **Bedrock context**: Complete conversation context sent to the AI
         - **Image references**: URLs to all images uploaded during the conversation
         - **Timestamps**: When the conversation was last updated
+        - **Concurrent support**: Multiple users can chat simultaneously without interference
         
         ### 🔍 Viewing Your Logs
         
-        {"**HF Dataset**: [https://huggingface.co/datasets/" + REPO_ID + "](https://huggingface.co/datasets/" + REPO_ID + ")" if HF_TOKEN else "**HF Dataset**: Not configured - set HF_TOKEN to enable"}
-        - Each record contains a complete conversation
+        **HF Dataset**: """ + (f"[https://huggingface.co/datasets/{REPO_ID}](https://huggingface.co/datasets/{REPO_ID})" if HF_TOKEN else "Not configured - set HF_TOKEN to enable") + """
+        - Each record contains a complete conversation from one user session
         - Images stored in `images/` folder with session and message indexing
         - Download as JSON, CSV, or Parquet for analysis
+        - Search/filter by session_id to find specific conversations
         
-        **Local logs**: Saved as `{session_id}_full.jsonl` (temporary until Space restarts)
+        **Local logs**: Saved as `[session_id]_full.jsonl` (temporary until Space restarts)
+        
+        ### 🔧 Performance Notes
+        
+        - Image uploads and dataset updates happen in background threads
+        - UI remains responsive during logging operations
+        - Each browser tab gets its own isolated conversation
+        - Session IDs are generated client-side for immediate isolation
         """
     )
 
@@ -387,5 +430,7 @@ if __name__ == "__main__":
         print(f"✅ HF logging enabled. Full conversations will be saved to: {REPO_ID}")
         print(f"📸 Images will be stored in: {REPO_ID}/images/")
     
+    print("🚀 Multi-user chat with full conversation logging ready!")
+    
     demo.queue(max_size=100)
     demo.launch(share=True)