Update main.py

main.py

@@ -1,252 +1,816 @@
-# main.py - FIXED VERSION
+# integrated_medical_system.py - Production Ready Single File
+"""
+Integrated Medical Information System
+Combines disease queries, medicine information, and image analysis
+"""
 
 import os
+import io
 import uuid
 import json
-import requests
+import time
+from datetime import datetime, timedelta
 from flask import Flask, request, jsonify, send_from_directory
 from werkzeug.utils import secure_filename
+from PIL import Image
 import google.generativeai as genai
 from dotenv import load_dotenv
 from flask_cors import CORS
+import logging
+from functools import wraps
+import threading
 
-# Step 1: API Key and Environment Setup
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
+
+# Load environment variables
 load_dotenv()
+
+# Configuration
+CONFIG = {
+    'GOOGLE_API_KEY': os.getenv("GOOGLE_API_KEY"),
+    'MODEL': os.getenv("MODEL", "gemini-2.0-flash-exp"),
+    'DISEASE_FACT_SHEETS_DIR': "Text_Files",
+    'MEDICINE_KNOWLEDGE_DIR': "MEDICINE_TXT",
+    'ALLOWED_EXTENSIONS': {'png', 'jpg', 'jpeg', 'gif', 'bmp', 'webp', 'txt'},
+    'MAX_FILE_SIZE': 16 * 1024 * 1024,  # 16MB
+    'SESSION_TIMEOUT': 300,  # 5 minutes
+    'RATE_LIMIT_PER_MINUTE': 60
+}
+
+# Validate API key
+if not CONFIG['GOOGLE_API_KEY']:
+    logger.error("GOOGLE_API_KEY not found in environment variables")
+    raise ValueError("GOOGLE_API_KEY is required. Please set it in your .env file")
+
+# Configure Gemini AI
 try:
-    genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))
-except TypeError:
-    print("ERROR: Google API Key not found.")
-    print("Create a .env file and add 'GOOGLE_API_KEY=your_key_here'")
-    exit()
+    genai.configure(api_key=CONFIG['GOOGLE_API_KEY'])
+    logger.info("✅ Gemini AI configured successfully")
+except Exception as e:
+    logger.error(f"❌ Error configuring Gemini AI: {e}")
+    raise
 
+# Initialize Flask app
 app = Flask(__name__, static_folder='.', static_url_path='')
 CORS(app)
+app.config['MAX_CONTENT_LENGTH'] = CONFIG['MAX_FILE_SIZE']
 
-# Serve static files
-@app.route('/<path:path>')
-def serve_static(path):
-    return send_from_directory('.', path)
+# Global storage
+SESSIONS = {}
+RATE_LIMITS = {}
 
