web_app.py

#!/usr/bin/env python3
"""
Web Interface for Enhanced Research Assistant
Provides a modern, interactive web UI for the research assistant
"""

import asyncio
import json
import os
from datetime import datetime
from flask import Flask, render_template, request, jsonify, session
from flask_socketio import SocketIO, emit
import uuid
from enhanced_research_assistant import EnhancedResearchAssistant

app = Flask(__name__)
app.config['SECRET_KEY'] = 'research_assistant_secret_key'
socketio = SocketIO(app, cors_allowed_origins="*", async_mode='threading')

# Global research assistant instance
research_assistant = None
active_sessions = {}

class WebResearchSession:
    def __init__(self, session_id):
        self.session_id = session_id
        # Get OpenAI API key from environment
        openai_api_key = os.getenv('OPENAI_API_KEY')
        if not openai_api_key:
            raise ValueError("OPENAI_API_KEY environment variable is required")
        self.assistant = EnhancedResearchAssistant(openai_api_key)
        self.is_processing = False
        
    async def initialize(self):
        """Initialize the research assistant"""
        self.assistant.setup_agents()
        
    async def process_command(self, command, socketio_instance, session_id):
        """Process a command and emit updates via WebSocket"""
        self.is_processing = True
        
        try:
            # Emit status update
            socketio_instance.emit('status_update', {
                'status': 'processing',
                'message': f'Processing: {command[:50]}...'
            }, room=session_id)
            
            # Process the command
            if command.lower().startswith('add '):
                paper_input = command[4:].strip()
                result = await self.assistant.add_new_paper(paper_input)
            elif command.lower().startswith('search '):
                query = command[7:].strip()
                result = await self.assistant.search_related_papers(query)
            elif command.lower().startswith('analyze '):
                aspect = command[8:].strip()
                result = await self.assistant.analyze_aspect(aspect)
            elif command.lower() == 'summary':
                result = await self.assistant.get_research_summary()
            elif command.lower() == 'context':
                result = self.assistant.get_context_summary()
            elif command.lower().startswith('authors '):
                author = command[8:].strip()
                result = await self.assistant.investigate_authors(author)
            elif command.lower() == 'refine proposal':
                result = await self.assistant.refine_proposal()
            elif command.lower().startswith('explore '):
                direction = command[8:].strip()
                result = await self.assistant.explore_direction(direction)
            else:
                result = await self.assistant.general_query(command)
            
            # Emit the result
            socketio_instance.emit('command_result', {
                'command': command,
                'result': result,
                'timestamp': datetime.now().isoformat()
            }, room=session_id)
            
            # Update context summary
            context_summary = self.assistant.get_context_summary()
            socketio_instance.emit('context_update', {
                'summary': context_summary
            }, room=session_id)
            
        except Exception as e:
            socketio_instance.emit('error', {
                'message': f'Error processing command: {str(e)}'
            }, room=session_id)
        
        finally:
            self.is_processing = False
            socketio_instance.emit('status_update', {
                'status': 'ready',
                'message': 'Ready for next command'
            }, room=session_id)

@app.route('/')
def index():
    """Main page"""
    return render_template('index.html')

@app.route('/api/session/new', methods=['POST'])
def new_session():
    """Create a new research session"""
    session_id = str(uuid.uuid4())
    session['session_id'] = session_id
    
    # Create new research session
    web_session = WebResearchSession(session_id)
    active_sessions[session_id] = web_session
    
    return jsonify({
        'session_id': session_id,
        'status': 'created'
    })

