Update app.py

app.py

@@ -19,7 +19,7 @@ Update REPO_ID to your HF dataset repository.
 from __future__ import annotations
 import base64, datetime, io, json, pathlib, uuid, os, threading
 from typing import Dict, List, Optional, Tuple
-from concurrent.futures import ThreadPoolExecutor
+from concurrent.futures import ThreadPoolExecutor, Future
 
 import boto3
 from botocore.config import Config
@@ -45,7 +45,7 @@ Never mention Amazon or Nova.
 
 # HuggingFace Configuration
 HF_TOKEN = os.getenv("HF_TOKEN")  # Set this in your Space's secrets
-REPO_ID = "collinear-ai/nova-premier-redteaming-external"  # Change this to your actual HF dataset repo
+REPO_ID = "collinear-ai/nova-premier-redteaming-external"  # Updated to your actual repo
 HF_API = HfApi()
 
 # Local directories (for temporary storage)
@@ -107,22 +107,21 @@ def call_bedrock(
     messages.append({"role": "assistant", "content": [{"text": reply}]})
     return reply, messages
 
-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."""
+def upload_image_to_hf_repo_sync(session_id: str, pil_img: Image.Image, message_index: int) -> Optional[str]:
+    """Upload image to HF repo synchronously and return the actual repo URL."""
     if not HF_TOKEN:
         return None
         
     try:
-        # Create unique filename with message index and random component to avoid conflicts
+        # Create unique filename with message index and timestamp
         ts = datetime.datetime.utcnow().strftime("%Y%m%dT%H%M%S")
-        random_suffix = str(uuid.uuid4())[:8]
-        filename = f"images/{session_id}_{message_index:03d}_{ts}_{random_suffix}.png"
+        filename = f"images/{session_id}_{message_index:03d}_{ts}.png"
         
         # Save locally first
         local_path = IMG_DIR / f"{session_id}_{message_index:03d}_{ts}.png"
         pil_img.save(local_path, format="PNG")
         
-        # Upload to HF repo
+        # Upload to HF repo and wait for completion
         upload_file(
             path_or_fileobj=str(local_path),
             path_in_repo=filename,
@@ -132,12 +131,13 @@ def upload_image_to_hf_repo_async(session_id: str, pil_img: Image.Image, message
             commit_message=f"Add image for session {session_id[:8]} message {message_index}"
         )
         
-        # Return the HF repo URL
+        # Return the actual HF repo URL
         hf_image_url = f"https://huggingface.co/datasets/{REPO_ID}/resolve/main/{filename}"
+        print(f"✅ Image uploaded successfully: {hf_image_url}")
         return hf_image_url
         
     except Exception as e:
-        print(f"Failed to upload image to HF repo: {e}")
+        print(f"❌ Failed to upload image to HF repo: {e}")
         return None
 
 def save_full_conversation_to_hf_dataset_async(session_id: str, chat_history: List[Tuple], br_history: List[Dict], image_urls: Dict[int, str]):
@@ -207,10 +207,10 @@ def save_full_conversation_to_hf_dataset_async(session_id: str, chat_history: Li
                     commit_message=f"Update conversation {session_id[:8]} ({len(chat_history)} exchanges)"
                 )
                 
-                print(f"Conversation {session_id[:8]} saved to HF dataset")
+                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}")
+            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)
@@ -238,27 +238,30 @@ def save_local_conversation_log(session_id: str, chat_history: List[Tuple], imag
         f.write(json.dumps(conversation_record, ensure_ascii=False) + "\n")
 
 # ====== Gradio UI ====== #
-with gr.Blocks(title="Multimodal Chat with Full Conversation Logging") as demo:
+with gr.Blocks(title="Nova Premier Red Team Chat") as demo:
     gr.Markdown(
         f"""
-        ## Multimodal Chat with Full Conversation Logging 📝
+        ## Nova Premier Red Team Chat 🔴🤖
         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
+        - ⚡ Real-time image upload with direct URLs
         
         **Logging Features:**
         - 💾 **Full conversation history** saved after each message
         - 🔄 **Context preservation** - entire chat context maintained
-        - 📸 **Image tracking** - all images linked to specific messages
+        - 📸 **Image tracking** - direct links to uploaded images
         - 🏷️ **Session management** - unique session ID per conversation
         
+        **Storage:**
+        - 🤗 HF Dataset: {"✅ Enabled" if HF_TOKEN else "❌ Disabled (set HF_TOKEN)"} - Repo: `{REPO_ID}`
+        - ��️ Images: {"✅ Uploaded with direct URLs" if HF_TOKEN else "❌ Local only"}
         """
     )
 
-    chatbot = gr.Chatbot(height=420, type="tuples")  # Fix the warning
+    chatbot = gr.Chatbot(height=420, type="tuples")
     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
@@ -284,7 +287,7 @@ with gr.Blocks(title="Multimodal Chat with Full Conversation Logging") as demo:
             label="Conversation Logging Status", 
             value="Ready to start logging full conversations...",
             interactive=False,
-            max_lines=3
+            max_lines=4
         )
 
     # Session info display
@@ -307,7 +310,24 @@ with gr.Blocks(title="Multimodal Chat with Full Conversation Logging") as demo:
         # Create new session if needed
         if not sess_id:
             sess_id = str(uuid.uuid4())
-            print(f"Created new session: {sess_id}")
+            print(f"🆕 Created new session: {sess_id}")
+
+        # Handle image upload FIRST (synchronously to get real URL)
+        hf_img_url = None
+        upload_status = ""
+        if image:
+            message_index = len(chat_log)  # Current message index
+            if HF_TOKEN:
+                upload_status = "🔄 Uploading image..."
+                # Upload synchronously to get the actual URL
+                hf_img_url = upload_image_to_hf_repo_sync(sess_id, image, message_index)
+                if hf_img_url:
+                    img_urls_dict[message_index] = hf_img_url
+                    upload_status = f"✅ Image uploaded: {hf_img_url}"
+                else:
+                    upload_status = "❌ Image upload failed"
+            else:
+                upload_status = "⚠️ Image saved locally only (no HF token)"
 
         # Call Bedrock with full conversation context
         try:
@@ -319,18 +339,6 @@ with gr.Blocks(title="Multimodal Chat with Full Conversation Logging") as demo:
         except Exception as e:
             raise gr.Error(f"Bedrock API error: {str(e)}")
 
-        # Handle image upload if present
-        hf_img_url = None
-        if image:
-            message_index = len(chat_log)  # Current message index
-            # 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))
@@ -339,21 +347,24 @@ with gr.Blocks(title="Multimodal Chat with Full Conversation Logging") as demo:
         save_local_conversation_log(sess_id, chat_log, img_urls_dict)
         save_full_conversation_to_hf_dataset_async(sess_id, chat_log, new_br_history, img_urls_dict)
         
-        # Update status message
+        # Update status message with real image URL
         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)}"
+        status_msg += f"📊 Total exchanges: {len(chat_log)} | Messages in context: {len(new_br_history)}\n"
         
-        if image:
-            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 image and hf_img_url:
+            status_msg += f"🖼️ Image URL: {hf_img_url}\n"
+        elif image and not hf_img_url:
+            status_msg += f"⚠️ Image upload failed - check logs\n"
+        
+        if upload_status:
+            status_msg += upload_status
         
         if not HF_TOKEN:
             status_msg += "\n❌ HF logging disabled (no token) - local only"
 
         # 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)])}"
+        image_count = len([k for k in img_urls_dict.keys() if isinstance(k, int)])
+        session_display = f"Session: {sess_id[:8]}... | Messages: {len(chat_log)} | Images: {image_count}"
 
         return (chat_log, chat_log, new_br_history, sess_id, img_urls_dict, 
                 None, "", status_msg, session_display)
@@ -384,33 +395,40 @@ with gr.Blocks(title="Multimodal Chat with Full Conversation Logging") as demo:
 
     # Add info about the logging structure
     gr.Markdown(
-        """
-        ### 📊 Multi-User Conversation Log Structure
+        f"""
+        ### 📊 Red Team Conversation Log Structure
         
         Each conversation is saved with:
         - **Session ID**: Unique identifier for each conversation (8-char prefix shown in UI)
         - **Full message history**: All user and assistant messages in chronological order
+        - **Image URLs**: Direct links to uploaded images in HF repo
         - **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
+        
+        ### 🖼️ Image Storage Format
+        
+        Images are stored as: `images/{{session_id}}_{{message_index:03d}}_{{timestamp}}.png`
+        - Example: `images/d8771c05_001_20240123T143022.png`
+        - **session_id**: First 8 chars of session UUID
+        - **message_index**: 3-digit message number (000, 001, 002...)
+        - **timestamp**: UTC timestamp when uploaded
         
         ### 🔍 Viewing Your Logs
         
-        **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") + """
+        **HF Dataset**: [https://huggingface.co/datasets/{REPO_ID}](https://huggingface.co/datasets/{REPO_ID})
         - 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
+        - **Images folder**: [https://huggingface.co/datasets/{REPO_ID}/tree/main/images](https://huggingface.co/datasets/{REPO_ID}/tree/main/images)
         - Search/filter by session_id to find specific conversations
+        - Image URLs in the conversation data link directly to the files
         
         **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
+        - Images are uploaded synchronously to get real URLs immediately
+        - Dataset updates happen in background threads for performance
         - Each browser tab gets its own isolated conversation
-        - Session IDs are generated client-side for immediate isolation
+        - Direct image URLs allow easy correlation with conversation data
         """
     )
 
@@ -422,12 +440,11 @@ if __name__ == "__main__":
         print("   To enable persistent logging:")
         print("   1. Go to your Space settings → Repository secrets")
         print("   2. Add HF_TOKEN with your HuggingFace token (write permissions)")
-        print("   3. Create a dataset repository and update REPO_ID in the code")
     else:
         print(f"✅ HF logging enabled. Full conversations will be saved to: {REPO_ID}")
-        print(f"📸 Images will be stored in: {REPO_ID}/images/")
+        print(f"📸 Images will be stored at: https://huggingface.co/datasets/{REPO_ID}/tree/main/images")
     
-    print("🚀 Multi-user chat with full conversation logging ready!")
+    print("🔴 Nova Premier Red Team Chat with full conversation logging ready!")
     
     demo.queue(max_size=100)
     demo.launch(share=True)