-@app.route('/')
-def serve_index():
-    return send_from_directory('.', 'index.html')
-
-# Step 2: FIXED API Endpoints
-API_ENDPOINTS = {
-    "skin_disease": "http://localhost:5003/api/skin",  # Update when ready
-    "medicine_info": "http://localhost:5002/api/query",
-    "report_reading": "http://localhost:5004/api/report",  # Update when ready
-    "disease_query": "http://localhost:5001/ask"
+# Create necessary directories
+for directory in [CONFIG['DISEASE_FACT_SHEETS_DIR'], CONFIG['MEDICINE_KNOWLEDGE_DIR']]:
+    os.makedirs(directory, exist_ok=True)
+
+# System Instructions
+SYSTEM_INSTRUCTIONS = {
+    'disease': """
+    You are a helpful Health Fact Sheet Assistant. Answer questions about diseases 
+    based on provided fact sheets or general medical knowledge. 
+    - Use fact sheet content when available
+    - Provide clear, accurate medical information
+    - Include disclaimers when using general knowledge
+    - Keep responses concise and helpful
+    - Respond in the same language as the user query
+    """,
+    
+    'medicine': """
+    You are a Medicine Information Assistant. Provide accurate information about 
+    medications, their uses, dosages, and side effects.
+    - Use knowledge base when available
+    - Provide dosage and usage instructions
+    - Include important warnings and side effects
+    - Add medical disclaimers
+    - Keep responses under 300 words
+    - Respond in the same language as the user query
+    """,
+    
+    'classifier': """
+    You are a medical query classifier. Classify queries into these categories:
+    - disease_query: General questions about diseases, symptoms, causes, treatments
+    - medicine_info: Questions about medicines, drugs, medications, pills
+    - skin_disease: Questions about skin conditions, rashes, moles, visible skin issues
+    - report_reading: Questions about interpreting medical reports, lab results, test results
+    
+    Also determine if an image is essential for accurate diagnosis/analysis.
+    """
 }
 
-# Session storage
-SESSIONS = {}
+# Utility Functions
+def cleanup_expired_sessions():
+    """Remove expired sessions"""
+    current_time = time.time()
+    expired_sessions = [
+        session_id for session_id, data in SESSIONS.items()
+        if current_time - data.get('created', 0) > CONFIG['SESSION_TIMEOUT']
+    ]
+    for session_id in expired_sessions:
+        del SESSIONS[session_id]
+    if expired_sessions:
+        logger.info(f"Cleaned up {len(expired_sessions)} expired sessions")
 
-# Step 3: Query Classification Function
-def classify_query_with_gemini(query: str):
-    """Classify user query using Gemini API"""
-    model = genai.GenerativeModel(os.getenv("MODEL","gemini-2.5-flash-lite"))
+def rate_limit_check(client_ip):
+    """Simple rate limiting"""
+    current_time = time.time()
+    minute_ago = current_time - 60
     
-    prompt = f"""
-    Analyze the user's medical query and classify it into one of the following categories:
-    - skin_disease: For queries about skin conditions, rashes, moles, spots, or any visible symptoms on the skin.
-    - medicine_info: For queries about medicines (usage, side effects, identification from image).
-    - report_reading: For queries asking to interpret medical reports, lab tests, or blood work from images.
-    - disease_query: For general questions about diseases, symptoms, causes, or treatments.
-
-    Determine if an image is essential to answer the query accurately.
+    if client_ip not in RATE_LIMITS:
+        RATE_LIMITS[client_ip] = []
+    
+    # Clean old requests
+    RATE_LIMITS[client_ip] = [
+        req_time for req_time in RATE_LIMITS[client_ip] 
+        if req_time > minute_ago
+    ]
     
-    User query: "{query}"
+    # Check limit
+    if len(RATE_LIMITS[client_ip]) >= CONFIG['RATE_LIMIT_PER_MINUTE']:
+        return False
+    
+    RATE_LIMITS[client_ip].append(current_time)
+    return True
 
-    Output ONLY valid JSON with "category" (string) and "image_required" (boolean).
+def allowed_file(filename, file_type='image'):
+    """Check if file extension is allowed"""
+    if not filename or '.' not in filename:
+        return False
+    extension = filename.rsplit('.', 1)[1].lower()
     
-    Example 1:
-    Query: "what are the symptoms of typhoid"
-    Output: {{"category": "disease_query", "image_required": false}}
+    if file_type == 'image':
+        return extension in {'png', 'jpg', 'jpeg', 'gif', 'bmp', 'webp'}
+    elif file_type == 'text':
+        return extension == 'txt'
+    return extension in CONFIG['ALLOWED_EXTENSIONS']
 
-    Example 2:
-    Query: "I have a red circular rash on my arm, what is it?"
-    Output: {{"category": "skin_disease", "image_required": true}}
+def safe_gemini_call(model_name, prompt, image=None, max_retries=3):
+    """Make safe Gemini API calls with retries"""
+    for attempt in range(max_retries):
+        try:
+            model = genai.GenerativeModel(model_name)
+            if image:
+                response = model.generate_content([prompt, image])
+            else:
+                response = model.generate_content(prompt)
+            return response.text.strip()
+        except Exception as e:
+            logger.warning(f"Gemini API attempt {attempt + 1} failed: {e}")
+            if attempt == max_retries - 1:
+                raise
+            time.sleep(1 * (attempt + 1))  # Exponential backoff
 
-    Example 3:
-    Query: "Can you tell me what this lab report says?"
-    Output: {{"category": "report_reading", "image_required": true}}
+# Core Functions
+def get_available_diseases():
+    """Get list of available disease fact sheets"""
+    try:
+        if not os.path.isdir(CONFIG['DISEASE_FACT_SHEETS_DIR']):
+            return []
+        return [
+            os.path.splitext(f)[0].replace('_', ' ')
+            for f in os.listdir(CONFIG['DISEASE_FACT_SHEETS_DIR'])
+            if f.endswith('.txt')
+        ]
+    except Exception as e:
+        logger.error(f"Error getting available diseases: {e}")
+        return []
 
-    Example 4:
-    Query: "What is this white pill with 'IP 204' written on it?"
-    Output: {{"category": "medicine_info", "image_required": true}}
-    """
+def get_disease_fact_sheet(disease_name):
+    """Read disease fact sheet content"""
+    try:
+        filename = disease_name.replace(' ', '_') + '.txt'
+        filepath = os.path.join(CONFIG['DISEASE_FACT_SHEETS_DIR'], filename)
+        
+        if os.path.exists(filepath):
+            with open(filepath, 'r', encoding='utf-8') as f:
+                content = f.read()
+            logger.info(f"Found fact sheet for: {disease_name}")
+            return {"disease": disease_name, "content": content}
+        else:
+            logger.warning(f"No fact sheet found for: {disease_name}")
+            return {"error": f"Fact sheet not found for: {disease_name}"}
+    except Exception as e:
+        logger.error(f"Error reading fact sheet for {disease_name}: {e}")
+        return {"error": f"Error reading fact sheet: {str(e)}"}
+
+def get_available_medicine_files():
+    """Get available medicine knowledge base files"""
+    try:
+        if not os.path.isdir(CONFIG['MEDICINE_KNOWLEDGE_DIR']):
+            return []
+        return [
+            f for f in os.listdir(CONFIG['MEDICINE_KNOWLEDGE_DIR'])
+            if f.endswith('.txt')
+        ]
+    except Exception as e:
+        logger.error(f"Error getting medicine files: {e}")
+        return []
+
+def find_relevant_medicine_file(topic):
+    """Find most relevant medicine file for topic"""
+    available_files = get_available_medicine_files()
+    if not available_files:
+        return None
+    
+    try:
+        prompt = f"""
+        From these files, which is most relevant for: "{topic}"?
+        Respond with ONLY the filename, nothing else.
+        Files: {', '.join(available_files)}
+        """
+        response = safe_gemini_call(CONFIG['MODEL'], prompt)
+        filename = response.replace('`', '').replace('"', '').strip()
+        
+        return filename if filename in available_files else None
+    except Exception as e:
+        logger.error(f"Error finding relevant medicine file: {e}")
+        return None
+
+def get_medicine_context(filename):
+    """Read medicine knowledge base file"""
+    if not filename:
+        return None
     
     try:
-        response = model.generate_content(prompt)
-        cleaned_text = response.text.strip().replace('```json', '').replace('```', '')
+        filepath = os.path.join(CONFIG['MEDICINE_KNOWLEDGE_DIR'], filename)
+        with open(filepath, 'r', encoding='utf-8') as f:
+            return f.read()
+    except Exception as e:
+        logger.error(f"Error reading medicine file {filename}: {e}")
+        return None
+
+def classify_query(query):
+    """Classify user query using Gemini"""
+    try:
+        prompt = f"""
+        {SYSTEM_INSTRUCTIONS['classifier']}
+        
+        User query: "{query}"
+        
+        Output ONLY valid JSON with "category" (string) and "image_required" (boolean).
+        
+        Examples:
+        {{"category": "disease_query", "image_required": false}}
+        {{"category": "medicine_info", "image_required": true}}
+        {{"category": "skin_disease", "image_required": true}}
+        {{"category": "report_reading", "image_required": true}}
+        """
+        
+        response = safe_gemini_call(CONFIG['MODEL'], prompt)
+        cleaned_text = response.replace('```json', '').replace('```', '').strip()
         result = json.loads(cleaned_text)
-        print(f"Classification result: {result}")
+        
+        # Validate result
+        valid_categories = ['disease_query', 'medicine_info', 'skin_disease', 'report_reading']
+        if result.get('category') not in valid_categories:
+            result['category'] = 'disease_query'
+        
+        logger.info(f"Query classified as: {result}")
         return result
+        
     except Exception as e:
-        print(f"Classification error: {e}")
-        # Default fallback
+        logger.error(f"Classification error: {e}")
         return {"category": "disease_query", "image_required": False}
 
-# Step 4: API Routes
+def process_disease_query(query):
+    """Process disease-related queries"""
+    try:
+        available_diseases = get_available_diseases()
+        
+        # Try to find relevant disease in query
+        relevant_disease = None
+        for disease in available_diseases:
+            if disease.lower() in query.lower():
+                relevant_disease = disease
+                break
+        
+        context = ""
+        source = "general_knowledge"
+        
+        if relevant_disease:
+            fact_sheet = get_disease_fact_sheet(relevant_disease)
+            if 'content' in fact_sheet:
+                context = f"FACT SHEET FOR {relevant_disease}:\n{fact_sheet['content']}"
+                source = f"fact_sheet_{relevant_disease}"
+        
+        prompt = f"""
+        {SYSTEM_INSTRUCTIONS['disease']}
+        
+        {context}
+        
+        User question: "{query}"
+        
+        Available diseases with fact sheets: {', '.join(available_diseases) if available_diseases else 'None'}
+        
+        Provide a helpful, accurate response. If using general knowledge, include appropriate disclaimers.
+        """
+        
+        response = safe_gemini_call(CONFIG['MODEL'], prompt)
+        
+        return {
+            "status": "success",
+            "response": response,
+            "source": source,
+            "available_diseases": len(available_diseases)
+        }
+        
+    except Exception as e:
+        logger.error(f"Error processing disease query: {e}")
+        return {"error": f"Failed to process disease query: {str(e)}"}
+
+def process_medicine_query(query, image=None):
+    """Process medicine-related queries"""
+    try:
+        medicine_topic = None
+        
+        # If image provided, analyze it first
+        if image:
+            vision_prompt = """
+            Identify the medicine from this image. Look for:
+            - Medicine name or brand
+            - Active ingredients  
+            - Any text on packaging or pills
+            Respond with just the medicine name or main component.
+            """
+            medicine_topic = safe_gemini_call(CONFIG['MODEL'], vision_prompt, image)
+            logger.info(f"Medicine identified from image: {medicine_topic}")
+        
+        # If no medicine from image, extract from query
+        if not medicine_topic:
+            extract_prompt = f"""
+            From this query: "{query}"
+            Extract the main medicine or medical topic being asked about.
+            Respond with ONLY the medicine/topic name.
+            """
+            medicine_topic = safe_gemini_call(CONFIG['MODEL'], extract_prompt)
+        
+        # Find relevant knowledge base file
+        context = None
+        source_file = find_relevant_medicine_file(medicine_topic)
+        if source_file:
+            context = get_medicine_context(source_file)
+        
+        # Build prompt
+        if context:
+            prompt = f"""
+            {SYSTEM_INSTRUCTIONS['medicine']}
+            
+            KNOWLEDGE BASE CONTEXT:
+            {context}
+            
+            IDENTIFIED MEDICINE/TOPIC: {medicine_topic}
+            USER QUESTION: {query}
+            
+            Answer based on the knowledge base context when available.
+            """
+        else:
+            prompt = f"""
+            {SYSTEM_INSTRUCTIONS['medicine']}
+            
+            MEDICINE/TOPIC: {medicine_topic}
+            USER QUESTION: {query}
+            
+            Provide accurate medical information about this medicine/topic.
+            """
+        
+        response = safe_gemini_call(CONFIG['MODEL'], prompt)
+        
+        return {
+            "status": "success",
+            "response": response,
+            "identified_topic": medicine_topic,
+            "source_file": source_file or "general_knowledge",
+            "knowledge_base_files": len(get_available_medicine_files())
+        }
+        
+    except Exception as e:
+        logger.error(f"Error processing medicine query: {e}")
+        return {"error": f"Failed to process medicine query: {str(e)}"}
+
+def process_skin_disease_query(query, image=None):
+    """Process skin disease queries (placeholder for future implementation)"""
+    try:
+        if image:
+            prompt = f"""
+            You are a dermatology assistant. Analyze this skin image and the user's query: "{query}"
+            
+            Provide information about possible skin conditions, but always include:
+            - This is not a medical diagnosis
+            - Recommend seeing a dermatologist
+            - General skin care advice
+            
+            Keep response under 250 words.
+            """
+            response = safe_gemini_call(CONFIG['MODEL'], prompt, image)
+        else:
+            response = f"""
+            Regarding your skin concern: "{query}"
+            
+            For accurate diagnosis of skin conditions, a visual examination is usually necessary. 
+            I recommend:
+            - Consulting with a dermatologist
+            - Taking clear photos in good lighting if seeking online consultation
+            - Noting any changes, symptoms, or triggers
+            
+            If this is urgent (rapid changes, pain, bleeding), please seek immediate medical attention.
+            """
+        
+        return {
+            "status": "success", 
+            "response": response,
+            "category": "skin_disease",
+            "disclaimer": "This is not medical diagnosis. Consult a dermatologist."
+        }
+        
+    except Exception as e:
+        logger.error(f"Error processing skin disease query: {e}")
+        return {"error": f"Failed to process skin disease query: {str(e)}"}
+
+def process_report_reading_query(query, image=None):
+    """Process medical report reading queries (placeholder for future implementation)"""
+    try:
+        if image:
+            prompt = f"""
+            You are a medical report analysis assistant. The user asks: "{query}"
+            
+            Analyze this medical report/lab result image and provide:
+            - Key findings in simple language
+            - What the values typically indicate
+            - Important notes or abnormalities
+            
+            Always include:
+            - This is not a medical interpretation
+            - Results should be discussed with healthcare provider
+            - Context and medical history are important for interpretation
+            
+            Keep response under 300 words.
+            """
+            response = safe_gemini_call(CONFIG['MODEL'], prompt, image)
+        else:
+            response = f"""
+            To help interpret medical reports or lab results, I would need to see the actual report image.
+            
+            However, please remember:
+            - Medical reports should always be discussed with your healthcare provider
+            - Lab values can vary by laboratory and individual circumstances
+            - Context, symptoms, and medical history are crucial for proper interpretation
+            
+            If you have urgent concerns about your results, contact your healthcare provider immediately.
+            """
+        
+        return {
+            "status": "success",
+            "response": response,
+            "category": "report_reading", 
+            "disclaimer": "This is not medical interpretation. Consult your healthcare provider."
+        }
+        
+    except Exception as e:
+        logger.error(f"Error processing report reading query: {e}")
+        return {"error": f"Failed to process report reading query: {str(e)}"}
+
+# Background cleanup task
+def background_cleanup():
+    """Background task to clean up expired sessions"""
+    while True:
+        try:
+            cleanup_expired_sessions()
+            time.sleep(60)  # Run every minute
+        except Exception as e:
+            logger.error(f"Cleanup task error: {e}")
+
+# Start background thread
+cleanup_thread = threading.Thread(target=background_cleanup, daemon=True)
+cleanup_thread.start()
+
+# API Routes
+@app.route('/')
+def serve_index():
+    """Serve main page"""
+    try:
+        return send_from_directory('.', 'index.html')
+    except:
+        return jsonify({
+            "service": "Integrated Medical Information System",
+            "status": "running",
+            "endpoints": {
+                "/start_session": "POST - Start new session",
+                "/process_query": "POST - Process text query", 
+                "/process_with_image": "POST - Process query with image",
+                "/upload_fact_sheet": "POST - Upload disease fact sheet",
+                "/upload_medicine_info": "POST - Upload medicine info",
+                "/health": "GET - Health check",
+                "/stats": "GET - System statistics"
+            }
+        })
+
+@app.route('/<path:path>')
+def serve_static(path):
+    """Serve static files"""
+    return send_from_directory('.', path)
+
 @app.route('/start_session', methods=['POST'])
 def start_session():
+    """Start a new session"""
+    # Rate limiting
+    client_ip = request.remote_addr
+    if not rate_limit_check(client_ip):
+        return jsonify({"error": "Rate limit exceeded"}), 429
+    
     session_id = str(uuid.uuid4())
-    SESSIONS[session_id] = {"status": "started"}
-    print(f"Session started: {session_id}")
+    SESSIONS[session_id] = {
+        "status": "started",
+        "created": time.time(),
+        "ip": client_ip
+    }
+    
+    logger.info(f"Session started: {session_id} from {client_ip}")
     return jsonify({"session_id": session_id}), 200
 
 @app.route('/process_query', methods=['POST'])
 def process_query():
-    data = request.get_json()
-    session_id = data.get('session_id')
-    query = data.get('query')
-
-    if not session_id or session_id not in SESSIONS:
-        return jsonify({"error": "Invalid or missing session_id"}), 400
-    if not query:
-        return jsonify({"error": "Query is required"}), 400
-
-    print(f"Session {session_id}: Query received: '{query}'")
-    classification = classify_query_with_gemini(query)
-    
-    SESSIONS[session_id]['classification'] = classification
-    SESSIONS[session_id]['query'] = query
-    
-    if classification.get('image_required'):
-        print(f"Session {session_id}: Image required for '{classification.get('category')}'")
-        return jsonify({
-            "status": "image_required",
-            "message": "Please upload an image via /process_with_image endpoint",
-            "category": classification.get('category')
-        }), 200
-    else:
-        # Process without image
+    """Process text-only queries"""
+    try:
+        data = request.get_json()
+        session_id = data.get('session_id')
+        query = data.get('query')
+        
+        if not session_id or session_id not in SESSIONS:
+            return jsonify({"error": "Invalid or missing session_id"}), 400
+        
+        if not query:
+            return jsonify({"error": "Query is required"}), 400
+        
+        logger.info(f"Session {session_id}: Processing query: '{query}'")
+        
+        # Classify query
+        classification = classify_query(query)
         category = classification['category']
-        endpoint_url = API_ENDPOINTS.get(category)
+        image_required = classification.get('image_required', False)
         
-        if not endpoint_url:
-            return jsonify({"error": f"No endpoint configured for category: {category}"}), 500
+        # Store classification in session
+        SESSIONS[session_id].update({
+            'classification': classification,
+            'query': query,
+            'last_activity': time.time()
+        })
         
-        try:
-            print(f"Forwarding to {category} API at {endpoint_url}")
-            
-            # FIXED: Proper request format for each service
-            if category == "disease_query":
-                # disease.py expects 'query' field
-                response = requests.post(endpoint_url, json={"query": query})
-            elif category == "medicine_info":
-                # medicine.py expects 'main_query' field
-                response = requests.post(endpoint_url, json={"main_query": query})
-            else:
-                # Default format for future services
-                response = requests.post(endpoint_url, json={"query": query})
-            
-            response.raise_for_status()
-            result = response.json()
-            
-            # Clean up session
-            del SESSIONS[session_id]
-            print(f"Session {session_id} closed successfully")
-            
+        if image_required:
             return jsonify({
-                "status": "success",
+                "status": "image_required",
+                "message": "Please upload an image for better analysis",
                 "category": category,
-                "response": result
-            })
-            
-        except requests.exceptions.RequestException as e:
-            print(f"Error calling {category} API: {e}")
-            del SESSIONS[session_id]
-            return jsonify({
-                "error": f"Failed to reach {category} service",
-                "details": str(e)
-            }), 503
+                "session_id": session_id
+            }), 200
+        
+        # Process query based on category
+        if category == 'disease_query':
+            result = process_disease_query(query)
+        elif category == 'medicine_info':
+            result = process_medicine_query(query)
+        elif category == 'skin_disease':
+            result = process_skin_disease_query(query)
+        elif category == 'report_reading':
+            result = process_report_reading_query(query)
+        else:
+            result = {"error": f"Unknown category: {category}"}
+        
+        # Clean up session
+        del SESSIONS[session_id]
+        logger.info(f"Session {session_id} completed successfully")
+        
+        result['category'] = category
+        return jsonify(result)
+        
+    except Exception as e:
+        logger.error(f"Error processing query: {e}")
+        return jsonify({"error": f"Failed to process query: {str(e)}"}), 500
 
 @app.route('/process_with_image', methods=['POST'])
 def process_with_image():
