gradio_interface.py

# session_isolated_interface.py - Session-isolated Gradio interface with Edit Prompts tab

import os
import gradio as gr
import json
import uuid
from datetime import datetime
from dataclasses import asdict
from typing import Dict, Any, Optional

from lifestyle_app import ExtendedLifestyleJourneyApp
from core_classes import SessionState, ChatMessage
from prompts import SYSTEM_PROMPT_MAIN_LIFESTYLE

try:
    from app_config import GRADIO_CONFIG
except ImportError:
    GRADIO_CONFIG = {"theme": "soft", "show_api": False}

class SessionData:
    """Container for user session data"""
    def __init__(self, session_id: str = None):
        self.session_id = session_id or str(uuid.uuid4())
        self.app_instance = ExtendedLifestyleJourneyApp()
        self.created_at = datetime.now().isoformat()
        self.last_activity = datetime.now().isoformat()
        # NEW: Custom prompts storage
        self.custom_prompts = {
            "main_lifestyle": SYSTEM_PROMPT_MAIN_LIFESTYLE  # Default prompt
        }
        self.prompts_modified = False
    
    def to_dict(self) -> Dict[str, Any]:
        """Serialize session for storage"""
        return {
            "session_id": self.session_id,
            "created_at": self.created_at,
            "last_activity": self.last_activity,
            "chat_history": [asdict(msg) for msg in self.app_instance.chat_history],
            "session_state": asdict(self.app_instance.session_state),
            "test_mode_active": self.app_instance.test_mode_active,
            "current_test_patient": self.app_instance.current_test_patient,
            "custom_prompts": self.custom_prompts,
            "prompts_modified": self.prompts_modified
        }
    
    def update_activity(self):
        """Update last activity timestamp"""
        self.last_activity = datetime.now().isoformat()
    
    def set_custom_prompt(self, prompt_name: str, prompt_text: str):
        """Set custom prompt for this session"""
        self.custom_prompts[prompt_name] = prompt_text
        self.prompts_modified = True
        # Update the app instance to use custom prompt
        if hasattr(self.app_instance, 'main_lifestyle_assistant'):
            self.app_instance.main_lifestyle_assistant.set_custom_system_prompt(prompt_text)
    
    def reset_prompt_to_default(self, prompt_name: str):
        """Reset prompt to default"""
        if prompt_name == "main_lifestyle":
            self.custom_prompts[prompt_name] = SYSTEM_PROMPT_MAIN_LIFESTYLE
            self.prompts_modified = False
            # Update the app instance
            if hasattr(self.app_instance, 'main_lifestyle_assistant'):
                self.app_instance.main_lifestyle_assistant.reset_to_default_prompt()

def load_instructions() -> str:
    """Load instructions from INSTRUCTION.md file"""
    try:
        with open("INSTRUCTION.md", "r", encoding="utf-8") as f:
            content = f.read()
        return content
    except FileNotFoundError:
        return """# 📖 Instructions Unavailable

❌ **File INSTRUCTION.md not found**

To view the full instructions, please ensure the `INSTRUCTION.md` file is in the application's root folder.

## 🚀 Quick Start

1. **For medical questions:** "I have a headache"
2. **For lifestyle coaching:** "I want to start exercising"  
3. **For testing:** Go to the "🧪 Testing Lab" tab

## ⚠️ Important
This application is not a substitute for professional medical advice. In case of serious symptoms, please consult a doctor.
"""
    except Exception as e:
        return f"""# ❌ Error Loading Instructions

An error occurred while reading the instructions file: `{str(e)}`

## 🔧 Recommendations
- Check that the INSTRUCTION.md file exists
- Ensure the file has the correct UTF-8 encoding
- Restart the application

## 🆘 Basic Help
For help, type "help" or "how to use" in the chat.
"""

