pages/ai_assistant.py

import gradio as gr
import datetime
from typing import Dict, List, Any, Union, Optional
import random
import os
import json
import numpy as np
from pathlib import Path

# Import utilities
from utils.storage import load_data, save_data
from utils.state import generate_id, get_timestamp, record_activity
from utils.ai_models import (
    generate_text, answer_question, analyze_image, transcribe_speech,
    translate_text, analyze_sentiment, summarize_text, generate_code
)
from utils.config import AI_MODELS, DATA_DIR
from utils.logging import get_logger
from utils.error_handling import handle_ai_model_exceptions, AIModelError

# Initialize logger
logger = get_logger(__name__)

# Define AI assistant types and their descriptions
AI_ASSISTANT_TYPES = {
    "General Chat": {
        "description": "Have natural conversations on any topic",
        "icon": "💬",
        "model": "microsoft/DialoGPT-medium",
        "task": "text_generation",
        "placeholder": "Chat with me about anything...",
        "examples": [
            "Tell me about the benefits of meditation",
            "What are some good productivity habits?",
            "Can you recommend some books on personal growth?"
        ]
    },
    "Task Assistant": {
        "description": "Get help with planning and organizing tasks",
        "icon": "📋",
        "model": "microsoft/DialoGPT-medium",
        "task": "text_generation",
        "placeholder": "Ask for help with your tasks and planning...",
        "examples": [
            "Help me break down this project into smaller tasks",
            "How can I prioritize my workload better?",
            "Create a schedule for my day"
        ]
    },
    "Writing Helper": {
        "description": "Assistance with writing and content creation",
        "icon": "✍️",
        "model": "microsoft/DialoGPT-medium",
        "task": "text_generation",
        "placeholder": "What would you like help writing?",
        "examples": [
            "Help me draft an email to my team about the project delay",
            "Give me ideas for a blog post about productivity",
            "Improve this paragraph: [your text here]"
        ]
    },
    "Code Assistant": {
        "description": "Get help with programming and coding",
        "icon": "💻",
        "model": "microsoft/CodeBERT-base",
        "task": "code_generation",
        "placeholder": "Describe what code you need help with...",
        "examples": [
            "Write a Python function to sort a list of dictionaries by a specific key",
            "How do I create a responsive navbar with CSS?",
            "Debug this code: [your code here]"
        ]
    },
    "Research Agent": {
        "description": "Help with gathering and organizing information",
        "icon": "🔎",
        "model": "distilbert-base-uncased-distilled-squad",
        "task": "question_answering",
        "placeholder": "What topic would you like to research?",
        "examples": [
            "Summarize the key points about climate change",
            "What are the main theories of motivation?",
            "Compare different project management methodologies"
        ]
    },
    "Learning Tutor": {
        "description": "Educational support and explanations",
        "icon": "🎓",
        "model": "microsoft/DialoGPT-medium",
        "task": "text_generation",
        "placeholder": "What would you like to learn about?",
        "examples": [
            "Explain quantum computing in simple terms",
            "Help me understand the concept of compound interest",
            "What are the key events of World War II?"
        ]
    },
    "Wellness Coach": {
        "description": "Guidance on health, fitness, and wellbeing",
        "icon": "🧘",
        "model": "microsoft/DialoGPT-medium",
        "task": "text_generation",
        "placeholder": "Ask about health, fitness, or wellbeing...",
        "examples": [
            "What are some good exercises for stress relief?",
            "Give me a simple meditation routine for beginners",
            "How can I improve my sleep quality?"
        ]
    }
}

