Update medicine.py

medicine.py

@@ -1,3 +1,5 @@
+# medicine.py - FIXED VERSION
+
 import os
 import io
 from flask import Flask, request, jsonify
@@ -7,190 +9,236 @@ import google.generativeai as genai
 import json
 
 # --- INITIAL SETUP ---
-
-# Load environment variables from the .env file
 load_dotenv()
 
-# Configure the Gemini API with your key
 api_key = os.getenv("GOOGLE_API_KEY")
 if not api_key:
     raise ValueError("GOOGLE_API_KEY not found. Please set it in your .env file.")
 genai.configure(api_key=api_key)
 
-# Initialize the Flask application
 app = Flask(__name__)
 
 # --- CONFIGURATION ---
 TEXT_FILES_DIR = "MEDICINE_TXT"
-# Allowed file extensions for image uploads
-ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif'}
+ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif', 'bmp', 'webp'}
 
-# Get a list of available knowledge base files
+# Check for available knowledge base files
 try:
+    os.makedirs(TEXT_FILES_DIR, exist_ok=True)
     AVAILABLE_FILES = [f for f in os.listdir(TEXT_FILES_DIR) if f.endswith('.txt')]
     if not AVAILABLE_FILES:
-        raise FileNotFoundError("No .txt files found in the 'Text_Files' directory.")
-except FileNotFoundError:
-    print("Warning: 'Text_Files' directory not found. The API will not have a knowledge base.")
+        print(f"Warning: No .txt files found in '{TEXT_FILES_DIR}'. Running without knowledge base.")
+except Exception as e:
+    print(f"Warning: Error accessing '{TEXT_FILES_DIR}': {e}")
     AVAILABLE_FILES = []
 
 # --- HELPER FUNCTIONS ---
-
 def allowed_file(filename):
-    return '.' in filename and \
-           filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
+    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
 
-def find_relevant_file(topic: str) -> str | None:
-    """
-    Uses Gemini to determine the most relevant file for a given topic.
-    This is more robust than simple keyword matching.
-    """
+def find_relevant_file(topic: str) -> str:
+    """Find the most relevant file for a given topic using Gemini"""
     if not AVAILABLE_FILES:
         return None
         
     try:
-        model = genai.GenerativeModel('gemini-2.5-flash-lite')
+        model = genai.GenerativeModel('gemini-2.0-flash-exp')
         prompt = f"""
-        From the following list of files, which one is the most relevant for a query about "{topic}"?
-        Respond with only the single, most relevant filename dont include any other text.
+        From the following list of files, which one is most relevant for: "{topic}"?
+        Respond with ONLY the filename, nothing else.
 
-        File List:
-        {', '.join(AVAILABLE_FILES)}
+        Files: {', '.join(AVAILABLE_FILES)}
         """
         response = model.generate_content(prompt)
-        # Clean up the response to get just the filename
-        filename = response.text.strip().replace("`", "")
+        filename = response.text.strip().replace("`", "").replace('"', '')
         
         if filename in AVAILABLE_FILES:
-            print(f"Gemini identified relevant file: {filename} for topic: {topic}")
+            print(f"Found relevant file: {filename} for topic: {topic}")
             return filename
         else:
-            print(f"Warning: Gemini suggested a file that doesn't exist: {filename}")
+            print(f"No matching file found for: {topic}")
             return None
     except Exception as e:
-        print(f"Error in find_relevant_file: {e}")
+        print(f"Error finding relevant file: {e}")
         return None
 
-def get_context_from_file(filename: str) -> str | None:
-    """Reads and returns the content of a specified text file."""
+def get_context_from_file(filename: str) -> str:
+    """Read content from a text file"""
+    if not filename:
+        return None
     filepath = os.path.join(TEXT_FILES_DIR, filename)
     try:
         with open(filepath, 'r', encoding='utf-8') as f:
-            return f.read()
-    except FileNotFoundError:
+            content = f.read()
+            print(f"Successfully loaded context from {filename}")
+            return content
+    except Exception as e:
+        print(f"Error reading file {filename}: {e}")
         return None
 
-# --- CORE API LOGIC ---
-
+# --- MAIN API ENDPOINT ---
 @app.route('/api/query', methods=['POST'])
 def handle_query():
     """
-    Main API endpoint to handle user queries.
-    Accepts JSON data with 'query' (required) and 'file' (optional image upload).
+    Handle medicine queries with optional image upload
+    Accepts both JSON and FormData
     """
