ui/lifeadmin_coach_ui.py

"""
LifeAdmin Coach UI - AI Chatbot Assistant with Tool Access
Provides conversational interface with RAG context, memory, and tool calling
"""

import gradio as gr
import asyncio
from datetime import datetime
import json
import os


def create_lifeadmin_coach_ui(agent):
    """Create LifeAdmin Coach chatbot interface with tool access"""
    
    async def chat_with_coach(message, history):
        """Process chat message with RAG, memory, and tool execution"""
        if not message or not message.strip():
            return history, ""
        
        try:
            # Get relevant context from RAG
            rag_context = ""
            uploaded_files = []
            try:
                rag_results = await agent.rag_engine.search(message, k=3)
                if rag_results:
                    rag_context = "\n".join([
                        f"- {doc.get('metadata', {}).get('filename', 'Unknown')}: {doc.get('text', '')[:200]}..." 
                        for doc in rag_results
                    ])
                    # Get list of uploaded files
                    for doc in rag_results:
                        filename = doc.get('metadata', {}).get('filename')
                        if filename and filename not in uploaded_files:
                            uploaded_files.append(filename)
            except Exception:
                pass
            
            # Check for files in upload directory
            upload_dir = "data/uploads"
            if os.path.exists(upload_dir):
                try:
                    files_in_dir = [f for f in os.listdir(upload_dir) if os.path.isfile(os.path.join(upload_dir, f))]
                    for f in files_in_dir:
                        if f not in uploaded_files:
                            uploaded_files.append(f)
                except Exception:
                    pass
            
            # Check for calendar events
            calendar_files = []
            output_dir = "data/outputs"
            if os.path.exists(output_dir):
                try:
                    calendar_files = [f for f in os.listdir(output_dir) if f.endswith('.ics')]
                except Exception:
                    pass
            
            # Check for drafted emails
            email_files = []
            if os.path.exists(output_dir):
                try:
                    email_files = [f for f in os.listdir(output_dir) if 'email' in f.lower() and (f.endswith('.txt') or f.endswith('.html'))]
                except Exception:
                    pass
            
            # Get relevant memories
            memory_context = ""
            try:
                memory_context = agent.memory.get_relevant_memories(message, k=3)
            except Exception:
                pass
            
            # Detect if user is asking about specific tools/actions
            message_lower = message.lower()
            needs_tool = False
            tool_action = None
            
            # Detect tool requests
            if any(word in message_lower for word in ['summarize', 'summary', 'key points', 'main points']):
                if uploaded_files:
                    needs_tool = True
                    tool_action = "summarize"
            elif any(word in message_lower for word in ['extract text', 'read', 'show content']):
                if uploaded_files:
                    needs_tool = True
                    tool_action = "extract"
            elif any(word in message_lower for word in ['draft email', 'write email', 'compose email']):
                needs_tool = True
                tool_action = "draft_email"
            elif any(word in message_lower for word in ['calendar', 'events', 'schedule', 'meeting']):
                needs_tool = True
                tool_action = "calendar"
            elif any(word in message_lower for word in ['uploaded', 'files', 'documents']):
                tool_action = "list_files"
            
            # Build enhanced prompt
            system_context = f"""You are LifeAdmin Coach, a helpful AI assistant that helps users manage their life admin tasks.

You have access to:
- User's uploaded documents: {len(uploaded_files)} files
- Calendar events: {len(calendar_files)} events
- Drafted emails: {len(email_files)} drafts
- Past conversation history (via memory)
- Various automation tools

Your role:
- Answer questions about uploaded documents
- Provide information about calendar events and emails
- Provide advice on life admin tasks
- Help organize and plan tasks
- Be friendly, concise, and actionable

"""
            
            if uploaded_files:
                system_context += f"\n📂 **Uploaded Files:**\n" + "\n".join([f"  • {f}" for f in uploaded_files[:10]]) + "\n"
            
            if rag_context:
                system_context += f"\n📚 **Document Content:**\n{rag_context}\n"
            
            if calendar_files:
                system_context += f"\n📅 **Calendar Events:**\n" + "\n".join([f"  • {f}" for f in calendar_files[:5]]) + "\n"
            
            if email_files:
                system_context += f"\n✉️ **Drafted Emails:**\n" + "\n".join([f"  • {f}" for f in email_files[:5]]) + "\n"
            
            if memory_context:
                system_context += f"\n💭 **Past Context:**\n{memory_context}\n"
            
            # Execute tools if needed
            tool_result = ""
            if needs_tool and tool_action:
                if tool_action == "summarize" and uploaded_files:
                    # Use the first PDF file
                    pdf_file = next((f for f in uploaded_files if f.lower().endswith('.pdf')), uploaded_files[0])
                    try:
                        from tools import summarize_pdf
                        result = await summarize_pdf(f"data/uploads/{pdf_file}", max_length=300)
                        if isinstance(result, dict) and 'summary' in result:
                            tool_result = f"\n**Summary of {pdf_file}:**\n{result['summary']}\n"
                    except Exception as e:
                        tool_result = f"\nCouldn't summarize: {str(e)}\n"
                
                elif tool_action == "extract" and uploaded_files:
                    # Extract text from first file
                    file = uploaded_files[0]
                    try:
                        if file.lower().endswith('.pdf'):
                            from utils.pdf_utils import extract_text_from_pdf
                            text = extract_text_from_pdf(f"data/uploads/{file}")
                            tool_result = f"\n**Text from {file}:**\n{text[:500]}...\n"
                        elif file.lower().endswith(('.png', '.jpg', '.jpeg')):
                            from tools import extract_text_ocr
                            result = await extract_text_ocr(f"data/uploads/{file}", 'en')
                            tool_result = f"\n**Text from {file}:**\n{result.get('text', '')[:500]}...\n"
                    except Exception as e:
                        tool_result = f"\nCouldn't extract text: {str(e)}\n"
                
                elif tool_action == "calendar":
                    if calendar_files:
                        tool_result = f"\n**Calendar Events:**\n" + "\n".join([f"  • {f}" for f in calendar_files]) + "\n"
                    else:
                        tool_result = "\nNo calendar events created yet. You can create events in the Manual Dashboard.\n"
            
            # Build full prompt
            full_prompt = f"""{system_context}

{tool_result}

**User Question:** {message}

Provide a helpful, concise response. Reference the documents and information available above.
"""
            
            # Get LLM response
            from utils.llm_utils import get_llm_response
            response = await get_llm_response(full_prompt, temperature=0.7, max_tokens=1000)
            
            # Save to memory
            try:
                agent.memory.add_memory(
                    content=f"User: {message}\nCoach: {response}",
                    memory_type='conversation',
                    metadata={'timestamp': datetime.now().isoformat()},
                    importance=5
                )
            except Exception:
                pass
            
            # Update history
            history.append({"role": "user", "content": message})
            history.append({"role": "assistant", "content": response})
            
            return history, ""
            
        except Exception as e:
            import traceback
            error_msg = f"⚠️ Error: {str(e)}\n\n{traceback.format_exc()}"
            print(error_msg)
            history.append({"role": "user", "content": message})
            history.append({"role": "assistant", "content": f"⚠️ Sorry, I encountered an error: {str(e)}"})
            return history, ""
    
    async def quick_action(action_type, history):
        """Handle quick action buttons"""
        actions = {
            "summarize_docs": "Can you summarize all the documents I've uploaded?",
            "list_tasks": "What tasks or action items can you find in my documents?",
            "deadlines": "Are there any deadlines or important dates I should know about?",
            "help": "What can you help me with?"
        }
        
        message = actions.get(action_type, "")
        if message:
            return await chat_with_coach(message, history)
        return history, ""
    
    def clear_chat():
        """Clear chat history"""
        return [], ""
    
    # Build UI
    with gr.Column():
        gr.Markdown("""
        ## 🤖 LifeAdmin Coach
        
        Your AI assistant for managing life admin tasks. Ask questions about your documents,
        get advice on tasks, or chat about anything related to organizing your life!
        
        **I can help you with:**
        - Answering questions about uploaded documents (from Manual Dashboard)
        - Summarizing PDFs and extracting key information
        - Finding deadlines and important dates
        - Checking your calendar events and drafted emails
        - Organizing tasks and creating plans
        - General life admin advice
        
        **💡 Tip:** Upload files in the Manual Dashboard first, then ask me questions here!
        """)
        
        # Chat interface
        chatbot = gr.Chatbot(
            label="💬 Chat with LifeAdmin Coach",
            height=500
        )
        
        # Input area
        with gr.Row():
            msg_input = gr.Textbox(
                label="",
                placeholder="Ask me anything about your documents or life admin tasks...",
                scale=4,
                lines=2
            )
            send_btn = gr.Button("Send", variant="primary", scale=1)
        
        # Quick action buttons
        with gr.Row():
            gr.Markdown("**Quick Actions:**")
        
        with gr.Row():
            summarize_btn = gr.Button("📄 Summarize Docs", size="sm")
            tasks_btn = gr.Button("✅ List Tasks", size="sm")
            deadlines_btn = gr.Button("📅 Find Deadlines", size="sm")
            help_btn = gr.Button("❓ Help", size="sm")
            clear_btn = gr.Button("🗑️ Clear Chat", size="sm", variant="stop")
        
        # Examples
        with gr.Accordion("💡 Example Questions", open=False):
            gr.Examples(
                examples=[
                    "What documents have I uploaded?",
                    "Summarize my PDF document",
                    "What calendar events do I have?",
                    "Have I drafted any emails?",
                    "Extract text from my uploaded image",
                    "What are the key points from my documents?",
                    "Find all phone numbers and emails in my documents",
                    "When is my next deadline?",
                ],
                inputs=msg_input
            )
        
        # Event handlers
        def chat_wrapper(message, history):
            """Sync wrapper for async chat function"""
            return asyncio.run(chat_with_coach(message, history))
        
        def quick_action_wrapper(action_type, history):
            """Sync wrapper for quick actions"""
            return asyncio.run(quick_action(action_type, history))
        
        # Connect events
        msg_input.submit(
            fn=chat_wrapper,
            inputs=[msg_input, chatbot],
            outputs=[chatbot, msg_input]
        )
        
        send_btn.click(
            fn=chat_wrapper,
            inputs=[msg_input, chatbot],
            outputs=[chatbot, msg_input]
        )
        
        summarize_btn.click(
            fn=lambda h: quick_action_wrapper("summarize_docs", h),
            inputs=[chatbot],
            outputs=[chatbot, msg_input]
        )
        
        tasks_btn.click(
            fn=lambda h: quick_action_wrapper("list_tasks", h),
            inputs=[chatbot],
            outputs=[chatbot, msg_input]
        )
        
        deadlines_btn.click(
            fn=lambda h: quick_action_wrapper("deadlines", h),
            inputs=[chatbot],
            outputs=[chatbot, msg_input]
        )
        
        help_btn.click(
            fn=lambda h: quick_action_wrapper("help", h),
            inputs=[chatbot],
            outputs=[chatbot, msg_input]
        )
        
        clear_btn.click(
            fn=clear_chat,
            outputs=[chatbot, msg_input]
        )
    
    return gr.Column()