app.py

# MoA Chatbot v2.1.0 - Flask Application with SSE Streaming
# Fixed for Flask-Uploads compatibility with newer Werkzeug versions
import werkzeug
if not hasattr(werkzeug, 'secure_filename'):
    from werkzeug.utils import secure_filename as _secure_filename
    werkzeug.secure_filename = _secure_filename
if not hasattr(werkzeug, 'FileStorage'):
    from werkzeug.datastructures import FileStorage
    werkzeug.FileStorage = FileStorage

from flask import Flask, request, render_template, redirect, url_for, flash, jsonify, Response, stream_with_context
from flask_sqlalchemy import SQLAlchemy
from flask_uploads import UploadSet, configure_uploads, DOCUMENTS
from werkzeug.utils import secure_filename
from datetime import datetime
from flask_migrate import Migrate
import os
import json
import traceback
import logging
import re
import queue
import threading
from logging.handlers import RotatingFileHandler
from typing import Generator, Optional

# Ensure directories exist
for directory in ['uploads', 'instance', 'logs']:
    if not os.path.exists(directory):
        os.makedirs(directory)

from utils import generate_with_references, generate_together_stream
import config
from update_checker import UpdateChecker, check_updates, update_application

app = Flask(__name__)

# ── Context Processor for Branding ─────────────────────────────────────
@app.context_processor
def inject_branding():
    return {
        'required_developer_name': config.REQUIRED_DEVELOPER_NAME,
        'required_company_name': config.REQUIRED_COMPANY_NAME,
        'required_company_url': config.REQUIRED_COMPANY_URL,
        'required_github_username': config.REQUIRED_GITHUB_USERNAME,
        'required_github_repo': config.REQUIRED_GITHUB_REPO,
        'required_company_logo': config.COMPANY_LOGO,
        'custom_developer_name': config.CUSTOM_DEVELOPER_NAME,
        'custom_company_name': config.CUSTOM_COMPANY_NAME,
        'custom_company_url': config.CUSTOM_COMPANY_URL,
        'custom_company_logo': config.CUSTOM_COMPANY_LOGO,
        'developer_name': config.DEVELOPER_NAME,
        'company_name': config.COMPANY_NAME,
        'company_url': config.COMPANY_URL,
        'github_username': config.GITHUB_USERNAME,
        'github_repo': config.GITHUB_REPO,
        'company_logo': config.COMPANY_LOGO,
        'app_name': config.APP_NAME,
        'app_description': config.APP_DESCRIPTION,
        'app_version': config.APP_VERSION,
    }

app.config['SECRET_KEY'] = os.environ.get('SECRET_KEY', os.urandom(24).hex())

# SQLite database path
instance_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'instance')
if not os.path.exists(instance_path):
    os.makedirs(instance_path)
db_path = os.path.join(instance_path, 'conversations.db')
db_path_uri = db_path.replace('\\', '/')
app.config['SQLALCHEMY_DATABASE_URI'] = f'sqlite:///{db_path_uri}'
app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = False
app.config['UPLOADED_FILES_DEST'] = 'uploads'
app.config['UPLOADED_FILES_ALLOW'] = DOCUMENTS
app.config['MAX_CONTENT_LENGTH'] = 32 * 1024 * 1024  # 32MB max file size

db = SQLAlchemy(app)
migrate = Migrate(app, db)

files = UploadSet('files', DOCUMENTS)
configure_uploads(app, files)

# ── Logging Configuration ──────────────────────────────────────────────
if not app.debug:
    log_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'logs')
    if not os.path.exists(log_dir):
        os.makedirs(log_dir)
    
    file_handler = RotatingFileHandler(
        os.path.join(log_dir, 'moa_app.log'),
        maxBytes=10240000,
        backupCount=10
    )
    file_handler.setFormatter(logging.Formatter(
        '%(asctime)s %(levelname)s: %(message)s [in %(pathname)s:%(lineno)d]'
    ))
    file_handler.setLevel(logging.INFO)
    app.logger.addHandler(file_handler)
    app.logger.setLevel(logging.INFO)
    app.logger.info('MoA Chatbot v2.1.0 startup')