-    # 1. Get and validate the request data
-    if not request.is_json:
-        return jsonify({"error": "Request must be in JSON format with 'query' field"}), 400
     
-    data = request.get_json()
-    user_query = data.get('query')
+    # FIXED: Handle both JSON and FormData
+    user_query = None
+    medicine_topic = None
+    
+    # Check if request is JSON
+    if request.is_json:
+        data = request.get_json()
+        user_query = data.get('main_query')
+    else:
+        # Handle FormData (when image is uploaded)
+        user_query = request.form.get('main_query')
     
     if not user_query:
-        return jsonify({"error": "Missing 'query' in request"}), 400
+        return jsonify({"error": "Missing 'main_query' in request"}), 400
     
-    # 2. Handle File Upload (if provided)
+    print(f"Received query: {user_query}")
+    
+    # Handle image upload if present
     if 'file' in request.files:
         file = request.files['file']
-        if file.filename == '':
-            return jsonify({"error": "No selected file"}), 400
-            
-        if file and allowed_file(file.filename):
+        
+        if file and file.filename != '' and allowed_file(file.filename):
             try:
-                print("Image file received. Identifying medicine from image...")
-                # Read the uploaded file directly
+                print(f"Processing uploaded image: {file.filename}")
+                
+                # Process image with Gemini Vision
                 img = Image.open(file.stream)
+                vision_model = genai.GenerativeModel('gemini-2.0-flash-exp')
                 
-                # Use the vision model to identify the medicine
-                vision_model = genai.GenerativeModel('gemini-2.5-flash')
-                prompt = ["""Identify the specific formula or Rx or medicine name or  primary subject from this image.""", img]
-                response = vision_model.generate_content(prompt)
+                vision_prompt = [
+                    """Identify the medicine from this image. Look for:
+                    - Medicine name or brand
+                    - Active ingredients
+                    - Rx number or formula
+                    - Any text on packaging or pills
+                    Respond with just the medicine name or main component.""",
+                    img
+                ]
                 
+                response = vision_model.generate_content(vision_prompt)
                 medicine_topic = response.text.strip()
-                print(f"Medicine identified from image: {medicine_topic}")
-
+                print(f"Identified from image: {medicine_topic}")
+                
             except Exception as e:
                 print(f"Error processing image: {e}")
-                return jsonify({"error": "Failed to process the uploaded image."}), 500
-        else:
-            return jsonify({"error": f"Invalid file type. Allowed types: {', '.join(ALLOWED_EXTENSIONS)}"}), 400
+                # Continue without image data
     
-    # 3. Handle Text-Only Input (or use the topic identified from the image)
-    medicine_topic = None
+    # If no medicine identified from image, extract from query text
     if not medicine_topic:
-        print("No image provided. Identifying topic from text query...")
         try:
-            model = genai.GenerativeModel('gemini-2.5-flash')
-            prompt = f"""
-            From the user query '{user_query}', identify the main medicine or medical topic.
-            Respond with only the name of the topic or medicine (e.g., 'Ibuprofen', 'Antacids', 'Cough Suppressants').
-
+            model = genai.GenerativeModel('gemini-2.0-flash-exp')
+            extract_prompt = f"""
+            From this query: "{user_query}"
+            Extract the main medicine or medical topic being asked about.
+            Respond with ONLY the medicine/topic name (e.g., 'Paracetamol', 'Antibiotics')
             """
-            response = model.generate_content(prompt)
+            response = model.generate_content(extract_prompt)
             medicine_topic = response.text.strip()
-            print(f"Topic identified from query: {medicine_topic}")
+            print(f"Extracted topic from query: {medicine_topic}")
         except Exception as e:
-            print(f"Error identifying topic from query: {e}")
-            return jsonify({"error": "Failed to understand the query topic."}), 500
-
-    # 4. Find the Relevant Knowledge Base File
-    relevant_filename = find_relevant_file(medicine_topic)
-    if not relevant_filename:
-        return jsonify({"error": f"Could not find a relevant information file for '{medicine_topic}'."}), 404
-
-    # 5. Get the Context from the File
-    context = get_context_from_file(relevant_filename)
-    if not context:
-        return jsonify({"error": "Failed to read the content of the relevant file."}), 500
-
-    # 6. Generate the Final Response Using the Context
+            print(f"Error extracting topic: {e}")
+            medicine_topic = "general medicine"
+    
+    # Find relevant knowledge base file
+    context = None
+    source_file = None
+    
+    if AVAILABLE_FILES:
+        relevant_file = find_relevant_file(medicine_topic)
+        if relevant_file:
+            context = get_context_from_file(relevant_file)
+            source_file = relevant_file
+    
+    # Generate response
     try:
-        model = genai.GenerativeModel('gemini-2.5-flash-lite')
-        final_prompt = f"""
-        You are a helpful medical information assistant.
-        Your task is to answer the user's question based ONLY on the provided context from the guide.
-        Generate response in same language as user query.
-        If there have no information about any medicine then prepare response using given context and your knowlage base make sure there have satisfied answer.
-        if there have any relevent medicine of provided medicine in context then prepare answer using that context.
-        Answer should be in simple language and short not more than 200 words.
-        If the answer cannot be found in the provided context, then you have to prepare response using your knowlage base make sure there have satisfied answer.
-        ---important---
-        Dont tell user to i have no information about that medicine. inplace of that prepare answer using given context and your knowlage base make sure there have satisfied answer.
-        user is also provide the medicine name and description of the medicine.
-        name:{medicine_topic}
-        ---important---
+        model = genai.GenerativeModel('gemini-2.0-flash-exp')
         
