disease.py

# disease.py - FIXED VERSION

import os
import json
from flask import Flask, request, jsonify
import google.generativeai as genai
import dotenv

dotenv.load_dotenv()

# --- Configuration ---
FACT_SHEET_DIR = "Text_Files"

# --- System Instruction ---
SYSTEM_INSTRUCTION = """
You are a helpful Health Fact Sheet Assistant. Your role is to answer questions
about specific diseases based ONLY on the information contained in the fact sheets
provided to you through the get_disease_fact_sheet tool. 
Generate response in same language as user query.

Follow these rules strictly:
1. Use the fact sheet for answers whenever possible.
2. To get information, call the get_disease_fact_sheet function.
3. If the fact sheet doesn't cover the answer, reply using general knowledge and include a disclaimer.
4. First, check the user's query to see which disease it refers to, then fetch that fact sheet.
5. If the query isn't about a specific disease, reply using general knowledge with a disclaimer.
6. Keep responses clear, short, and simple. Don't mention the source of the information.
"""

# Configure Gemini
try:
    api_key = os.getenv("GOOGLE_API_KEY")
    if not api_key:
        raise ValueError("GOOGLE_API_KEY not found in environment")
    
    genai.configure(api_key=api_key)
    model = genai.GenerativeModel(
        os.getenv("MODEL","gemini-2.5-flash-lite"),
        system_instruction=SYSTEM_INSTRUCTION
    )
    print("✅ Gemini AI configured successfully")
except Exception as e:
    print(f"❌ Error configuring Google Generative AI: {e}")
    print("Please ensure GOOGLE_API_KEY is set in your .env file")
    model = None

# Flask App
app = Flask(__name__)

# Create fact sheet directory if it doesn't exist
os.makedirs(FACT_SHEET_DIR, exist_ok=True)

# --- Helper Functions ---
def get_available_diseases():
    """Scans the directory for available disease fact sheets."""
    if not os.path.isdir(FACT_SHEET_DIR):
        return []
    return [
        os.path.splitext(f)[0].replace('_', ' ') 
        for f in os.listdir(FACT_SHEET_DIR) 
        if f.endswith(".txt")
    ]

def get_disease_fact_sheet(disease_name: str):
    """
    Reads the content of a specific disease's text file.
    """
    print(f"📂 Looking for fact sheet: '{disease_name}'")
    
    # Convert to filename format
    filename = disease_name.replace(' ', '_') + ".txt"
    filepath = os.path.join(FACT_SHEET_DIR, filename)

    if os.path.exists(filepath):
        try:
            with open(filepath, 'r', encoding='utf-8') as f:
                content = f.read()
            print(f"✅ Found and read '{filename}'")
            return {"disease": disease_name, "content": content}
        except Exception as e:
            print(f"❌ Error reading file: {e}")
            return {"error": f"Could not read fact sheet for {disease_name}"}
    else:
        print(f"⚠️ Fact sheet not found for '{disease_name}'")
        return {"error": f"Fact sheet not found for: {disease_name}"}

# --- API Endpoints ---
@app.route('/ask', methods=['POST'])
def ask_question():
    """
    Handle disease-related queries.
    Expects JSON with 'query' field.
    """
    if not model:
        return jsonify({
            "error": "Gemini AI not configured. Check your API key.",
            "status": "error"
        }), 500

    # Handle both JSON and form data
    if request.is_json:
        data = request.get_json()
        user_query = data.get('query')
    else:
        user_query = request.form.get('query')
    
    if not user_query:
        return jsonify({
            "error": "Missing 'query' in request",
            "status": "error"
        }), 400

    print(f"\n{'='*60}")
    print(f"📝 New query: '{user_query}'")
    print(f"{'='*60}")

    available_diseases = get_available_diseases()
    
    if not available_diseases:
        print("⚠️ No fact sheets available, using general knowledge")
        # Continue with general knowledge if no fact sheets
        available_diseases_str = "No specific disease fact sheets available"
    else:
        available_diseases_str = ', '.join(available_diseases)
        print(f"📚 Available fact sheets: {len(available_diseases)} files")

    # Generate response
    try:
        # Start chat with automatic function calling
        chat = model.start_chat(enable_automatic_function_calling=True)

        # Create prompt
        prompt = f"""
        User question: '{user_query}'

        Available disease fact sheets: {available_diseases_str}
        
        Please provide a helpful response. If a specific disease is mentioned and 
        we have a fact sheet for it, use the get_disease_fact_sheet tool to retrieve it.
        """

        print("🤖 Sending to Gemini AI...")
        
        # Send message with tool
        response = chat.send_message(
            prompt,
            tools=[get_disease_fact_sheet]
        )

        final_answer = response.text
        print(f"✅ Response generated")
        
        return jsonify({
            "status": "success",
            "response": final_answer,
            "available_diseases": len(available_diseases)
        })

    except Exception as e:
        print(f"❌ Error: {e}")
        return jsonify({
            "error": f"Failed to generate response: {str(e)}",
            "status": "error"
        }), 500

@app.route('/health', methods=['GET'])
def health_check():
    """Health check endpoint"""
    diseases = get_available_diseases()
    return jsonify({
        "status": "running",
        "service": "disease_query",
        "gemini_configured": model is not None,
        "fact_sheets_available": len(diseases),
        "available_diseases": diseases[:5] if diseases else [],
        "port": 5001
    })

@app.route('/', methods=['GET'])
def index():
    """Basic info endpoint"""
    diseases = get_available_diseases()
    return jsonify({
        "service": "Disease Information API",
        "endpoint": "/ask",
        "methods": ["POST"],
        "accepts": "JSON with 'query' field",
        "fact_sheets": f"{len(diseases)} disease fact sheets available",
        "example_request": {
            "query": "What are the symptoms of diabetes?"
        }
    })

@app.route('/diseases', methods=['GET'])
def list_diseases():
    """List all available diseases"""
    diseases = get_available_diseases()
    return jsonify({
        "available_diseases": diseases,
        "count": len(diseases)
    })

# --- File Upload Endpoint (optional) ---
@app.route('/upload_fact_sheet', methods=['POST'])
def upload_fact_sheet():
    """Upload a new disease fact sheet"""
    if 'file' not in request.files:
        return jsonify({"error": "No file provided"}), 400
    
    file = request.files['file']
    if file.filename == '' or not file.filename.endswith('.txt'):
        return jsonify({"error": "Invalid file. Must be a .txt file"}), 400
    
    try:
        # Save the file
        filename = file.filename.replace(' ', '_')
        filepath = os.path.join(FACT_SHEET_DIR, filename)
        file.save(filepath)
        
        return jsonify({
            "status": "success",
            "message": f"Fact sheet '{filename}' uploaded successfully",
            "total_fact_sheets": len(get_available_diseases())
        })
    except Exception as e:
        return jsonify({"error": f"Failed to upload: {str(e)}"}), 500

if __name__ == '__main__':
    print("="*60)
    print("Starting Disease Query Service")
    print(f"Fact sheets directory: {FACT_SHEET_DIR}")
    
    # Check and list available fact sheets
    diseases = get_available_diseases()
    if diseases:
        print(f"✅ Loaded {len(diseases)} fact sheets:")
        for disease in diseases[:5]:
            print(f"  - {disease}")
        if len(diseases) > 5:
            print(f"  ... and {len(diseases) - 5} more")
    else:
        print("⚠️ No fact sheets found. Service will use general knowledge only.")
        print(f"💡 Add .txt files to '{FACT_SHEET_DIR}/' directory for specific disease info")
    
    print("="*60)
    app.run(debug=True, port=5001)