-    session_id = request.form.get('session_id')
-    
-    if not session_id or session_id not in SESSIONS:
-        return jsonify({"error": "Invalid or missing session_id"}), 400
-    
-    if 'photo' not in request.files:
-        return jsonify({"error": "No photo file found in request"}), 400
+    """Process queries with image upload"""
+    try:
+        session_id = request.form.get('session_id')
         
-    file = request.files['photo']
-    if file.filename == '':
-        return jsonify({"error": "No file selected"}), 400
-
-    # Get session data
-    query = SESSIONS[session_id].get('query')
-    classification = SESSIONS[session_id].get('classification')
-    category = classification.get('category')
-    endpoint_url = API_ENDPOINTS.get(category)
-    
-    if not endpoint_url:
+        if not session_id or session_id not in SESSIONS:
+            return jsonify({"error": "Invalid or missing session_id"}), 400
+        
+        if 'photo' not in request.files:
+            return jsonify({"error": "No photo file found"}), 400
+        
+        file = request.files['photo']
+        if file.filename == '' or not allowed_file(file.filename, 'image'):
+            return jsonify({"error": "Invalid image file"}), 400
+        
+        # Get session data
+        session = SESSIONS[session_id]
+        query = session.get('query')
+        classification = session.get('classification', {})
+        category = classification.get('category', 'disease_query')
+        
+        logger.info(f"Session {session_id}: Processing image for category '{category}'")
+        
+        # Process image
+        try:
+            image = Image.open(file.stream)
+            # Convert to RGB if needed
+            if image.mode != 'RGB':
+                image = image.convert('RGB')
+        except Exception as e:
+            logger.error(f"Error processing image: {e}")
+            return jsonify({"error": "Invalid image format"}), 400
+        
+        # Process based on category
+        if category == 'medicine_info':
+            result = process_medicine_query(query, image)
+        elif category == 'skin_disease':
+            result = process_skin_disease_query(query, image)
+        elif category == 'report_reading':
+            result = process_report_reading_query(query, image)
+        else:
+            # Fallback to disease query
+            result = process_disease_query(query)
+        
+        # Clean up session
         del SESSIONS[session_id]
