gradio_app.py

#!/usr/bin/env python3
"""
Final Enhanced Gradio Web UI with Auto-Loading User Conversations
"""

import os
import gradio as gr
from dotenv import load_dotenv
from graph import ArchitectureAssistant
from user_state_manager import user_state_manager
import json
import hashlib
from datetime import datetime


class GradioArchitectureApp:
    def __init__(self):
        load_dotenv()
        api_key = os.getenv("OPENAI_API_KEY")
        if not api_key:
            raise ValueError("Please set OPENAI_API_KEY in .env file")
        
        self.api_key = api_key
        self.assistant = None
        self.current_user_id = None
        self.conversation_history = []
        self.state_manager = user_state_manager
    
    def initialize_user_session(self, user_identifier: str = None):
        """Initialize user session with persistent state"""
        if not user_identifier:
            # Generate user ID from timestamp for demo
            user_identifier = hashlib.md5(str(datetime.now()).encode()).hexdigest()[:12]
        
        self.current_user_id = user_identifier
        self.assistant = ArchitectureAssistant(self.api_key, user_identifier)
        
        return user_identifier
    
    def change_user_id(self, new_user_id: str):
        """Change user ID and auto-load their conversation history"""
        if not new_user_id.strip():
            return [], "❌ Please enter a valid User ID", "", "Start chatting to see your requirements here!"
        
        try:
            # Set new user ID
            self.current_user_id = new_user_id.strip()
            
            # Create new assistant instance for this user
            self.assistant = ArchitectureAssistant(self.api_key, self.current_user_id)
            
            # Try to load their latest conversation
            history = self.state_manager.get_user_history(self.current_user_id)
            
            if history:
                # Load most recent conversation
                latest_session = history[0]["session_id"]
                success = self.assistant.load_previous_state(latest_session)
                
                if success:
                    # Convert state messages to chat history
                    chat_history = []
                    messages = self.assistant.state.get("messages", [])
                    
                    current_pair = [None, None]
                    for msg in messages:
                        if msg["role"] == "user":
                            current_pair = [msg["content"], None]
                        elif msg["role"] == "assistant":
                            if current_pair[0] is not None:
                                current_pair[1] = msg["content"]
                                chat_history.append(current_pair)
                                current_pair = [None, None]
                    
                    status_msg = f"✅ Loaded {len(history)} previous conversation(s) for user {new_user_id}"
                    summary = self.get_conversation_summary()
                    
                    return chat_history, status_msg, new_user_id, summary
                else:
                    status_msg = f"✅ User {new_user_id} found but couldn't load conversation. Starting fresh."
                    summary = self.get_conversation_summary()
                    return [], status_msg, new_user_id, summary
            else:
                # New user
                status_msg = f"✅ New user {new_user_id} - starting fresh conversation"
                summary = self.get_conversation_summary()
                return [], status_msg, new_user_id, summary
        
        except Exception as e:
            return [], f"❌ Error loading user {new_user_id}: {str(e)}", "", "Error loading user"
    
    def chat_with_assistant(self, message, history, user_id_input=""):
        """Process user message and return response with state saving"""
        if not message.strip():
            return history, "", user_id_input, self.get_conversation_summary()
        
        # Ensure we have an assistant instance
        if not self.assistant:
            if user_id_input.strip():
                self.current_user_id = user_id_input.strip()
                self.assistant = ArchitectureAssistant(self.api_key, self.current_user_id)
            else:
                user_id = self.initialize_user_session()
                return history, "", user_id, self.get_conversation_summary()
        
        # Handle special commands
        if message.lower() == 'reset':
            self.assistant.reset_conversation()
            self.conversation_history = []
            status_msg = "🔄 Conversation reset! New session started."
            return [[status_msg, ""]], "", self.current_user_id or "", self.get_conversation_summary()
        
        # Get response from assistant
        try:
            response = self.assistant.chat(message)
            
            # Add response to history
            history.append([message, response])
            
            # Add to conversation history
            self.conversation_history.append({"user": message, "assistant": response})
            
            # Get updated summary
            summary = self.get_conversation_summary()
            
            return history, "", self.current_user_id or "", summary
            
        except Exception as e:
            error_msg = f"❌ Error: {str(e)}"
            history.append([message, error_msg])
            return history, "", self.current_user_id or "", self.get_conversation_summary()
    
    def get_conversation_summary(self):
        """Get formatted conversation summary"""
        if not self.assistant:
            return "Enter a User ID above to start or resume a conversation!"
        
        summary = self.assistant.get_conversation_summary()
        
        summary_text = f"📋 **CONVERSATION SUMMARY**\n"
        summary_text += f"👤 **User ID:** {self.current_user_id or 'Not set'}\n\n"
        
        # User requirements
        reqs = summary["user_requirements"]
        summary_text += "**USER REQUIREMENTS:**\n"
        if reqs["budget"]:
            summary_text += f"• Budget: ${reqs['budget']:,.0f}\n"
        if reqs["location"]:
            summary_text += f"• Location: {reqs['location']}\n"
        if reqs["family_size"]:
            summary_text += f"• Family size: {reqs['family_size']}\n"
        if reqs["lifestyle_preferences"]:
            summary_text += f"• Preferences: {', '.join(reqs['lifestyle_preferences'])}\n"
        
        # Floorplan requirements
        floor_reqs = summary["floorplan_requirements"]
        summary_text += "\n**FLOORPLAN REQUIREMENTS:**\n"
        if floor_reqs["num_floors"]:
            summary_text += f"• Floors: {floor_reqs['num_floors']}\n"
        if floor_reqs["total_sqft"]:
            summary_text += f"• Total sq ft: {floor_reqs['total_sqft']}\n"
        if floor_reqs["lot_shape"]:
            summary_text += f"• Lot shape: {floor_reqs['lot_shape']}\n"
        if floor_reqs["lot_dimensions"]:
            summary_text += f"• Lot dimensions: {floor_reqs['lot_dimensions']}\n"
        if floor_reqs["rooms"]:
            rooms_str = ", ".join([f"{r['count']}x {r['type']}" for r in floor_reqs["rooms"]])
            summary_text += f"• Rooms: {rooms_str}\n"
        
        # Project progress
        agent_memory = self.assistant.state.get("agent_memory", {})
        completed_phases = []
        if self.assistant.state.get("detailed_floorplan", {}).get("detailed_rooms"):
            completed_phases.append("Architectural Design")
        if self.assistant.state.get("budget_breakdown", {}).get("total_construction_cost"):
            completed_phases.append("Budget Analysis")
        if agent_memory.get("structural_analysis"):
            completed_phases.append("Structural Analysis")
        if agent_memory.get("sustainability"):
            completed_phases.append("Sustainability Review")
        if agent_memory.get("permits"):
            completed_phases.append("Permit Planning")
        if agent_memory.get("interior_design"):
            completed_phases.append("Interior Design")
        
        summary_text += f"\n**PROJECT PROGRESS:**\n"
        summary_text += f"• Completed: {len(completed_phases)}/6 phases\n"
        if completed_phases:
            summary_text += f"• Phases: {', '.join(completed_phases)}\n"
        summary_text += f"• Total messages: {summary['total_messages']}\n"
        
        return summary_text
    
    def get_user_history_display(self, user_id: str):
        """Get formatted user history for display"""
        if not user_id.strip():
            return "Please enter a User ID to view history."
        
        try:
            history = self.state_manager.get_user_history(user_id.strip())
            
            if not history:
                return f"No conversation history found for user: {user_id}"
            
            display_text = f"📚 **CONVERSATION HISTORY FOR USER: {user_id}**\n\n"
            
            for i, conv in enumerate(history, 1):
                timestamp = conv["timestamp"]
                session_id = conv["session_id"]
                summary = conv.get("summary", {})
                
                display_text += f"**#{i} - {timestamp}**\n"
                display_text += f"Session ID: `{session_id}`\n"
                
                if summary.get("total_messages", 0) > 0:
                    display_text += f"Messages: {summary['total_messages']}\n"
                
                if summary.get("user_requirements"):
                    reqs = summary["user_requirements"]
                    if reqs.get("budget"):
                        display_text += f"Budget: {reqs['budget']}\n"
                    if reqs.get("location"):
                        display_text += f"Location: {reqs['location']}\n"
                    if reqs.get("family_size"):
                        display_text += f"Family: {reqs['family_size']} people\n"
                
                if summary.get("floorplan_status"):
                    floor_status = summary["floorplan_status"]
                    if floor_status.get("size"):
                        display_text += f"House Size: {floor_status['size']}\n"
                    if floor_status.get("rooms"):
                        display_text += f"Rooms: {floor_status['rooms']}\n"
                
                if summary.get("project_progress"):
                    progress = summary["project_progress"]
                    completed = progress.get("completed_phases", [])
                    percentage = progress.get("completion_percentage", 0)
                    display_text += f"Progress: {percentage}% - {', '.join(completed)}\n"
                
                display_text += "\n---\n\n"
            
            return display_text
            
        except Exception as e:
            return f"❌ Error retrieving history: {str(e)}"
    
    def get_all_users_summary(self):
        """Get summary of all users"""
        try:
            users = self.state_manager.get_all_users()
            
            if not users:
                return "No users found in the system."
            
            display_text = f"👥 **ALL USERS SUMMARY** ({len(users)} users)\n\n"
            
            for i, user in enumerate(users, 1):
                user_id = user["user_id"]
                total_conversations = user["total_conversations"]
                last_activity = user["last_activity"]
                latest_summary = user.get("latest_summary", {})
                
                display_text += f"**#{i} User: {user_id}**\n"
                display_text += f"Conversations: {total_conversations}\n"
                display_text += f"Last Activity: {last_activity}\n"
                
                if latest_summary.get("project_progress"):
                    progress = latest_summary["project_progress"]
                    percentage = progress.get("completion_percentage", 0)
                    display_text += f"Latest Progress: {percentage}%\n"
                
                display_text += "\n---\n\n"
            
            return display_text
            
        except Exception as e:
            return f"❌ Error retrieving users summary: {str(e)}"