-        --- CONTEXT FROM THE GUIDE ---
-        {context}
-        --- END OF CONTEXT ---
-
-        USER'S QUESTION: {user_query}
-
-        YOUR ANSWER:
-        """
+        # Build prompt based on available context
+        if context:
+            final_prompt = f"""
+            You are a medical information assistant.
+            
+            CONTEXT FROM KNOWLEDGE BASE:
+            {context}
+            
+            IDENTIFIED MEDICINE/TOPIC: {medicine_topic}
+            
+            USER QUESTION: {user_query}
+            
+            Instructions:
+            - Answer based on the context if available
+            - Use simple language
+            - Keep response under 200 words
+            - Include dosage, usage, and warnings if relevant
+            - If context doesn't have the info, use general knowledge
+            - Respond in the same language as the user query
+            """
+        else:
+            final_prompt = f"""
+            You are a medical information assistant.
+            
+            MEDICINE/TOPIC: {medicine_topic}
+            USER QUESTION: {user_query}
+            
+            Provide accurate medical information about this topic.
+            - Use simple language
+            - Keep response under 200 words
+            - Include dosage, usage, side effects if relevant
+            - Add standard medical disclaimers
+            - Respond in the same language as the user query
+            """
         
-        final_response = model.generate_content(final_prompt)
+        response = model.generate_content(final_prompt)
         
-        # 7. Return the final, context-aware response
         return jsonify({
-            "response": final_response.text.strip(),
+            "status": "success",
+            "response": response.text.strip(),
             "identified_topic": medicine_topic,
-            "source_file": relevant_filename
+            "source_file": source_file if source_file else "general_knowledge",
+            "knowledge_base_available": len(AVAILABLE_FILES) > 0
         })
-
+        
     except Exception as e:
-        print(f"Error generating final response: {e}")
-        return jsonify({"error": "An error occurred while generating the response."}), 500
+        print(f"Error generating response: {e}")
+        return jsonify({
+            "error": "Failed to generate response",
+            "details": str(e)
+        }), 500
+
+@app.route('/health', methods=['GET'])
+def health_check():
+    """Health check endpoint"""
+    return jsonify({
+        "status": "running",
+        "service": "medicine_info",
+        "knowledge_base_files": len(AVAILABLE_FILES),
+        "port": 5002
+    })
+
+@app.route('/', methods=['GET'])
+def index():
+    """Basic info endpoint"""
+    return jsonify({
+        "service": "Medicine Information API",
+        "endpoint": "/api/query",
+        "methods": ["POST"],
+        "accepts": "JSON or FormData with 'main_query' and optional 'file'",
+        "knowledge_base": f"{len(AVAILABLE_FILES)} files available"
+    })
 
 if __name__ == '__main__':
-    # Runs the Flask server
+    print("="*50)
+    print("Starting Medicine Information Service")
+    print(f"Knowledge base: {len(AVAILABLE_FILES)} files in '{TEXT_FILES_DIR}'")
+    if AVAILABLE_FILES:
+        print(f"Available files: {', '.join(AVAILABLE_FILES[:5])}")
+    print("="*50)
     app.run(host='0.0.0.0', port=5002, debug=True)