-        return jsonify({"error": f"No endpoint configured for category: {category}"}), 500
-    
-    print(f"Session {session_id}: Processing image for '{category}'")
+        logger.info(f"Session {session_id} with image completed successfully")
+        
+        result['category'] = category
+        return jsonify(result)
+        
+    except Exception as e:
+        logger.error(f"Error processing query with image: {e}")
+        return jsonify({"error": f"Failed to process query with image: {str(e)}"}), 500
 
+@app.route('/upload_fact_sheet', methods=['POST'])
+def upload_fact_sheet():
+    """Upload disease fact sheet"""
     try:
-        # FIXED: Proper request format based on category
-        if category == "medicine_info":
-            # medicine.py expects 'file' and 'main_query'
-            files = {'file': (file.filename, file.stream, file.mimetype)}
-            data = {'main_query': query}
-            response = requests.post(endpoint_url, files=files, data=data)
-            
-        elif category == "disease_query":
-            # disease.py doesn't typically need images, but if it did:
-            files = {'file': (file.filename, file.stream, file.mimetype)}
-            data = {'query': query}
-            response = requests.post(endpoint_url, files=files, data=data)
-            
-        else:
-            # Default format for skin_disease, report_reading
-            files = {'file': (file.filename, file.stream, file.mimetype)}
-            data = {'query': query}
-            response = requests.post(endpoint_url, files=files, data=data)
+        if 'file' not in request.files:
+            return jsonify({"error": "No file provided"}), 400
         
