Update disease.py

disease.py

@@ -1,3 +1,5 @@
+# disease.py - FIXED VERSION
+
 import os
 import json
 from flask import Flask, request, jsonify
@@ -9,9 +11,7 @@ dotenv.load_dotenv()
 # --- Configuration ---
 FACT_SHEET_DIR = "Text_Files"
 
-# --- System Instruction for the Gemini Model ---
-# This instruction guides the model's behavior, ensuring it stays on task
-# and uses only the tools and information we provide.
+# --- 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
@@ -27,120 +27,220 @@ Follow these rules strictly:
 6. Keep responses clear, short, and simple. Don't mention the source of the information.
 """
 
-# Configure the Google Generative AI SDK.
+# Configure Gemini
 try:
-    # This will automatically look for the GOOGLE_API_KEY environment variable.
-    genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))
-
-    # Initialize the Gemini model with the system instruction.
-    # We recommend using a model that is highly optimized for tool use.
+    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(
-        'gemini-1.5-flash',
+        'gemini-2.0-flash-exp',
         system_instruction=SYSTEM_INSTRUCTION
     )
+    print("✅ Gemini AI configured successfully")
 except Exception as e:
-    print(f"Error configuring Google Generative AI: {e}")
-    print("Please make sure your GOOGLE_API_KEY environment variable is set.")
+    print(f"❌ Error configuring Google Generative AI: {e}")
+    print("Please ensure GOOGLE_API_KEY is set in your .env file")
     model = None
 
-# --- Flask App Initialization ---
+# Flask App
 app = Flask(__name__)
 
+# Create fact sheet directory if it doesn't exist
+os.makedirs(FACT_SHEET_DIR, exist_ok=True)
 
-# --- Helper Functions (Tool Implementation) ---
-
+# --- Helper Functions ---
 def get_available_diseases():
     """Scans the directory for available disease fact sheets."""
     if not os.path.isdir(FACT_SHEET_DIR):
         return []
-    # Create a clean list of names from filenames (e.g., "Chickenpox_and_Shingles.txt" -> "Chickenpox and Shingles")
-    return [os.path.splitext(f)[0].replace('_', ' ') for f in os.listdir(FACT_SHEET_DIR) if f.endswith(".txt")]
+    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):
     """
-    This is the actual Python function that gets executed by the model.
-    It reads the content of a specific disease's text file from the local directory.
+    Reads the content of a specific disease's text file.
     """
-    print(f"--- TOOL EXECUTION: Running get_disease_fact_sheet(disease_name='{disease_name}') ---")
-    # Convert the friendly name back to a filename format
+    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):
-        with open(filepath, 'r', encoding='utf-8') as f:
-            content = f.read()
-        print(f"--- SUCCESS: Found and read '{filename}' ---")
-        # Return a dictionary, which will be implicitly handled by the Gemini SDK
-        return {"disease": disease_name, "content": content}
+        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"--- ERROR: Fact sheet not found for '{disease_name}' ---")
-        return {"error": f"Fact sheet not found for the disease: {disease_name}."}
-
-
-# --- Main API Endpoint ---
+        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():
     """
-    Handles user queries by orchestrating the interaction with the Gemini model,
-    which is guided by the system instruction to use the provided tools and context.
+    Handle disease-related queries.
+    Expects JSON with 'query' field.
     """
     if not model:
-        return jsonify({"error": "Gemini client is not configured. Check your API key."}), 500
-
-    data = request.get_json()
-    form_data = request.form
-    if not form_data and 'query' not in data:
-        return jsonify({"error": "Request must be JSON and contain a 'query' field."}), 400
-
-    user_query = data['query'] or form_data['query']
-    print(f"\n=================================================")
-    print(f"Received new query: '{user_query}'")
-    print(f"=================================================")
-
+        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:
-        return jsonify({"error": f"No fact sheets found in the '{FACT_SHEET_DIR}' directory."}), 500
+        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")
 
-    # === Orchestration with Gemini ===
+    # Generate response
     try:
-        # We start a chat session. The model will automatically handle calling the
-        # function and using its output to generate a final answer, thanks to
-        # automatic function calling and the system instruction.
+        # Start chat with automatic function calling
         chat = model.start_chat(enable_automatic_function_calling=True)
 
-        # Construct a more informative prompt for the model.
+        # Create prompt
         prompt = f"""
-        Here is the user's question: '{user_query}'
+        User question: '{user_query}'
 
-        Please use your tools to answer it. The available diseases you can look up are:
-        {', '.join(available_diseases)}
+        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 request to Gemini... ---")
-        # Send the user's query and the definitions of the available tools
+        print("🤖 Sending to Gemini AI...")
+        
+        # Send message with tool
         response = chat.send_message(
             prompt,
-            tools=[get_disease_fact_sheet] # Pass the actual function reference
+            tools=[get_disease_fact_sheet]
         )
 
         final_answer = response.text
-        print(f"--- Gemini's Final Answer: ---\n{final_answer}\n")
-        return jsonify({"response": final_answer})
+        print(f"✅ Response generated")
+        
+        return jsonify({
+            "status": "success",
+            "response": final_answer,
+            "available_diseases": len(available_diseases)
+        })
 
     except Exception as e:
-        print(f"--- An unexpected error occurred: {e} ---")
-        return jsonify({"error": f"Gemini API Error: {e}"}), 500
-
+        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
 
-# --- To run the app ---
 if __name__ == '__main__':
-    # Make sure the Text_Files directory exists before starting
-    if not os.path.isdir(FACT_SHEET_DIR):
-        print(f"CRITICAL ERROR: The directory '{FACT_SHEET_DIR}' does not exist.")
-        print("Please create it and populate it with the disease .txt files.")
+    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("Starting Flask server...")
-        print(f"Fact sheets loaded for: {', '.join(get_available_diseases())}")
-        app.run(debug=True, port=5001)
+        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)