WanIrfan/Atlas
0
1from flask import Flask, request, render_template, session, url_for, redirect, jsonify2# from flask_session import Session <--- REMOVED3from langchain_core.messages import HumanMessage, AIMessage4from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder5import os6import logging7import re8import traceback9import base6410import shutil11import zipfile12from dotenv import load_dotenv13from huggingface_hub import hf_hub_download14from PIL import Image15 16# --- Core Application Imports ---17from src.medical_swarm import run_medical_swarm18from src.utils import load_rag_system, standardize_query, get_standalone_question, parse_agent_response, markdown_bold_to_html19from langchain_google_genai import ChatGoogleGenerativeAI20 21# Setup logging22logging.basicConfig(level=logging.DEBUG)23logger = logging.getLogger(__name__)24 25# Load environment variables26load_dotenv()27 28# These are your "customers". You give them a key.29# In a real app, this would be in a database.30VALID_API_KEYS = {31 "amelia_key_123": "Dr. Amelia (Premium Plan, Unlimited)",32 "irfan_key_456": "Irfan (Admin, Unlimited)",33 "sistem_gelap": "Demo User (Free Tier, 10 requests)"34}35 36# --- 1. NEW HELPER FUNCTIONS TO FIX 'TypeError' ---37def hydrate_history(raw_history_list: list) -> list:38 """Converts a list of dicts from session back into LangChain Message objects."""39 history = []40 if not raw_history_list:41 return history42 for item in raw_history_list:43 if item.get('type') == 'human':44 history.append(HumanMessage(content=item.get('content', '')))45 elif item.get('type') == 'ai':46 history.append(AIMessage(content=item.get('content', '')))47 return history48 49def dehydrate_history(history_messages: list) -> list:50 """Converts LangChain Message objects into a JSON-serializable list of dicts."""51 raw_list = []52 for msg in history_messages:53 if isinstance(msg, HumanMessage):54 raw_list.append({'type': 'human', 'content': msg.content})55 elif isinstance(msg, AIMessage):56 raw_list.append({'type': 'ai', 'content': msg.content})57 return raw_list58 59# --- 2. DATABASE SETUP FUNCTION (For Deployment) ---60def setup_database():61 """Downloads and unzips the ChromaDB folder from Hugging Face Datasets."""62 DATASET_REPO_ID = "WanIrfan/atlast-db" 63 ZIP_FILENAME = "chroma_db.zip"64 DB_DIR = "chroma_db"65 if os.path.exists(DB_DIR) and os.listdir(DB_DIR):66 logger.info("✅ Database directory already exists. Skipping download.")67 return68 logger.info(f"📥 Downloading database from HF Hub: {DATASET_REPO_ID}")69 try:70 zip_path = hf_hub_download(repo_id=DATASET_REPO_ID, filename=ZIP_FILENAME, repo_type="dataset")71 logger.info(f"📦 Unzipping database from {zip_path}...")72 with zipfile.ZipFile(zip_path, 'r') as zip_ref:73 zip_ref.extractall(".")74 logger.info("✅ Database setup complete!")75 if os.path.exists(zip_path):76 os.remove(zip_path)77 except Exception as e:78 logger.error(f"❌ CRITICAL ERROR setting up database: {e}", exc_info=True)79 80# --- RUN DATABASE SETUP *BEFORE* INITIALIZING THE APP ---81setup_database()82 83# --- STANDARD FLASK APP INITIALIZATION ---84app = Flask(__name__)85app.secret_key = "a_really_strong_static_secret_key_12345" 86# --- REMOVED flask_session CONFIG ---87 88google_api_key = os.getenv("GOOGLE_API_KEY")89if not google_api_key:90 logger.warning("⚠️ GOOGLE_API_KEY not found.")91else:92 logger.info("GOOGLE_API_KEY loaded successfully.")93 94llm = ChatGoogleGenerativeAI(model="gemini-2.5-flash", temperature=0.05, google_api_key=google_api_key)95 96# --- LOAD RAG SYSTEMS (AFTER DB SETUP) ---97logger.info("🌟 Starting Multi-Domain AI Assistant...")98try:99 rag_systems = {100 'medical': load_rag_system(collection_name="medical_csv_Agentic_retrieval", domain="medical"),101 'islamic': load_rag_system(collection_name="islamic_texts_Agentic_retrieval", domain="islamic"),102 'insurance': load_rag_system(collection_name="etiqa_Agentic_retrieval", domain="insurance")103 }104except Exception as e:105 logger.error(f"❌ FAILED to load RAG systems. Error: {e}", exc_info=True)106 rag_systems = {'medical': None, 'islamic': None, 'insurance': None}107 108app.rag_systems = rag_systems109app.llm = llm110 111logger.info("\n📊 SYSTEM STATUS:")112for domain, system in rag_systems.items():113 status = "✅ Ready" if system else "❌ Failed (DB missing?)" 114 logger.info(f" {domain}: {status}")115 116# --- FLASK WEB UI ROUTES ---117@app.route("/")118def homePage():119 session.clear() # Clear all keys120 return render_template("homePage.html")121 122# --- MEDICAL PAGE ---123@app.route("/medical", methods=["GET", "POST"])124def medical_page():125 if request.method == "GET":126 latest_response = session.pop('latest_medical_response', {}) 127 return render_template("medical_page.html", 128 history=session.get('medical_history', []),129 answer=latest_response.get('answer', ""),130 thoughts=latest_response.get('thoughts', ""),131 validation=latest_response.get('validation', ""),132 source=latest_response.get('source', ""))133 134 answer, thoughts, validation, source = "", "", "", ""135 raw_history_list = session.get('medical_history', [])136 history_for_agent = hydrate_history(raw_history_list)137 current_medical_document = session.get('current_medical_document', "")138 query = ""139 140 try:141 query=standardize_query(request.form.get("query", ""))142 has_image = 'image' in request.files and request.files['image'].filename143 has_document = 'document' in request.files and request.files['document'].filename144 145 if not (query or has_image or has_document):146 raise ValueError("No query or file provided.")147 148 if has_document:149 logger.info("Processing Document with Medical Swarm")150 file = request.files['document']151 document_text = file.read().decode("utf-8")152 session['current_medical_document'] = document_text153 current_medical_document = document_text154 swarm_answer = run_medical_swarm(current_medical_document, query)155 answer = markdown_bold_to_html(swarm_answer)156 thoughts = "Swarm analysis complete."157 validation = (True, "Swarm output generated.")158 source = "Medical Swarm"159 history_for_agent.append(HumanMessage(content=f"[Document Uploaded] Query: '{query}'"))160 history_for_agent.append(AIMessage(content=answer))161 162 elif has_image :163 logger.info("Processing Multimodal RAG: Query + Image")164 file = request.files['image']165 upload_dir = "Uploads"166 os.makedirs(upload_dir, exist_ok=True)167 image_path = os.path.join(upload_dir, file.filename)168 try:169 file.save(image_path); file.close()170 with open(image_path, "rb") as img_file:171 img_data = base64.b64encode(img_file.read()).decode("utf-8")172 vision_prompt = f"Analyze image. Query: '{query}'"173 message = HumanMessage(content=[{"type": "text", "text": vision_prompt}, {"type": "image_url", "image_url": f"data:image/jpeg;base64,{img_data}"}])174 visual_prediction = llm.invoke([message]).content175 enhanced_query = (f'User Query: "{query}" Context from Image: "{visual_prediction}"')176 agent = rag_systems['medical']177 if not agent: raise Exception("Medical RAG system not loaded.")178 response_dict = agent.answer(enhanced_query, chat_history=history_for_agent)179 answer, thoughts, validation, source = parse_agent_response(response_dict)180 history_for_agent.append(HumanMessage(content=query + " [Image Attached]"))181 history_for_agent.append(AIMessage(content=answer))182 finally:183 if os.path.exists(image_path):184 try: os.remove(image_path)185 except Exception as e: logger.warning(f"Could not remove {image_path}. Error: {e}")186 187 elif query:188 history_doc_context = history_for_agent189 if current_medical_document:190 history_doc_context = [HumanMessage(content=f"Document Context:\n{current_medical_document}")] + history_for_agent191 else:192 logger.info("Processing Text RAG query for Medical domain")193 194 standalone_query = get_standalone_question(query, history_doc_context, llm)195 logger.info(f"Standalone Query : {standalone_query}")196 agent = rag_systems['medical']197 if not agent: raise Exception("Medical RAG system not loaded.")198 response_dict = agent.answer(standalone_query, chat_history=history_doc_context)199 answer, thoughts, validation, source = parse_agent_response(response_dict)200 history_for_agent.append(HumanMessage(content=query))201 history_for_agent.append(AIMessage(content=answer))202 203 except Exception as e:204 logger.error(f"Error on /medical page: {e}", exc_info=True)205 answer = f"An error occurred: {e}"206 thoughts = traceback.format_exc()207 validation = (False, "Exception")208 source = "Application Error"209 history_for_agent.append(HumanMessage(content=query if query else "Failed request"))210 history_for_agent.append(AIMessage(content=answer))211 212 session['medical_history'] = dehydrate_history(history_for_agent)213 session['latest_medical_response'] = {'answer': answer, 'thoughts': thoughts, 'validation': validation, 'source': source}214 session.modified = True215 216 logger.info(f"DEBUG: Saving to session: ANSWER='{answer[:50]}...'")217 return redirect(url_for('medical_page'))218 219@app.route("/medical/clear")220def clear_medical_chat():221 session.pop('medical_history', None)222 session.pop('current_medical_document', None)223 return redirect(url_for('medical_page'))224 225# --- ISLAMIC PAGE ---226@app.route("/islamic", methods=["GET", "POST"])227def islamic_page():228 if request.method == "GET":229 latest_response = session.pop('latest_islamic_response', {})230 return render_template("islamic_page.html",231 history=session.get('islamic_history', []),232 answer=latest_response.get('answer', ""),233 thoughts=latest_response.get('thoughts', ""),234 validation=latest_response.get('validation', ""),235 source=latest_response.get('source', ""))236 237 answer, thoughts, validation, source = "", "", "", ""238 raw_history_list = session.get('islamic_history', [])239 history_for_agent = hydrate_history(raw_history_list)240 query = ""241 try:242 query = standardize_query(request.form.get("query", ""))243 has_image = 'image' in request.files and request.files['image'].filename244 if not (query or has_image):245 raise ValueError("No query or file provided.")246 final_query = query247 248 if has_image:249 logger.info("Processing Multimodal RAG query for Islamic domain")250 file = request.files['image']251 upload_dir = "Uploads"252 os.makedirs(upload_dir, exist_ok=True)253 image_path = os.path.join(upload_dir, file.filename)254 try:255 file.save(image_path); file.close() 256 with open(image_path, "rb") as img_file:257 img_base64 = base64.b64encode(img_file.read()).decode("utf-8")258 vision_prompt = f"Analyze image. Query: '{query}'"259 message = HumanMessage(content=[{"type": "text", "text": vision_prompt}, {"type": "image_url", "image_url": f"data:image/jpeg;base64,{img_base64}"}])260 visual_prediction = llm.invoke([message]).content261 final_query = (f'User Query: "{query}" Context from Image: "{visual_prediction}"')262 finally:263 if os.path.exists(image_path):264 try: os.remove(image_path)265 except Exception as e: logger.warning(f"Could not remove {image_path}. Error: {e}")266 history_for_agent.append(HumanMessage(content=query + " [Image Attached]"))267 268 elif query:269 logger.info("Processing Text RAG query for Islamic domain")270 final_query = get_standalone_question(query, history_for_agent, llm)271 history_for_agent.append(HumanMessage(content=query))272 273 agent = rag_systems['islamic']274 if not agent: raise Exception("Islamic RAG system is not loaded.")275 response_dict = agent.answer(final_query, chat_history=history_for_agent[:-1])276 answer, thoughts, validation, source = parse_agent_response(response_dict)277 history_for_agent.append(AIMessage(content=answer))278 279 except Exception as e:280 logger.error(f"Error on /islamic page: {e}", exc_info=True)281 answer = f"An error occurred: {e}"; thoughts = traceback.format_exc(); validation = (False, "Exception"); source = "Application Error"282 if not (has_image or query): history_for_agent.append(HumanMessage(content="Failed request"))283 else: history_for_agent.append(HumanMessage(content=query))284 history_for_agent.append(AIMessage(content=answer))285 286 session['islamic_history'] = dehydrate_history(history_for_agent)287 session['latest_islamic_response'] = {'answer': answer, 'thoughts': thoughts, 'validation': validation, 'source': source}288 session.modified = True289 logger.info(f"DEBUG: Saving to session: ANSWER='{answer[:50]}...'")290 return redirect(url_for('islamic_page'))291 292@app.route("/islamic/clear")293def clear_islamic_chat():294 session.pop('islamic_history', None)295 return redirect(url_for('islamic_page'))296 297# --- INSURANCE PAGE ---298@app.route("/insurance", methods=["GET", "POST"])299def insurance_page():300 if request.method == "GET" :301 latest_response = session.pop('latest_insurance_response',{})302 return render_template("insurance_page.html",303 history=session.get('insurance_history', []),304 answer=latest_response.get('answer', ""),305 thoughts=latest_response.get('thoughts', ""),306 validation=latest_response.get('validation', ""),307 source=latest_response.get('source', ""))308 309 answer, thoughts, validation, source = "", "", "", ""310 raw_history_list = session.get('insurance_history', [])311 history_for_agent = hydrate_history(raw_history_list)312 query = ""313 try:314 query = standardize_query(request.form.get("query", ""))315 if not query:316 raise ValueError("No query provided.")317 318 standalone_query = get_standalone_question(query, history_for_agent, llm)319 agent = rag_systems['insurance']320 if not agent: raise Exception("Insurance RAG system is not loaded.")321 322 response_dict = agent.answer(standalone_query, chat_history=history_for_agent)323 answer, thoughts, validation, source = parse_agent_response(response_dict)324 history_for_agent.append(HumanMessage(content=query))325 history_for_agent.append(AIMessage(content=answer))326 327 except Exception as e:328 logger.error(f"Error on /insurance page: {e}", exc_info=True)329 answer = f"An error occurred: {e}"; thoughts = traceback.format_exc(); validation = (False, "Exception"); source = "Application Error"330 history_for_agent.append(HumanMessage(content=query))331 history_for_agent.append(AIMessage(content=answer))332 333 session['insurance_history'] = dehydrate_history(history_for_agent)334 session['latest_insurance_response'] = {'answer': answer, 'thoughts': thoughts, 'validation': validation, 'source': source}335 session.modified = True336 logger.debug(f"Redirecting after saving latest response.")337 return redirect(url_for('insurance_page'))338 339@app.route("/insurance/clear")340def clear_insurance_chat():341 session.pop('insurance_history', None)342 return redirect(url_for('insurance_page'))343 344@app.route("/about", methods=["GET"])345def about():346 return render_template("about.html")347 348# --- (Metrics routes remain unchanged) ---349@app.route('/metrics/<domain>')350def get_metrics(domain):351 try:352 if domain == "medical" and rag_systems['medical']:353 stats = rag_systems['medical'].metrics_tracker.get_stats()354 elif domain == "islamic" and rag_systems['islamic']:355 stats = rag_systems['islamic'].metrics_tracker.get_stats()356 elif domain == "insurance" and rag_systems['insurance']:357 stats = rag_systems['insurance'].metrics_tracker.get_stats()358 elif not rag_systems.get(domain):359 return jsonify({"error": f"{domain} RAG system not loaded"}), 500360 else:361 return jsonify({"error": "Invalid domain"}), 400362 return jsonify(stats)363 except Exception as e:364 return jsonify({"error": str(e)}), 500365 366@app.route('/metrics/reset/<domain>', methods=['POST'])367def reset_metrics(domain):368 try:369 if domain == "medical" and rag_systems['medical']:370 rag_systems['medical'].metrics_tracker.reset_metrics()371 elif domain == "islamic" and rag_systems['islamic']:372 rag_systems['islamic'].metrics_tracker.reset_metrics()373 elif domain == "insurance" and rag_systems['insurance']:374 rag_systems['insurance'].metrics_tracker.reset_metrics()375 elif not rag_systems.get(domain):376 return jsonify({"error": f"{domain} RAG system not loaded"}), 500377 else:378 return jsonify({"error": "Invalid domain"}), 400379 return jsonify({"success": True, "message": f"Metrics reset for {domain}"})380 except Exception as e:381 return jsonify({"error": str(e)}), 500382 383# Helper function to check API key384API_USAGE = {}385 386def check_api_key(request_data):387 api_key = request_data.get("api_key")388 389 # 1. Check if key exists390 if not api_key or api_key not in VALID_API_KEYS:391 return False, {"error": "Invalid API key"}, 401392 393 # 2. Initialize counter for this key if new394 if api_key not in API_USAGE:395 API_USAGE[api_key] = 0396 397 # 3. Check Quota (The "Selling" Logic)398 if api_key == "sistem_gelap" and API_USAGE[api_key] >= 10:399 logger.warning(f"Quota exceeded for user: {VALID_API_KEYS[api_key]}")400 return False, {"error": "Quota exceeded. Free tier is limited to 10 requests."}, 429401 402 # 4. Increment Counter403 API_USAGE[api_key] += 1404 logger.info(f"User {VALID_API_KEYS[api_key]} used {API_USAGE[api_key]} requests.")405 406 return True, None, None407 408# Helper function to save and process uploaded files (Base64)409def process_base64_file(base64_string, file_type):410 try:411 # Decode the base64 string412 file_bytes = base64.b64decode(base64_string)413 414 # Save to a temporary file415 upload_dir = "Uploads"416 os.makedirs(upload_dir, exist_ok=True)417 # Use a unique filename418 temp_filename = f"{file_type}_{int(time.time())}.tmp"419 temp_path = os.path.join(upload_dir, temp_filename)420 421 with open(temp_path, 'wb') as f:422 f.write(file_bytes)423 424 logger.info(f"Saved temporary {file_type} to {temp_path}")425 return temp_path426 except Exception as e:427 logger.error(f"Error decoding/saving base64 file: {e}")428 return None429# --- 3. NEW API-ONLY ROUTES ---430 431@app.route("/api/medical", methods=["POST"])432def medical_api():433 try:434 data = request.json435 is_valid, error_response, status_code = check_api_key(data)436 if not is_valid:437 return jsonify(error_response), status_code438 439 query = data.get("query")440 if not query:441 return jsonify({"error": "No query provided"}), 400442 443 # Hydrate history from the JSON payload444 raw_history = data.get("history", [])445 history_for_agent = hydrate_history(raw_history)446 447 agent = rag_systems['medical']448 if not agent:449 return jsonify({"error": "Medical RAG system not loaded"}), 500450 451 # --- Handle File Uploads (Base64) ---452 enhanced_query = query453 temp_file_path = None454 455 if data.get("document_base64"):456 logger.info("API: Processing base64 document for Swarm")457 doc_text = base64.b64decode(data.get("document_base64")).decode('utf-8')458 swarm_answer = run_medical_swarm(doc_text, query)459 response_dict = {460 "answer": markdown_bold_to_html(swarm_answer),461 "thoughts": "Swarm analysis complete.",462 "validation": (True, "Swarm output generated."),463 "source": "Medical Swarm",464 "response_time": 0 # Not tracked for swarm in this path465 }466 return jsonify(response_dict)467 468 elif data.get("image_base64"):469 logger.info("API: Processing base64 image")470 temp_file_path = process_base64_file(data.get("image_base64"), "image")471 if not temp_file_path:472 return jsonify({"error": "Invalid base64 image data"}), 400473 474 with open(temp_file_path, "rb") as img_file:475 img_data = base64.b64encode(img_file.read()).decode("utf-8")476 477 vision_prompt = f"Analyze image. Query: '{query}'"478 message = HumanMessage(content=[{"type": "text", "text": vision_prompt}, {"type": "image_url", "image_url": f"data:image/jpeg;base64,{img_data}"}])479 visual_prediction = llm.invoke([message]).content480 enhanced_query = (f'User Query: "{query}" Context from Image: "{visual_prediction}"')481 482 # Run the agent483 response_dict = agent.answer(enhanced_query, chat_history=history_for_agent)484 485 # Clean up temp file486 if temp_file_path and os.path.exists(temp_file_path):487 os.remove(temp_file_path)488 489 # Return the full, clean JSON response490 return jsonify(response_dict)491 492 except Exception as e:493 logger.error(f"Error on /api/medical: {e}", exc_info=True)494 return jsonify({"error": str(e)}), 500495 496@app.route("/api/islamic", methods=["POST"])497def islamic_api():498 try:499 data = request.json500 is_valid, error_response, status_code = check_api_key(data)501 if not is_valid: return jsonify(error_response), status_code502 503 query = data.get("query")504 if not query: return jsonify({"error": "No query provided"}), 400505 506 raw_history = data.get("history", [])507 history_for_agent = hydrate_history(raw_history)508 509 agent = rag_systems['islamic']510 if not agent: return jsonify({"error": "Islamic RAG system not loaded"}), 500511 512 enhanced_query = query513 temp_file_path = None514 515 if data.get("image_base64"):516 logger.info("API: Processing base64 image")517 temp_file_path = process_base64_file(data.get("image_base64"), "image")518 if not temp_file_path:519 return jsonify({"error": "Invalid base64 image data"}), 400520 521 with open(temp_file_path, "rb") as img_file:522 img_data = base64.b64encode(img_file.read()).decode("utf-8")523 524 vision_prompt = f"Analyze image. Query: '{query}'"525 message = HumanMessage(content=[{"type": "text", "text": vision_prompt}, {"type": "image_url", "image_url": f"data:image/jpeg;base64,{img_data}"}])526 visual_prediction = llm.invoke([message]).content527 enhanced_query = (f'User Query: "{query}" Context from Image: "{visual_prediction}"')528 529 response_dict = agent.answer(enhanced_query, chat_history=history_for_agent)530 531 if temp_file_path and os.path.exists(temp_file_path):532 os.remove(temp_file_path)533 534 return jsonify(response_dict)535 536 except Exception as e:537 logger.error(f"Error on /api/islamic: {e}", exc_info=True)538 return jsonify({"error": str(e)}), 500539 540@app.route("/api/insurance", methods=["POST"])541def insurance_api():542 try:543 data = request.json544 is_valid, error_response, status_code = check_api_key(data)545 if not is_valid: return jsonify(error_response), status_code546 547 query = data.get("query")548 if not query: return jsonify({"error": "No query provided"}), 400549 550 raw_history = data.get("history", [])551 history_for_agent = hydrate_history(raw_history)552 553 agent = rag_systems['insurance']554 if not agent: return jsonify({"error": "Insurance RAG system not loaded"}), 500555 556 response_dict = agent.answer(query, chat_history=history_for_agent)557 return jsonify(response_dict)558 559 except Exception as e:560 logger.error(f"Error on /api/insurance: {e}", exc_info=True)561 return jsonify({"error": str(e)}), 500562 563if __name__ == "__main__":564 logger.info("Starting Flask app for deployment testing...")565 app.run(host="0.0.0.0", port=7860, debug=False)