-        response.raise_for_status()
-        result = response.json()
+        file = request.files['file']
+        if not file.filename or not allowed_file(file.filename, 'text'):
+            return jsonify({"error": "Invalid file. Must be a .txt file"}), 400
         
-        # Clean up session
-        del SESSIONS[session_id]
-        print(f"Session {session_id} closed successfully")
+        # Secure filename
+        filename = secure_filename(file.filename)
+        filepath = os.path.join(CONFIG['DISEASE_FACT_SHEETS_DIR'], filename)
+        
+        file.save(filepath)
+        logger.info(f"Disease fact sheet uploaded: {filename}")
         
         return jsonify({
             "status": "success",
-            "category": category,
-            "response": result
+            "message": f"Fact sheet '{filename}' uploaded successfully",
+            "total_fact_sheets": len(get_available_diseases())
         })
+        
+    except Exception as e:
+        logger.error(f"Error uploading fact sheet: {e}")
+        return jsonify({"error": f"Failed to upload: {str(e)}"}), 500
 
-    except requests.exceptions.RequestException as e:
-        print(f"Error calling {category} API with image: {e}")
-        del SESSIONS[session_id]
+@app.route('/upload_medicine_info', methods=['POST'])
+def upload_medicine_info():
+    """Upload medicine knowledge base file"""
+    try:
+        if 'file' not in request.files:
+            return jsonify({"error": "No file provided"}), 400
+        
+        file = request.files['file']
+        if not file.filename or not allowed_file(file.filename, 'text'):
+            return jsonify({"error": "Invalid file. Must be a .txt file"}), 400
+        
+        # Secure filename
+        filename = secure_filename(file.filename)
+        filepath = os.path.join(CONFIG['MEDICINE_KNOWLEDGE_DIR'], filename)
+        
+        file.save(filepath)
+        logger.info(f"Medicine info file uploaded: {filename}")
+        
         return jsonify({
-            "error": f"Failed to reach {category} service",
-            "details": str(e)
-        }), 503
+            "status": "success",
+            "message": f"Medicine info '{filename}' uploaded successfully",
+            "total_medicine_files": len(get_available_medicine_files())
+        })
+        
+    except Exception as e:
+        logger.error(f"Error uploading medicine info: {e}")
+        return jsonify({"error": f"Failed to upload: {str(e)}"}), 500
 
 @app.route('/health', methods=['GET'])
 def health_check():