def create_session_isolated_interface():
    """Create session-isolated Gradio interface with Edit Prompts tab"""
    
    log_prompts_enabled = os.getenv("LOG_PROMPTS", "false").lower() == "true"
    
    theme_name = GRADIO_CONFIG.get("theme", "soft")
    if theme_name.lower() == "soft":
        theme = gr.themes.Soft()
    elif theme_name.lower() == "default":
        theme = gr.themes.Default()
    else:
        theme = gr.themes.Soft()
    
    with gr.Blocks(
        title=GRADIO_CONFIG.get("title", "Lifestyle Journey MVP + Testing Lab"), 
        theme=theme,
        analytics_enabled=False
    ) as demo:
        # Session state - CRITICAL: Each user gets isolated state
        session_data = gr.State(value=None)
        
        # Header
        if log_prompts_enabled:
            gr.Markdown("# 🏥 Lifestyle Journey MVP + 🧪 Testing Lab + 🔧 Prompt Editor 📝")
            gr.Markdown("⚠️ **DEBUG MODE:** LLM prompts and responses are saved to `lifestyle_journey.log`")
        else:
            gr.Markdown("# 🏥 Lifestyle Journey MVP + 🧪 Testing Lab + 🔧 Prompt Editor")
            
        gr.Markdown("Medical chatbot with lifestyle coaching, testing system, and prompt customization")
        
        # Session info
        with gr.Row():
            session_info = gr.Markdown("🔄 **Initializing session...**")
        
        # Initialize session on load
        def initialize_session():
            """Initialize new user session"""
            new_session = SessionData()
            session_info_text = f"""
✅ **Session Initialized**
🆔 **Session ID:** `{new_session.session_id[:8]}...`
🕒 **Started:** {new_session.created_at[:19]}
👤 **Isolated Instance:** Each user has separate data
🔧 **Custom Prompts:** {'✅ Modified' if new_session.prompts_modified else '🔄 Default'}
            """
            return new_session, session_info_text
        
        # Main tabs
        with gr.Tabs():
            # Main chat tab
            with gr.TabItem("💬 Patient Chat", id="main_chat"):
                with gr.Row():
                    with gr.Column(scale=2):
                        chatbot = gr.Chatbot(
                            label="💬 Conversation with Assistant",
                            height=400,
                            show_copy_button=True,
                            type="messages"
                        )
                        
                        with gr.Row():
                            msg = gr.Textbox(
                                label="Your message",
                                placeholder="Type your question...",
                                scale=4
                            )
                            send_btn = gr.Button("📤 Send", scale=1)
                        
                        with gr.Row():
                            clear_btn = gr.Button("🗑️ Clear Chat", scale=1)
                            end_conversation_btn = gr.Button("🏁 End Conversation", scale=1, variant="secondary")
                        
                        # Quick start examples
                        gr.Markdown("### ⚡ Quick Start:")
                        with gr.Row():
                            example_medical_btn = gr.Button("🩺 I have a headache", size="sm")
                            example_lifestyle_btn = gr.Button("💚 I want to start exercising", size="sm")
                            example_help_btn = gr.Button("❓ Help", size="sm")
                        
                    with gr.Column(scale=1):
                        status_box = gr.Markdown(
                            value="🔄 Loading status...",
                            label="📊 System Status"
                        )
                        
                        refresh_status_btn = gr.Button("🔄 Refresh Status", size="sm")
                        
                        end_conversation_result = gr.Markdown(value="", visible=False)

            # NEW: Edit Prompts tab
            with gr.TabItem("🔧 Edit Prompts", id="edit_prompts"):
                gr.Markdown("## 🔧 Customize AI Assistant Prompts")
                gr.Markdown("⚠️ **Note:** Changes apply only to your current session and will be lost when you close the browser.")
                
                with gr.Row():
                    with gr.Column(scale=3):
                        gr.Markdown("### 💚 Main Lifestyle Assistant Prompt")
                        
                        main_lifestyle_prompt = gr.Textbox(
                            label="System Prompt for Lifestyle Coaching",
                            value=SYSTEM_PROMPT_MAIN_LIFESTYLE,
                            lines=20,
                            max_lines=30,
                            placeholder="Enter your custom system prompt here...",
                            info="This prompt defines how the AI behaves during lifestyle coaching sessions."
                        )
                        
                        with gr.Row():
                            apply_prompt_btn = gr.Button("✅ Apply Changes", variant="primary", scale=2)
                            reset_prompt_btn = gr.Button("🔄 Reset to Default", variant="secondary", scale=1)
                            preview_prompt_btn = gr.Button("👁️ Preview", size="sm", scale=1)
                        
                        prompt_status = gr.Markdown(value="", visible=True)
                    
                    with gr.Column(scale=1):
                        gr.Markdown("### 📋 Prompt Guidelines")
                        gr.Markdown("""
**🎯 Key Elements to Include:**
- **Role definition** (lifestyle coach)
- **Safety principles** (medical limitations)
- **Action logic** (gather_info/lifestyle_dialog/close)
- **Output format** (JSON with message/action/reasoning)

**⚠️ Important:**
- Keep JSON format for actions
- Maintain safety guidelines
- Consider patient's medical conditions
- Use same language as patient

**🔧 Actions:**
- `gather_info` - collect more details
- `lifestyle_dialog` - provide coaching
- `close` - end session safely

**💡 Tips:**
- Test changes with simple questions
- Use "🔄 Reset" if issues occur
- Check JSON format carefully
                        """)
                        
                        gr.Markdown("### 📊 Current Settings")
                        prompt_info = gr.Markdown(value="🔄 Default prompt active")
            
            # Testing Lab tab
            with gr.TabItem("🧪 Testing Lab", id="testing_lab"):
                gr.Markdown("## 📁 Load Test Patient")
                
                with gr.Row():
                    with gr.Column():
                        clinical_file = gr.File(
                            label="🏥 Clinical Background JSON",
                            file_types=[".json"],
                            type="filepath"
                        )
                        lifestyle_file = gr.File(
                            label="💚 Lifestyle Profile JSON", 
                            file_types=[".json"],
                            type="filepath"
                        )
                        
                        load_patient_btn = gr.Button("📋 Load Patient", variant="primary")
                    
                    with gr.Column():
                        load_result = gr.Markdown(value="Select files to load")
                
                # Quick test buttons
                gr.Markdown("## ⚡ Quick Testing (Built-in Data)")
                with gr.Row():
                    quick_elderly_btn = gr.Button("👵 Elderly Mary", size="sm")
                    quick_athlete_btn = gr.Button("🏃 Athletic John", size="sm") 
                    quick_pregnant_btn = gr.Button("🤰 Pregnant Sarah", size="sm")
                
                gr.Markdown("## 👤 Patient Preview")
                patient_preview = gr.Markdown(value="No patient loaded")
                
                gr.Markdown("## 🎯 Test Session Management")
                with gr.Row():
                    end_session_notes = gr.Textbox(
                        label="Session End Notes",
                        placeholder="Describe testing results...",
                        lines=3
                    )
                    with gr.Column():
                        end_session_btn = gr.Button("⏹️ End Test Session")
                        end_session_result = gr.Markdown(value="")
            
            # Test results tab
            with gr.TabItem("📊 Test Results", id="test_results"):
                gr.Markdown("## 📈 Test Session Analysis")
                
                refresh_results_btn = gr.Button("🔄 Refresh Results")
                
                with gr.Row():
                    with gr.Column(scale=2):
                        results_summary = gr.Markdown(value="Click 'Refresh Results'")
                    
                    with gr.Column(scale=1):
                        export_btn = gr.Button("💾 Export to CSV")
                        export_result = gr.Markdown(value="")
                
                gr.Markdown("## 📋 Recent Test Sessions")
                results_table = gr.Dataframe(
                    headers=["Patient", "Time", "Messages", "Medical", "Lifestyle", "Escalations", "Duration", "Notes"],
                    datatype=["str", "str", "number", "number", "number", "number", "str", "str"],
                    label="Session Details",
                    value=[]
                )

            # Instructions tab
            with gr.TabItem("📖 Instructions", id="instructions"):
                gr.Markdown("## 📚 User Guide")
                
                # Load and display instructions
                instructions_content = load_instructions()
                
                with gr.Row():
                    with gr.Column(scale=4):
                        instructions_display = gr.Markdown(
                            value=instructions_content,
                            label="📖 Instructions"
                        )
                    
                    with gr.Column(scale=1):
                        gr.Markdown("### 🔗 Quick Links")
                        
                        # Quick navigation buttons
                        medical_example_btn = gr.Button("🩺 Medical Example", size="sm")
                        lifestyle_example_btn = gr.Button("💚 Lifestyle Example", size="sm")
                        testing_example_btn = gr.Button("🧪 Testing", size="sm")
                        prompts_example_btn = gr.Button("🔧 Edit Prompts", size="sm")
                        
                        gr.Markdown("### 📞 Help")
                        refresh_instructions_btn = gr.Button("🔄 Refresh Instructions", size="sm")
                        
                        gr.Markdown("""
**💡 Quick Commands:**
- "help" - get assistance
- "example" - see examples
- "clear" - start over
                        """)
        
        # Session-isolated event handlers
        def handle_message_isolated(message: str, history, session: SessionData):
            """Session-isolated message handler"""
            if session is None:
                session = SessionData()
            
            session.update_activity()
            new_history, status = session.app_instance.process_message(message, history)
            return new_history, status, session
        
        def handle_clear_isolated(session: SessionData):
            """Session-isolated clear handler"""
            if session is None:
                session = SessionData()
            
            session.update_activity()
            new_history, status = session.app_instance.reset_session()
            return new_history, status, session
        
        def handle_load_patient_isolated(clinical_file, lifestyle_file, session: SessionData):
            """Session-isolated patient loading"""
            if session is None:
                session = SessionData()
            
            session.update_activity()
            result = session.app_instance.load_test_patient(clinical_file, lifestyle_file)
            return result + (session,)
        
        def handle_quick_test_isolated(patient_type: str, session: SessionData):
            """Session-isolated quick test loading"""
            if session is None:
                session = SessionData()
            
            session.update_activity()
            result = session.app_instance.load_quick_test_patient(patient_type)
            return result + (session,)
        
        def handle_end_conversation_isolated(session: SessionData):
            """Session-isolated conversation end"""
            if session is None:
                session = SessionData()
            
            session.update_activity()
            return session.app_instance.end_conversation_with_profile_update() + (session,)
        
        def get_status_isolated(session: SessionData):
            """Get session-isolated status"""
            if session is None:
                return "❌ Session not initialized"
            
            session.update_activity()
            base_status = session.app_instance._get_status_info()
            
            # Add prompt status
            prompt_status = ""
            if session.prompts_modified:
                prompt_status = f"""
🔧 **CUSTOM PROMPTS:**
• Main Lifestyle: ✅ Modified ({len(session.custom_prompts.get('main_lifestyle', ''))} chars)
• Status: Custom prompt active for this session
"""
            else:
                prompt_status = f"""
🔧 **CUSTOM PROMPTS:**
• Main Lifestyle: 🔄 Default prompt
• Status: Using original system prompts
"""
            
            session_status = f"""
🔐 **SESSION ISOLATION:**
• Session ID: {session.session_id[:8]}...
• Created: {session.created_at[:19]}
• Last Activity: {session.last_activity[:19]}
• Isolated: ✅ Your data is private
{prompt_status}
{base_status}
            """
            return session_status
        
        # NEW: Prompt editing handlers
        def apply_custom_prompt(prompt_text: str, session: SessionData):
            """Apply custom prompt to session"""
            if session is None:
                session = SessionData()
            
            session.update_activity()
            
            # Validate prompt (basic check)
            if not prompt_text.strip():
                return "❌ Prompt cannot be empty", session, "❌ Empty prompt"
            
            if len(prompt_text.strip()) < 50:
                return "⚠️ Prompt seems too short. Are you sure it's complete?", session, "⚠️ Short prompt"
            
            try:
                # Apply the custom prompt
                session.set_custom_prompt("main_lifestyle", prompt_text.strip())
                
                status_msg = f"""✅ **Custom prompt applied successfully!**

📊 **Details:**
• Length: {len(prompt_text.strip())} characters
• Applied to: Main Lifestyle Assistant
• Session: {session.session_id[:8]}...
• Status: Active for this session only

🔄 **Next steps:**
• Test the changes by starting a lifestyle conversation
• Use "Reset to Default" if you encounter issues
"""
                
                info_msg = f"✅ Custom prompt active ({len(prompt_text.strip())} chars)"
                
                return status_msg, session, info_msg
                
            except Exception as e:
                error_msg = f"❌ Error applying prompt: {str(e)}"
                return error_msg, session, "❌ Application failed"
        
        def reset_prompt_to_default(session: SessionData):
            """Reset prompt to default"""
            if session is None:
                session = SessionData()
            
            session.update_activity()
            session.reset_prompt_to_default("main_lifestyle")
            
            status_msg = f"""🔄 **Prompt reset to default**

📊 **Details:**
• Main Lifestyle Assistant prompt restored
• Session: {session.session_id[:8]}...
• All customizations removed

💡 You can edit and apply again at any time.
"""
            
            return SYSTEM_PROMPT_MAIN_LIFESTYLE, status_msg, session, "🔄 Default prompt active"
        
        def preview_prompt_changes(prompt_text: str):
            """Preview prompt changes"""
            if not prompt_text.strip():
                return "❌ No prompt text to preview"
            
            preview = f"""📋 **Prompt Preview:**

**Length:** {len(prompt_text.strip())} characters
**Lines:** {len(prompt_text.strip().split(chr(10)))} lines

**First 200 characters:**
```
{prompt_text.strip()[:200]}{'...' if len(prompt_text.strip()) > 200 else ''}
```

**Contains key elements:**
• JSON format mentioned: {'✅' if 'json' in prompt_text.lower() or 'JSON' in prompt_text else '❌'}
• Actions mentioned: {'✅' if 'gather_info' in prompt_text and 'lifestyle_dialog' in prompt_text and 'close' in prompt_text else '❌'}
• Safety guidelines: {'✅' if 'safety' in prompt_text.lower() or 'medical' in prompt_text.lower() else '❌'}

**Ready to apply:** {'✅ Yes' if len(prompt_text.strip()) > 50 else '❌ Too short'}
"""
            return preview
        
        # Helper functions for examples and instructions
        def send_example_message(example_text: str, history, session: SessionData):
            """Send example message to chat"""
            return handle_message_isolated(example_text, history, session)
        
        def refresh_instructions():
            """Refresh instructions content"""
            return load_instructions()
        
        def handle_end_session_isolated(notes: str, session: SessionData):
            """Session-isolated end session handler"""
            if session is None:
                session = SessionData()
            
            session.update_activity()
            result = session.app_instance.end_test_session(notes)
            return result, session
        
        def handle_refresh_results_isolated(session: SessionData):
            """Session-isolated refresh results handler"""
            if session is None:
                session = SessionData()
            
            session.update_activity()
            result = session.app_instance.get_test_results_summary()
            return result + (session,)
        
        def handle_export_isolated(session: SessionData):
            """Session-isolated export handler"""
            if session is None:
                session = SessionData()
            
            session.update_activity()
            result = session.app_instance.export_test_results()
            return result, session
        
        # Event binding with session isolation
        demo.load(
            initialize_session,
            outputs=[session_data, session_info]
        )
        
        # Main chat events
        send_btn.click(
            handle_message_isolated,
            inputs=[msg, chatbot, session_data],
            outputs=[chatbot, status_box, session_data]
        ).then(
            lambda: "",
            outputs=[msg]
        )
        
        msg.submit(
            handle_message_isolated,
            inputs=[msg, chatbot, session_data], 
            outputs=[chatbot, status_box, session_data]
        ).then(
            lambda: "",
            outputs=[msg]
        )
        
        clear_btn.click(
            handle_clear_isolated,
            inputs=[session_data],
            outputs=[chatbot, status_box, session_data]
        )
        
        end_conversation_btn.click(
            handle_end_conversation_isolated,
            inputs=[session_data],
            outputs=[chatbot, status_box, end_conversation_result, session_data]
        )
        
        # Status refresh
        refresh_status_btn.click(
            get_status_isolated,
            inputs=[session_data],
            outputs=[status_box]
        )
        
        # NEW: Prompt editing events
        apply_prompt_btn.click(
            apply_custom_prompt,
            inputs=[main_lifestyle_prompt, session_data],
            outputs=[prompt_status, session_data, prompt_info]
        )
        
        reset_prompt_btn.click(
            reset_prompt_to_default,
            inputs=[session_data],
            outputs=[main_lifestyle_prompt, prompt_status, session_data, prompt_info]
        )
        
        preview_prompt_btn.click(
            preview_prompt_changes,
            inputs=[main_lifestyle_prompt],
            outputs=[prompt_status]
        )
        
        # Quick example buttons in chat
        example_medical_btn.click(
            lambda history, session: send_example_message("I have a headache", history, session),
            inputs=[chatbot, session_data],
            outputs=[chatbot, status_box, session_data]
        )
        
        example_lifestyle_btn.click(
            lambda history, session: send_example_message("I want to start exercising", history, session),
            inputs=[chatbot, session_data],
            outputs=[chatbot, status_box, session_data]
        )
        
        example_help_btn.click(
            lambda history, session: send_example_message("Help - how do I use this application?", history, session),
            inputs=[chatbot, session_data],
            outputs=[chatbot, status_box, session_data]
        )
        
        # Instructions tab events
        refresh_instructions_btn.click(
            refresh_instructions,
            outputs=[instructions_display]
        )
        
        # Navigation from instructions to examples
        medical_example_btn.click(
            lambda: gr.update(selected="main_chat"),  # Switch to chat tab
            outputs=[]
        )
        
        lifestyle_example_btn.click(
            lambda: gr.update(selected="main_chat"),  # Switch to chat tab
            outputs=[]
        )
        
        testing_example_btn.click(
            lambda: gr.update(selected="testing_lab"),  # Switch to testing tab
            outputs=[]
        )
        
        prompts_example_btn.click(
            lambda: gr.update(selected="edit_prompts"),  # Switch to prompts tab
            outputs=[]
        )
        
        # Testing Lab handlers with session isolation
        load_patient_btn.click(
            handle_load_patient_isolated,
            inputs=[clinical_file, lifestyle_file, session_data],
            outputs=[load_result, patient_preview, chatbot, status_box, session_data]
        )
        
        quick_elderly_btn.click(
            lambda session: handle_quick_test_isolated("elderly", session),
            inputs=[session_data],
            outputs=[load_result, patient_preview, chatbot, status_box, session_data]
        )
        
        quick_athlete_btn.click(
            lambda session: handle_quick_test_isolated("athlete", session),
            inputs=[session_data],
            outputs=[load_result, patient_preview, chatbot, status_box, session_data]
        )
        
        quick_pregnant_btn.click(
            lambda session: handle_quick_test_isolated("pregnant", session),
            inputs=[session_data],
            outputs=[load_result, patient_preview, chatbot, status_box, session_data]
        )
        
        end_session_btn.click(
            handle_end_session_isolated,
            inputs=[end_session_notes, session_data],
            outputs=[end_session_result, session_data]
        )
        
        # Results handlers
        refresh_results_btn.click(
            handle_refresh_results_isolated,
            inputs=[session_data],
            outputs=[results_summary, results_table, session_data]
        )
        
        export_btn.click(
            handle_export_isolated,
            inputs=[session_data],
            outputs=[export_result, session_data]
        )
    
    return demo

# Create alias for backward compatibility
create_gradio_interface = create_session_isolated_interface

# Usage
if __name__ == "__main__":
    demo = create_session_isolated_interface()
    demo.launch()