def create_gradio_interface():
    """Create and return the final enhanced Gradio interface"""
    app = GradioArchitectureApp()
    
    # Custom CSS for better styling
    css = """
    .gradio-container {
        max-width: 1400px !important;
    }
    .chat-message {
        border-radius: 10px !important;
    }
    .user-id-input {
        background-color: #e3f2fd !important;
        font-weight: bold !important;
    }
    """
    
    with gr.Blocks(
        title="🏠 Architecture Assistant - Multi-User",
        theme=gr.themes.Soft(),
        css=css
    ) as interface:
        
        gr.HTML("""
        <div style="text-align: center; margin-bottom: 20px;">
            <h1>🏠 Residential Architecture Assistant</h1>
            <p>Multi-user system with automatic conversation loading and precise floorplan specifications</p>
        </div>
        """)
        
        # User ID input at the top
        with gr.Row():
            with gr.Column(scale=4):
                user_id_input = gr.Textbox(
                    placeholder="Enter your User ID (e.g., 'john_smith_house') - automatically loads your conversation history",
                    label="👤 User ID",
                    elem_classes=["user-id-input"],
                    interactive=True
                )
            with gr.Column(scale=1):
                load_user_btn = gr.Button("Load User", variant="primary")
        
        status_display = gr.Markdown(
            label="📊 Status",
            value="Enter a User ID above to start or resume your architectural consultation."
        )
        
        # Main interface
        with gr.Row():
            with gr.Column(scale=2):
                chatbot = gr.Chatbot(
                    label="💬 Chat with your Architecture Assistant",
                    height=500,
                    show_copy_button=True,
                    bubble_full_width=False
                )
                
                with gr.Row():
                    msg = gr.Textbox(
                        placeholder="Ask about home design, budget, or floorplan planning...",
                        container=False,
                        scale=4,
                        label="Your message"
                    )
                    send_btn = gr.Button("Send", variant="primary", scale=1)
                
                with gr.Row():
                    clear_btn = gr.Button("🔄 Reset Conversation", variant="secondary")
                    
                gr.HTML("""
                <div style="margin-top: 15px; padding: 10px; background-color: #f0f0f0; border-radius: 5px;">
                    <h4>💡 How to use:</h4>
                    <ul>
                        <li><strong>New User:</strong> Enter any User ID above and start chatting</li>
                        <li><strong>Returning User:</strong> Enter your previous User ID - conversation automatically loads</li>
                        <li><strong>Multiple Users:</strong> Each User ID has completely separate conversations</li>
                        <li><strong>Floorplans:</strong> Get precise room dimensions for drawing actual plans</li>
                    </ul>
                </div>
                """)
            
            with gr.Column(scale=1):
                summary_display = gr.Markdown(
                    label="📋 Current Project Summary",
                    value="Enter a User ID above to start or resume a conversation!"
                )
                
                gr.HTML("""
                <div style="margin-top: 20px; padding: 10px; background-color: #e8f4f8; border-radius: 5px;">
                    <h4>🤖 Available Core Specialists:</h4>
                    <p><strong>🧭 Router:</strong> Intelligent conversation routing</p>
                    <p><strong>🏛️ General Design:</strong> Architecture principles & design guidance</p>
                    <p><strong>💰 Budget Analysis:</strong> Montreal market costs & feasibility</p>
                    <p><strong>📐 Floorplan:</strong> Layout planning with exact dimensions</p>
                    <p><strong>📋 Regulation:</strong> Montreal building codes & permit requirements</p>
                    <br>
                </div>
                """)
        
        # Example prompts
        with gr.Row():
            gr.HTML("<h3>🚀 Try these examples:</h3>")
        
        with gr.Row():
            example1 = gr.Button("What are key home design principles?", size="sm")
            example2 = gr.Button("I have $800k budget for Montreal - realistic?", size="sm") 
            example3 = gr.Button("Need 3BR/2BA house, 2500 sqft - help floorplan design", size="sm")
        
        with gr.Row():
            example4 = gr.Button("Show me a detailed budget breakdown", size="sm")
            example5 = gr.Button("What permits do I need for my house design in Montreal?", size="sm")
            example6 = gr.Button("Help me with Montreal building codes", size="sm")
        
        # History tab
        with gr.Accordion("📚 User History & Management", open=False):
            with gr.Row():
                with gr.Column():
                    gr.HTML("<h4>📋 View User History</h4>")
                    
                    with gr.Row():
                        history_user_id = gr.Textbox(
                            placeholder="Enter User ID to view history",
                            label="User ID for History",
                            scale=3
                        )
                        view_history_btn = gr.Button("View History", variant="primary", scale=1)
                    
                    user_history_display = gr.Markdown(
                        label="User History",
                        value="Enter a User ID to view their conversation history."
                    )
                
                with gr.Column():
                    gr.HTML("<h4>👥 All Users Overview</h4>")
                    
                    view_all_btn = gr.Button("View All Users", variant="secondary")
                    
                    all_users_display = gr.Markdown(
                        label="All Users Summary",
                        value="Click 'View All Users' to see system overview."
                    )
        
        # Event handlers
        def send_message(message, history, user_id):
            return app.chat_with_assistant(message, history, user_id)
        
        def change_user(user_id):
            return app.change_user_id(user_id)
        
        def view_user_history(user_id):
            return app.get_user_history_display(user_id)
        
        def view_all_users():
            return app.get_all_users_summary()
        
        # Wire up the events
        load_user_btn.click(change_user, [user_id_input], [chatbot, status_display, user_id_input, summary_display])
        
        msg.submit(send_message, [msg, chatbot, user_id_input], [chatbot, msg, user_id_input, summary_display])
        send_btn.click(send_message, [msg, chatbot, user_id_input], [chatbot, msg, user_id_input, summary_display])
        
        clear_btn.click(lambda: app.assistant.reset_conversation() if app.assistant else None, outputs=[chatbot, summary_display])
        
        # History management events
        view_history_btn.click(view_user_history, [history_user_id], [user_history_display])
        view_all_btn.click(view_all_users, outputs=[all_users_display])
        
        # Example button handlers
        example1.click(lambda: "What are key home design principles?", outputs=msg)
        example2.click(lambda: "I have $800k budget for Montreal - is that realistic?", outputs=msg)
        example3.click(lambda: "I need a 3 bedroom, 2 bathroom house with about 2500 square feet. Can you help me plan the layout?", outputs=msg)
        example4.click(lambda: "Can you show me a detailed budget breakdown for my house design?", outputs=msg)
        example5.click(lambda: "What permits do I need for my house design in Montreal?", outputs=msg)
        example6.click(lambda: "Help me understand Montreal building codes and regulations", outputs=msg)
    
    return interface


def main():
    """Launch the final enhanced Gradio interface"""
    try:
        interface = create_gradio_interface()
        print("🏠 Starting Final Architecture Assistant Web UI...")
        print("👥 Features: Multi-user support, auto-loading conversations, precise floorplans")
        print("📝 Make sure you have set OPENAI_API_KEY in your .env file")
        print("🔧 Fixed issues: User conversation resumption, exact floorplan dimensions, SVG generation")
        
        interface.launch(
            server_name="0.0.0.0",
            server_port=7862,
            share=False,
            show_error=True,
            quiet=False
        )
        
    except ValueError as e:
        print(f"❌ Configuration Error: {e}")
        print("Please copy .env.example to .env and add your OpenAI API key")
    except Exception as e:
        print(f"❌ Error starting the application: {e}")


if __name__ == "__main__":
    main()