-    """Health check endpoint to verify services"""
-    service_status = {}
-    
-    for service, url in API_ENDPOINTS.items():
-        try:
-            # Try to connect to each service
-            if "localhost" in url:
-                test_url = url.replace('/ask', '').replace('/api/query', '')
-                requests.get(test_url, timeout=1)
-                service_status[service] = "✅ Running"
-        except:
-            service_status[service] = "❌ Not reachable"
+    """System health check"""
+    try:
+        diseases = get_available_diseases()
+        medicine_files = get_available_medicine_files()
+        
+        return jsonify({
+            "status": "✅ Running",
+            "service": "Integrated Medical Information System",
+            "timestamp": datetime.now().isoformat(),
+            "gemini_configured": True,
+            "active_sessions": len(SESSIONS),
+            "disease_fact_sheets": len(diseases),
+            "medicine_knowledge_files": len(medicine_files),
+            "rate_limits_active": len(RATE_LIMITS),
+            "system_info": {
+                "max_file_size_mb": CONFIG['MAX_FILE_SIZE'] // (1024 * 1024),
+                "session_timeout_minutes": CONFIG['SESSION_TIMEOUT'] // 60,
+                "rate_limit_per_minute": CONFIG['RATE_LIMIT_PER_MINUTE']
+            }
+        })
+    except Exception as e:
+        logger.error(f"Health check error: {e}")
+        return jsonify({"status": "❌ Error", "error": str(e)}), 500
+
+@app.route('/stats', methods=['GET'])
+def get_stats():
+    """Get system statistics"""
+    diseases = get_available_diseases()
+    medicine_files = get_available_medicine_files()
     
     return jsonify({
-        "main_service": "✅ Running",
-        "connected_services": service_status,
-        "sessions_active": len(SESSIONS)
+        "available_diseases": diseases,
+        "available_medicine_files": medicine_files,
+        "counts": {
+            "diseases": len(diseases),
+            "medicine_files": len(medicine_files),
+            "active_sessions": len(SESSIONS),
+            "rate_limited_ips": len(RATE_LIMITS)
+        },
+        "recent_sessions": len([
+            s for s in SESSIONS.values() 
+            if time.time() - s.get('created', 0) < 300  # Last 5 minutes
+        ])
+    })
+
+@app.route('/diseases', methods=['GET'])
+def list_diseases():
+    """List available diseases"""
+    diseases = get_available_diseases()
+    return jsonify({
+        "available_diseases": diseases,
+        "count": len(diseases)
+    })
+
+@app.route('/medicines', methods=['GET'])
+def list_medicines():
+    """List available medicine files"""
+    files = get_available_medicine_files()
+    return jsonify({
+        "available_medicine_files": files,
+        "count": len(files)
     })