@app.route('/api/context/load', methods=['POST'])
def load_context():
    """Load existing research context"""
    session_id = session.get('session_id')
    if not session_id or session_id not in active_sessions:
        return jsonify({'error': 'No active session'}), 400
    
    try:
        web_session = active_sessions[session_id]
        filename = request.json.get('filename', 'research_context.json')
        
        # Load context in the assistant
        if os.path.exists(filename):
            web_session.assistant.load_context(filename)
            context_summary = web_session.assistant.get_context_summary()
            
            return jsonify({
                'status': 'loaded',
                'context_summary': context_summary
            })
        else:
            return jsonify({'error': 'Context file not found'}), 404
            
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/api/context/save', methods=['POST'])
def save_context():
    """Save current research context"""
    session_id = session.get('session_id')
    if not session_id or session_id not in active_sessions:
        return jsonify({'error': 'No active session'}), 400
    
    try:
        web_session = active_sessions[session_id]
        filename = request.json.get('filename', 'research_context.json')
        
        web_session.assistant.save_context(filename)
        
        return jsonify({
            'status': 'saved',
            'filename': filename
        })
        
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@socketio.on('connect')
def handle_connect():
    """Handle WebSocket connection"""
    # Create a new session for this connection
    session_id = str(uuid.uuid4())
    session['session_id'] = session_id
    
    # Create new research session
    try:
        openai_api_key = os.getenv('OPENAI_API_KEY')
        if not openai_api_key:
            emit('error', {'message': 'OPENAI_API_KEY not configured'})
            return
            
        web_session = WebResearchSession(session_id)
        active_sessions[session_id] = web_session
        
        # Join the session room
        from flask_socketio import join_room
        join_room(session_id)
        
        emit('connected', {
            'session_id': session_id,
            'message': 'Connected to research assistant'
        })
        
        # Auto-initialize the session
        def init_sync():
            try:
                # Initialize the research assistant
                web_session.assistant.setup_agents()
                
                # Try to load existing context
                if os.path.exists('research_context.json'):
                    web_session.assistant.load_context('research_context.json')
                    context_summary = web_session.assistant.get_context_summary()
                    
                    socketio.emit('session_initialized', {
                        'status': 'ready',
                        'context_loaded': True,
                        'context_summary': context_summary
                    }, room=session_id)
                else:
                    socketio.emit('session_initialized', {
                        'status': 'ready',
                        'context_loaded': False,
                        'message': 'New research session started'
                    }, room=session_id)
                    
            except Exception as e:
                socketio.emit('error', {
                    'message': f'Failed to initialize session: {str(e)}'
                }, room=session_id)

        
        # Run initialization in background thread
        import threading
        threading.Thread(target=init_sync).start()
        
    except Exception as e:
        emit('error', {'message': f'Failed to create session: {str(e)}'})



@socketio.on('execute_command')
def handle_execute_command(data):
    """Execute a research command"""
    session_id = session.get('session_id')
    if not session_id or session_id not in active_sessions:
        emit('error', {'message': 'No active session'})
        return
    
    command = data.get('command', '').strip()
    if not command:
        emit('error', {'message': 'Empty command'})
        return
    
    web_session = active_sessions[session_id]
    
    if web_session.is_processing:
        emit('error', {'message': 'Another command is already being processed'})
        return
    
    def process_sync():
        try:
            # Run async command processing in new event loop
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            loop.run_until_complete(web_session.process_command(command, socketio, session_id))
        except Exception as e:
            socketio.emit('error', {
                'message': f'Error processing command: {str(e)}'
            }, room=session_id)
        finally:
            loop.close()
    
    # Process command in background thread
    import threading
    threading.Thread(target=process_sync).start()

@socketio.on('add_paper')
def handle_add_paper(data):
    """Handle adding a new paper"""
    session_id = session.get('session_id')
    if not session_id or session_id not in active_sessions:
        emit('error', {'message': 'No active session'})
        return
    
    paper_input = data.get('paper_input', '').strip()
    if not paper_input:
        emit('error', {'message': 'No paper input provided'})
        return
    
    # Execute add command
    command = f"add {paper_input}"
    handle_execute_command({'command': command})

@socketio.on('disconnect')
def handle_disconnect():
    """Handle WebSocket disconnection"""
    session_id = session.get('session_id')
    if session_id:
        from flask_socketio import leave_room
        leave_room(session_id)

def run_web_app(host='0.0.0.0', port=52609, debug=False):
    """Run the web application"""
    print(f"🌐 Starting Research Assistant Web Interface...")
    print(f"📍 Access at: http://localhost:{port}")
    print(f"🔧 Debug mode: {debug}")
    
    socketio.run(app, host=host, port=port, debug=debug, allow_unsafe_werkzeug=True)

if __name__ == '__main__':
    run_web_app(debug=True)