medicine.py

# medicine.py - FIXED VERSION

import os
import io
from flask import Flask, request, jsonify
from PIL import Image
from dotenv import load_dotenv
import google.generativeai as genai
import json

# --- INITIAL SETUP ---
load_dotenv()

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)

app = Flask(__name__)

# --- CONFIGURATION ---
TEXT_FILES_DIR = "MEDICINE_TXT"
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif', 'bmp', 'webp'}

# 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:
        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

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(os.getenv("MODEL","gemini-2.5-flash-lite"))
        prompt = f"""
        From the following list of files, which one is most relevant for: "{topic}"?
        Respond with ONLY the filename, nothing else.

        Files: {', '.join(AVAILABLE_FILES)}
        """
        response = model.generate_content(prompt)
        filename = response.text.strip().replace("`", "").replace('"', '')
        
        if filename in AVAILABLE_FILES:
            print(f"Found relevant file: {filename} for topic: {topic}")
            return filename
        else:
            print(f"No matching file found for: {topic}")
            return None
    except Exception as e:
        print(f"Error finding relevant file: {e}")
        return None

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:
            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

# --- MAIN API ENDPOINT ---
@app.route('/api/query', methods=['POST'])
def handle_query():
    """
    Handle medicine queries with optional image upload
    Accepts both JSON and FormData
    """
    
    # 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 'main_query' in request"}), 400
    
    print(f"Received query: {user_query}")
    
    # Handle image upload if present
    if 'file' in request.files:
        file = request.files['file']
        
        if file and file.filename != '' and allowed_file(file.filename):
            try:
                print(f"Processing uploaded image: {file.filename}")
                
                # Process image with Gemini Vision
                img = Image.open(file.stream)
                vision_model = genai.GenerativeModel(os.getenv("MODEL","gemini-2.5-flash-lite"))
                
                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"Identified from image: {medicine_topic}")
                
            except Exception as e:
                print(f"Error processing image: {e}")
                # Continue without image data
    
    # If no medicine identified from image, extract from query text
    if not medicine_topic:
        try:
            model = genai.GenerativeModel(os.getenv("MODEL","gemini-2.5-flash-lite"))
            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(extract_prompt)
            medicine_topic = response.text.strip()
            print(f"Extracted topic from query: {medicine_topic}")
        except Exception as e:
            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.0-flash-exp')
        
        # 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
            """
        
        response = model.generate_content(final_prompt)
        
        return jsonify({
            "status": "success",
            "response": response.text.strip(),
            "identified_topic": medicine_topic,
            "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 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__':
    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)