+
+# Error handlers
+@app.errorhandler(413)
+def too_large(e):
+    return jsonify({"error": "File too large. Maximum size is 16MB"}), 413
+
+@app.errorhandler(404)
+def not_found(e):
+    return jsonify({"error": "Endpoint not found"}), 404
+
+@app.errorhandler(500)
+def internal_error(e):
+    return jsonify({"error": "Internal server error"}), 500
+
+if __name__ == '__main__':
+    logger.info("=" * 60)
+    logger.info("🏥 Starting Integrated Medical Information System")
+    logger.info(f"📁 Disease fact sheets: {CONFIG['DISEASE_FACT_SHEETS_DIR']}/")
+    logger.info(f"💊 Medicine knowledge: {CONFIG['MEDICINE_KNOWLEDGE_DIR']}/")
+    
+    # Check available knowledge base
+    diseases = get_available_diseases()
+    medicine_files = get_available_medicine_files()
+    
+    if diseases:
+        logger.info(f"✅ {len(diseases)} disease fact sheets loaded")
+        for disease in diseases[:3]:
+            logger.info(f"  - {disease}")
+        if len(diseases) > 3:
+            logger.info(f"  ... and {len(diseases) - 3} more")
+    else:
+        logger.warning(f"⚠️ No disease fact sheets found in '{CONFIG['DISEASE_FACT_SHEETS_DIR']}'")
+    
+    if medicine_files:
+        logger.info(f"✅ {len(medicine_files)} medicine files loaded")
+        for file in medicine_files[:3]:
+            logger.info(f"  - {file}")
+        if len(medicine_files) > 3:
+            logger.info(f"  ... and {len(medicine_files) - 3} more")
+    else:
+        logger.warning(f"⚠️ No medicine files found in '{CONFIG['MEDICINE_KNOWLEDGE_DIR']}'")
+    
+    logger.info("=" * 60)
+    logger.info("🚀 Server starting on http://localhost:5000")
+    logger.info("📚 Upload .txt files to knowledge directories for enhanced responses")
+    logger.info("=" * 60)
+    
+    app.run(host='0.0.0.0', port=7860, debug=False)