Update disease.py

disease.py

@@ -1,145 +1,142 @@
-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 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 = """
-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 the Google Generative AI SDK.
-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.
-    model = genai.GenerativeModel(
-        'gemini-1.5-flash',
-        system_instruction=SYSTEM_INSTRUCTION
-    )
-except Exception as e:
-    print(f"Error configuring Google Generative AI: {e}")
-    print("Please make sure your GOOGLE_API_KEY environment variable is set.")
-    model = None
-
-# --- Flask App Initialization ---
-app = Flask(__name__)
-
-
-# --- Helper Functions (Tool Implementation) ---
-
-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")]
-
-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.
-    """
-    print(f"--- TOOL EXECUTION: Running get_disease_fact_sheet(disease_name='{disease_name}') ---")
-    # Convert the friendly name back to a 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}
-    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 ---
-
-@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.
-    """
-    if not model:
-        return jsonify({"error": "Gemini client is not configured. Check your API key."}), 500
-
-    data = request.get_json()
-    if not data or 'query' not in data:
-        return jsonify({"error": "Request must be JSON and contain a 'query' field."}), 400
-
-    user_query = data['query']
-    print(f"\n=================================================")
-    print(f"Received new query: '{user_query}'")
-    print(f"=================================================")
-
-
-    available_diseases = get_available_diseases()
-    if not available_diseases:
-        return jsonify({"error": f"No fact sheets found in the '{FACT_SHEET_DIR}' directory."}), 500
-
-    # === Orchestration with Gemini ===
-    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.
-        chat = model.start_chat(enable_automatic_function_calling=True)
-
-        # Construct a more informative prompt for the model.
-        prompt = f"""
-        Here is the user's question: '{user_query}'
-
-        Please use your tools to answer it. The available diseases you can look up are:
-        {', '.join(available_diseases)}
-        """
-
-        print("--- Sending request to Gemini... ---")
-        # Send the user's query and the definitions of the available tools
-        response = chat.send_message(
-            prompt,
-            tools=[get_disease_fact_sheet] # Pass the actual function reference
-        )
-
-        final_answer = response.text
-        print(f"--- Gemini's Final Answer: ---\n{final_answer}\n")
-        return jsonify({"response": final_answer})
-
-    except Exception as e:
-        print(f"--- An unexpected error occurred: {e} ---")
-        return jsonify({"error": f"Gemini API Error: {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.")
-    else:
-        print("Starting Flask server...")
-        print(f"Fact sheets loaded for: {', '.join(get_available_diseases())}")
+import os
+import json
+from flask import Flask, request, jsonify
+import google.generativeai as genai
+
+# --- 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 = """
+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 the Google Generative AI SDK.
+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.
+    model = genai.GenerativeModel(
+        'gemini-1.5-flash',
+        system_instruction=SYSTEM_INSTRUCTION
+    )
+except Exception as e:
+    print(f"Error configuring Google Generative AI: {e}")
+    print("Please make sure your GOOGLE_API_KEY environment variable is set.")
+    model = None
+
+# --- Flask App Initialization ---
+app = Flask(__name__)
+
+
+# --- Helper Functions (Tool Implementation) ---
+
+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")]
+
+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.
+    """
+    print(f"--- TOOL EXECUTION: Running get_disease_fact_sheet(disease_name='{disease_name}') ---")
+    # Convert the friendly name back to a 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}
+    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 ---
+
+@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.
+    """
+    if not model:
+        return jsonify({"error": "Gemini client is not configured. Check your API key."}), 500
+
+    data = request.get_json()
+    if not data or 'query' not in data:
+        return jsonify({"error": "Request must be JSON and contain a 'query' field."}), 400
+
+    user_query = data['query']
+    print(f"\n=================================================")
+    print(f"Received new query: '{user_query}'")
+    print(f"=================================================")
+
+
+    available_diseases = get_available_diseases()
+    if not available_diseases:
+        return jsonify({"error": f"No fact sheets found in the '{FACT_SHEET_DIR}' directory."}), 500
+
+    # === Orchestration with Gemini ===
+    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.
+        chat = model.start_chat(enable_automatic_function_calling=True)
+
+        # Construct a more informative prompt for the model.
+        prompt = f"""
+        Here is the user's question: '{user_query}'
+
+        Please use your tools to answer it. The available diseases you can look up are:
+        {', '.join(available_diseases)}
+        """
+
+        print("--- Sending request to Gemini... ---")
+        # Send the user's query and the definitions of the available tools
+        response = chat.send_message(
+            prompt,
+            tools=[get_disease_fact_sheet] # Pass the actual function reference
+        )
+
+        final_answer = response.text
+        print(f"--- Gemini's Final Answer: ---\n{final_answer}\n")
+        return jsonify({"response": final_answer})
+
+    except Exception as e:
+        print(f"--- An unexpected error occurred: {e} ---")
+        return jsonify({"error": f"Gemini API Error: {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.")
+    else:
+        print("Starting Flask server...")
+        print(f"Fact sheets loaded for: {', '.join(get_available_diseases())}")
         app.run(debug=True, port=5001)