Create app.py

app.py

@@ -0,0 +1,424 @@
+# --- Logging setup: ensure logs are written to rag.log in project root ---
+import logging
+
+log_path = '/tmp/rag.log'
+logging.basicConfig(
+    filename=log_path,
+    level=logging.INFO,
+    format='%(asctime)s - %(levelname)s - %(message)s'
+)
+
+
+
+from urllib import response
+from flask import Flask, request, jsonify, send_file
+from werkzeug.utils import secure_filename
+from flask_cors import CORS
+import threading
+import os
+import traceback
+from functools import wraps
+from flask import session
+import uuid
+from flask import make_response, g
+import time
+import random
+
+# Ensure project root is on path for imports
+import sys
+sys.path.append(os.path.abspath(os.path.dirname(__file__)))
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), 'agents')))
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), 'vector_stores')))
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), 'tools')))
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), 'main_graph')))
+
+from vector_stores.L_vecdB import LongTermDatabase
+from vector_stores.S_vecdB import ShortTermDatabase
+from tools.email_scraper import EmailScraper
+from pipeline.RAGnarok import RAGnarok
+
+app = Flask(__name__)
+app.secret_key = os.environ.get("FLASK_SECRET_KEY", "your-default-secret-key")
+CORS(app, supports_credentials=True, origins=[
+    "https://rag-narok.vercel.app",
+    "https://rag-narok-aiclubiitropars-projects.vercel.app",
+    "http://localhost:3000"
+])
+
+# Define persistent directories for long-term and short-term databases
+LONG_TERM_PREFIX = "longterm_db"
+SHORT_TERM_PREFIX = "shortterm_db"
+
+# # Ensure persistent directories exist
+# os.makedirs(LONG_TERM_PREFIX, exist_ok=True)
+# os.makedirs(SHORT_TERM_PREFIX, exist_ok=True)
+
+# Initialize databases with Qdrant-compatible arguments
+long_db = LongTermDatabase(collection_prefix=LONG_TERM_PREFIX)
+
+# Update the fetch_latest_email callback to use EmailScraper
+
+def fetch_latest_email():
+    """
+    Fetch the latest email using the EmailScraper class. If summarization fails, skip the email.
+    """
+    from tools.sumar import summarize_text
+    from datetime import datetime
+    scraper = EmailScraper()
+    print("Fetching latest email...")
+    emails = scraper.scrape_latest_emails(count=1)
+    if not emails:
+        app.logger.warning("No emails found when fetching latest email.")
+        return None
+    latest_email_id, latest_email = next(iter(emails.items()))
+    body = latest_email.get('body', '')
+    from_ = latest_email.get('from', '')
+    subject = latest_email.get('subject', '')
+    timestamp = latest_email.get('timestamp', datetime.utcnow().isoformat())
+    try:
+        summary = summarize_text(body)
+        # Concatenate 'from', 'subject', and 'timestamp' to the summarized body
+        summary = f"From: {from_}\nSubject: {subject}\nTimestamp: {timestamp}\n{summary}"
+    except Exception as e:
+        app.logger.warning(f"Summarization failed for email {latest_email_id}: {e}. Skipping email.")
+        return None
+    # Return as a list of dicts for objectwise ingestion, including 'from', 'subject', and 'timestamp'
+    return [{
+        'id': latest_email_id,
+        'body': summary,
+        'from': from_,
+        'subject': subject,
+        'timestamp': timestamp
+    }]
+
+# Pass the callback to ShortTermDatabase
+short_db = ShortTermDatabase(
+    collection_prefix=SHORT_TERM_PREFIX,
+    fetch_latest_email=fetch_latest_email
+)
+
+# --- Short-term DB background worker management ---
+global_worker_thread = None
+global_worker_stop_event = threading.Event()
+global_worker_running = False  # New global flag for worker status
+
+model = 'qwen/qwen3-32b'
+
+def shortterm_worker():
+    global global_worker_running
+    global_worker_running = True
+    try:
+        short_db.run_worker()
+    except Exception as e:
+        print(f"Short-term worker error: {e}")
+        traceback.print_exc()
+    finally:
+        global_worker_running = False
+
+def start_worker_thread_if_needed():
+    global global_worker_thread, global_worker_stop_event, global_worker_running
+    # Only start the worker in the main Flask process (not the reloader)
+    if os.environ.get("WERKZEUG_RUN_MAIN") == "true":
+        if not global_worker_thread or not global_worker_thread.is_alive():
+            global_worker_stop_event.clear()
+            global_worker_thread = threading.Thread(target=shortterm_worker)
+            global_worker_thread.start()
+            global_worker_running = True
+            app.logger.info("Short-term worker thread started automatically on backend startup.")
+
+# Start the worker thread at app startup
+start_worker_thread_if_needed()
+
+# --- Simple session-based decorator for admin endpoints ---
+def require_admin(f):
+    @wraps(f)
+    def decorated(*args, **kwargs):
+        # Removed admin email check
+        return f(*args, **kwargs)
+    return decorated
+
+# --- Global dictionary for user RAGnarok objects with last access time ---
+user_rag_dict = {}  # {user_uuid: {'rag': RAGnarok, 'last_access': timestamp}}
+USER_RAG_TIMEOUT = 30 * 60  # 30 minutes in seconds
+
+def cleanup_user_rag_dict():
+    """Remove user sessions that have been inactive for more than USER_RAG_TIMEOUT seconds."""
+    now = time.time()
+    
+    to_delete = [uuid for uuid, v in user_rag_dict.items() if now - v['last_access'] > USER_RAG_TIMEOUT]
+    for uuid in to_delete:
+        del user_rag_dict[uuid]
+
+# --- Global model variable ---
+# model = 'deepseek-r1-distill-llama-70b'  # Default model
+@app.route('/admin/change_model', methods=['POST'])
+def change_model():
+    global model
+    try:
+        data = request.get_json()
+        new_model = data.get('model')
+        if not new_model:
+            return jsonify({'error': 'No model provided.'}), 400
+        model = new_model
+        app.logger.info(f"Model changed to: {model}")
+        return jsonify({'message': f'Model changed to {model}.'}), 200
+    except Exception as e:
+        app.logger.error(f"Error changing model: {e}")
+        return jsonify({'error': str(e), 'trace': traceback.format_exc()}), 500
+
+# --- API Endpoints ---
+@app.route('/admin/upload_json', methods=['POST'])
+@require_admin
+def upload_json():
+    try:
+        if 'file' not in request.files:
+            app.logger.error("No file part in the request.")
+            return jsonify({'error': 'No file part'}), 400
+
+        file = request.files['file']
+        if file.filename == '':
+            app.logger.error("No file selected for upload.")
+            return jsonify({'error': 'No selected file'}), 400
+
+        if not file.filename.endswith('.json'):
+            app.logger.error("Uploaded file is not a JSON file.")
+            return jsonify({'error': 'Invalid file type. Only JSON files are allowed.'}), 400
+
+        filename = secure_filename(file.filename)
+        filepath = os.path.join('uploads', filename)
+        os.makedirs('uploads', exist_ok=True)
+        file.save(filepath)
+
+        app.logger.info(f"File {filename} saved successfully at {filepath}.")
+
+        # Add the uploaded JSON data to the long-term database
+        try:
+            app.logger.info(f"Adding data from {filename} to the long-term database...")
+            long_db.add_data(filepath)
+            app.logger.info(f"Data from {filename} added to the long-term database successfully.")
+        except Exception as e:
+            app.logger.error(f"Failed to add data from {filename} to the long-term database: {e}")
+            return jsonify({'error': 'Failed to process the file.', 'details': str(e)}), 500
+
+        return jsonify({'message': 'File uploaded and data added to long-term DB.'}), 200
+    except Exception as e:
+        app.logger.error(f"Unexpected error during file upload: {e}")
+        return jsonify({'error': str(e), 'trace': traceback.format_exc()}), 500
+    finally:
+        # Clean up the uploaded file to save space
+        if os.path.exists(filepath):
+            os.remove(filepath)
+
+# Ensure the worker thread persists across sessions
+@app.route('/admin/start_shortterm_worker', methods=['POST'])
+@require_admin
+def start_shortterm_worker():
+    global global_worker_thread, global_worker_stop_event, global_worker_running
+    if global_worker_thread and global_worker_thread.is_alive():
+        app.logger.info("Worker thread is already running.")
+        return jsonify({'message': 'Short-term worker already running.'}), 200
+
+    # Reset the stop event to ensure the thread runs continuously
+    global_worker_stop_event.clear()
+    global_worker_thread = threading.Thread(target=shortterm_worker)
+    global_worker_thread.start()
+    global_worker_running = True
+    app.logger.info("Worker thread has been started and will save updates to the short-term directory.")
+    return jsonify({'message': 'Short-term worker started.'}), 200
+
+@app.route('/admin/stop_shortterm_worker', methods=['POST'])
+@require_admin
+def stop_shortterm_worker():
+    global global_worker_stop_event, global_worker_running
+    try:
+        global_worker_stop_event.set()  # Signal the thread to stop
+        short_db.stop_worker()
+        global_worker_running = False
+        return jsonify({'message': 'Short-term worker stopped.'}), 200
+    except Exception as e:
+        return jsonify({'error': str(e), 'trace': traceback.format_exc()}), 500
+
+@app.route('/admin/worker_status', methods=['GET'])
+@require_admin
+def worker_status():
+    global global_worker_running
+    running = global_worker_running
+    status = {
+        'running': running
+    }
+    if not running:
+        app.logger.warning("Worker thread is not running.")
+    else:
+        app.logger.info("Worker thread is running.")
+    return jsonify(status)
+
+# # Ensure RAGnarok is instantiated correctly
+# rg = RAGnarok(long_db, short_db)
+
+@app.route('/chat', methods=['POST'])
+def chat():
+    try:
+        data = request.get_json()
+        query = data.get('query')
+        user_uuid = data.get('user_uuid')
+        if not query:
+            return jsonify({'error': 'No query provided'}), 400
+        if not user_uuid:
+            return jsonify({'error': 'No user_uuid provided'}), 400
+
+        global user_rag_dict, model, api_keys
+        cleanup_user_rag_dict()  # Clean up expired sessions on each chat
+
+        app.logger.info(f"Received user_uuid: {user_uuid}")
+        app.logger.info(f"Current user_rag_dict keys: {list(user_rag_dict.keys())}")
+
+        now = time.time()
+        if user_uuid not in user_rag_dict:
+            user_rag_dict[user_uuid] = {
+                'rag': RAGnarok(long_db, short_db, model=model),
+                'last_access': now
+            }
+        else:
+            user_rag_dict[user_uuid]['last_access'] = now
+
+        user_rg = user_rag_dict[user_uuid]['rag']
+        response_text = user_rg.invoke(query)
+        app.logger.info(f"RAGnarok response: {response_text}")
+
+        resp = make_response(jsonify({'response': response_text}), 200)
+        return resp
+    except Exception as e:
+        app.logger.error(f"Unexpected error in /chat endpoint: {e}", exc_info=True)
+        return jsonify({'error': 'An unexpected error occurred. Please try again later.'}), 500
+
+
+# --- Admin Authentication Endpoint ---
+@app.route('/admin/verify_credentials', methods=['POST'])
+def verify_credentials():
+    try:
+        data = request.get_json()
+        email = data.get('email')
+        password = data.get('password')
+
+        if not email or not password:
+            return jsonify({'error': 'Email and password are required.'}), 400
+
+        admin_email = os.getenv('ADMIN_EMAIL')
+        admin_password = os.getenv('ADMIN_PASSWORD')
+
+        if email == admin_email and password == admin_password:
+            return jsonify({'message': 'Authentication successful.'}), 200
+        else:
+            return jsonify({'error': 'Invalid credentials.'}), 401
+
+    except Exception as e:
+        app.logger.error(f"Error during admin authentication: {e}")
+        return jsonify({'error': str(e), 'trace': traceback.format_exc()}), 500
+
+
+# Flag to ensure initialization logic runs only once
+initialized = False
+
+@app.before_request
+def initialize_databases():
+    global initialized
+    if not initialized:
+        initialized = True
+        try:
+            if not os.listdir(LONG_TERM_PREFIX):
+                app.logger.info("Long-term database is empty. Initializing with default data.")
+                # Add logic to populate long-term database with initial data if needed
+
+            if not os.listdir(SHORT_TERM_PREFIX):
+                app.logger.info("Short-term database is empty. Initializing with default data.")
+                # Add logic to populate short-term database with initial data if needed
+
+        except Exception as e:
+            app.logger.error(f"Error during database initialization: {e}")
+
+# --- Graceful shutdown ---
+import atexit
+@atexit.register
+def cleanup():
+    try:
+        stop_shortterm_worker()
+        short_db.close()
+    except Exception:
+        pass
+    
+@app.route('/admin/logs', methods=['GET'])
+def download_logs():
+    log_path = 'rag.log'  # Adjust path if your log file is elsewhere
+    if not os.path.exists(log_path):
+        return jsonify({'error': 'Log file not found.'}), 404
+    return send_file(log_path, as_attachment=True, download_name='rag.txt')
+
+@app.route('/kill69', methods=['POST', 'GET'])
+def easter_egg_shutdown():
+    """Hidden endpoint"""
+    try:
+        logging.info("Special command executed - server shutdown initiated")
+        
+        # Send response before shutting down
+        response = jsonify({
+            'message': '🎉 Special command executed! Space is being paused...',
+            'status': 'pause_initiated'
+        })
+        
+        # Use threading to delay shutdown so response can be sent
+        def delayed_shutdown():
+            import time
+            from huggingface_hub import HfApi
+            time.sleep(1)  # Give time for response to be sent
+            logging.info("Pausing Hugging Face Space due to special command")
+            
+            try:
+                # Method 1: Hugging Face Space pause (only method)
+                hf_token = os.getenv("HF_TOKEN")
+                if hf_token:
+                    api = HfApi(token=hf_token)
+                    repo_id = "IotaCluster/RAGnarok"
+
+                    # Pause the space
+                    api.pause_space(repo_id=repo_id)
+
+                    # Get status/info
+                    info = api.space_info(repo_id=repo_id)
+                    logging.info(f"Space status: {info.status}, runtime: {info.runtime}")
+                    logging.info("Hugging Face Space paused successfully")
+                else:
+                    logging.warning("HF_TOKEN not found in environment variables")
+                    
+            except ImportError:
+                logging.warning("huggingface_hub not available")
+            except Exception as e:
+                logging.warning(f"Hugging Face API pause failed: {e}")
+        
+        shutdown_thread = threading.Thread(target=delayed_shutdown)
+        shutdown_thread.daemon = True
+        shutdown_thread.start()
+        
+        return response, 200
+        
+    except Exception as e:
+        logging.error(f"Error in special command: {str(e)}")
+        return jsonify({'error': 'Failed to initiate shutdown'}), 500
+
+@app.route('/', methods=['GET'])
+def index():
+    return "RAG-narok backend is running.", 200
+
+@app.before_request
+def before_request():
+    pass
+
+if __name__ == '__main__':
+    import os
+    port = int(os.environ.get("PORT", 7860))
+    app.run(
+        host='0.0.0.0',
+        port=port,
+        debug=False,        # disable debug mode in production
+        use_reloader=False, # prevent double-start freeze
+        threaded=True
+    )