@handle_ai_model_exceptions
def create_ai_assistant_page(state: Dict[str, Any]) -> None:
    """
    Create the AI Assistant Hub page with various AI assistants
    
    Args:
        state: Application state
    """
    logger.info("Creating AI Assistant Hub page")
    
    # Create the AI Assistant Hub layout
    with gr.Column(elem_id="ai-assistant-page"):
        gr.Markdown("# 🤖 AI Assistant Hub")
        
        # Assistant selector
        with gr.Row():
            assistant_selector = gr.Radio(
                choices=list(AI_ASSISTANT_TYPES.keys()),
                value=list(AI_ASSISTANT_TYPES.keys())[0],
                label="Select Assistant",
                elem_id="assistant-selector"
            )
        
        # Assistant description
        assistant_description = gr.Markdown(
            f"### {AI_ASSISTANT_TYPES[list(AI_ASSISTANT_TYPES.keys())[0]]['icon']} {list(AI_ASSISTANT_TYPES.keys())[0]}"
            f"\n{AI_ASSISTANT_TYPES[list(AI_ASSISTANT_TYPES.keys())[0]]['description']}"
        )
        
        # Chat interface
        with gr.Group(elem_id="chat-interface"):
            # Chat history
            chat_history = gr.Chatbot(
                elem_id="chat-history",
                height=400
            )
            
            # Input and send button
            with gr.Row():
                with gr.Column(scale=4):
                    chat_input = gr.Textbox(
                        placeholder=AI_ASSISTANT_TYPES[list(AI_ASSISTANT_TYPES.keys())[0]]['placeholder'],
                        label="",
                        elem_id="chat-input"
                    )
                with gr.Column(scale=1):
                    send_btn = gr.Button("Send", elem_id="send-btn")
            
            # Example queries
            with gr.Group(elem_id="example-queries"):
                gr.Markdown("### Example Queries")
                example_btns = []
                for example in AI_ASSISTANT_TYPES[list(AI_ASSISTANT_TYPES.keys())[0]]['examples']:
                    example_btns.append(gr.Button(example))
        
        # Function to update assistant description
        @handle_ai_model_exceptions
        def update_assistant_description(assistant_name):
            """
            Update the assistant description based on selection
            
            Args:
                assistant_name: Name of the selected assistant
                
            Returns:
                Updated description markdown
            """
            logger.debug(f"Updating assistant description for: {assistant_name}")
            assistant_info = AI_ASSISTANT_TYPES[assistant_name]
            
            # Update chat input placeholder
            chat_input.placeholder = assistant_info['placeholder']
            
            # Update example queries
            for i, example_btn in enumerate(example_btns):
                if i < len(assistant_info['examples']):
                    example_btn.value = assistant_info['examples'][i]
            
            return f"### {assistant_info['icon']} {assistant_name}\n{assistant_info['description']}"
        
        # Connect assistant selector to description update
        assistant_selector.change(
            update_assistant_description,
            inputs=[assistant_selector],
            outputs=[assistant_description]
        )
        
        # Function to handle chat messages
        @handle_ai_model_exceptions
        def chat_with_assistant(message, history, assistant_name):
            """
            Process chat messages and generate responses
            
            Args:
                message: User message
                history: Chat history
                assistant_name: Name of the selected assistant
                
            Returns:
                Updated chat history
            """
            if not message.strip():
                return history
            
            logger.info(f"Processing message for {assistant_name}: {message[:30]}...")
            
            # Get assistant info
            assistant_info = AI_ASSISTANT_TYPES[assistant_name]
            task = assistant_info['task']
            
            try:
                # Generate response based on assistant type
                if task == "text_generation":
                    # Prepare context from history
                    context = "\n".join([f"User: {h[0]}\nAssistant: {h[1]}" for h in history[-3:]])
                    context += f"\nUser: {message}\nAssistant:"
                    
                    response = generate_text(context)
                    
                elif task == "question_answering":
                    # For QA, we need a context, so we'll use the history as context
                    context = "\n".join([f"Q: {h[0]}\nA: {h[1]}" for h in history[-3:]])
                    response = answer_question(message, context)
                    
                elif task == "code_generation":
                    # For code generation, we'll use a specialized prompt
                    prompt = f"Generate code for: {message}"
                    response = generate_code(prompt)
                    
                else:
                    # Default to text generation
                    response = generate_text(message)
                
                # Record activity
                record_activity({
                    "type": "ai_assistant_used",
                    "assistant": assistant_name,
                    "message": message[:50] + ("..." if len(message) > 50 else ""),
                    "timestamp": datetime.datetime.now().isoformat()
                })
                
                # Update history
                history.append((message, response))
                return history
                
            except AIModelError as e:
                logger.error(f"AI model error: {str(e)}")
                return history + [(message, f"I'm sorry, I encountered an error: {e.message}")]
            except Exception as e:
                logger.error(f"Unexpected error in chat: {str(e)}")
                return history + [(message, "I'm sorry, I encountered an unexpected error. Please try again.")]
        
        # Connect send button to chat function
        send_btn.click(
            chat_with_assistant,
            inputs=[chat_input, chat_history, assistant_selector],
            outputs=[chat_history],
            clear_button=chat_input
        )
        
        # Connect chat input to chat function (for Enter key)
        chat_input.submit(
            chat_with_assistant,
            inputs=[chat_input, chat_history, assistant_selector],
            outputs=[chat_history],
            clear_button=chat_input
        )
        
        # Connect example buttons to chat input
        for example_btn in example_btns:
            example_btn.click(
                lambda example: example,
                inputs=[example_btn],
                outputs=[chat_input]
            )