main.py

"""
Main integration file for HF Spaces deployment
Wires together UI, agents, and orchestrator
"""

import gradio as gr
import logging
import os
from typing import Dict, Any

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Import components
try:
    from app import create_mobile_optimized_interface
    from src.agents.intent_agent import create_intent_agent
    from src.agents.synthesis_agent import create_synthesis_agent
    from src.agents.safety_agent import create_safety_agent
    from src.config import settings
    from src.llm_router import LLMRouter
    from src.orchestrator_engine import MVPOrchestrator
    from src.context_manager import EfficientContextManager
    from src.mobile_handlers import MobileUXHandlers
except ImportError as e:
    logger.warning(f"Some components not available: {e}")
    # Create mock components for basic functionality
    class MockComponent:
        async def execute(self, *args, **kwargs):
            return {"result": "Mock response", "status": "mock"}
    
    # Mock imports for deployment
    create_intent_agent = lambda x: MockComponent()
    create_synthesis_agent = lambda x: MockComponent()
    create_safety_agent = lambda x: MockComponent()
    settings = type('Settings', (), {'hf_token': os.getenv('HF_TOKEN', '')})()
    LLMRouter = type('MockRouter', (), {'__init__': lambda self, x: None})
    MVPOrchestrator = type('MockOrchestrator', (), {'__init__': lambda self, x, y, z: None})
    EfficientContextManager = type('MockContextManager', (), {'__init__': lambda self: None})
    MobileUXHandlers = type('MockHandlers', (), {'__init__': lambda self, x: None})

def initialize_components():
    """Initialize all system components with error handling"""
    components = {}
    
    try:
        # Initialize LLM Router
        logger.info("Initializing LLM Router...")
        llm_router = LLMRouter(settings.hf_token)
        components['llm_router'] = llm_router
        
        # Initialize Agents
        logger.info("Initializing Agents...")
        agents = {
            'intent_recognition': create_intent_agent(llm_router),
            'response_synthesis': create_synthesis_agent(llm_router),
            'safety_check': create_safety_agent(llm_router)
        }
        components['agents'] = agents
        
        # Initialize Context Manager
        logger.info("Initializing Context Manager...")
        context_manager = EfficientContextManager()
        components['context_manager'] = context_manager
        
        # Initialize Orchestrator
        logger.info("Initializing Orchestrator...")
        orchestrator = MVPOrchestrator(llm_router, context_manager, agents)
        components['orchestrator'] = orchestrator
        
        # Initialize Mobile Handlers
        logger.info("Initializing Mobile Handlers...")
        mobile_handlers = MobileUXHandlers(orchestrator)
        components['mobile_handlers'] = mobile_handlers
        
        logger.info("All components initialized successfully")
        
    except Exception as e:
        logger.error(f"Component initialization failed: {e}")
        logger.info("Falling back to mock mode for basic functionality")
        components['mock_mode'] = True
    
    return components

def create_event_handlers(demo, components):
    """Connect UI events to backend handlers"""
    
    def get_response_handler(message, chat_history, session_id, show_reasoning, show_agent_trace, request):
        """Handle user messages with proper error handling"""
        try:
            if components.get('mock_mode'):
                # Mock response for basic functionality
                response = f"Mock response to: {message}"
                chat_history.append((message, response))
                return "", chat_history, {}, {}
            
            # Use mobile handlers if available
            if 'mobile_handlers' in components:
                # This would be the real implementation
                result = components['mobile_handlers'].handle_mobile_submit(
                    message, chat_history, session_id, show_reasoning, show_agent_trace, request
                )
                return result
            else:
                # Fallback mock response
                response = f"System response to: {message}"
                chat_history.append((message, response))
                return "", chat_history, {"status": "processed"}, {"response_time": 0.5}
                
        except Exception as e:
            logger.error(f"Error handling message: {e}")
            # Graceful error response
            error_response = "I apologize, but I'm experiencing technical difficulties. Please try again."
            chat_history.append((message, error_response))
            return "", chat_history, {"error": str(e)}, {"status": "error"}
    
    return get_response_handler

def setup_application():
    """Setup and return the Gradio application"""
    logger.info("Starting application setup...")
    
    # Initialize components
    components = initialize_components()
    
    # Create the interface
    logger.info("Creating mobile-optimized interface...")
    demo = create_mobile_optimized_interface()
    
    # Setup event handlers
    logger.info("Setting up event handlers...")
    
    # For now, use a simple chat interface until full integration is ready
    try:
        # Get the chat function from the demo
        chat_interface = demo.get_blocks().children[0]  # Get first component
        
        # Simple message handling for MVP
        def simple_chat_fn(message, history):
            if components.get('mock_mode'):
                return f"I'm running in mock mode. You said: {message}"
            else:
                return f"System is processing: {message}"
        
        # Set the chat function
        if hasattr(chat_interface, 'chat_fn'):
            chat_interface.chat_fn = simple_chat_fn
        
    except Exception as e:
        logger.warning(f"Could not setup advanced handlers: {e}")
    
    logger.info("Application setup completed")
    return demo

def main():
    """Main entry point for HF Spaces"""
    logger.info("🚀 Starting AI Research Assistant MVP")
    
    # Check for HF Token
    hf_token = os.getenv('HF_TOKEN')
    if not hf_token:
        logger.warning("HF_TOKEN not found in environment. Some features may be limited.")
    
    # Create and launch application
    demo = setup_application()
    
    # Launch configuration for HF Spaces
    launch_config = {
        'server_name': '0.0.0.0',
        'server_port': 7860,
        'share': False,
        'debug': False
    }
    
    logger.info("✅ Application ready for launch")
    return demo.launch(**launch_config)

if __name__ == "__main__":
    main()