DevForML/Multi_Agent_System
0
1# app.py2 3import eventlet4eventlet.monkey_patch()5 6from flask import Flask, render_template, request, redirect, url_for, flash, session, send_from_directory7from flask_socketio import SocketIO8import traceback9import os10from werkzeug.utils import secure_filename11import json12import logging13import agent # your agent.py module14from agent import refresh_memory15from agent import run_stream 16from typing import List, Dict17import markdown218import re19import time20 21 22# ──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────23# ─────────────────────────────────────────────── Inialized VAR & FS ───────────────────────────────────────────────────────────────────24# ──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────25BASE_DIR = os.path.abspath(os.path.dirname(__file__))26 27UPLOAD_FOLDER = os.path.join(BASE_DIR, "uploads")28os.makedirs(UPLOAD_FOLDER, exist_ok=True)29 30CHAT_FOLDER = os.path.join(BASE_DIR, "chats")31os.makedirs(CHAT_FOLDER, exist_ok=True)32 33try:34 os.chmod(CHAT_FOLDER, 0o777)35except Exception:36 pass37 38app = Flask(__name__, template_folder="templates")39 40# For storing the processed files41app.config["UPLOAD_FOLDER"] = UPLOAD_FOLDER42# For storing the Chat history and other files43app.config["CHAT_FOLDER"] = CHAT_FOLDER44 45app.secret_key = os.getenv("FLASK_SECRET", "supersecret")46 47# Use eventlet for async SocketIO48socketio = SocketIO(app, cors_allowed_origins="*", async_mode="eventlet")49 50# Allowed file extensions51ALLOWED_EXTENSIONS = {52 ".db", ".sqlite", # SQLite databases53 ".pdf", ".txt", ".doc", ".docx", # Documents54 ".png", ".jpg", ".jpeg", ".gif" # Images55}56 57import config58DB_PATH = None # will be set when a .db is uploaded59DOC_PATH = None # will be set when a document is uploaded60IMG_PATH = None # will be set when an image is uploaded61OTH_PATH = None # will be set when an other file is uploaded62 63# import config64# IMG_PATH = "path/to/user_uploaded_file.jpg"65# DOC_PATH = "path/to/user_uploaded_file.pdf"66# DB_PATH = "path/to/user_uploaded_file.db"67# OTH_PATH = "path/to/user_uploaded_file.txt"68 69def allowed_file(filename: str) -> bool:70 ext = os.path.splitext(filename.lower())[1]71 return ext in ALLOWED_EXTENSIONS72 73def ensure_user_session():74 if "user_id" not in session:75 session["user_id"] = os.urandom(16).hex()76 session["uploads"] = []77 session["chat_history"] = []78 refresh_memory()79 80#text cleaning function81def clean_html_chunk(text):82 """83 Removes outer <p>...</p> tags and trims extra backticks or 'json' words.84 Similar to your 'format:' cleaning logic.85 """86 text = text.strip()87 88 # Pattern to match single <p>...</p> wrapping89 pattern = r'^<p>(.*?)</p>$'90 match = re.match(pattern, text, re.DOTALL)91 if match:92 text = match.group(1).strip()93 94 # Extra clean-up (optional, like your example)95 text = text.strip('`').strip('json').strip()96 97 return text98 99 100# ──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────101# ─────────────────────────────────────────────── AGENT Defination ────────────────────────────────────────────────────────────────────────────102# ──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────103 104def run_agent_thread(105 prompt: str,106 user_id: str,107 chat_history: List[Dict], # ← new parameter108):109 """110 Launches the agent in a background thread, streaming results back over SocketIO.111 `data` is the uploaded file path (image, document, or DB) that will be injected112 into config before the agent runs.113 """114 global DB_PATH, DOC_PATH, IMG_PATH, OTH_PATH # force re-initializing agent115 116 # Build the list of all paths, skipping None or empty117 data_paths = [DB_PATH, DOC_PATH, IMG_PATH, OTH_PATH]118 data_paths = [p for p in data_paths if p]119 print(f"Data paths----------------->: {data_paths}")120 text_accum = ""121 try:122 # Stream using run_stream (which will also pick up the same globals)123 for piece in run_stream(prompt, data_paths):124 html_chunk = markdown2.markdown(piece, extras=["fenced-code-blocks", "tables", "strike", "task_list", "break-on-newline"])125 print(f"HTML chunk: {html_chunk}") # Debugging output126 html_chunk = clean_html_chunk(html_chunk)127 text_accum += html_chunk # accumulate HTML128 socketio.emit("final_stream", {"message": html_chunk})129 print(f"Streaming chunk: {html_chunk}") # Debugging output130 except Exception as e:131 socketio.emit("error", {"message": f"Streaming error: {e}"})132 traceback.print_exc()133 return134 135 # Fallback / finalize136 try:137 if not text_accum:138 text_accum = markdown2.markdown(agent.agent.executor.run(prompt), extras=["fenced-code-blocks", "tables", "strike", "task_list", "break-on-newline"])139 print(f"Final HTML chunk: {text_accum}") # Debugging output140 text_accum = clean_html_chunk(text_accum)141 print(f"Final text: {text_accum}") # Debugging output142 143 socketio.emit("stream_complete", {"message": text_accum})144 socketio.emit("final", {"message": text_accum})145 chat_history.append({"user": prompt, "assistant": text_accum})146 147 output_path = os.path.join(148 app.config["CHAT_FOLDER"],149 f"user_chat_no_{user_id}.json"150 )151 with open(output_path, "w", encoding="utf-8") as f:152 json.dump(chat_history, f, ensure_ascii=False, indent=4)153 154 except Exception as e:155 socketio.emit("error", {"message": f"Final generation error: {e}"})156 traceback.print_exc()157 158# ──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────159# ─────────────────────────────────────────────── Main Page ────────────────────────────────────────────────────────────────────────────────160# ──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────161 162@app.route("/")163def index():164 if "user_id" not in session:165 session["user_id"] = os.urandom(16).hex()166 session["uploads"] = []167 session["chat_history"] = []168 refresh_memory()169 else:170 # Load previous chat history from JSON file if exists171 user_id = session["user_id"]172 chat_file = os.path.join(app.config.get("CHAT_FOLDER", "chat_history"), f"user_chat_no_{user_id}.json")173 if os.path.exists(chat_file):174 with open(chat_file, "r", encoding="utf-8") as f:175 session["chat_history"] = json.load(f)176 177 return render_template("index.html", chat_history=session.get("chat_history", []))178 179# ──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────180# ─────────────────────────────────────────────── Upload section ──────────────────────────────────────────────────────────────────────────181# ──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────182 183@app.route("/upload", methods=["GET", "POST"])184def upload():185 ensure_user_session()186 global DB_PATH, DOC_PATH, IMG_PATH, OTH_PATH # force re-initializing agent187 188 if request.method == "POST":189 f = request.files.get("file")190 if not f or f.filename == "":191 flash("No file selected", "error")192 return render_template("upload.html")193 194 filename = secure_filename(f.filename)195 if not allowed_file(filename):196 flash("File type not supported", "error")197 return render_template("upload.html")198 199 ext = os.path.splitext(filename.lower())[1]200 201 if ext in (".db", ".sqlite"):202 save_path = os.path.join(app.config["UPLOAD_FOLDER"], "databases", filename)203 os.makedirs(os.path.dirname(save_path), exist_ok=True)204 global DB_PATH205 DB_PATH = save_path206 # DB_PATH = f"http://127.0.0.1:5000/uploads/databases/{filename}"207 print(f"Database path: {save_path}")208 f.save(save_path)209 elif ext in (".pdf", ".txt", ".doc", ".docx"):210 save_path = os.path.join(app.config["UPLOAD_FOLDER"], "documents", filename)211 os.makedirs(os.path.dirname(save_path), exist_ok=True)212 global DOC_PATH213 DOC_PATH = save_path214 # DOC_PATH = f"http://127.0.0.1:5000/uploads/documents/{filename}"215 print(f"Document path: {save_path}")216 f.save(save_path)217 elif ext in (".png", ".jpg", ".jpeg", ".gif"):218 save_path = os.path.join(app.config["UPLOAD_FOLDER"], "images", filename)219 os.makedirs(os.path.dirname(save_path), exist_ok=True)220 global IMG_PATH221 IMG_PATH = save_path222 # IMG_PATH = f"http://127.0.0.1:5000/uploads/images/{filename}"223 print(f"Image path: {save_path}")224 f.save(save_path)225 else:226 save_path = os.path.join(app.config["UPLOAD_FOLDER"], "others", filename)227 os.makedirs(os.path.dirname(save_path), exist_ok=True)228 global OTH_PATH229 OTH_PATH = save_path230 # OTH_PATH = f"http://127.0.0.1:5000/uploads/others/{filename}"231 print(f"Other file path: {save_path}")232 f.save(save_path)233 234 #f.save(save_path)235 236 # Add the uploaded file to the session237 session["uploads"].append(filename)238 239 # — Database files —240 if ext in (".db", ".sqlite"):241 DB_PATH = save_path242 agent.GLOBAL_DB_PATH = DB_PATH243 flash(f"Database uploaded and set: {filename}", "success")244 245 # — Documents for RAG indexing —246 elif ext in (".pdf", ".txt", ".doc", ".docx"):247 agent.rag_index_document(save_path)248 flash(f"Document indexed for RAG: {filename}", "success")249 250 # — Images or other files —251 else:252 flash(f"File uploaded: {filename}", "success")253 254 return redirect(url_for("index"))255 256 # GET257 return render_template("upload.html")258 259# ──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────260# ─────────────────────────────────────────────── Static Upload ───────────────────────────────────────────────────────────────────────────261# ──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────262 263@app.route('/uploads/databases/<filename>')264def serve_database(filename):265 """266 Serve an database file from the uploads/databases folder.267 """268 return send_from_directory(os.path.join(app.root_path, 'uploads', 'databases'), filename)269 270@app.route('/uploads/images/<filename>')271def serve_image(filename):272 """273 Serve an document file from the uploads/images folder.274 """275 return send_from_directory(os.path.join(app.root_path, 'uploads', 'images'), filename)276 277@app.route('/uploads/documents/<filename>')278def serve_document(filename):279 """280 Serve an image file from the uploads/documents folder.281 """282 return send_from_directory(os.path.join(app.root_path, 'uploads', 'documents'), filename)283 284@app.route('/uploads/others/<filename>')285def serve_other(filename):286 """287 Serve an other file from the uploads/others folder.288 """289 return send_from_directory(os.path.join(app.root_path, 'uploads', 'others'), filename)290 291# ──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────292# ─────────────────────────────────────────────── AGENT calling ───────────────────────────────────────────────────────────────────────────293# ──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────294 295@app.route("/generate", methods=["POST"])296def generate():297 prompt = request.json.get("prompt", "").strip()298 ensure_user_session()299 if not prompt:300 return "No prompt provided", 400301 302 socketio.start_background_task(303 run_agent_thread,304 prompt,305 session["user_id"],306 session["chat_history"]307 # ← pass it here308 )309 return "OK", 200310 311# ──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────312# ─────────────────────────────────────────────── SESSION handling ───────────────────────────────────────────────────────────────────────313# ──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────314 315@app.route("/session-info")316def session_info():317 # Endpoint to view session details (for debugging purposes)318 return {319 "user_id": session.get("user_id"),320 "uploads": session.get("uploads", [])321 }322 323# ──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────324# ─────────────────────────────────────────────── SESSION clearing ───────────────────────────────────────────────────────────────────────325# ──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────326 327@app.route("/clear_chat", methods=["POST"])328def clear_chat():329 ensure_user_session()330 user_id = session.get("user_id")331 332 # Remove saved JSON chat file333 if user_id:334 chat_file = os.path.join(app.config["CHAT_FOLDER"], f"user_chat_no_{user_id}.json")335 if os.path.exists(chat_file):336 os.remove(chat_file)337 338 # Reset session339 session.clear()340 341 # Generate new session id and chat history342 session["user_id"] = os.urandom(16).hex()343 session["uploads"] = []344 session["chat_history"] = []345 346 # Refresh agent memory347 refresh_memory()348 349 return {"message": "Chat history and session cleared!"}, 200350 351if __name__ == "__main__":352 socketio.run(app, debug=True, port=5000)353