# ── Database Models ─────────────────────────────────────────────────────
class Conversation(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    topic = db.Column(db.String(200), nullable=False)
    messages = db.Column(db.Text, nullable=False, default='[]')
    created_at = db.Column(db.DateTime, default=datetime.utcnow)
    updated_at = db.Column(db.DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
    
    def __repr__(self):
        return f'<Conversation {self.id}: {self.topic}>'

    def to_dict(self):
        return {
            'id': self.id,
            'topic': self.topic,
            'created_at': self.created_at.isoformat() if self.created_at else None,
            'updated_at': self.updated_at.isoformat() if self.updated_at else None,
            'message_count': len(json.loads(self.messages)) if self.messages else 0
        }


class Document(db.Model):
    """Store uploaded documents for RAG"""
    id = db.Column(db.Integer, primary_key=True)
    conversation_id = db.Column(db.Integer, db.ForeignKey('conversation.id'), nullable=True)
    filename = db.Column(db.String(255), nullable=False)
    original_name = db.Column(db.String(255), nullable=False)
    content = db.Column(db.Text, nullable=False)
    chunks = db.Column(db.Text, nullable=True)  # JSON array of chunks
    created_at = db.Column(db.DateTime, default=datetime.utcnow)
    
    def to_dict(self):
        return {
            'id': self.id,
            'filename': self.original_name,
            'created_at': self.created_at.isoformat() if self.created_at else None
        }


# ── Routes ──────────────────────────────────────────────────────────────
@app.route('/')
def home():
    try:
        conversations = Conversation.query.order_by(Conversation.updated_at.desc()).all()
        return render_template('index.html', conversations=conversations)
    except Exception as e:
        app.logger.error(f"Error in home route: {str(e)}")
        flash('An error occurred while loading conversations.', 'error')
        return render_template('index.html', conversations=[])


@app.route('/chat/<int:conv_id>', methods=['GET'])
def chat(conv_id):
    try:
        conversation = Conversation.query.get_or_404(conv_id)
        conversations = Conversation.query.order_by(Conversation.updated_at.desc()).all()
        documents = Document.query.filter_by(conversation_id=conv_id).all()
        
        try:
            messages = json.loads(conversation.messages) if conversation.messages else []
        except (json.JSONDecodeError, TypeError):
            messages = []
        
        return render_template('chat.html', 
                             conversation=conversation, 
                             messages=messages, 
                             conversations=conversations,
                             documents=documents)
    
    except Exception as e:
        app.logger.error(f"Error in chat route: {str(e)}\n{traceback.format_exc()}")
        flash('An error occurred.', 'error')
        return redirect(url_for('home'))


@app.route('/new', methods=['POST'])
def new_conversation():
    try:
        topic = request.form.get('topic', '').strip()
        if not topic:
            flash('Please provide a conversation topic.', 'error')
            return redirect(url_for('home'))
        
        if len(topic) > 200:
            topic = topic[:200]
        
        new_conv = Conversation(topic=topic, messages='[]')
        db.session.add(new_conv)
        db.session.commit()
        
        return redirect(url_for('chat', conv_id=new_conv.id))
    
    except Exception as e:
        app.logger.error(f"Error creating conversation: {str(e)}")
        flash('Error creating conversation.', 'error')
        return redirect(url_for('home'))


# ── SSE Streaming Endpoint ───────────────────────────────────────────────
@app.route('/api/chat/stream', methods=['POST'])
def chat_stream():
    """Server-Sent Events streaming endpoint for real-time responses"""
    try:
        data = request.get_json()
        
        conv_id = data.get('conversation_id')
        instruction = data.get('instruction', '').strip()
        model = data.get('model', 'llama-3.1-70b-versatile')
        temperature = float(data.get('temperature', 0.7))
        max_tokens = int(data.get('max_tokens', 2048))
        
        if not instruction:
            return jsonify({'error': 'No instruction provided'}), 400
        
        conversation = Conversation.query.get(conv_id)
        if not conversation:
            return jsonify({'error': 'Conversation not found'}), 404
        
        # Load existing messages
        try:
            messages = json.loads(conversation.messages) if conversation.messages else []
        except (json.JSONDecodeError, TypeError):
            messages = []
        
        # Add user message
        messages.append({"role": "user", "content": instruction})
        
        def generate() -> Generator[str, None, None]:
            """Generator for SSE streaming"""
            try:
                # Send initial event
                yield f"data: {json.dumps({'type': 'start', 'model': model})}\n\n"
                
                # Stream response from Groq
                response = generate_together_stream(
                    model=model,
                    messages=messages,
                    temperature=temperature,
                    max_tokens=max_tokens,
                )
                
                full_response = ""
                for chunk in response:
                    if hasattr(chunk, 'choices') and chunk.choices:
                        delta = chunk.choices[0].delta
                        if hasattr(delta, 'content') and delta.content:
                            full_response += delta.content
                            # Send chunk to client
                            yield f"data: {json.dumps({'type': 'token', 'content': delta.content})}\n\n"
                
                # Save the complete conversation
                messages.append({"role": "assistant", "content": full_response})
                conversation.messages = json.dumps(messages)
                conversation.updated_at = datetime.utcnow()
                db.session.commit()
                
                # Send completion event
                yield f"data: {json.dumps({'type': 'done', 'full_response': full_response})}\n\n"
                
            except Exception as e:
                app.logger.error(f"Error in stream: {str(e)}")
                yield f"data: {json.dumps({'type': 'error', 'message': str(e)})}\n\n"
        
        return Response(
            stream_with_context(generate()),
            mimetype='text/event-stream',
            headers={
                'Cache-Control': 'no-cache',
                'X-Accel-Buffering': 'no',
                'Connection': 'keep-alive',
            }
        )
        
    except Exception as e:
        app.logger.error(f"Error in chat_stream: {str(e)}\n{traceback.format_exc()}")
        return jsonify({'error': str(e)}), 500


# ── Non-streaming fallback (for backward compatibility) ────────────────
@app.route('/chat/<int:conv_id>/send', methods=['POST'])
def chat_send(conv_id):
    """Non-streaming endpoint for clients without SSE support"""
    try:
        conversation = Conversation.query.get_or_404(conv_id)
        
        instruction = None
        # Handle file upload
        if 'file' in request.files:
            file = request.files['file']
            if file and file.filename:
                try:
                    filename = secure_filename(file.filename)
                    if not filename:
                        flash('Invalid file name.', 'error')
                        return redirect(url_for('chat', conv_id=conv_id))
                    
                    file_path = os.path.join(app.config['UPLOADED_FILES_DEST'], filename)
                    file.save(file_path)
                    
                    content = extract_content(file_path)
                    if not content.strip():
                        flash('Could not extract content from the file.', 'error')
                        return redirect(url_for('chat', conv_id=conv_id))
                    
                    user_instruction = request.form.get('instruction', '').strip()
                    if user_instruction:
                        instruction = f"{user_instruction}\n\nDocument content:\n{content}"
                    else:
                        instruction = f"Please analyze and summarize the following document:\n\n{content}"
                    
                    # Store document for RAG
                    doc = Document(
                        conversation_id=conv_id,
                        filename=filename,
                        original_name=file.filename,
                        content=content
                    )
                    db.session.add(doc)
                    
                    try:
                        os.remove(file_path)
                    except Exception:
                        pass
                        
                except Exception as e:
                    app.logger.error(f"Error processing file: {str(e)}")
                    flash(f'Error processing file: {str(e)}', 'error')
                    return redirect(url_for('chat', conv_id=conv_id))
        
        if not instruction:
            instruction = request.form.get('instruction', '').strip()
            if not instruction:
                flash('Please provide an instruction or upload a file.', 'error')
                return redirect(url_for('chat', conv_id=conv_id))
        
        model = request.form.get('model', 'llama-3.1-70b-versatile')
        try:
            temperature = float(request.form.get('temperature', 0.7))
            temperature = max(0.0, min(2.0, temperature))
        except (ValueError, TypeError):
            temperature = 0.7
        
        try:
            max_tokens = int(request.form.get('max_tokens', 2048))
            max_tokens = max(1, min(32768, max_tokens))
        except (ValueError, TypeError):
            max_tokens = 2048
        
        try:
            messages = json.loads(conversation.messages) if conversation.messages else []
        except (json.JSONDecodeError, TypeError):
            messages = []
        
        messages.append({"role": "user", "content": instruction})
        
        try:
            output = generate_with_references(
                model=model,
                temperature=temperature,
                max_tokens=max_tokens,
                messages=messages,
                references=[],
                generate_fn=generate_together_stream,
            )
            
            final_output = ""
            for chunk in output:
                if hasattr(chunk, 'choices') and chunk.choices:
                    delta = chunk.choices[0].delta
                    if hasattr(delta, 'content') and delta.content:
                        final_output += delta.content
            
            if not final_output.strip():
                flash('No response generated. Please try again.', 'error')
                return redirect(url_for('chat', conv_id=conv_id))
            
            messages.append({"role": "assistant", "content": final_output})
            conversation.messages = json.dumps(messages)
            conversation.updated_at = datetime.utcnow()
            db.session.commit()
            
        except Exception as e:
            app.logger.error(f"Error generating response: {str(e)}\n{traceback.format_exc()}")
            flash(f'Error generating response: {str(e)}', 'error')
            return redirect(url_for('chat', conv_id=conv_id))
        
        return redirect(url_for('chat', conv_id=conv_id))
    
    except Exception as e:
        app.logger.error(f"Error in chat_send: {str(e)}\n{traceback.format_exc()}")
        flash('An error occurred.', 'error')
        return redirect(url_for('home'))


# ── Search API ──────────────────────────────────────────────────────────
@app.route('/api/search')
def search():
    """Full-text search across conversations"""
    query = request.args.get('q', '').strip()
    if not query or len(query) < 2:
        return jsonify({'results': []})
    
    try:
        # Search in topics and messages
        results = Conversation.query.filter(
            db.or_(
                Conversation.topic.ilike(f'%{query}%'),
                Conversation.messages.ilike(f'%{query}%')
            )
        ).order_by(Conversation.updated_at.desc()).limit(20).all()
        
        # Highlight matches
        matched = []
        for conv in results:
            messages = json.loads(conv.messages) if conv.messages else []
            # Find matching message snippets
            snippets = []
            for msg in messages[-10:]:  # Last 10 messages
                if query.lower() in msg.get('content', '').lower():
                    content = msg['content']
                    # Extract snippet around match
                    idx = content.lower().find(query.lower())
                    start = max(0, idx - 50)
                    end = min(len(content), idx + len(query) + 50)
                    snippet = content[start:end]
                    if start > 0:
                        snippet = '...' + snippet
                    if end < len(content):
                        snippet = snippet + '...'
                    snippets.append({
                        'role': msg['role'],
                        'snippet': snippet
                    })
            
            matched.append({
                'id': conv.id,
                'topic': conv.topic,
                'updated_at': conv.updated_at.isoformat() if conv.updated_at else None,
                'snippets': snippets[:3]  # Max 3 snippets per conversation
            })
        
        return jsonify({'results': matched, 'count': len(matched)})
        
    except Exception as e:
        app.logger.error(f"Search error: {str(e)}")
        return jsonify({'error': str(e)}), 500


# ── Export Endpoints ────────────────────────────────────────────────────
@app.route('/export/<int:conv_id>/markdown')
def export_markdown(conv_id):
    """Export conversation as Markdown"""
    conversation = Conversation.query.get_or_404(conv_id)
    messages = json.loads(conversation.messages) if conversation.messages else []
    
    lines = [
        f"# {conversation.topic}",
        f"*Exported from {config.APP_NAME} on {datetime.now().strftime('%Y-%m-%d %H:%M')}*",
        "",
        "---",
        ""
    ]
    
    for msg in messages:
        role = "**You**" if msg['role'] == 'user' else f"**{config.APP_NAME}**"
        lines.append(f"### {role}")
        lines.append("")
        lines.append(msg['content'])
        lines.append("")
        lines.append("---")
        lines.append("")
    
    content = "\n".join(lines)
    filename = f"{conversation.topic.replace(' ', '_')[:50]}.md"
    
    return Response(
        content,
        mimetype='text/markdown',
        headers={'Content-Disposition': f'attachment; filename="{filename}"'}
    )


@app.route('/export/<int:conv_id>/json')
def export_json(conv_id):
    """Export conversation as JSON"""
    conversation = Conversation.query.get_or_404(conv_id)
    
    data = {
        'conversation': conversation.to_dict(),
        'messages': json.loads(conversation.messages) if conversation.messages else [],
        'exported_at': datetime.now().isoformat(),
        'app': config.APP_NAME,
        'version': config.APP_VERSION
    }
    
    filename = f"{conversation.topic.replace(' ', '_')[:50]}.json"
    
    return Response(
        json.dumps(data, indent=2),
        mimetype='application/json',
        headers={'Content-Disposition': f'attachment; filename="{filename}"'}
    )


# ── Document Management ────────────────────────────────────────────────
@app.route('/api/documents/<int:conv_id>')
def list_documents(conv_id):
    """List documents attached to a conversation"""
    documents = Document.query.filter_by(conversation_id=conv_id).all()
    return jsonify({'documents': [d.to_dict() for d in documents]})


@app.route('/api/documents/<int:doc_id>', methods=['DELETE'])
def delete_document(doc_id):
    """Delete a document"""
    doc = Document.query.get_or_404(doc_id)
    db.session.delete(doc)
    db.session.commit()
    return jsonify({'success': True})


# ── Conversation Management ─────────────────────────────────────────────
@app.route('/reset', methods=['GET', 'POST'])
def reset():
    try:
        db.drop_all()
        db.create_all()
        flash('All conversations have been reset.', 'info')
        return redirect(url_for('home'))
    except Exception as e:
        app.logger.error(f"Error resetting database: {str(e)}")
    flash('Error resetting conversations.', 'error')
    return redirect(url_for('home'))


@app.route('/api/conversations/<int:conv_id>', methods=['DELETE'])
def delete_conversation(conv_id):
    """Delete a single conversation"""
    try:
        conv = Conversation.query.get_or_404(conv_id)
        # Delete associated documents
        Document.query.filter_by(conversation_id=conv_id).delete()
        db.session.delete(conv)
        db.session.commit()
        return jsonify({'success': True})
    except Exception as e:
        return jsonify({'error': str(e)}), 500


# ── Update API ──────────────────────────────────────────────────────────
@app.route('/api/check-updates', methods=['GET'])
def api_check_updates():
    try:
        info = check_updates()
        return jsonify({'success': True, 'data': info})
    except Exception as e:
        return jsonify({'success': False, 'error': str(e)}), 500


@app.route('/api/update', methods=['POST'])
def api_update():
    try:
        success, message = update_application()
        return jsonify({'success': success, 'message': message})
    except Exception as e:
        return jsonify({'success': False, 'error': str(e)}), 500


# ── Document Extraction ────────────────────────────────────────────────
def extract_content(file_path: str) -> str:
    """Extract text content from various file formats."""
    content = ""
    try:
        if file_path.endswith('.pdf'):
            import fitz
            doc = fitz.open(file_path)
            for page in doc:
                content += page.get_text()
            doc.close()
        elif file_path.endswith('.docx'):
            import docx
            doc = docx.Document(file_path)
            for para in doc.paragraphs:
                content += para.text + "\n"
        elif file_path.endswith('.txt'):
            with open(file_path, 'r', encoding='utf-8') as f:
                content = f.read()
        else:
            try:
                with open(file_path, 'r', encoding='utf-8') as f:
                    content = f.read()
            except Exception:
                pass
    except Exception as e:
        app.logger.error(f"Error extracting content from {file_path}: {str(e)}")
        raise
    
    return content


def chunk_text(text: str, chunk_size: int = 1000, overlap: int = 200) -> list:
    """Simple text chunking with overlap for RAG"""
    chunks = []
    start = 0
    while start < len(text):
        end = start + chunk_size
        chunk = text[start:end]
        # Try to end on a sentence
        last_period = chunk.rfind('.')
        if last_period > chunk_size // 2:
            chunk = chunk[:last_period + 1]
            end = start + last_period + 1
        chunks.append(chunk.strip())
        start = end - overlap
    return [c for c in chunks if c]


# ── Error Handlers ──────────────────────────────────────────────────────
@app.errorhandler(404)
def not_found(error):
    flash('Page not found.', 'error')
    return redirect(url_for('home'))


@app.errorhandler(500)
def internal_error(error):
    db.session.rollback()
    app.logger.error(f"Internal server error: {str(error)}")
    flash('An internal error occurred.', 'error')
    return redirect(url_for('home'))


# ── CLI Commands ────────────────────────────────────────────────────────
@app.cli.command('init-db')
def init_db():
    """Initialize the database."""
    db.create_all()
    print('Database initialized.')


if __name__ == "__main__":
    with app.app_context():
        db.create_all()
    app.run(debug=True, host='0.0.0.0', port=5000)